What is the best practice for choosing between CPU and GPU dispatch for custom C++ extensions?

mdoyle: <p>Hi everyone,</p>
<p>I am currently writing C++ extensions which should sit on top of Torch that implement custom layers. I have both CPU and GPU implementations of my layers which I developed following the guide for developing C++ and CUDA extensions to Torch (<a href="https://pytorch.org/tutorials/advanced/cpp_extension.html" class="inline-onebox" rel="noopener nofollow ugc">Custom C++ and CUDA Extensions — PyTorch Tutorials 1.9.0+cu102 documentation</a>).</p>
<p>The guide is very helpful, however I have a remaining question of what is the best way of selecting which back-end (CPU or GPU) implementation of the layer should be executed at runtime? The example guide defines the steps to provide a C++ extension for either an explicit CPU <strong>or</strong> GPU implementation, but not for both.</p>
<p>This can be seen here: <a href="https://pytorch.org/tutorials/advanced/cpp_extension.html#writing-a-mixed-c-cuda-extension" class="inline-onebox" rel="noopener nofollow ugc">Custom C++ and CUDA Extensions — PyTorch Tutorials 1.9.0+cu102 documentation</a>. Let’s consider only the forward pass of the layer in the example (see below for simplified code excerpt from the link);</p>
<pre><code class="lang-auto">#include &lt;torch/extension.h&gt;
#include &lt;vector&gt;

std::vector&lt;torch::Tensor&gt; lltm_cuda_forward(
    torch::Tensor input,
    torch::Tensor weights,
    torch::Tensor bias,
    torch::Tensor old_h,
    torch::Tensor old_cell);

#define CHECK_CUDA(x) TORCH_CHECK(x.type().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)

std::vector&lt;torch::Tensor&gt; lltm_forward(
    torch::Tensor input,
    torch::Tensor weights,
    torch::Tensor bias,
    torch::Tensor old_h,
    torch::Tensor old_cell) {
  CHECK_INPUT(input);
  CHECK_INPUT(weights);
  CHECK_INPUT(bias);
  CHECK_INPUT(old_h);
  CHECK_INPUT(old_cell);

  return lltm_cuda_forward(input, weights, bias, old_h, old_cell);
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("forward", &amp;lltm_forward, "LLTM forward (CUDA)");
}
</code></pre>
<ul>
<li>
<code>lltm_cuda_forward</code> - Function which launches the CUDA kernel is declared.</li>
<li>
<code>lltm_forward</code> - Wrapper function is defined which calls <code>lltm_cuda_forward</code>. However, <code>lltm_forward</code> has it’s inputs checked to ensure the inputs reside on GPU memory (using the macro <code>CHECK_INPUT</code>).</li>
<li>The <code>lltm_forward</code> routine is then bound to Python with PyBind11.</li>
</ul>
<p>To me it makes sense to have a condition which chooses the CPU or GPU implementation in <code>lltm_forward</code>, by just checking where the input tensors reside (on CPU or GPU memory) and select the appropriate back-end. However, is this good practice? Or have you any better suggestions?</p>
<p>The only other option I can see is exposing both the CPU and GPU back-ends to Python through the bindings, and again checking the input types to choose which back-end to choose. This seems a bit less desirable to me.</p>
<p>Thanks for reading this, and for any suggestions/advice!</p>

googlebot: <p>There is a newer registration mechanism with backend dispatching, described <a href="https://pytorch.org/cppdocs/library.html" rel="noopener nofollow ugc">here</a>. Though I’m not sure if this is the current suggested approach; if you’re not creating a library with more than 2 backend implementations, cuda/cpu branch dispatch seems good enough to me.</p>

mdoyle: <p>Thank you Alex for your reply.</p>
<p>I haven’t come across that link yet, so that is very helpful and seems to be what I was looking for. Yes, I suppose ideally my plan is to then tie them in with TorchScript but the link you posted details that process to some extent too.</p>