Extending Numba for CUDA inside Cppyy

Numba Support: How do I do ...?
7

Numba has a built-in (albeit slightly limited) interface to NVRTC which you might be able to use to compile CUDA C++ to PTX, which you can then add to the link - here’s a simple example:

from numba.cuda.cudadrv import nvrtc

code = '''
extern "C" __device__ int foo(int* return_value){
  *return_value = 42;
  return 0;
}
'''

# Returns a tuple of the PTX code and the compiler output messages
# Arguments are (<source code>, 
#                <source filename>, 
#                <compute capability as a tuple>)
ptx, log = nvrtc.compile(code, "test.cu", (8, 9))

print(ptx)

which prints (slightly abridged):

//
// Generated by NVIDIA NVVM Compiler
//
// Compiler Build ID: CL-34177558
// Cuda compilation tools, release 12.5, V12.5.40
// Based on NVVM 7.0.1
//

.version 8.5
.target sm_89
.address_size 64

	// .globl	foo

.visible .func  (.param .b32 func_retval0) foo(
	.param .b64 foo_param_0
)
{
	.reg .b32 	%r<3>;
	.reg .b64 	%rd<2>;

	ld.param.u64 	%rd1, [foo_param_0];
	mov.u32 	%r1, 42;
	st.u32 	[%rd1], %r1;
	mov.u32 	%r2, 0;
	st.param.b32 	[func_retval0+0], %r2;
	ret;
}

There are a couple of points to be aware of:

  • Arbitrary function pointers are not supported in CUDA, because of the way register allocation works (kernels can use different numbers of registers, but their register usage needs to be declared and is static for the kernel, so a device function that uses some unknown number of registers is not possible)
  • Therefore, in cppyy you will need an alternative mechanism to function pointers to call the C++ functions. The existing built-in CUDA C/C++ support does this with the declare_device() function to insert Numba typing and lowering that implements the call using the name of the C/C++ device function - you may be able to do the same in the cppyy implementation.

If you need to be able to have more flexibility in the way that NVRTC is called (e.g. using arbitrary flags, include paths, etc) then you can use the CUDA Python bindings to access it more directly:

1 Like