Summary
- With
torch.export(), an EXIR (export intermediate representation) is generated with ATen dialect. All AOT compilations are based on this EXIR, but can have multiple dialects along the lowering path as detailed below
example_args = (torch.randn(1, 3, 256, 256),)
aten_dialect: ExportedProgram = export(SimpleConv(), example_args)-
ATen Dialect. PyTorch Edge is based on PyTorch’s Tensor library ATen, which has clear contracts for efficient execution. dw
- ATen Dialect is a graph represented by ATen nodes which are fully ATen compliant.
- Custom operators are allowed, but must be registered with the dispatcher
- It’s flat with no module hierarchy (submodules in a bigger module), but the source code and module hierarchy are preserved in the metadata.
- This representation is also autograd safe
- (quantization, either QAT (quantization-aware training) or PTQ (post training quantization) can be applied to the whole ATen graph before converting to Core ATen.)
-
Core ATen Dialect. ATen has thousands of operators. It’s not ideal for some fundamental transforms and kernel library implementation.
- The operators from the ATen Dialect graph are decomposed into fundamental operators so that the operator set (op set) is smaller and more fundamental transforms can be applied.
- The Core ATen dialect is also serializable and convertible to Edge Dialect.