xref: /llvm-project/mlir/docs/Dialects/Affine.md (revision d28a4f1fc02dc34a87fa22af0a053e8f1e7f6cea)
1# 'affine' Dialect
2
3This dialect provides a powerful abstraction for affine operations and analyses.
4
5[TOC]
6
7## Polyhedral Structures
8
9MLIR uses techniques from polyhedral compilation to make dependence analysis and
10loop transformations efficient and reliable. This section introduces some of the
11core concepts that are used throughout the document.
12
13### Dimensions and Symbols
14
15Dimensions and symbols are the two kinds of identifiers that can appear in the
16polyhedral structures, and are always of [`index`](Builtin.md/#indextype) type.
17Dimensions are declared in parentheses and symbols are declared in square
18brackets.
19
20Examples:
21
22```mlir
23// A 2d to 3d affine mapping.
24// d0/d1 are dimensions, s0 is a symbol
25#affine_map2to3 = affine_map<(d0, d1)[s0] -> (d0, d1 + s0, d1 - s0)>
26```
27
28Dimensional identifiers correspond to the dimensions of the underlying structure
29being represented (a map, set, or more concretely a loop nest or a tensor); for
30example, a three-dimensional loop nest has three dimensional identifiers. Symbol
31identifiers represent an unknown quantity that can be treated as constant for a
32region of interest.
33
34Dimensions and symbols are bound to SSA values by various operations in MLIR and
35use the same parenthesized vs square bracket list to distinguish the two.
36
37Syntax:
38
39```
40// Uses of SSA values that are passed to dimensional identifiers.
41dim-use-list ::= `(` ssa-use-list? `)`
42
43// Uses of SSA values that are used to bind symbols.
44symbol-use-list ::= `[` ssa-use-list? `]`
45
46// Most things that bind SSA values bind dimensions and symbols.
47dim-and-symbol-use-list ::= dim-use-list symbol-use-list?
48```
49
50SSA values bound to dimensions and symbols must always have 'index' type.
51
52Example:
53
54```mlir
55#affine_map2to3 = affine_map<(d0, d1)[s0] -> (d0, d1 + s0, d1 - s0)>
56// Binds %N to the s0 symbol in affine_map2to3.
57%x = memref.alloc()[%N] : memref<40x50xf32, #affine_map2to3>
58```
59
60### Restrictions on Dimensions and Symbols
61
62The affine dialect imposes certain restrictions on dimension and symbolic
63identifiers to enable powerful analysis and transformation. An SSA value's use
64can be bound to a symbolic identifier if that SSA value is either:
65
661. a region argument for an op with trait `AffineScope` (eg. `FuncOp`),
672. a value defined at the top level of an `AffineScope` op (i.e.,
68immediately enclosed by the latter),
693. a value that dominates the `AffineScope` op enclosing the value's
70use,
714. the result of a constant operation,
725. the result of a `Pure` operation whose operands are valid symbolic identifiers.
736. the result of a
74[`dim` operation](MemRef.md/#memrefdim-mlirmemrefdimop) on either a memref that
75is an argument to a `AffineScope` op or a memref where the corresponding
76dimension is either static or a dynamic one in turn bound to a valid symbol.
77
78*Note:* if the use of an SSA value is not contained in any op with the
79`AffineScope` trait, only the rules 4-6 can be applied.
80
81Note that as a result of rule (3) above, symbol validity is sensitive to the
82location of the SSA use. Dimensions may be bound not only to anything that a
83symbol is bound to, but also to induction variables of enclosing
84[`affine.for`](#affinefor-mliraffineforop) and
85[`affine.parallel`](#affineparallel-mliraffineparallelop) operations, and the result
86of an [`affine.apply` operation](#affineapply-mliraffineapplyop) (which recursively
87may use other dimensions and symbols).
88
89### Affine Expressions
90
91Syntax:
92
93```
94affine-expr ::= `(` affine-expr `)`
95              | affine-expr `+` affine-expr
96              | affine-expr `-` affine-expr
97              | `-`? integer-literal `*` affine-expr
98              | affine-expr `ceildiv` integer-literal
99              | affine-expr `floordiv` integer-literal
100              | affine-expr `mod` integer-literal
101              | `-`affine-expr
102              | bare-id
103              | `-`? integer-literal
104
105multi-dim-affine-expr ::= `(` `)`
106                        | `(` affine-expr (`,` affine-expr)* `)`
107```
108
109`ceildiv` is the ceiling function which maps the result of the division of its
110first argument by its second argument to the smallest integer greater than or
111equal to that result. `floordiv` is a function which maps the result of the
112division of its first argument by its second argument to the largest integer
113less than or equal to that result. `mod` is the modulo operation: since its
114second argument is always positive, its results are always positive in our
115usage. The `integer-literal` operand for ceildiv, floordiv, and mod is always
116expected to be positive. `bare-id` is an identifier which must have type
117[index](Builtin.md/#indextype). The precedence of operations in an affine
118expression are ordered from highest to lowest in the order: (1)
119parenthesization, (2) negation, (3) modulo, multiplication, floordiv, and
120ceildiv, and (4) addition and subtraction. All of these operators associate from
121left to right.
122
123A *multidimensional affine expression* is a comma separated list of
124one-dimensional affine expressions, with the entire list enclosed in
125parentheses.
126
127**Context:** An affine function, informally, is a linear function plus a
128constant. More formally, a function f defined on a vector $\vec{v} \in
129\mathbb{Z}^n$ is a multidimensional affine function of $\vec{v}$ if $f(\vec{v})$
130can be expressed in the form $M \vec{v} + \vec{c}$ where $M$ is a constant
131matrix from $\mathbb{Z}^{m \times n}$ and $\vec{c}$ is a constant vector from
132$\mathbb{Z}$. $m$ is the dimensionality of such an affine function. MLIR further
133extends the definition of an affine function to allow 'floordiv', 'ceildiv', and
134'mod' with respect to positive integer constants. Such extensions to affine
135functions have often been referred to as quasi-affine functions by the
136polyhedral compiler community. MLIR uses the term 'affine map' to refer to these
137multidimensional quasi-affine functions. As examples, $(i+j+1, j)$, $(i \mod 2,
138j+i)$, $(j, i/4, i \mod 4)$, $(2i+1, j)$ are two-dimensional affine functions of
139$(i, j)$, but $(i \cdot j, i^2)$, $(i \mod j, i/j)$ are not affine functions of
140$(i, j)$.
141
142### Affine Maps
143
144Syntax:
145
146```
147affine-map-inline
148   ::= dim-and-symbol-value-lists `->` multi-dim-affine-expr
149```
150
151The identifiers in the dimensions and symbols lists must be unique. These are
152the only identifiers that may appear in 'multi-dim-affine-expr'. Affine maps
153with one or more symbols in its specification are known as "symbolic affine
154maps", and those with no symbols as "non-symbolic affine maps".
155
156**Context:** Affine maps are mathematical functions that transform a list of
157dimension indices and symbols into a list of results, with affine expressions
158combining the indices and symbols. Affine maps distinguish between
159[indices and symbols](#dimensions-and-symbols) because indices are inputs to the
160affine map when the map is called (through an operation such as
161[affine.apply](#affineapply-mliraffineapplyop)), whereas symbols are bound when the
162map is established (e.g. when a memref is formed, establishing a memory
163[layout map](Builtin.md/#layout)).
164
165Affine maps are used for various core structures in MLIR. The restrictions we
166impose on their form allows powerful analysis and transformation, while keeping
167the representation closed with respect to several operations of interest.
168
169#### Named affine mappings
170
171Syntax:
172
173```
174affine-map-id ::= `#` suffix-id
175
176// Definitions of affine maps are at the top of the file.
177affine-map-def    ::= affine-map-id `=` affine-map-inline
178module-header-def ::= affine-map-def
179
180// Uses of affine maps may use the inline form or the named form.
181affine-map ::= affine-map-id | affine-map-inline
182```
183
184Affine mappings may be defined inline at the point of use, or may be hoisted to
185the top of the file and given a name with an affine map definition, and used by
186name.
187
188Examples:
189
190```mlir
191// Affine map out-of-line definition and usage example.
192#affine_map42 = affine_map<(d0, d1)[s0] -> (d0, d0 + d1 + s0 floordiv 2)>
193
194// Use an affine mapping definition in an alloc operation, binding the
195// SSA value %N to the symbol s0.
196%a = memref.alloc()[%N] : memref<4x4xf32, #affine_map42>
197
198// Same thing with an inline affine mapping definition.
199%b = memref.alloc()[%N] : memref<4x4xf32, affine_map<(d0, d1)[s0] -> (d0, d0 + d1 + s0 floordiv 2)>>
200```
201
202### Semi-affine maps
203
204Semi-affine maps are extensions of affine maps to allow multiplication,
205`floordiv`, `ceildiv`, and `mod` with respect to symbolic identifiers.
206Semi-affine maps are thus a strict superset of affine maps.
207
208Syntax of semi-affine expressions:
209
210```
211semi-affine-expr ::= `(` semi-affine-expr `)`
212                   | semi-affine-expr `+` semi-affine-expr
213                   | semi-affine-expr `-` semi-affine-expr
214                   | symbol-or-const `*` semi-affine-expr
215                   | semi-affine-expr `ceildiv` symbol-or-const
216                   | semi-affine-expr `floordiv` symbol-or-const
217                   | semi-affine-expr `mod` symbol-or-const
218                   | bare-id
219                   | `-`? integer-literal
220
221symbol-or-const ::= `-`? integer-literal | symbol-id
222
223multi-dim-semi-affine-expr ::= `(` semi-affine-expr (`,` semi-affine-expr)* `)`
224```
225
226The precedence and associativity of operations in the syntax above is the same
227as that for [affine expressions](#affine-expressions).
228
229Syntax of semi-affine maps:
230
231```
232semi-affine-map-inline
233   ::= dim-and-symbol-value-lists `->` multi-dim-semi-affine-expr
234```
235
236Semi-affine maps may be defined inline at the point of use, or may be hoisted to
237the top of the file and given a name with a semi-affine map definition, and used
238by name.
239
240```
241semi-affine-map-id ::= `#` suffix-id
242
243// Definitions of semi-affine maps are at the top of file.
244semi-affine-map-def ::= semi-affine-map-id `=` semi-affine-map-inline
245module-header-def ::= semi-affine-map-def
246
247// Uses of semi-affine maps may use the inline form or the named form.
248semi-affine-map ::= semi-affine-map-id | semi-affine-map-inline
249```
250
251### Integer Sets
252
253An integer set is a conjunction of affine constraints on a list of identifiers.
254The identifiers associated with the integer set are separated out into two
255classes: the set's dimension identifiers, and the set's symbolic identifiers.
256The set is viewed as being parametric on its symbolic identifiers. In the
257syntax, the list of set's dimension identifiers are enclosed in parentheses
258while its symbols are enclosed in square brackets.
259
260Syntax of affine constraints:
261
262```
263affine-constraint ::= affine-expr `>=` `affine-expr`
264                    | affine-expr `<=` `affine-expr`
265                    | affine-expr `==` `affine-expr`
266affine-constraint-conjunction ::= affine-constraint (`,` affine-constraint)*
267```
268
269Integer sets may be defined inline at the point of use, or may be hoisted to the
270top of the file and given a name with an integer set definition, and used by
271name.
272
273```
274integer-set-id ::= `#` suffix-id
275
276integer-set-inline
277   ::= dim-and-symbol-value-lists `:` '(' affine-constraint-conjunction? ')'
278
279// Declarations of integer sets are at the top of the file.
280integer-set-decl ::= integer-set-id `=` integer-set-inline
281
282// Uses of integer sets may use the inline form or the named form.
283integer-set ::= integer-set-id | integer-set-inline
284```
285
286The dimensionality of an integer set is the number of identifiers appearing in
287dimension list of the set. The affine-constraint non-terminals appearing in the
288syntax above are only allowed to contain identifiers from dims and symbols. A
289set with no constraints is a set that is unbounded along all of the set's
290dimensions.
291
292Example:
293
294```mlir
295// A example two-dimensional integer set with two symbols.
296#set42 = affine_set<(d0, d1)[s0, s1]
297   : (d0 >= 0, -d0 + s0 - 1 >= 0, d1 >= 0, -d1 + s1 - 1 >= 0)>
298
299// Inside a Region
300affine.if #set42(%i, %j)[%M, %N] {
301  ...
302}
303```
304
305`d0` and `d1` correspond to dimensional identifiers of the set, while `s0` and
306`s1` are symbol identifiers.
307
308## Operations
309
310[include "Dialects/AffineOps.md"]
311
312### `affine.dma_start` (mlir::AffineDmaStartOp)
313
314Syntax:
315
316```
317operation ::= `affine.dma_start` ssa-use `[` multi-dim-affine-map-of-ssa-ids `]`, `[` multi-dim-affine-map-of-ssa-ids `]`, `[` multi-dim-affine-map-of-ssa-ids `]`, ssa-use `:` memref-type
318```
319
320The `affine.dma_start` op starts a non-blocking DMA operation that transfers
321data from a source memref to a destination memref. The source and destination
322memref need not be of the same dimensionality, but need to have the same
323elemental type. The operands include the source and destination memref's each
324followed by its indices, size of the data transfer in terms of the number of
325elements (of the elemental type of the memref), a tag memref with its indices,
326and optionally at the end, a stride and a number_of_elements_per_stride
327arguments. The tag location is used by an AffineDmaWaitOp to check for
328completion. The indices of the source memref, destination memref, and the tag
329memref have the same restrictions as any affine.load/store. In particular, index
330for each memref dimension must be an affine expression of loop induction
331variables and symbols. The optional stride arguments should be of 'index' type,
332and specify a stride for the slower memory space (memory space with a lower
333memory space id), transferring chunks of number_of_elements_per_stride every
334stride until %num_elements are transferred. Either both or no stride arguments
335should be specified. The value of 'num_elements' must be a multiple of
336'number_of_elements_per_stride'.
337
338Example 1:
339
340For example, a `DmaStartOp` operation that transfers 256 elements of a memref
341`%src` in memory space 0 at indices `[%i + 3, %j]` to memref `%dst` in memory
342space 1 at indices `[%k + 7, %l]`, would be specified as follows:
343
344
345```mlir
346%num_elements = arith.constant 256
347%idx = arith.constant 0 : index
348%tag = memref.alloc() : memref<1xi32, 4>
349affine.dma_start %src[%i + 3, %j], %dst[%k + 7, %l], %tag[%idx],
350  %num_elements :
351    memref<40x128xf32, 0>, memref<2x1024xf32, 1>, memref<1xi32, 2>
352```
353
354Example 2:
355
356If `%stride` and `%num_elt_per_stride` are specified, the DMA is expected to
357transfer `%num_elt_per_stride` elements every `%stride elements` apart from
358memory space 0 until `%num_elements` are transferred.
359
360```mlir
361affine.dma_start %src[%i, %j], %dst[%k, %l], %tag[%idx], %num_elements,
362  %stride, %num_elt_per_stride : ...
363```
364
365### `affine.dma_wait` (mlir::AffineDmaWaitOp)
366
367Syntax:
368
369```
370operation ::= `affine.dma_wait` ssa-use `[` multi-dim-affine-map-of-ssa-ids `]`, ssa-use `:` memref-type
371```
372
373The `affine.dma_wait` op blocks until the completion of a DMA operation
374associated with the tag element `%tag[%index]`. `%tag` is a memref, and `%index` has
375to be an index with the same restrictions as any load/store index. In
376particular, index for each memref dimension must be an affine expression of loop
377induction variables and symbols. `%num_elements` is the number of elements
378associated with the DMA operation.
379
380Example:
381
382```mlir
383affine.dma_start %src[%i, %j], %dst[%k, %l], %tag[%index], %num_elements :
384  memref<2048xf32, 0>, memref<256xf32, 1>, memref<1xi32, 2>
385...
386...
387affine.dma_wait %tag[%index], %num_elements : memref<1xi32, 2>
388```
389