How to tell if calculations are happening on GPU?

Karthik_Ganesan: <p>I am writing a custom linear layer and I want to execute it on GPU. The tensors are correctly located on the GPU ( *.is_cuda returns true for input &amp; weight). I am performing multiplications directly using two loops and I am getting the same answer as the in-built PyTorch Linear layer, but I’m not sure if this calculation is actually being done on the GPU. Is there a way to check?</p>
<blockquote>
<pre><code>def forward(ctx, X, weight, bias):
    ctx.save_for_backward(X, weight, bias)

    print(f"Input:: GPU: {X.is_cuda}, Weight:: GPU: {weight.is_cuda})

    (m, n_W) = X.shape
    (A, _) = weight.shape

    # Transpose the weights
    weight = torch.transpose(weight, 0, 1)

    output = torch.cuda.FloatTensor(m, A).fill_(0)

    for k in range(m):
        for l in range(A):
            accumulation = 0.0
            for j in range(n_W):
                accumulation += X[k, j] * weight[j, l]
            output[k, l] = accumulation + bias[l]

    return output
</code></pre>
</blockquote>
<p>Thanks!</p>

eqy: <p>If everything is indeed a <code>cuda</code> tensor, then the computation should be happening on GPU. You can run <code>nvprof</code> e.g., <code>nvprof python myscript.py</code> to get a sense of which kernels are being executed on GPU.</p>
<p>Alternatively if you are running the script interactively (or it is long-running enough), you can do something like <code>watch nvidia-smi</code> to get a sense of GPU usage as the script is running.</p>

Karthik_Ganesan: <p>Hi <a class="mention" href="/u/eqy">@eqy</a>, thank you. Yes, the script does run for a long time, so I can try those commands you mentioned. But what I really wanted to check is whether this simple python code gets converted to NV PTX code to run on the GPU automatically? Since I am doing the multiply accumulate manually and until I came across a tutorial, I imagined I would have to write CUDA C++ code to use the GPU. That’s why I wanted to ask and double check.</p>

eqy: <p>Normally when you are doing common operations such as matrix multiplication, matrix-vector multiplication, elementwise multiply-add, convolution, etc., these are dispatched to precompiled kernels in ATen (PyTorch’s tensor operator library), cuBLAS, or cuDNN. No compilation to PTX etc., will happen unless you are doing some kind of JIT compilation (e.g., <a href="https://pytorch.org/docs/stable/jit.html#id9" rel="noopener nofollow ugc">TorchScript — PyTorch 1.9.1 documentation</a>).</p>
<p>Rather than Python code being “converted,” usually what is happening is the Python code is calling into C++ code that calls the relevant CUDA kernels.</p>

Karthik_Ganesan: <p>Thank you. I tried nvprof as you suggested and it seems to be spending a lot of time in this function:<br>
<code>void at::native::vectorized_elementwise_kernel</code></p>
<p>How can I find where this function is located? I also have the entire PyTorch source downloaded and when I grep for this function there, I just get hits in a bunch of .json files but no source files.</p>
<p>Alternatively, is there somewhere layers are defined in CUDA code directly? If I could get that, I could modify that and plug that in via PyBind. It’s a little hard to parse through the various backends supported, like iDeep, MKLDNN etc. So not sure where I might just find the base CUDA code.</p>

eqy: <p>For elementwise kernels, these are typically implemented with PyTorch’s TensorIterator to generate kernels rather than with handwritten CUDA, so there isn’t really any CUDA code to directly inspect in this case.</p>
<aside class="onebox allowlistedgeneric" data-onebox-src="https://labs.quansight.org/blog/2021/04/pytorch-tensoriterator-internals-update/">
  <header class="source">
      <img src="https://labs.quansight.org/blog/2021/04/pytorch-tensoriterator-internals-update/../../../../favicon.ico" class="site-icon" width="16" height="16">

      <a href="https://labs.quansight.org/blog/2021/04/pytorch-tensoriterator-internals-update/" target="_blank" rel="noopener nofollow ugc" title="08:00AM - 09 April 2021">Quansight Labs – 9 Apr 21</a>
  </header>

  <article class="onebox-body">
    <img src="https://labs.quansight.org/images/QuansightLabs_logo_V2_rect.png" class="thumbnail onebox-avatar" width="500" height="500">

<h3><a href="https://labs.quansight.org/blog/2021/04/pytorch-tensoriterator-internals-update/" target="_blank" rel="noopener nofollow ugc">PyTorch TensorIterator Internals - 2021 Update</a></h3>

  <p>For contributors to the PyTorch codebase, one of the most commonly encountered
C++ classes is
TensorIterator.
TensorIterator offers a standardized way to iterate over elements of
a tensor, automatical</p>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


Karthik_Ganesan: <p>Thank you! I get it now. I’ve also been following this tutorial to see how to write my own efficient gpu version of the layer. Since I need to change the multiplication operation, I need quite low level access to the operation, without sacrificing too much performance. I hadn’t seen that page about the tensor iterator so I’ll be sure to check that out as well!</p>