How can I use the function at::cuda::blas::gemm<float>()?

wangps: <p>I follow the official <a href="https://pytorch.org/tutorials/advanced/cpp_extension.html" rel="nofollow noopener">tutorial</a> to build custom CUDA extensions. And I would like to use the function <code>at::cuda::blas::gemm&lt;float&gt;()</code> to do the matrix product, which is defined in  <code>#include &lt;ATen/cuda/CUDABlas.h&gt;</code>. But the g++ compiler seems to fail to link this function according to current configurations. Could anyone give me some help?</p>
<p>Steps to reproduce the behavior with a toy example:</p>
<ol>
<li>
<p>The cpp file <code>pytorch_gemm_gpu.cpp</code>:</p>
<pre><code class="lang-cpp"> #include &lt;ATen/cuda/CUDABlas.h&gt;
 #include &lt;torch/extension.h&gt;

 torch::Tensor gemm(torch::Tensor A, torch::Tensor B) {
   int64_t N = A.size(0);
   torch::Tensor C = torch::zeros_like(A);
   at::cuda::blas::gemm&lt;float&gt;('n', 'n',                       // transa, transb
                                N, N, N,                       // M, N, K
                                1.0f, A.data_ptr&lt;float&gt;(), N,  // alpha, a, lda
                                B.data_ptr&lt;float&gt;(), N, 0.0f,  // b, ldb, beta
                                C.data_ptr&lt;float&gt;(), N);       // c, ldc
 }

 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 
   m.def("gemm", &amp;gemm, "gemm"); 
 }
</code></pre>
<p>The setup file <code>setup.py</code>:</p>
<pre><code class="lang-python"> from setuptools import setup
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension

 setup(
     name='pytorch_gemm_gpu',
     ext_modules=[
         CUDAExtension('pytorch_gemm_gpu', ['pytorch_gemm_gpu.cpp']),
     ],
     cmdclass={
         'build_ext': BuildExtension
     })
</code></pre>
</li>
<li>
<p>Run the following command to install the package. This procedure finished without errors.</p>
<pre><code class="lang-auto">python setup.py install
</code></pre>
</li>
<li>
<p>Run the following command to import the package.</p>
<pre><code class="lang-auto">python -c "import torch; import pytorch_gemm_gpu"
</code></pre>
<p>The error message is :</p>
<pre><code class="lang-auto">pytorch_gemm_gpu.cpython-37m-x86_64-linux-gnu.so:undefined symbol: _ZN2at4cuda4blas4gemmIfEEvcclllT_PKS3_lS5_lS3_PS3_l
</code></pre>
</li>
</ol>
<h2>Expected behavior</h2>
<p>I expect the command <code>python -c "import torch; import pytorch_gemm_gpu"</code> succeeds without errors.<br>
And I did find the symbol <code>_ZN2at4cuda4blas4gemmIfEEvcclllT_PKS3_lS5_lS3_PS3_l</code> in the <code>libtorch_cuda.so</code>.  I do not know why this symbol can not be found at runtime.</p>
<h2>Environment</h2>
<pre><code class="lang-auto">PyTorch version: 1.6.0
Is debug build: False
CUDA used to build PyTorch: 10.1

OS: Ubuntu 16.04.6 LTS (x86_64)
GCC version: (Ubuntu 5.4.0-6ubuntu1~16.04.12) 5.4.0 20160609
Clang version: Could not collect
CMake version: version 3.14.5

Python version: 3.7 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: 10.1.168
GPU models and configuration:
GPU 0: GeForce RTX 2080 Ti
GPU 1: GeForce RTX 2080 Ti

Nvidia driver version: 418.67
cuDNN version: Could not collect

Versions of relevant libraries:
[pip3] numpy==1.19.1
[pip3] pytorch-gemm-gpu==0.0.0
[pip3] torch==1.6.0
[pip3] torchvision==0.7.0
[conda] blas                      1.0                         mkl
[conda] cudatoolkit               10.1.243             h6bb024c_0
[conda] mkl                       2020.1                      217
[conda] mkl-service               2.3.0            py37he904b0f_0
[conda] mkl_fft                   1.1.0            py37h23d657b_0
[conda] mkl_random                1.1.1            py37h0da4684_0    conda-forge
[conda] numpy                     1.19.1           py37hbc911f0_0
[conda] numpy-base                1.19.1           py37hfa32c7d_0
[conda] pytorch                   1.6.0           py3.7_cuda10.1.243_cudnn7.6.3_0    pytorch
[conda] pytorch-gemm-gpu          0.0.0                    pypi_0    pypi
[conda] torchvision               0.7.0                py37_cu101    pytorch
</code></pre>

albanD: <p>Hi,</p>
<p>In the extensions, you should not try to access low level functions like this one directly.<br>
You can use <code>at::mm()</code> or <code>at::bmm()</code> directly with Tensor arguments to get that result (also these function will have all the nice autograd support while low level ones won’t).</p>

wangps: <p>Thank you so much for your reply!</p>

creiser: <p>Hello, for my use case it is absolutely nessecary to use this “low level” function <code>at::cuda::blas::gemm&lt;scalar_t&gt;</code> and I am getting the same linker error like OP. Otherwise I suppose I could also directly link against cuBLAS and take the original function from there, but it would be nicer to have access to this wrapper since it elegantly handles different floating types.</p>
<p>I have made sure that PyTorch and the extensions are built with the same compiler and ABI flag (which seems to be another source for this linker error)</p>

maktukmak: <p><a class="mention" href="/u/creiser">@creiser</a>, could you solve this problem? I need to do the exact thing like you but having the same error.</p>