nvrtc_runtime_cuda_kernel_compilation_to_ptx_with_driver_api.py

1import cuda.cuda as cuda
2import cuda.nvrtc as nvrtc
3import numpy as np
4
5def check_cuda_error(res):
6    if isinstance(res, cuda.CUresult):
7        if res != cuda.CUresult.CUDA_SUCCESS:
8            raise RuntimeError(f"CUDA Error: {res}")
9    elif isinstance(res, nvrtc.nvrtcResult):
10        if res != nvrtc.nvrtcResult.NVRTC_SUCCESS:
11            raise RuntimeError(f"NVRTC Error: {res}")
12
13# Kernel source code
14saxpy = """\
15extern "C" __global__
16void saxpy(float a, float *x, float *y, int n)
17{
18    int i = blockIdx.x * blockDim.x + threadIdx.x;
19    if (i < n) y[i] = a * x[i] + y[i];
20}
21"""
22
23# 1. Compile kernel to PTX with NVRTC
24err, program = nvrtc.nvrtcCreateProgram(str.encode(saxpy), b"saxpy.cu", 0, [], [])
25check_cuda_error(err)
26
27err, = nvrtc.nvrtcCompileProgram(program, 0, [])
28check_cuda_error(err)
29
30err, ptx_size = nvrtc.nvrtcGetPTXSize(program)
31check_cuda_error(err)
32
33ptx = b" " * ptx_size
34err, = nvrtc.nvrtcGetPTX(program, ptx)
35check_cuda_error(err)
36
37# 2. Initialize CUDA Driver API
38err, = cuda.cuInit(0)
39check_cuda_error(err)
40
41err, device = cuda.cuDeviceGet(0)
42check_cuda_error(err)
43
44err, context = cuda.cuCtxCreate(0, device)
45check_cuda_error(err)
46
47# 3. Load PTX and setup data
48err, module = cuda.cuModuleLoadData(np.char.array(ptx))
49check_cuda_error(err)
50
51err, kernel = cuda.cuModuleGetFunction(module, b"saxpy")
52check_cuda_error(err)
53
54n = 1024
55a = np.array([2.0], dtype=np.float32)
56host_x = np.arange(n, dtype=np.float32)
57host_y = np.arange(n, dtype=np.float32)
58
59err, device_x = cuda.cuMemAlloc(host_x.nbytes)
60check_cuda_error(err)
61err, device_y = cuda.cuMemAlloc(host_y.nbytes)
62check_cuda_error(err)
63
64err, = cuda.cuMemcpyHtoD(device_x, host_x, host_x.nbytes)
65check_cuda_error(err)
66err, = cuda.cuMemcpyHtoD(device_y, host_y, host_y.nbytes)
67check_cuda_error(err)
68
69# 4. Launch Kernel
70args = [a, device_x, device_y, np.int32(n)]
71arg_types = [None, None, None, None] # Use default size-based inference
72err, = cuda.cuLaunchKernel(kernel, 
73                           n // 64, 1, 1,  # grid dim
74                           64, 1, 1,       # block dim
75                           0, None,        # shared mem and stream
76                           args, 0)        # arguments
77check_cuda_error(err)
78
79# 5. Retrieve result
80err, = cuda.cuMemcpyDtoH(host_y, device_y, host_y.nbytes)
81check_cuda_error(err)
82
83print(f"Result (first 5): {host_y[:5]}")
84
85# Cleanup
86cuda.cuMemFree(device_x)
87cuda.cuMemFree(device_y)
88cuda.cuCtxDestroy(context)
89nvrtc.nvrtcDestroyProgram(program)