render.py

import numpy as np
from PIL import Image

import pycuda.autoinit
import pycuda.driver as cuda
from pycuda.compiler import SourceModule

# Image size
WIDTH = 800
HEIGHT = 600

cuda_code = r"""
__global__ void render(unsigned char *img, int width, int height)
{
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    int y = blockIdx.y * blockDim.y + threadIdx.y;

    if (x >= width || y >= height) return;

    int idx = (y * width + x) * 3;

    // Normalized screen coordinates
    float u = (x / (float)width) * 2.0f - 1.0f;
    float v = (y / (float)height) * 2.0f - 1.0f;
    v *= height / (float)width;

    // Ray origin and direction
    float3 ro = make_float3(0, 0, -3);
    float3 rd = normalize(make_float3(u, v, 1));

    // Sphere
    float3 center = make_float3(0, 0, 0);
    float radius = 1.0f;

    float3 oc = ro - center;
    float b = dot(oc, rd);
    float c = dot(oc, oc) - radius * radius;
    float h = b*b - c;

    if (h > 0.0f) {
        h = sqrtf(h);
        float t = -b - h;
        float3 p = ro + rd * t;
        float3 n = normalize(p - center);

        // Light
        float3 light = normalize(make_float3(1, 1, -1));
        float diff = fmaxf(dot(n, light), 0.0f);

        img[idx + 0] = (unsigned char)(diff * 255);
        img[idx + 1] = (unsigned char)(diff * 180);
        img[idx + 2] = (unsigned char)(diff * 120);
    } else {
        img[idx + 0] = 20;
        img[idx + 1] = 20;
        img[idx + 2] = 40;
    }
}
"""

# Compile CUDA
mod = SourceModule(cuda_code)
render = mod.get_function("render")

# Output buffer
img = np.zeros((HEIGHT, WIDTH, 3), dtype=np.uint8)
img_gpu = cuda.mem_alloc(img.nbytes)

# Launch kernel
block = (16, 16, 1)
grid = (
    (WIDTH + block[0] - 1) // block[0],
    (HEIGHT + block[1] - 1) // block[1],
    1,
)

render(
    img_gpu,
    np.int32(WIDTH),
    np.int32(HEIGHT),
    block=block,
    grid=grid,
)

# Copy back
cuda.memcpy_dtoh(img, img_gpu)

# Save image
Image.fromarray(img).save("cuda_render.png")
print("Rendered image saved as cuda_render.png")