Cuda extension with multiple gpus

pickandroll3: <p>Hey everyone!</p>
<p>I have a problem with my custom cuda kernel with multiple gpus.<br>
I wrote the following kernel to essentially do <code>torch.logical_not(torch.isclose(X, scalar, rtol=radius))</code>. My kernel code looks like this:</p>
<p>(file <code>custom_ops_kernel.cu</code>)</p>
<pre><code class="lang-auto">#include &lt;cuda.h&gt;
#include &lt;cuda_runtime.h&gt;
#include &lt;c10/cuda/CUDAFunctions.h&gt;

// C++ PROJECT INCLUDES


template&lt;typename scalar_t&gt;
__global__
void
cuda_notnear_center_kernel(const scalar_t* __restrict__ x,
                           bool* __restrict__ z,
                           const double lthreshold,
                           const double rthreshold,
                           size_t N)
{
    const int index = blockIdx.x * blockDim.x + threadIdx.x;

    if (index &lt; N)
    {
        z[index] = x[index] &lt; lthreshold || x[index] &gt; rthreshold;
    }
}



torch::Tensor
cuda_notnear_center(torch::Tensor x,
                    torch::Tensor out,
                    double center,
                    double radius)
{
    const double lthreshold = center - radius;
    const double rthreshold = center + radius;

    // auto z = torch::empty_like(x, torch::kBool);
    const int num_threads = 1024;
    const int num_blocks = (x.size(0) / num_threads) + 1;
    // printf("num_blocks: %i\n", num_blocks);

    const auto current_device = c10::cuda::current_device();
    c10::cuda::set_device(x.get_device());
    AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), "cuda_notnear_center", ([&amp;] {
        cuda_notnear_center_kernel&lt;scalar_t&gt;&lt;&lt;&lt;num_blocks, num_threads&gt;&gt;&gt;(
            x.data_ptr&lt;scalar_t&gt;(),
            out.data_ptr&lt;bool&gt;(),
            lthreshold,
            rthreshold,
            x.size(0));
    }));
    c10::cuda::set_device(current_device);

    // printf("Error: %s", cudaGetLastError());

    return out;
}
</code></pre>
<p>and is managed by a cpp file <code>custom_ops.cc</code> which the relevant parts look like:</p>
<pre><code class="lang-auto">#include &lt;torch/extension.h&gt;
#include &lt;vector&gt;


// C++ PROJECT INCLUDES

// taken from https://pytorch.org/tutorials/advanced/cpp_extension.html
#define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
#define CHECK_SAME_SIZE(x, y) TORCH_CHECK(torch::is_same_size(x, y), #x " must have save size as " #y)
#define CHECK_TYPE(x, t) TORCH_CHECK(x.scalar_type() == t, #x " must have same type as " #t)

torch::Tensor
cuda_notnear_center(torch::Tensor x,
                    torch::Tensor out,
                    double center,
                    double radius);

torch::Tensor&amp;
notnear_center_out(torch::Tensor&amp; out,
                   const torch::Tensor&amp; x,
                   const double center,
                   const double radius)
{
    CHECK_INPUT(x);
    CHECK_INPUT(out);
    CHECK_SAME_SIZE(x, out);
    CHECK_TYPE(out, torch::kBool);

    cuda_notnear_center(x.view({-1}), out.view({-1}), center, radius);
    return out;
}


torch::Tensor
notnear_center(const torch::Tensor&amp; x,
               const double center,
               const double radius)
{
    CHECK_INPUT(x);
    torch::Tensor out = torch::empty_like(x, torch::kBool);
    cuda_notnear_center(x.view({-1}), out.view({-1}), center, radius);
    return out;
}
</code></pre>
<p>My problem is that my machine has multiple gpus. When I launch this kernel on a gpu other than <code>gpu0</code> I end up allocated with memory on <code>gpu0</code> from this kernel. I have done some testing with my code and have narrowed it down to the kernel. I thought that I could set the device with <code>c10::cuda::set_device()</code> but is this the proper way?</p>
<p>Here is the output of <code>nvidia-smi</code> on my 4-gpu machine. The main processes are running on <code>gpu1</code> and <code>gpu3</code> but both processes are allocating memory on <code>gpu0</code>:</p>
<pre><code class="lang-auto">+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A    128024      C   python3                          1481MiB |
|    0   N/A  N/A    129010      C   python3                          1481MiB |
|    1   N/A  N/A    129010      C   python3                         10157MiB |
|    3   N/A  N/A    128024      C   python3                         10195MiB |
+-----------------------------------------------------------------------------+
</code></pre>

tom: <p>Are you looking for a <a href="https://github.com/pytorch/pytorch/blob/master/c10/cuda/CUDAGuard.h" rel="noopener nofollow ugc"><code>torch::CUDAGuard</code></a> RAII variable?</p>
<p>Best regards</p>
<p>Thomas</p>

pickandroll3: <p>I’m not sure. This is my first attempt at using a kernel (that I wrote) on a gpu that isnt <code>gpu0</code>, so I don’t know much at all about specifying the device.</p>
<p>For instance, if I were to use a <code>CUDAGuard</code>, does it go in the kernel itself or is my placement correct?</p>

tom: <p>Take these with a grain of salt as from someone who does single GPU more often than multi, but</p>
<ul>
<li>I’d use the guard instead of set device,</li>
<li>I’d set the device based on the input tensors, not the local state of PyTorch,</li>
<li>I’d do it right at the top of the function taking the tensor (it also affects new tensors that you might create).</li>
</ul>
<p>Best regards</p>
<p>Thomas</p>

pickandroll3: <p>Thanks a ton, I really appreciate it <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

pickandroll3: <p>update. I tried using <code>c10::cuda::CUDAGuard guard(x.get_device())</code> but I still have the same problem. I did however move <code>c10::cuda::set_device(x.get_device())</code> to the top of the function which worked for some reason.</p>
<p>To be more clear, I have the following lines at the top of my function:</p>
<pre><code class="lang-auto">torch::Tensor
cuda_notnear_center(torch::Tensor x,
                    torch::Tensor out,
                    double center,
                    double radius)
{
    const auto current_device = c10::cuda::current_device(); // will use to replace device after call
    c10::cuda::set_device(x.get_device()); // set device to be device of tensor X

    ... // function body

    c10::cuda::set_device(current_device); // put state back after call 
    return out;
}
</code></pre>
<p>This seemed to work, but I have no idea why:</p>
<pre><code class="lang-auto">+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    1   N/A  N/A    104618      C   python3                         10155MiB |
|    3   N/A  N/A    103802      C   python3                         10195MiB |
+-----------------------------------------------------------------------------+

</code></pre>

tom: <p>Glad you figured it out!</p>