Developer documentation for PyTorch 2.x compiler

anubane: <p>Could someone please point me to the <strong>developer</strong> docs for <em>dynamo</em>, <em>inductor</em>?</p>
<p>Please DO NOT give me the <strong>user</strong> doc links currently available on the official PyTorch docs page. I need to understand the complete flow of dynamo so as to extend it for my own purpose.</p>
<p>Specifically, I need to understand how does dynamo capture the forward and backward graphs when training any random model - how it captures <em>conditional branches</em>, <em>loops</em> and <em>dynamic tensor shapes</em>.</p>
<p>Any help is welcome.</p>

anubane: <p>I was following the example at <a href="https://pytorch.org/docs/stable/torch.compiler_get_started.html" rel="noopener nofollow ugc">this PyTorch tutorial page</a></p>
<p>There is some discrepency: where as the tutorial states that upon executing the command:</p>
<pre><code class="lang-auto">TORCH_COMPILE_DEBUG=1 python example.py
</code></pre>
<p>the generated <code>output_code.py</code> file must contain the Triton IR, but I am seeing C++ code instead:</p>
<pre><code class="lang-python">cpp_fused_add_cos_sin_0 = async_compile.cpp('''
#include "/tmp/torchinductor_&lt;username&gt;/ib/&lt;a_random_hash_like_string&gt;.h"
extern "C" void kernel(const float* in_ptr0,
                       float* out_ptr0)
{
    {
        for(long i0=static_cast&lt;long&gt;(0L); i0&lt;static_cast&lt;long&gt;(10000L); i0+=static_cast&lt;long&gt;(16L))
        {
            auto tmp0 = at::vec::Vectorized&lt;float&gt;::loadu(in_ptr0 + static_cast&lt;long&gt;(i0));
            auto tmp1 = tmp0.cos();
            auto tmp2 = tmp0.sin();
            auto tmp3 = tmp1 + tmp2;
            tmp3.store(out_ptr0 + static_cast&lt;long&gt;(i0));
        }
    }
}
''')
</code></pre>
<p>Embedded within other generated Python code.</p>
<p>What am I missing?</p>

bdhirsh: <p><a class="mention" href="/u/anubane">@anubane</a> inductor (the default compiler) will emit cpp code if its given cpu tensor inputs, and it will emit triton if it’s given cuda tensor inputs</p>

xgbj: <p>support with <a class="mention" href="/u/anubane">@anubane</a>’s suggestion. I want to know if there are any good ways to learn torch.compile. It’s not easy to analyze and debug after encountering problems with compilation.</p>

marksaroufim: <p>We’ve tried to recently aggregate some</p>
<ol>
<li>dev docs here <a href="https://pytorch.org/docs/main/torch.compiler.html" class="inline-onebox">torch.compiler — PyTorch main documentation</a></li>
<li>There’s also the 2.0 series here <a href="https://www.youtube.com/watch?v=v4nDZTK_eJg&amp;list=PL_lsbAsL_o2CQr8oh5sNWt96yWQphNEzM">https://www.youtube.com/watch?v=v4nDZTK_eJg&amp;list=PL_lsbAsL_o2CQr8oh5sNWt96yWQphNEzM</a></li>
<li>The PyTorch dev podcast <a href="https://pytorch-dev-podcast.simplecast.com/">https://pytorch-dev-podcast.simplecast.com/</a></li>
</ol>
<p>For some community tutorials I personally really liked</p>
<ul>
<li><a href="https://x.com/tarantulae/status/1734192354716725401?s=20">https://x.com/tarantulae/status/1734192354716725401?s=20</a></li>
<li><a href="https://towardsdatascience.com/how-pytorch-2-0-accelerates-deep-learning-with-operator-fusion-and-cpu-gpu-code-generation-35132a85bd26?gi=846f935de46e" class="inline-onebox">How Pytorch 2.0 Accelerates Deep Learning with Operator Fusion and CPU/GPU Code-Generation | by Shashank Prasanna | Towards Data Science</a></li>
</ul>
<p>I understand the sentiment of <code>torch.compile</code> is complicated but then again compilers are generally complicated but torch.compile has going for it that it’s mostly written in python. If there’s concrete parts of the stack you find hard to parse and would like more tutorials written about please let me know</p>

xgbj: <p>Thanks a lot!  this is exactly what I wanted.</p>

anubane: <p>Thanks a lot for the pointers!</p>
<p>I would specifically like to understand the codegen part of torch-inductor along with an overall code structure understanding for the Inductor module.</p>

marksaroufim: <p>Here you go <a href="https://www.youtube.com/watch?v=p13HpZv2S3Q">https://www.youtube.com/watch?v=p13HpZv2S3Q</a></p>

bdhirsh: <p>There are also two (excellent) podcasts on Inductor internals here:</p>
<p>Inductor IR: <a href="https://open.spotify.com/episode/3c3FjUcQLUCCsTvEznKcBZ?si=jLNEhskySSiAuvq703vUPg" class="inline-onebox" rel="noopener nofollow ugc">Spotify</a></p>
<p>Details on inductor’s define-by-run IR semantics: <a href="https://open.spotify.com/episode/7hFLcgZim1G6T9gN5SdNNB?si=PTTPBgadQ5q20pM8IA4Qrw" class="inline-onebox" rel="noopener nofollow ugc">Spotify</a></p>

anubane: <p>I believe I have been unable to explain my question:</p>
<p>I am <strong>NOT</strong> looking for developer docs from the perspective of model developers who <strong>USE</strong> PyTorch.</p>
<p>I am looking for developer docs, design docs from the perspective of <strong>PyTorch core developers</strong> - <em>the logic and working of the source code</em> - that is what I need to understand and modify.</p>
<p>Most of the links here explain things from the pespective of PyTorch <strong>users</strong>.</p>
<p>(<em>Further, I prefer docs/videos over podcasts since podcasts make it hard to follow the actual source code</em>)</p>
<p>Currently, the only method to understand this seems to build PyTorch in debug mode and step through the inductor module to understand the flow. This can be done but is tedious, that is why I was looking for some documentation.</p>
<p>(<em>the question was cross posted on the <a href="https://dev-discuss.pytorch.org/t/developer-docs-for-pytorch-inductor/1714/3?u=anubane" rel="noopener nofollow ugc">developer forum</a></em>)</p>

marksaroufim: <p>The videos and podcasts we shared were indeed targeted to new PyTorch developer. Honestly I only started to learn how PT2 worked after filing and fixing some random easy bugs <a href="https://github.com/pytorch/pytorch/issues?q=is%3Aissue+is%3Aopen+sort%3Aupdated-desc++label%3A%22oncall%3A+pt2%22+" class="inline-onebox">Issues · pytorch/pytorch · GitHub</a> otherwise a debugger sounds great, that’s how I ramp up on new projects</p>

anubane: <p>I understand that for a beginner like me, there will always be a learning curve. However, since PT2.x is new, and the components in question are different from PT1.x, I was hoping to catch up sooner. The resources shared here are great and give a holistic overview, but, since they move on to performance comparisons and benefits to end-users of PT, the purpose to me is lost.</p>
<p>I am trying to extend PT to a custom backend; while a public API is available for Dynamo, similar API is not publicly available for Inductor. Therefore I needed to understand the module itself.</p>
<p><em>For example</em>, in the video:</p>
<aside class="quote no-group" data-username="marksaroufim" data-post="8" data-topic="191625">
<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/marksaroufim/48/36747_2.png" class="avatar"> marksaroufim:</div>
<blockquote>
<p>Here you go <a href="https://www.youtube.com/watch?v=p13HpZv2S3Q" rel="noopener nofollow ugc">https://www.youtube.com/watch?v=p13HpZv2S3Q</a></p>
</blockquote>
</aside>
<p>The logic of inductor is explained, but, that is not enough for me to understand the actual source code, it’s structure - what goes in the codegen module, what goes in the scheduler module, how to extend that, etc.</p>

anubane: <p>Found a relevant post: <a href="https://dev-discuss.pytorch.org/t/what-should-we-do-about-developer-documentation/491" rel="noopener nofollow ugc">What should we do about developer documentation?</a></p>