How to use CUDA Dynamic Parallelism in PyTorch CPP extension?

[qinjian623]

I found discussions at discuss.pytorch.org . But there is still no solution now.

Here is the error message:

error: kernel launch from __device__ or __global__ functions requires separate compilation mode

and

error: a __device__ function call cannot be configured

Thanks.

cc @ngimel

[ajinkyaambatwar]

You can remove that error by using -rdc=True flag in nvcc. But that solves only the compile issue. I am facing the error
undefined symbol: __cudaRegisterLinkedBinary_50
which has something to do with device linking. More can be read here. I am not sure how to implement this in torch.utils.cpp_extension.load

[coreylammie]

I can confirm that this is still an issue in 2021. Defining the arguments extra_compile_args=['-rdc=true', '-arch=sm_61', '-lcudart', '-Xcompiler', '-lcudadevrt'] when instantiating torch.utils.cpp_extension.CUDAExtension raises the compilation error error LNK2001: unresolved external symbol __cudaRegisterLinkedBinary_54_.

The current CUDA toolkit documentation suggests using the following to compile cuda files with dynamic parallelism:

nvcc --gpu-architecture=sm_50 --device-c a.cu b.cu
nvcc --gpu-architecture=sm_50 --device-link a.o b.o --output-file link.o
nvcc --lib --output-file libgpu.a a.o b.o link.o
g++ host.o --library=gpu --library-path=<path> \
    --library=cudadevrt --library=cudart

The automatically generated build.ninja file (from PyTorch) contains the following rule for compiling cuda files:

rule cuda_compile
  command = $nvcc $cuda_cflags -c $in -o $out $cuda_post_cflags 

Perhaps a boolean input argument could be added to torch.utils.cpp_extension.CUDAExtension to indicate that dynamic parallelism is to be used, and the cuda_compile rule in build.ninja could be modified accordingly? cl can easily be used in-place of g++, and --gpu-architecture could be set using extra_compile_args for sake of modularity.

The cuda_compile rule appears to be defined/generated in lines 1916-1932 of torch.utils.cpp_extension.