test.py

import pycuda.autoinit
import pycuda.driver as cuda
from pycuda.compiler import SourceModule
import numpy as np
from PIL import Image

width, height = 512, 512

# CUDA kernel
mod = SourceModule("""
__global__ void render(unsigned char *img, int w, int h)
{
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    int y = blockIdx.y * blockDim.y + threadIdx.y;

    if (x >= w || y >= h) return;

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

    img[idx + 0] = (unsigned char)(255.0f * x / w); // R
    img[idx + 1] = (unsigned char)(255.0f * y / h); // G
    img[idx + 2] = 128;                             // B
}
""")

render = mod.get_function("render")

# Allocate GPU memory
img = np.zeros((height, width, 3), dtype=np.uint8)
img_gpu = cuda.mem_alloc(img.nbytes)

# Launch kernel
block = (16, 16, 1)
grid = ((width + 15) // 16, (height + 15) // 16)

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, "RGB").save("cuda_render.png")

print("Image rendered on GPU and saved as cuda_render.png")