We’ve mentioned the shape dialect in our documents and on the mailing list a few times, but mostly in the context of shape inference instead of as a dialect on its own. The goal of the shape dialect is to develop in MLIR core a dialect for shaped type constraints specification/solving. In particular interest it will be used for
ShapedType (for example, to determine the output shape and element type of a result) as representation of shape functions.
I’ll describe the dialect following the MLIR developer guide’s guidelines for contributing a new dialect, here are the elements:
What is the overall goal of the dialect? What is the first implementation milestone?
The goal of the dialect is to describe the shape constraints of operations in a manner that could be consumed by different backends (either during compilation or for a runtime). Shape functions should be expressible without being tied to a specific use case or obscured by lowering considerations (e.g., returning early when an error encountered should be possible to do without being expressly part of shape function description). The first milestone would be to be able to represent all operation shape constraints in OpBase and to automatically generate both verification and build methods for operations fully constrained by their shape functions.
How does it fit into the MLIR dialect ecosystem? Connection: how does it connect to the existing dialects in a compilation pipeline(s)?
The dialect is intended to be used by ODS/DRR and other dialects to define & solve shape constraints for operations defined in other dialects. But in turn functions in this dialect could be generated from other dialects (e.g., a shape function for a structured ops could be generated from the iteration descriptions of the op).
Consolidation: is there already a dialect with a similar goal or matching abstractions; if so, can it be improved instead of adding a new one?
None in MLIR core, there are existing library functions and different users have different shape inference implementations, but no dialect with the same goal. The standard and affine arithmetic operations could be used to specify shape function too, and for static shaped computations might be sufficient but for dynamic shapes would result in unnecessarily complicating/obscuring the functions (e.g., having
shape.add operate on two dimensions which may be dynamic should result in the correct addition result vs needing to insert additional conditionals to handle dynamic cases inside the shape function if one were to represent this using standard dialect).
Reuse: how does it generalize to similar but slightly different use-cases? What is the community of users that it is serving?
A successful implementation should allow specifying shape functions for compilers, runtimes and different solvers:
- It enables automatically generating build methods for operations whose constraints fully define an operation. Functions in this dialect would be used to generate new build methods, making it easier to add new operations where the verification and inference is based off of the same constraints.
- It could be used to generate constraint for symbolic solvers.
- Shape functions could be reified into dynamic allocation/checking calls for runtime memory checking/allocation.
In particular this would be used to represent TensorFlow, TFLite and XLA shape functions. nGraph, the tensor compute dialect and others could have similar needs.
Who are the future contributors/maintainers beyond those who propose the dialect?
Folks who need shape inference for their ops and want a common dialect to perform optimization, reification and constraint solving with.