CUDA Extension error (an illegal memory access was encountered)

Bokyeong1001: <p>Hello</p>
<p>I am trying to build matrix multiplication cuda extension using cublas.</p>
<p>Here is my code</p>
<pre><code class="lang-auto">#include &lt;torch/extension.h&gt;

void matmul_cuda_forward(
    float* A,
    float* B,
    float* C,
    int N);

void matmul_forward(
    float* A,
    float* B,
    float* C,
    int N) {

  matmul_cuda_forward(A, B, C, N);
}


PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("forward", &amp;matmul_forward, "matmul forward (CUDA)");
}
</code></pre>
<pre><code class="lang-auto">#include &lt;torch/extension.h&gt;

#include &lt;cuda_runtime_api.h&gt; 
#include &lt;cublas_v2.h&gt;


void matmul_cuda_forward(
  float* A,
  float* B,
  float* C,
  int N) {

    int lda=N,ldb=N,ldc=N;
    const float alf = 1;
    const float bet = 0;
    const float *alpha = &amp;alf;
    const float *beta = &amp;bet;
 
    // Create a handle for CUBLAS
    cublasHandle_t handle;
    cublasCreate(&amp;handle);
 
    // Do the actual multiplication
    cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, N, N, N, alpha, A, lda, B, ldb, beta, C, ldc);
 
    // Destroy the handle
    cublasDestroy(handle);

}
</code></pre>
<pre><code class="lang-auto">from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension

setup(
    name='matmul_cuda',
    ext_modules=[
        CUDAExtension('matmul_cuda', ['matmul_cuda.cpp', 'matmul_cublas_kernel.cu'],
        extra_compile_args={'cxx':[], 'nvcc':['-arch=sm_80', '-lcublas']})
    ],
    cmdclass={
        'build_ext': BuildExtension
    })
</code></pre>
<pre><code class="lang-auto">import math
from torch import nn
from torch.autograd import Function
import torch

import matmul_cuda

torch.manual_seed(42)
device = torch.device("cuda")

N = 4

A = torch.zeros((N*N),dtype=torch.float32).to(device)
B = torch.zeros((N*N),dtype=torch.float32).to(device)
C = torch.zeros((N*N),dtype=torch.float32).to(device)

A_p = A.contiguous().data_ptr()
B_p = B.contiguous().data_ptr()
C_p = C.contiguous().data_ptr()

for i in range(N*N):
  A[i] = (i%4) + 1
  B[i] = (i%4) + 5

matmul_cuda.forward(A_p, B_p, C_p, N)

print(C)
</code></pre>
<p>It builds cuda extension successfully but when I run the python code<br>
I got this error<br>
RuntimeError: CUDA error: an illegal memory access was encountered<br>
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.<br>
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.</p>
<p>Can anyone tell me what is wrong with this code?<br>
Does torch cuda extension cannot get the float* value?</p>
<p>Thank you</p>

eqy: <p>I’m not sure about the convention of a CUDA extension taking raw pointers via <code>.data_ptr()</code>, but it sounds like it could work in theory. Alternatively you could take a look at e.g., the APEX normalization extension and the convention for passing tensors there: <a href="https://github.com/NVIDIA/apex/blob/master/csrc/layer_norm_cuda.cpp" class="inline-onebox" rel="noopener nofollow ugc">apex/layer_norm_cuda.cpp at master · NVIDIA/apex · GitHub</a></p>
<p>If the pointer passing is not the issue, I would try the following debugging steps:</p>
<ol>
<li>
<p>Check the status of each cublas API call to make sure that you are getting status success (0) before continuing to the next API call.</p>
</li>
<li>
<p>Consider running your script with <code>compute-sanitizer</code> which may help you locate where the IMA is occurring.</p>
</li>
<li>
<p>Try removing operations from your extension to isolate the IMA.</p>
</li>
</ol>