affine.max are operations that take operands and return a value of type
index and an attribute of builtin type
The particularity of the builtin type
AffineMap is that it decomposes its operands between so-called “dimensions” and “symbols”. This decomposition was originally motivated by polyhedral compilation and Presburger arithmetic in which we distinguish between:
- runtime quantities that follow an arithmetic progression (called “dimensions” in MLIR parlance)
- runtime quantities that are constant but statically unknown at compile-time (called “symbols” in MLIR parlance)
For those who remember the Polylib tools , these were called “domain iterators” and “parameters”.
One interesting property is that the definition of a “dimension” and a “symbol” is not universal: depending on which “op with a region” is considered the root of a particular analysis, the classification of which variable can be interpreted as a “dimension” or a “symbol” may change.
affine.apply was originally introduced as a means to implement classical analysis subject to polyhedral and affine restrictions in MLIR.
affine.apply has emerged as a good abstraction to represent more general notions of indexing that revolve around the
AffineMap::compose method that do not have restrictions about what can be considered a dim and what must be considered as a symbol.
As the evolution of usages evolved in MLIR, it now appears clear that the implementation of
affine.apply conflates 2 independent things:
- the machinery and rewrites based on mathematical properties of
builtin::AffineMapthat are universal and reusable.
- the polyhedral-era restrictions on what is allowed to be a “dimension” and a “symbol” in the context of an enclosing op with a region; also known as
AffineScopein MLIR (not to be mistaken with a traditional polyhedral SCoP (static control part)).
As a result of this conflation,
affine.apply depends on
This makes it impossible to have a
tensor.expand/collapse_shape that matches
memref.expand/collapse_shape due to cyclic dependencies.
More generally, such an
apply operation is a better representation than chains of AddIOp/MulIOp to manipulate in the IR and generally composes better, irrespective of polyhedral-inspired analyses.
Additionally, in the longer term, we anticipate to benefit from canonicalizations and rewrites that mix affine and non-affine indexing. In particular, I have always been a fan of the conciseness of the Halide folding rules once a more general notion of
IndexingExpr is available; see for instance Halide/Simplify_And.cpp at master · halide/Halide · GitHub. We hope we can also look at recursive indexing for FFT/sort and other recursive quantities such as
Decouple the machinery and rewrites based on mathematical properties of
arith dialect operations and keep AffineScope-related considerations to the
affine dialect where they belong. Since
AffineMap is already
builtin, this should be relatively mechanical.
I am not too attached to the specifics but here are a few questions I expect others will have strong opinions on:
- start a new
indexdialect vs plop ops into
arithand split out later on a per-need basis
affine_apply/affine_min/affine_maxvs other names
- how to actually carve out the AffineScope related constraints and make them compose properly. This may be tricky as
affinehas historically not been very composable with other things.
Thanks for reading !