What's the difference between .cuda() and .to(device)

username: Eta_C
cooked: <p>What’s the difference between <code>tensor.cuda()</code> and <code>tensor.to(0)</code>?</p>
<p>I copy function <code>CUDA_tensor_apply2</code> from <code>ATen/cuda/CUDAApplyUtils.cuh</code> and use it as a PyTorch extension.</p>
<p>When I run</p>
<pre><code class="lang-python">import torch
import my_extension.run as run

x = torch.rand(3, 4)
y = x.cuda()
print(run(y))  # all is well
print(y)  # all is well
print(x)  # all is well
</code></pre>
<p>But if I run</p>
<pre><code class="lang-python">import torch
import my_extension.run as run

x = torch.rand(3, 4)
y = x.to(0)
print(run(y))  # incorrect result
print(y)  # RuntimeError (700)
print(x)  # no problem

z = x.cuda()  # RuntimeError (700)
</code></pre>
<p>I got the error <code>RuntimeError (700) : an illegal memory access was encountered at /pytorch/aten/src/THC/THCReduceAll.cuh:327</code></p>
<p>Why? Is this error related to <code>at::cuda::current_device()</code>?</p>

username: ptrblck
cooked: <p>Both calls should just transfer the tensor to the <code>cuda:0</code> device.<br>
Do you have a small reproducible code snippet, so that we can have a look?</p>

username: Eta_C
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a><br>
Thank you for replying, and I am so sorry for my deviation from the above description.<br>
I mean, my code can run when using <code>.cuda()</code>, but got error when using <code>.to(1)</code>.</p>
<p>Here’s my code, consist of 3 files, <code>my_extension.cpp</code>, <code>my_extension_kernel.cu</code> and <code>setup.py</code>.<br>
File: my_extension.cpp</p>
<pre><code class="lang-cpp">// file my_extension.cpp
#include &lt;torch/torch.h&gt;
#include &lt;cmath&gt;

void run_cpu(at::Tensor&amp; out, const at::Tensor&amp; x);
void run_gpu(at::Tensor&amp; out, const at::Tensor&amp; x);

void run_cpu(at::Tensor&amp; out, const at::Tensor&amp; x) {
    out = x + 1;
}

#define CHECK_CONTIGUOUS(x) AT_ASSERT(x.is_contiguous(), #x " must be contiguous.")
#define CHECK_INPUT(x) CHECK_CONTIGUOUS(x)

at::Tensor run(const at::Tensor x) {
    CHECK_INPUT(x);
    at::Tensor out = at::empty(x.sizes(), x.options());
    if(x.type().is_cuda()) {
        run_gpu(out, x);
    } else {
        run_cpu(out, x);
    }
    return out;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &amp;run, "My Extension");
}
</code></pre>
<p>File: my_extension_kernel.cu</p>
<pre><code class="lang-cpp">// file my_extension_kernel.cu
#include &lt;ATen/ATen.h&gt;
#include &lt;cuda.h&gt;
#include &lt;cuda_runtime.h&gt;
#include &lt;cuda_runtime_api.h&gt;
#include &lt;ATen/cuda/CUDAApplyUtils.cuh&gt;
#include &lt;THC/generic/THCTensorCopy.h&gt;
#include &lt;ATen/LegacyTHFunctionsCUDA.h&gt;
#include &lt;cmath&gt;


template &lt;typename scalar1, typename scalar2, int step, typename Op&gt;
inline bool CUDA_tensor_apply22(at::Tensor a,
                               at::Tensor b,
                               const Op op,
                               at::cuda::TensorArgType aType = at::cuda::TensorArgType::ReadWrite,
                               at::cuda::TensorArgType bType = at::cuda::TensorArgType::ReadOnly) {
    int64_t totalElements = a.numel();

    if (a.numel() == 0) { return true; }
    constexpr uint32_t AT_APPLY_THREADS_PER_BLOCK = 512;
    const dim3 block = dim3(AT_APPLY_THREADS_PER_BLOCK);
 
    dim3 grid;
    int64_t curDevice = at::cuda::current_device();    

    if (curDevice == -1) return false;
    if (!at::cuda::getApplyGrid&lt;step&gt;(totalElements, grid, curDevice)) {
      return false;
    }

#define HANDLE_CASE(TYPE, A, B)                                        \
  at::cuda::kernelPointwiseApply2&lt;Op,                                  \
                        scalar1,                                       \
                        scalar2,                                       \
                        TYPE, A, B, step&gt;                              \
   &lt;&lt;&lt;grid, block, 0, at::cuda::getCurrentCUDAStream(curDevice)&gt;&gt;&gt;(    \
       aInfo, bInfo, static_cast&lt;TYPE&gt;(totalElements), op);

#define HANDLE_B_CASE(TYPE, A, B) {         \
  switch (B) {                              \
    case 1:                                 \
      HANDLE_CASE(TYPE, A, 1);              \
      break;                                \
    case 2:                                 \
      HANDLE_CASE(TYPE, A, 2);              \
      break;                                \
    default:                                \
      HANDLE_CASE(TYPE, A, -1);             \
      break;                                \
  }                                         \
}

#define HANDLE_A_CASE(TYPE, A, B) {         \
  switch (A) {                              \
    case 1:                                 \
      HANDLE_B_CASE(TYPE, 1, B);            \
      break;                                \
    case 2:                                 \
      HANDLE_B_CASE(TYPE, 2, B);            \
      break;                                \
    default:                                \
      HANDLE_B_CASE(TYPE, -1, B);           \
      break;                                \
  }                                         \
}

    if (at::cuda::detail::canUse32BitIndexMath(a) &amp;&amp;
        at::cuda::detail::canUse32BitIndexMath(b)) {
      at::cuda::detail::TensorInfo&lt;scalar1, unsigned int&gt; aInfo =
        at::cuda::detail::getTensorInfo&lt;scalar1, unsigned int&gt;(a);

      at::cuda::detail::TensorInfo&lt;scalar2, unsigned int&gt; bInfo =
        at::cuda::detail::getTensorInfo&lt;scalar2, unsigned int&gt;(b);
      at::cuda::rearrangeDims(&amp;aInfo, &amp;bInfo);
      aInfo.collapseDims();
      bInfo.collapseDims();

      HANDLE_A_CASE(unsigned int, aInfo.dims, bInfo.dims);
    } else {
      at::cuda::detail::TensorInfo&lt;scalar1, uint64_t&gt; aInfo =
        at::cuda::detail::getTensorInfo&lt;scalar1, uint64_t&gt;(a);

      at::cuda::detail::TensorInfo&lt;scalar2, uint64_t&gt; bInfo =
        at::cuda::detail::getTensorInfo&lt;scalar2, uint64_t&gt;(b);
      at::cuda::rearrangeDims(&amp;aInfo, &amp;bInfo);
      aInfo.collapseDims();
      bInfo.collapseDims();

      if (aInfo.dims == 1 &amp;&amp; bInfo.dims == 1) {
        HANDLE_CASE(uint64_t, 1, 1);
      } else {
        HANDLE_CASE(uint64_t, -1, -1);
      }
    }
#undef HANDLE_CASE
#undef HANDLE_B_CASE
#undef HANDLE_A_CASE

  return true;
}


template &lt;typename scalar1, typename scalar2, typename Op&gt;
inline bool CUDA_tensor_apply22(at::Tensor a,
                               at::Tensor b,
                               const Op op,
                               at::cuda::TensorArgType aType = at::cuda::TensorArgType::ReadWrite,
                               at::cuda::TensorArgType bType = at::cuda::TensorArgType::ReadOnly) {
  return CUDA_tensor_apply22&lt;scalar1, scalar2, 1, Op&gt;(a, b, op, aType, bType);
}


void run_gpu(at::Tensor&amp; out, const at::Tensor&amp; x) {
    CUDA_tensor_apply22&lt;float, float&gt;(
        x, out, [=] __device__(const float&amp; val, float&amp; res) {
            res = val + 1;
        }
    );
}
</code></pre>
<p>File: setup.py</p>
<pre><code class="lang-python"># file setup.py
from setuptools import setup
from torch.utils.cpp_extension import CUDAExtension, BuildExtension


setup(
  name="my_extension",
  ext_modules=[
    CUDAExtension("my_extension", 
                  ["my_extension.cpp", "my_extension_kernel.cu"],
                  extra_compile_args=["-expt-extended-lambda"])
    ],
  cmdclass={"build_ext": BuildExtension}    
)
</code></pre>
<p>Then run <code>python setup.py install</code>.<br>
And run the test as follow:</p>
<pre><code class="lang-python">import torch
import my_extension

x = torch.rand(3, 4)
y = x.cuda()
print(my_extension.run(y))
print(y)

z = x.to(1)
print(my_extension.run(z))
print(z) 
</code></pre>
<p>I do some simple check. The function <code>inline bool CUDA_tensor_apply22</code> in <code>my_extension_kernel.cu</code> returns <code>true</code>.</p>

username: ptrblck
cooked: <p>Could you try to get the current device from the passed tensors instead of</p>
<pre><code class="lang-auto">int64_t curDevice = at::cuda::current_device();  
</code></pre>
<p>I haven’t tested the code yet, but if I’m not mistaken, this would use the current device specified by a device guard.</p>

username: Eta_C
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a>  Just now, I change that line to</p>
<pre><code class="lang-cpp">int64_t curDevice = at::device_of(a).value().index();
int64_t bDevice = at::device_of(b).value().index();
</code></pre>
<p>and add</p>
<pre><code class="lang-auto">std::cout &lt;&lt; curDevice &lt;&lt; bDevice &lt;&lt; std::endl;
</code></pre>
<p>before <code>#define HANDLE_CASE</code>.</p>
<p>Then run this test code,</p>
<pre><code class="lang-python">x = torch.rand(3, 4)
y = x.cuda()
z = x.to(1)
r = my_extension(y) # correct result and print 0 0
k = my_extension(z)  # incorrect result and print 1 1
print(y)  # RuntimeError
</code></pre>
<p>PS: If I print <code>y</code> without running <code>my_extension(z)</code>, that’s no problem. <code>my_extension(z)</code> always return a tensor which all elements are 0.</p>
<p>PS: my function <code>CUDA_tensor_apply22</code> is modified from <code>CUDA_tensor_apply2</code> in PyTorch source file <code>include/ATen/cuda/CUDAApplyUtils.cuh</code>.</p>