Custom C++ / CUDA Extension - device placement in Python?

sebastian-sz: <p>I’m following the <a href="https://pytorch.org/tutorials/advanced/cpp_extension.html" rel="noopener nofollow ugc">tutorial on Custom C++ / CUDA Extensions</a>.</p>
<p>In the end there are two implementations of the module:</p>
<ol>
<li>In C++, imported as <code>lltm_cpp</code>
</li>
<li>In CUDA, imported as <code>lltm_cuda</code>.</li>
</ol>
<p>(Note that although <code>lltm_cpu</code> also runs on the GPU, CUDA is still a lot faster)</p>
<p>Is there a recommended way to use LLTM (or any other extension) inside my Python’s <code>torch.nn.Module</code> so that <code>to(device)</code> would execute <code>lltm_cpp</code> in running on CPU and <code>lltm_cuda</code> when running on GPU?</p>
<p>In more detail:</p>
<pre><code class="lang-python">class FinalModel(torch.nn.Module):
    def __init__(self):
        self.lltm = LLTM()

model = FinalModel()
model.to("cpu")  # executes lltm_cpp
model.to("cuda")  # executes lltm_cuda
</code></pre>
<p>Any help would be greatly appreciated.</p>

ptrblck: <p>The dispatching mechanism should already abstract away the device and should execute the corresponding code as described in the tutorial. Is this not working for you and do you see any errors after implementing the custom CUDA extension?</p>

sebastian-sz: <p>The dispatching mechanism seems to work only for C++ extension, but not for CUDA extension (unless I’m missing something).</p>
<p>My <code>setup.py</code> looks as follows:</p>
<pre><code class="lang-python">ext_modules = [
      cpp_extension.CppExtension("lltm_cpp", sources=['lltm.cpp']),
      cpp_extension.CUDAExtension(
          "lltm_cuda", sources=["lltm_cuda.cpp", "lltm_cuda_kernel.cu"],
      )
]

setup(name='lltm_fast',
      ext_modules=ext_modules,
      cmdclass={'build_ext': cpp_extension.BuildExtension})
</code></pre>
<p>I can successfully run <code>lltm_cpp</code> on CPU and GPU. But when running <code>lltm_cuda</code> on CPU i get:</p>
<pre><code class="lang-python">Traceback (most recent call last):
  File "&lt;stdin&gt;", line 1, in &lt;module&gt;
RuntimeError: x must be a CUDA tensor
</code></pre>
<p>Therefore I would like to automatically run <code>lltm_cpp</code> on CPU and <code>lltm_cuda</code> on GPU.</p>
<p>My question is: is there a recommended way to dispatch these ops on the Python side?</p>

ptrblck: <p>Ah OK, I see. In this case, this <a href="https://pytorch.org/tutorials/advanced/dispatcher.html">Dispatching tutorial</a> might be useful, which seems to target your use case:</p>
<blockquote>
<p>The dispatcher is an internal component of PyTorch which is responsible for figuring out what code should actually get run when you call a function like torch::add. This can be nontrivial, because PyTorch operations need to handle a lot of cross-cutting concerns that are “layered” on top of one of another. Here is a sampling of some of the things it handles:</p>
</blockquote>
<ul>
<li>
<blockquote>
<p>Switching between the CPU and CUDA implementations of an operator, depending on the devices of the input tensors. …</p>
</blockquote>
</li>
</ul>

sebastian-sz: <p>This is exactly what I needed, thank you!</p>