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
-
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.
Detailed
Writing exportable modeling files
Dynamic shapes
- https://pytorch.org/docs/main/export.html#expressing-dynamism
- export will trace through the exact shapes it’s given if not specified