PrimTorch: How backend/compiler writers interact with various IRs

[gilfree]

@Minerva_Yu

For Core ATen IR is the subset of aten operators, the way I use it is the same with former, i.e. through “native_functions.yaml”. ( But how could I distinguish who are the “Core ATen” ops and who are the “ATen” ones? )

Hi @Minerva_Yu,

No answers, just hints - I am at the same situation as you are - trying to understand how should I integrate my compiler with the torch IR. I got some pointers here:

functorch.compile.aot_function fails on reshape in TorchRefsMode

You can play with the backend and get only torch.xxx with no aten, or only aten if you run the aot compiler. It seems that to get the core aten and prims - you can use what @ezyang pointed there - using proxy_tensor, make_fx & TorchRefsMode, but I am not sure how these levels relate to IRs — PyTorch master documentation.

Also see the post here: PrimTorch: could we get pure core-aten-ops or prims-ops after aot_autograd - #5 by SherlockNoMad, on how to get the core aten ir.

Seems also that you can define the decompositions you want, in python, and get varying levels of IRs from torch.xxx to “core atens and prims” - See the notebook I linked to in that issue.