I am writing my own CUDA implementation of a CNN and to do so, I’m compiling it as a CPP extension file, using Ninja. All the base CUDA code (e.g., malloc, memcopy, kernel launch etc) works fine, but I am having an issue using CUDNN in my custom extension.
I’m running PyTorch 1.10 using a Conda environment. PyTorch is using CUDNN, ( torch.backends.cudnn.enabled == True and torch.backends.cudnn.version() returns 8200).
I also have CUDNN installed in my /usr/local/ folder, since I think this is CUDA code that Ninja uses?
When I try to call CUDNN functions, it builds fine, but when I call it from the Python script, I get an error:
undefined symbol: cudnnCreateTensorDescriptor
<cudnn.h> is included in my .cu file. Also, when I compile this code using nvcc directly (and I pass -lcudnn flag), it compiles and runs fine. But if I do the same in my setup.py file when I build the extension, it still gives me the same error.
Specifically, I added the extra_compile_args line to my setup.py file to build the extension:
CUDAExtension(
              name='<cuda code>',
              sources=['<cuda file>.cu'],
              extra_compile_args={'nvcc': ['-lcudnn']}
          )

I also confirmed in my build log that -lcudnn is being passed to nvcc. But I still get the same error.
Is there something else I need to do to get nvcc to link the cudnn library in when it builds the extension?