nvrtc_cuda_kernel_string_to_ptx_compilation.py

1import numpy as np
2from cuda import nvrtc
3
4def check_nvrtc_errors(result):
5    if result.value != nvrtc.nvrtcResult.NVRTC_SUCCESS:
6        raise RuntimeError(f"NVRTC Error: {result}")
7
8# 1. Define the CUDA kernel as a string
9sak_kernel = """\
10extern "C" __global__
11void saxpy(float a, float *x, float *y, int n)
12{
13    int i = blockIdx.x * blockDim.x + threadIdx.x;
14    if (i < n) y[i] = a * x[i] + y[i];
15}
16"""
17
18# 2. Create a program with NVRTC
19# Arguments: (program_name, source_code, num_headers, headers, include_names)
20err, prog = nvrtc.nvrtcCreateProgram(
21    str.encode(sak_kernel), 
22    b"saxpy.cu", 
23    0, [], []
24)
25check_nvrtc_errors(err)
26
27# 3. Compile the program
28# Options can include include paths, fast-math, etc.
29opts = [b"--gpu-architecture=compute_75"]
30err, = nvrtc.nvrtcCompileProgram(prog, len(opts), opts)
31
32# 4. Check compilation log (optional but recommended)
33err, log_size = nvrtc.nvrtcGetProgramLogSize(prog)
34check_nvrtc_errors(err)
35log = b" " * log_size
36err, = nvrtc.nvrtcGetProgramLog(prog, log)
37print(f"Compilation Log: {log.decode('utf-8')}")
38
39# 5. Extract the PTX
40err, ptx_size = nvrtc.nvrtcGetPTXSize(prog)
41check_nvrtc_errors(err)
42ptx = b" " * ptx_size
43err, = nvrtc.nvrtcGetPTX(prog, ptx)
44
45print("Successfully compiled CUDA kernel to PTX!")
46print(ptx.decode('utf-8')[:100] + "...")
47
48# 6. Cleanup
49err, = nvrtc.nvrtcDestroyProgram(prog)
50check_nvrtc_errors(err)