@jino_rohit: understanding the torch compile stack torch.compile is a technique to speed up your pytorch code. torch.compile makes t…

understanding the torch compile stack

torch.compile is a technique to speed up your pytorch code. torch.compile makes torch code run faster by JIT-compiling PyTorch code into optimized kernels, while requiring minimal code changes.

this is the overall flow for the torch compile stack from my understanding -

1. pytorch api - this is your regular nn.module that you write in torch.
2. dynamo - dynamo intercepts the regular python flow and captures these pytorch specific operations into a graph. you can think of them like DAGs.
3. fx graph - fx graph is pytorch’s internal graph representation. this IR is pretty easy to work with and debug since its just graphs and it has only 6 main instructions. with fx graph, all the operations are represented as nodes and connections.
4. aten ops - all the operation captured in the graph have to be lowered to the primitives written in C++ in torch, for instance cos, sin etc. all of them are present in the aten/ library.
5. torch inductor - this is the actual compiler backend that takes these aten ops, and finally lowers them into triton kernels and ptx and so on.

oh nice, thaats very useful, thanks!

ah thanks!

aten ops are just the primitive ops written in C++ all the fundamental ones youll need , peek into the ATen/ folder youll see