Developer documentation for PyTorch 2.x compiler

torch.compile torch._dynamo
1

Could someone please point me to the developer docs for dynamo, inductor?

Please DO NOT give me the user 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.

Specifically, I need to understand how does dynamo capture the forward and backward graphs when training any random model - how it captures conditional branches, loops and dynamic tensor shapes.

Any help is welcome.

2

I was following the example at this PyTorch tutorial page

There is some discrepency: where as the tutorial states that upon executing the command:

TORCH_COMPILE_DEBUG=1 python example.py

the generated output_code.py file must contain the Triton IR, but I am seeing C++ code instead:

cpp_fused_add_cos_sin_0 = async_compile.cpp('''
#include "/tmp/torchinductor_<username>/ib/<a_random_hash_like_string>.h"
extern "C" void kernel(const float* in_ptr0,
                       float* out_ptr0)
{
    {
        for(long i0=static_cast<long>(0L); i0<static_cast<long>(10000L); i0+=static_cast<long>(16L))
        {
            auto tmp0 = at::vec::Vectorized<float>::loadu(in_ptr0 + static_cast<long>(i0));
            auto tmp1 = tmp0.cos();
            auto tmp2 = tmp0.sin();
            auto tmp3 = tmp1 + tmp2;
            tmp3.store(out_ptr0 + static_cast<long>(i0));
        }
    }
}
''')

Embedded within other generated Python code.

What am I missing?

3

@anubane 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

1 Like