Declarative specification of the Torchscript IR

jit
1

Hi - We are trying to implement a lowering of Torchscript to MLIR/ATen dialect as part of the LLVM npcomp incubator effort [1]. Is there a declarative specification of the Torchscript IR we could leverage for this work ? Thanks

[1] https://llvm.discourse.group/t/npcomp-next-steps-for-torch-ir-aten-dialect/1777

1 Like
2

I think the closest thing to what you are looking for is the core program representation section of the JIT overview documentation.

3

Thanks - that is a really helpful and nicely written doc.

Are there any facilities for systematically enumerating the ops and how they are materialized in the IR? (i.e. I’ve got a harness that traces each op and then extracts the corresponding IR to make such a mapping, but it would be better if this was in some machine-usable source form somewhere)

4

If by ops, you mean things like convolutions, not really. Most ops of that sort show up in the IR as aten::op, but they are not a formally defined part of the JIT IR. There is no formal restriction on which ops are “allowed” and which ops are not. They map to operators in ATen, some of which can be user-defined and registered outside of the core PyTorch library.

5

Thanks - that was my understanding but was just double checking that I wasn’t missing something obvious.