CUDA extension with `TORCH_CUDABLAS_CHECK` throws undefined symbol error at import

[xwang233]

🐛 Bug

CUDA extension with TORCH_CUDABLAS_CHECK throws undefined symbol error

To Reproduce

Prepare two files to build a CUDA extension

• cuda_ext.cpp

#include <iostream>

#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>

#include <cublas_v2.h>

void a_cublas_function() {
    printf("hello world\n");

    cublasHandle_t handle;
    TORCH_CUDABLAS_CHECK(cublasCreate(&handle));

    TORCH_CUDABLAS_CHECK(cublasDestroy(handle));
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("a_cublas_function", &a_cublas_function, "a cublas function");
}

• setup.py

from setuptools import setup, Extension
from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CppExtension

cublas_module = CUDAExtension(
            name='cublas_ext',
            sources=['cublas_ext.cpp']
        )

setup(
    name='cublas_ext_root',
    version='0.1',
    ext_modules=[cublas_module],
    cmdclass={
        'build_ext': BuildExtension.with_options(use_ninja=False)
    }
)

build cuda extension with

pip install -v --no-cache-dir .   

run test with

python -c 'import torch; import cublas_ext; cublas_ext.a_cublas_function();'

got error message

Traceback (most recent call last):
  File "<string>", line 1, in <module>
ImportError: /home/xwang/.local/lib/python3.9/site-packages/cublas_ext.cpython-39-x86_64-linux-gnu.so: undefined symbol: _ZN2at4cuda4blas19_cublasGetErrorEnumE14cublasStatus_t

Expected behavior

No error for TORCH_CUDABLAS_CHECK in cuda extensions.

Environment

pytorch is source build using gcc 10.3 from latest master commit

Collecting environment information...
PyTorch version: 1.11.0a0+gitf56a1a5
Is debug build: False
CUDA used to build PyTorch: 11.4
ROCM used to build PyTorch: N/A

OS: Manjaro Linux (x86_64)
GCC version: (GCC) 11.1.0
Clang version: Could not collect
CMake version: version 3.21.1
Libc version: glibc-2.33

Python version: 3.9.6 (default, Jun 30 2021, 10:22:16)  [GCC 11.1.0] (64-bit runtime)
Python platform: Linux-5.10.60-1-MANJARO-x86_64-with-glibc2.33
Is CUDA available: True
CUDA runtime version: 11.4.100
GPU models and configuration: 
GPU 0: NVIDIA GeForce RTX 2070 SUPER
GPU 1: NVIDIA GeForce GTX 1070 Ti

Nvidia driver version: 470.63.01
cuDNN version: Probably one of the following:
/usr/lib/libcudnn.so.8.2.2
/usr/lib/libcudnn_adv_infer.so.8.2.2
/usr/lib/libcudnn_adv_train.so.8.2.2
/usr/lib/libcudnn_cnn_infer.so.8.2.2
/usr/lib/libcudnn_cnn_train.so.8.2.2
/usr/lib/libcudnn_ops_infer.so.8.2.2
/usr/lib/libcudnn_ops_train.so.8.2.2
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] mypy==0.812
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.19.5
[pip3] torch==1.11.0a0+gitf56a1a5
[pip3] torch-tb-profiler==0.2.0
[conda] Could not collect

Additional context

N/A

cc @malfet @zou3519 @crcrpar @ptrblck @ngimel

[ptrblck]

Could this be failing because TORCH_CUDABLAS_CHECK might be missing the TORCH_CUDA_CPP_API export?

[crcrpar]

TORCH_CUSOLVER_CHECK and TORCH_CUDASPARSE_CHECK have the same issue.
I could fix the former, cusolver check by exporting one internal function, but not the latter; I encountered ImportError: dynamic module does not define module export function on my workstation

EDIT: I saw different behaviors of TORCH_CUDASPARSE_CHECK in different environments. Let me revisit this one later.

[byronyi]

Does this affect version 1.10?

[crcrpar]

@byronyi it seems like the macro has not been available since the beginning (of course, including 1.10) in custom extensions.

[crcrpar]

Retried TORCH_CUDASPARSE_CHECK in a custom extension and it worked without any changes.

The code used are as follows:

#include <iostream>

#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>

#include <cusparse.h>

void a_custom_function() {
    printf("hello world\n");
    cusparseHandle_t sparse_handle;
    TORCH_CUDASPARSE_CHECK(cusparseCreate(&sparse_handle));
    TORCH_CUDASPARSE_CHECK(cusparseDestroy(sparse_handle));
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("a_custom_function", &a_custom_function, "a custom function");
}

# setup.py
from setuptools import setup, Extension
from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CppExtension

cublas_module = CUDAExtension(
            name='cuda_ext',
            sources=['cuda_ext.cpp']
        )

setup(
    name='cuda_ext_root',
    version='0.1',
    ext_modules=[cublas_module],
    cmdclass={
        'build_ext': BuildExtension.with_options(use_ninja=False)
    }
)

my environment is:

PyTorch version: 1.11.0a0+gitcb2a41e
Is debug build: False
CUDA used to build PyTorch: 11.4
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.2 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: Could not collect
CMake version: version 3.19.6
Libc version: glibc-2.31

Python version: 3.8.10 (default, Jun  4 2021, 15:09:15)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.8.0-55-generic-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: 11.4.100
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3080
Nvidia driver version: 470.63.01
cuDNN version: Probably one of the following:
/usr/local/cuda-11.4/targets/x86_64-linux/lib/libcudnn.so.8.2.2
/usr/local/cuda-11.4/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.2.2
/usr/local/cuda-11.4/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.2.2
/usr/local/cuda-11.4/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.2.2
/usr/local/cuda-11.4/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.2.2
/usr/local/cuda-11.4/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.2.2
/usr/local/cuda-11.4/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.2.2
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.20.2
[pip3] pytorch-lightning==1.5.0.dev0
[pip3] torch==1.11.0a0+gitcb2a41e
[pip3] torchmetrics==0.5.1
[conda] blas                      1.0                         mkl  
[conda] magma-cuda111             2.5.2                         1    pytorch
[conda] mkl                       2021.2.0           h06a4308_296  
[conda] mkl-include               2021.2.0           h06a4308_296  
[conda] mkl-service               2.3.0            py38h27cfd23_1  
[conda] mkl_fft                   1.3.0            py38h42c9631_2  
[conda] mkl_random                1.2.1            py38ha9443f7_2  
[conda] mypy_extensions           0.4.3                    py38_0  
[conda] numpy                     1.20.2           py38h2d18471_0  
[conda] numpy-base                1.20.2           py38hfae3a4d_0  
[conda] pytorch-lightning         1.5.0.dev0               pypi_0    pypi
[conda] torch                     1.11.0a0+gitcb2a41e           dev_0    <develop>
[conda] torchmetrics              0.5.1                    pypi_0    pypi