Aten Native CUDA Op Tracing

Jerome_Ku: <p>Hi!</p>
<p>I’m profiling pytorch code using NVIDIA <code>nsys</code> and noticed calls to kernels such as<br>
<code>at::native::::CatArrayBatchedCopy</code> and <code>at::native::elementwise_kernel</code>.  I can see the implementations of these kernels <a href="https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/cuda/Shape.cu" rel="noopener nofollow ugc">here</a> and <a href="https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/cuda/CUDALoops.cuh" rel="noopener nofollow ugc">here</a>.</p>
<p>However, I can’t seem to find where these kernels are actually called when searching the repo – searching for these identifiers only <a href="https://github.com/search?q=repo%3Apytorch%2Fpytorch%20batchedcopy&amp;type=code" rel="noopener nofollow ugc">results</a> in the implementation files.</p>
<p>I’d like the trace the chain of calls leading to these kernel launches from higher-level operators (e.g., <code>Conv2d</code>), and more generally, understand the internal plumbing of <code>pytorch</code> in greater detail.</p>
<p>Would greatly appreciate if anybody could explain how <code>Aten native</code> operators are connected to higher-level functions and eventually connected to Python!</p>

bdhirsh: <p>at::native kernels are dynamically registered to the pytorch dispatcher through code-generated files (which has been a very useful abstraction, but can make grepping the codebase for them difficult).</p>
<p>Some helpful links:</p>
<p>intro to the dispatcher: <a href="https://blog.ezyang.com/2020/09/lets-talk-about-the-pytorch-dispatcher/" rel="noopener nofollow ugc">https://blog.ezyang.com/2020/09/lets-talk-about-the-pytorch-dispatcher/</a></p>
<p>deeper dive into the full set of function calls from torch* to a native kernel: <a href="https://github.com/pytorch/pytorch/wiki/PyTorch-dispatcher-walkthrough" class="inline-onebox" rel="noopener nofollow ugc">PyTorch dispatcher walkthrough · pytorch/pytorch Wiki · GitHub</a></p>

Jerome_Ku: <p>Thanks!  Those are great resources – been a big fan of EZYang’s in-depth podcasts and posts on pytorch internals.  The wiki is also a goldmine.</p>
<p>FWIW, running <code>nsys</code> with sampling enabled (<code>--sample cpu</code>) lets you see full function traces.  Here is a sample screenshot:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/d/0d60769e10e3e91ba8ac132504b14ee2293258b9.png" data-download-href="https://discuss.pytorch.org/uploads/default/0d60769e10e3e91ba8ac132504b14ee2293258b9" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/0/d/0d60769e10e3e91ba8ac132504b14ee2293258b9_2_345x172.png" alt="image" data-base62-sha1="1UkSurNAiA3DoVKTFaHpBVI1woh" width="345" height="172" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/0/d/0d60769e10e3e91ba8ac132504b14ee2293258b9_2_345x172.png, https://discuss.pytorch.org/uploads/default/optimized/3X/0/d/0d60769e10e3e91ba8ac132504b14ee2293258b9_2_517x258.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/0/d/0d60769e10e3e91ba8ac132504b14ee2293258b9_2_690x344.png 2x" data-dominant-color="EAECDB"><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">image</span><span class="informations">1787×894 177 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>

saurabh-singh-rajput: <p>I’m sorry to bother you, but is there a way to run these kernels standalone in exclusion to benchmark?</p>