Why indexing a tensor on GPU involves data transactions to CPU?

marcortiz11: <p>Hey community!</p>
<p>I’ve recently been studying a very simple case where I index a CUDA tensor in GPU. As far as I know, the indexing operation is adapted for GPU execution with potential speedups regarding CPU.</p>
<p>In the small example below, I access elements in tensor <code>a</code> according to the mask tensor <code>b</code>. I have both the indexed tensor and the tensor of indices on  GPU, and after running the script, I would expect no CPU involvement at all. I study the behavior of the code with a profiler:</p>
<pre><code class="lang-auto">import torch
import torch.autograd.profiler as profiler

a = torch.rand((10000)).to('cuda:0')
b = torch.ones_like(a).bool()

with profiler.profile() as prof:
       with profiler.record_function("CHAIN_FORWARD"):
               for i in range(10):
                     a[b]

print(prof.key_averages().table(sort_by='self_cpu_time_total', row_limit=5))
</code></pre>
<p>Here is the output of the profiling:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/d/bd953fe8f02722b1dcf4ac2fe8b80e5f758ec5d9.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/bd953fe8f02722b1dcf4ac2fe8b80e5f758ec5d9" title="indexing"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/d/bd953fe8f02722b1dcf4ac2fe8b80e5f758ec5d9_2_690x93.jpeg" alt="indexing" data-base62-sha1="r380u1CJBfXH7dLxJq7tS1bbo9j" width="690" height="93" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/d/bd953fe8f02722b1dcf4ac2fe8b80e5f758ec5d9_2_690x93.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/b/d/bd953fe8f02722b1dcf4ac2fe8b80e5f758ec5d9_2_1035x139.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/b/d/bd953fe8f02722b1dcf4ac2fe8b80e5f758ec5d9_2_1380x186.jpeg 2x" data-dominant-color="181818"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">indexing</span><span class="informations">1530×208 53.5 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>As you can see, the operation <code>to()</code>, which transfers data to and from different devices, has been called 60 times. How comes? I am using PyTorch version 1.5.0</p>

marcortiz11: <p>One of the possible reasons is that GPU may be sending back to CPU profiling information. Is that true?</p>

googlebot: <p>.to() doesn’t necessarily involve cross-device copying, and I wouldn’t expect it here. I’d guess it is either to(dtype=int) form or no-op protective calls (like x=x.to(y)). Chrome trace or nvtx profiling mode would tell you more about these calls.</p>
<aside class="quote no-group" data-username="marcortiz11" data-post="2" data-topic="111569" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/marcortiz11/48/34553_2.png" class="avatar"> marcortiz11:</div>
<blockquote>
<p>One of the possible reasons is that GPU may be sending back to CPU profiling information. Is that true?</p>
</blockquote>
</aside>
<p>I don’t think this would be directly recorded, this would imply messed up profiling report.</p>

ptrblck: <p>I agree with <a class="mention" href="/u/googlebot">@googlebot</a> what capturing the profiling information would be confusing and bad.</p>
<p>Note that your current code snippet uses a <code>BoolTensor</code> to index <code>a</code>, which will yield a variable sized output tensor (in your example you are using <code>torch.ones_like</code>, so all values would be returned).<br>
This would call into <code>nonzero</code>, which needs to synchronize as seen <a href="https://github.com/pytorch/pytorch/blob/594a66d778a660faed0b0fbbe1dd8c2c318707ff/aten/src/ATen/native/cuda/Indexing.cu#L873-L876">here</a>. Besides that, the origin of the <code>to()</code> op could be found in a profiler as already explained.<br>
If you are using Nsight Systems, you could have a look at <a href="https://dev-discuss.pytorch.org/t/using-nsight-systems-to-profile-gpu-workload/59?u=ptrblck">this post</a> to see how to enable backtraces.</p>

Ramansh_Sharma: <p>Hi, does this mean that indexing a CUDA PyTorch tensor with say a python list of indices on the CPU does not carry information back and forth from GPU to CPU?</p>

ptrblck: <p>A good indicator would be to check for synchronizations or profile the actual code. To do so you could use <code>torch.cuda.set_sync_debug_mode(debug_mode)</code> and the profiling steps mentioned before.</p>