Why didn't this cuda kernel implement?

Mon: <p>I’m new about cuda programming but I need to make some pytorch cuda programming now.<br>
So I looked <a href="https://pytorch.org/tutorials/advanced/cpp_extension.html" rel="nofollow noopener">the pytorch cuda extension tutorial</a> and had a try. First I’d like to try the <code>add</code> example, so I tried:</p>
<pre><code class="lang-cpp">// add_cuda.h
#ifndef _ADD_CUDA
#define _ADD_CUDA

#include &lt;torch/extension.h&gt;
void add(torch::Tensor a, torch::Tensor b, torch::Tensor c);
#endif
</code></pre>
<pre><code class="lang-cpp">// add_wrapper.cpp
#include "add_cuda.h"

#define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)

void add(torch::Tensor a, torch::Tensor b, torch::Tensor c){
    CHECK_INPUT(a);
    CHECK_INPUT(b);
    CHECK_INPUT(c);
    add_cu(a, b, c);
}

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

template &lt;typename scalar_t&gt;
__global__ void add_kernel(
    scalar_t* __restrict__ a, 
    scalar_t* __restrict__ b, 
    scalar_t* __restrict__ c, 
    size_t size
){
    const int index = blockIdx.x * blockDim.x + threadIdx.x;
    const int stride = blockIdx.x * gridDim.x;
    for (int i = index; i &lt; size; i += stride){
        a[i] = b[i] + c[i];
    }
}


void add_cu(torch::Tensor a, torch::Tensor b, torch::Tensor c){
    const auto size = a.size(0);

    const int threads = 8;
    const dim3 blocks((size + threads - 1) / threads);

    AT_DISPATCH_FLOATING_TYPES(a.type(), "add cuda", ([&amp;] {
        add_kernel&lt;scalar_t&gt;&lt;&lt;&lt;blocks, threads&gt;&gt;&gt;(
            a.data&lt;scalar_t&gt;(), 
            b.data&lt;scalar_t&gt;(), 
            c.data&lt;scalar_t&gt;(), 
            size
        );
    }));
}
</code></pre>
<p>And the test py file:</p>
<pre><code class="lang-python"># test.py
import torch
import add

if __name__ == "__main__":
    a = torch.zeros((100, ))
    b = torch.ones((100, )) * 10
    c = torch.ones((100, ))
    a = a.cuda(1)
    b = b.cuda(1)
    c = c.cuda(1)
    add.add(a, b, c)
    print(a)
</code></pre>
<p>Theoretically the print output should be a torch tensor with 100 <code>11</code>s, but actually the output is</p>
<pre><code class="lang-auto">tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
        0., 0., 0., 0.], device='cuda:1')
</code></pre>
<p>It seems that the <code>add_kernel</code> in <code>add_cuda.cu</code> didn’t implement at all, but I don’t know why it didn’t implement.<br>
Any help would be appreciated. Thank you!</p>

ptrblck: <p>Could you check the definition of <code>stride</code>?<br>
It seems it should be defined as the total number of threads in the grid:</p>
<pre><code class="lang-auto">const int stride = blockDim.x * gridDim.x
</code></pre>

Mon: <p>Thank you for your reply, and I did misunderstand it.<br>
But it seems that the most essential problem is not there, because the front elements are 0s after <code>add_kernel</code>, and after I changed <code>stride</code>, the result was still all 0s.</p>

ptrblck: <p>That’s strange. I’ve just tested your code as an inline method and it seems to work:</p>
<pre><code class="lang-python">import torch
from torch.utils.cpp_extension import load_inline

cpp_src = """
void add_cu(torch::Tensor a, torch::Tensor b, torch::Tensor c);

#define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)

void add(torch::Tensor a, torch::Tensor b, torch::Tensor c){
    CHECK_INPUT(a);
    CHECK_INPUT(b);
    CHECK_INPUT(c);
    add_cu(a, b, c);
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("add", &amp;add, "add(CUDA)");
}
"""

cuda_src = """
#include &lt;torch/extension.h&gt;
#include &lt;cuda.h&gt;
#include &lt;cuda_runtime.h&gt;

void add_cu(torch::Tensor a, torch::Tensor b, torch::Tensor c);


template &lt;typename scalar_t&gt;
__global__ void add_kernel(
    scalar_t* __restrict__ a, 
    scalar_t* __restrict__ b, 
    scalar_t* __restrict__ c, 
    size_t size
){
    const int index = blockIdx.x * blockDim.x + threadIdx.x;
    const int stride = blockDim.x * gridDim.x;
    for (int i = index; i &lt; size; i += stride){
        a[i] = b[i] + c[i];
    }
}


void add_cu(torch::Tensor a, torch::Tensor b, torch::Tensor c){
    const auto size = a.size(0);

    const int threads = 8;
    const dim3 blocks((size + threads - 1) / threads);

    AT_DISPATCH_FLOATING_TYPES(a.type(), "add cuda", ([&amp;] {
        add_kernel&lt;scalar_t&gt;&lt;&lt;&lt;blocks, threads&gt;&gt;&gt;(
            a.data&lt;scalar_t&gt;(), 
            b.data&lt;scalar_t&gt;(), 
            c.data&lt;scalar_t&gt;(), 
            size
        );
    }));
}
"""

add = load_inline(name='add', cpp_sources=[cpp_src],
                   cuda_sources=[cuda_src])

a = torch.zeros((100, ))
b = torch.ones((100, )) * 10
c = torch.ones((100, ))
a = a.cuda(0)
b = b.cuda(0)
c = c.cuda(0)
add.add(a, b, c)
print(a)
</code></pre>
<p>Could you check, if it gives the right result and compare your code against mine?</p>

Mon: <p>I’ve just found the reason that the program must be used on <code>cuda:0</code>. I’m sorry that I used <code>cuda(1)</code> actually as <code>0</code> was used by others then but I changed it to <code>0</code> here.<br>
Now when I changed it to <code>cuda(0)</code>, everything goes right.<br>
But it is still very strange for me why the cuda kernel doesn’t implement at cudas except <code>cuda(0)</code>?  As our GPU usage is of high frequency, it is of importance that these pytorch cuda extensions can be used on any GPU, so <strong>how</strong> can it be implemented on other GPUs?</p>
<p>I’m sorry that I forgot to say the versions:<br>
OS: ubuntu 16.04 LTS<br>
python: 3.6.8<br>
pytorch: 1.0.0</p>

tom: <p>You probably want to</p>
<ul>
<li>use <code>const torch::OptionalDeviceGuard device_guard(a.device());</code> or so.<br>
This is similar to <a href="https://pytorch.org/docs/master/notes/cuda.html" rel="nofollow noopener"><code>with torch.cuda.device(...)</code> in Python</a>.</li>
<li>Check that <code>a, b, c</code> actually are on the same device.</li>
</ul>
<p>Best regards</p>
<p>Thomas</p>
<p>P.S.: If you take inspiration from PyTorch internals: The reason you don’t see this in <code>aten/src/Aten/native/cuda/*</code> is that there is a auto-generated wrapper doing it for you.<br>
P.P.S.: <a href="https://github.com/pytorch/vision/blob/master/torchvision/csrc/cuda/ROIAlign_cuda.cu#L316" rel="nofollow noopener">The new torchvision C++ extension</a> does it right and is a great example of a model PyTorch extension.</p>

Mon: <p>Thank you very much! It works!</p>