xref: /llvm-project/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td (revision 5d26a6d7590f13d21d78f7f0a443b92b04c80f98)
1//===- AffineOps.td - Affine operation definitions ---------*- tablegen -*-===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8//
9// Defines MLIR affine operations.
10//
11//===----------------------------------------------------------------------===//
12
13#ifndef AFFINE_OPS
14#define AFFINE_OPS
15
16include "mlir/Dialect/Arith/IR/ArithBase.td"
17include "mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.td"
18include "mlir/Interfaces/ControlFlowInterfaces.td"
19include "mlir/Interfaces/InferTypeOpInterface.td"
20include "mlir/Interfaces/LoopLikeInterface.td"
21include "mlir/Interfaces/SideEffectInterfaces.td"
22
23def Affine_Dialect : Dialect {
24  let name = "affine";
25  let cppNamespace = "::mlir::affine";
26  let hasConstantMaterializer = 1;
27  let dependentDialects = ["arith::ArithDialect", "ub::UBDialect"];
28}
29
30// Base class for Affine dialect ops.
31class Affine_Op<string mnemonic, list<Trait> traits = []> :
32    Op<Affine_Dialect, mnemonic, traits>;
33
34// Require regions to have affine.yield.
35def ImplicitAffineTerminator
36    : SingleBlockImplicitTerminator<"AffineYieldOp">;
37
38def AffineApplyOp : Affine_Op<"apply", [Pure]> {
39  let summary = "affine apply operation";
40  let description = [{
41    The `affine.apply` operation applies an [affine mapping](#affine-maps)
42    to a list of SSA values, yielding a single SSA value. The number of
43    dimension and symbol arguments to `affine.apply` must be equal to the
44    respective number of dimensional and symbolic inputs to the affine mapping;
45    the affine mapping has to be one-dimensional, and so the `affine.apply`
46    operation always returns one value. The input operands and result must all
47    have ‘index’ type.
48
49    Example:
50
51    ```mlir
52    #map10 = affine_map<(d0, d1) -> (d0 floordiv 8 + d1 floordiv 128)>
53    ...
54    %1 = affine.apply #map10 (%s, %t)
55
56    // Inline example.
57    %2 = affine.apply affine_map<(i)[s0] -> (i+s0)> (%42)[%n]
58    ```
59  }];
60  let arguments = (ins AffineMapAttr:$map, Variadic<Index>:$mapOperands);
61  let results = (outs Index);
62
63  // TODO: The auto-generated builders should check to see if the return type
64  // has a constant builder. That way we wouldn't need to explicitly specify the
65  // result types here.
66  let builders = [
67    OpBuilder<(ins "ArrayRef<AffineExpr> ":$exprList,"ValueRange":$mapOperands),
68    [{
69      build($_builder, $_state, $_builder.getIndexType(),
70            AffineMap::inferFromExprList(exprList, $_builder.getContext())
71                                        .front(), mapOperands);
72    }]>
73  ];
74
75  let extraClassDeclaration = [{
76    /// Returns the affine map to be applied by this operation.
77    AffineMap getAffineMap() { return getMap(); }
78
79    /// Returns the affine value map computed from this operation.
80    AffineValueMap getAffineValueMap();
81
82    /// Returns true if the result of this operation can be used as dimension id
83    /// in the region of the closest surrounding op with trait AffineScope.
84    bool isValidDim();
85
86    /// Returns true if the result of this operation can be used as dimension id
87    /// within 'region', i.e., for all its uses with `region`.
88    bool isValidDim(Region *region);
89
90    /// Returns true if the result of this operation is a symbol in the region
91    /// of the closest surrounding op that has the trait AffineScope.
92    bool isValidSymbol();
93
94    /// Returns true if the result of this operation is a symbol for all its
95    /// uses in `region`.
96    bool isValidSymbol(Region *region);
97
98    /// Returns all dimension operands.
99    ValueRange getDimOperands() {
100      return OperandRange{getOperands().begin(),
101                          getOperands().begin() + getMap().getNumDims()};
102    }
103
104    /// Returns all symbol operands.
105    ValueRange getSymbolOperands() {
106      return OperandRange{getOperands().begin() + getMap().getNumDims(),
107                          getOperands().end()};
108    }
109  }];
110
111  let hasCanonicalizer = 1;
112  let hasCustomAssemblyFormat = 1;
113  let hasFolder = 1;
114  let hasVerifier = 1;
115}
116
117def AffineForOp : Affine_Op<"for",
118    [AttrSizedOperandSegments, AutomaticAllocationScope,
119     ImplicitAffineTerminator, ConditionallySpeculatable,
120     RecursiveMemoryEffects, DeclareOpInterfaceMethods<LoopLikeOpInterface,
121     ["getLoopInductionVars", "getLoopLowerBounds", "getLoopSteps",
122      "getLoopUpperBounds", "getYieldedValuesMutable",
123      "replaceWithAdditionalYields"]>,
124     DeclareOpInterfaceMethods<RegionBranchOpInterface,
125     ["getEntrySuccessorOperands"]>]> {
126  let summary = "for operation";
127  let description = [{
128    Syntax:
129
130    ```
131    operation   ::= `affine.for` ssa-id `=` lower-bound `to` upper-bound
132                    (`step` integer-literal)? `{` op* `}`
133
134    lower-bound ::= `max`? affine-map-attribute dim-and-symbol-use-list | shorthand-bound
135    upper-bound ::= `min`? affine-map-attribute dim-and-symbol-use-list | shorthand-bound
136    shorthand-bound ::= ssa-id | `-`? integer-literal
137    ```
138
139    The `affine.for` operation represents an affine loop nest. It has one region
140    containing its body. This region must contain one block that terminates with
141    [`affine.yield`](#affineyield-mliraffineyieldop). *Note:* when
142    `affine.for` is printed in custom format, the terminator is omitted. The
143    block has one argument of [`index`](Builtin.md/#indextype) type that
144    represents the induction variable of the loop.
145
146    The `affine.for` operation executes its body a number of times iterating
147    from a lower bound to an upper bound by a stride. The stride, represented by
148    `step`, is a positive constant integer which defaults to "1" if not present.
149    The lower and upper bounds specify a half-open range: the range includes the
150    lower bound but does not include the upper bound.
151
152    The lower and upper bounds of a `affine.for` operation are represented as an
153    application of an affine mapping to a list of SSA values passed to the map.
154    The [same restrictions](#restrictions-on-dimensions-and-symbols) hold for
155    these SSA values as for all bindings of SSA values to dimensions and
156    symbols.
157
158    The affine mappings for the bounds may return multiple results, in which
159    case the `max`/`min` keywords are required (for the lower/upper bound
160    respectively), and the bound is the maximum/minimum of the returned values.
161    There is no semantic ambiguity, but MLIR syntax requires the use of these
162    keywords to make things more obvious to human readers.
163
164    Many upper and lower bounds are simple, so MLIR accepts two custom form
165    syntaxes: the form that accepts a single 'ssa-id' (e.g. `%N`) is shorthand
166    for applying that SSA value to a function that maps a single symbol to
167    itself, e.g., `()[s]->(s)()[%N]`. The integer literal form (e.g. `-42`) is
168    shorthand for a nullary mapping function that returns the constant value
169    (e.g. `()->(-42)()`).
170
171    Example showing reverse iteration of the inner loop:
172
173    ```mlir
174    #map57 = affine_map<(d0)[s0] -> (s0 - d0 - 1)>
175
176    func.func @simple_example(%A: memref<?x?xf32>, %B: memref<?x?xf32>) {
177      %N = dim %A, 0 : memref<?x?xf32>
178      affine.for %i = 0 to %N step 1 {
179        affine.for %j = 0 to %N {   // implicitly steps by 1
180          %0 = affine.apply #map57(%j)[%N]
181          %tmp = call @F1(%A, %i, %0) : (memref<?x?xf32>, index, index)->(f32)
182          call @F2(%tmp, %B, %i, %0) : (f32, memref<?x?xf32>, index, index)->()
183        }
184      }
185      return
186    }
187    ```
188    `affine.for` can also operate on loop-carried variables (`iter_args`) and
189    return the final values after loop termination. The initial values of the
190    variables are passed as additional SSA operands to the `affine.for`
191    following the operands for the loop's lower and upper bounds. The
192    operation's region has equivalent arguments for each variable representing
193    the value of the variable at the current iteration.
194
195    The region must terminate with an `affine.yield` that passes all the current
196    iteration variables to the next iteration, or to the `affine.for`'s results
197    if at the last iteration. For `affine.for`'s that execute zero iterations, the
198    initial values of the loop-carried variables (corresponding to the SSA
199    operands) will be the op's results.
200
201    For example, to sum-reduce a memref:
202
203     ```mlir
204    func.func @reduce(%buffer: memref<1024xf32>) -> (f32) {
205      // Initial sum set to 0.
206      %sum_0 = arith.constant 0.0 : f32
207      // iter_args binds initial values to the loop's region arguments.
208      %sum = affine.for %i = 0 to 10 step 2
209          iter_args(%sum_iter = %sum_0) -> (f32) {
210        %t = affine.load %buffer[%i] : memref<1024xf32>
211        %sum_next = arith.addf %sum_iter, %t : f32
212        // Yield current iteration sum to next iteration %sum_iter or to %sum
213        // if final iteration.
214        affine.yield %sum_next : f32
215      }
216      return %sum : f32
217    }
218    ```
219
220    ```mlir
221    %res:2 = affine.for %i = 0 to 128 iter_args(%arg0 = %init0, %arg1 = %init1)
222               -> (index, index) {
223      %y0 = arith.addi %arg0, %c1 : index
224      %y1 = arith.addi %arg1, %c2 : index
225      affine.yield %y0, %y1 : index, index
226    }
227    ```
228    If the `affine.for` defines any values, a yield terminator must be
229    explicitly present. The number and types of the "affine.for" results must
230    match the initial values in the `iter_args` binding and the yield operands.
231  }];
232  let arguments = (ins Variadic<Index>:$lowerBoundOperands,
233                       Variadic<Index>:$upperBoundOperands,
234                       Variadic<AnyType>:$inits,
235                       AffineMapAttr:$lowerBoundMap,
236                       AffineMapAttr:$upperBoundMap,
237                       IndexAttr:$step);
238  let results = (outs Variadic<AnyType>:$results);
239  let regions = (region SizedRegion<1>:$region);
240
241  let skipDefaultBuilders = 1;
242  let builders = [
243    OpBuilder<(ins "int64_t":$lowerBound, "int64_t":$upperBound,
244      CArg<"int64_t", "1">:$step, CArg<"ValueRange", "std::nullopt">:$iterArgs,
245      CArg<"function_ref<void(OpBuilder &, Location, Value, ValueRange)>",
246           "nullptr">:$bodyBuilder)>,
247    OpBuilder<(ins "ValueRange":$lbOperands, "AffineMap":$lbMap,
248      "ValueRange":$ubOperands, "AffineMap":$ubMap, CArg<"int64_t", "1">:$step,
249      CArg<"ValueRange", "std::nullopt">:$iterArgs,
250      CArg<"function_ref<void(OpBuilder &, Location, Value, ValueRange)>",
251           "nullptr">:$bodyBuilder)>
252  ];
253
254  let extraClassDeclaration = [{
255    /// Defining the function type we use for building the body of affine.for.
256    using BodyBuilderFn =
257        function_ref<void(OpBuilder &, Location, Value, ValueRange)>;
258
259    BlockArgument getInductionVar() { return getBody()->getArgument(0); }
260    Block::BlockArgListType getRegionIterArgs() {
261      return getBody()->getArguments().drop_front();
262    }
263
264    /// Returns operands for the lower and upper bound maps with the operands
265    /// for the lower bound map in front of those for the upper bound map.
266    operand_range getControlOperands();
267
268    /// Returns information about the lower bound as a single object.
269    AffineBound getLowerBound();
270
271    /// Returns information about the upper bound as a single object.
272    AffineBound getUpperBound();
273
274    /// Returns loop step.
275    int64_t getStepAsInt() { return getStep().getSExtValue(); }
276
277    /// Set lower bound. The new bound must have the same number of operands as
278    /// the current bound map. Otherwise, 'replaceForLowerBound' should be used.
279    void setLowerBound(ValueRange operands, AffineMap map);
280    /// Set upper bound. The new bound must not have more operands than the
281    /// current bound map. Otherwise, 'replaceForUpperBound' should be used.
282    void setUpperBound(ValueRange operands, AffineMap map);
283
284    /// Set loop step.
285    void setStep(int64_t step) {
286      assert(step > 0 && "step has to be a positive integer constant");
287      setStep(APInt(/*numBits=*/64, step, /*isSigned=*/true));
288    }
289
290    /// Returns number of region arguments for loop-carried values.
291    unsigned getNumRegionIterArgs() {
292      return getBody()->getNumArguments() - 1;
293    }
294
295    /// Number of operands controlling the loop: lb and ub.
296    unsigned getNumControlOperands() {
297      return getOperation()->getNumOperands() - getNumIterOperands();
298    }
299
300    /// Get the number of loop-carried values.
301    unsigned getNumIterOperands();
302
303    /// Returns true if the lower bound is constant.
304    bool hasConstantLowerBound();
305    /// Returns true if the upper bound is constant.
306    bool hasConstantUpperBound();
307    /// Returns true if both bounds are constant.
308    bool hasConstantBounds() {
309      return hasConstantLowerBound() && hasConstantUpperBound();
310    }
311    /// Returns the value of the constant lower bound.
312    /// Fails assertion if the bound is non-constant.
313    int64_t getConstantLowerBound();
314    /// Returns the value of the constant upper bound. The upper bound is
315    /// exclusive. Fails assertion if the bound is non-constant.
316    int64_t getConstantUpperBound();
317    /// Sets the lower bound to the given constant value.
318    void setConstantLowerBound(int64_t value);
319    /// Sets the upper bound to the given constant value.
320    void setConstantUpperBound(int64_t value);
321
322    /// Returns true if both the lower and upper bound have the same operand
323    /// lists (same operands in the same order).
324    bool matchingBoundOperandList();
325
326    /// Interface method for ConditionallySpeculatable.
327    Speculation::Speculatability getSpeculatability();
328  }];
329
330  let hasCanonicalizer = 1;
331  let hasCustomAssemblyFormat = 1;
332  let hasFolder = 1;
333  let hasRegionVerifier = 1;
334}
335
336def AffineIfOp : Affine_Op<"if",
337                           [ImplicitAffineTerminator, RecursivelySpeculatable,
338                            RecursiveMemoryEffects, NoRegionArguments,
339                            DeclareOpInterfaceMethods<RegionBranchOpInterface>
340                           ]> {
341  let summary = "if-then-else operation";
342  let description = [{
343    Syntax:
344
345    ```
346    operation  ::= `affine.if` if-op-cond `{` op* `}` (`else` `{` op* `}`)?
347    if-op-cond ::= integer-set-attr dim-and-symbol-use-list
348    ```
349
350    The `affine.if` operation restricts execution to a subset of the loop
351    iteration space defined by an integer set (a conjunction of affine
352    constraints). A single `affine.if` may end with an optional `else` clause.
353
354    The condition of the `affine.if` is represented by an
355    [integer set](#integer-sets) (a conjunction of affine constraints),
356    and the SSA values bound to the dimensions and symbols in the integer set.
357    The [same restrictions](#restrictions-on-dimensions-and-symbols) hold for
358    these SSA values as for all bindings of SSA values to dimensions and
359    symbols.
360
361    The `affine.if` operation contains two regions for the "then" and "else"
362    clauses.  `affine.if` may return results that are defined in its regions.
363    The values defined are determined by which execution path is taken.  Each
364    region of the `affine.if` must contain a single block with no arguments,
365    and be terminated by `affine.yield`.  If `affine.if` defines no values,
366    the `affine.yield` can be left out, and will be inserted implicitly.
367    Otherwise, it must be explicit.  If no values are defined, the else block
368    may be empty (i.e. contain no blocks).
369
370    Example:
371
372    ```mlir
373    #set = affine_set<(d0, d1)[s0]: (d0 - 10 >= 0, s0 - d0 - 9 >= 0,
374                                     d1 - 10 >= 0, s0 - d1 - 9 >= 0)>
375    func.func @reduced_domain_example(%A, %X, %N) : (memref<10xi32>, i32, i32) {
376      affine.for %i = 0 to %N {
377         affine.for %j = 0 to %N {
378           %0 = affine.apply #map42(%j)
379           %tmp = call @S1(%X, %i, %0)
380           affine.if #set(%i, %j)[%N] {
381              %1 = affine.apply #map43(%i, %j)
382              call @S2(%tmp, %A, %i, %1)
383           }
384        }
385      }
386      return
387    }
388    ```
389
390    Example with an explicit yield (initialization with edge padding):
391
392    ```mlir
393    #interior = affine_set<(i, j) : (i - 1 >= 0, j - 1 >= 0,  10 - i >= 0, 10 - j >= 0)> (%i, %j)
394    func.func @pad_edges(%I : memref<10x10xf32>) -> (memref<12x12xf32) {
395      %O = alloc memref<12x12xf32>
396      affine.parallel (%i, %j) = (0, 0) to (12, 12) {
397        %1 = affine.if #interior (%i, %j) {
398          %2 = load %I[%i - 1, %j - 1] : memref<10x10xf32>
399          affine.yield %2
400        } else {
401          %2 = arith.constant 0.0 : f32
402          affine.yield %2 : f32
403        }
404        affine.store %1, %O[%i, %j] : memref<12x12xf32>
405      }
406      return %O
407    }
408    ```
409  }];
410  let arguments = (ins Variadic<AnyType>,
411                       IntegerSetAttr:$condition);
412  let results = (outs Variadic<AnyType>:$results);
413  let regions = (region SizedRegion<1>:$thenRegion, AnyRegion:$elseRegion);
414
415  let skipDefaultBuilders = 1;
416
417  let builders = [
418    OpBuilder<(ins "IntegerSet":$set, "ValueRange":$args,
419      "bool":$withElseRegion)>,
420    OpBuilder<(ins "TypeRange":$resultTypes, "IntegerSet":$set,
421      "ValueRange":$args, "bool":$withElseRegion)>,
422  ];
423
424  let extraClassDeclaration = [{
425    static StringRef getConditionAttrStrName() { return "condition"; }
426
427    IntegerSet getIntegerSet();
428    void setIntegerSet(IntegerSet newSet);
429
430    /// Sets the integer set with its operands.
431    void setConditional(IntegerSet set, ValueRange operands);
432
433    /// Returns true if an else block exists.
434    bool hasElse() { return !getElseRegion().empty(); }
435
436    Block *getThenBlock() {
437      assert(!getThenRegion().empty() && "Unexpected empty 'then' region.");
438      return &getThenRegion().front();
439    }
440
441    Block *getElseBlock() {
442      assert(hasElse() && "Empty 'else' region.");
443      return &getElseRegion().front();
444    }
445
446    OpBuilder getThenBodyBuilder() {
447      assert(!getThenRegion().empty() && "Unexpected empty 'then' region.");
448      Block &body = getThenRegion().front();
449      return OpBuilder(&body, std::prev(body.end()));
450    }
451    OpBuilder getElseBodyBuilder() {
452      assert(hasElse() && "No 'else' block");
453      Block &body = getElseRegion().front();
454      return OpBuilder(&body, std::prev(body.end()));
455    }
456  }];
457
458  let hasCanonicalizer = 1;
459  let hasCustomAssemblyFormat = 1;
460  let hasFolder = 1;
461  let hasVerifier = 1;
462}
463
464class AffineLoadOpBase<string mnemonic, list<Trait> traits = []> :
465    Affine_Op<mnemonic, !listconcat(traits,
466        [DeclareOpInterfaceMethods<AffineReadOpInterface>,
467        DeclareOpInterfaceMethods<AffineMapAccessInterface>,
468        MemRefsNormalizable])> {
469  let arguments = (ins Arg<AnyMemRef, "the reference to load from",
470      [MemRead]>:$memref,
471      Variadic<Index>:$indices,
472      AffineMapAttr:$map);
473
474  code extraClassDeclarationBase = [{
475    /// Returns the operand index of the memref.
476    unsigned getMemRefOperandIndex() { return 0; }
477
478    void setMemRef(Value value) { setOperand(getMemRefOperandIndex(), value); }
479
480    /// Returns the affine map used to index the memref for this operation.
481    AffineMapAttr getAffineMapAttr() {
482      return getProperties().map;
483    }
484
485    static StringRef getMapAttrStrName() { return "map"; }
486  }];
487}
488
489def AffineLoadOp : AffineLoadOpBase<"load"> {
490  let summary = "affine load operation";
491  let description = [{
492    Syntax:
493
494    ```
495    operation ::= ssa-id `=` `affine.load` ssa-use `[` multi-dim-affine-map-of-ssa-ids `]` `:` memref-type
496    ```
497
498    The `affine.load` op reads an element from a memref, where the index
499    for each memref dimension is an affine expression of loop induction
500    variables and symbols. The output of `affine.load` is a new value with the
501    same type as the elements of the memref. An affine expression of loop IVs
502    and symbols must be specified for each dimension of the memref. The keyword
503    `symbol` can be used to indicate SSA identifiers which are symbolic.
504
505    Example 1:
506
507    ```mlir
508    %1 = affine.load %0[%i0 + 3, %i1 + 7] : memref<100x100xf32>
509    ```
510
511    Example 2: Uses `symbol` keyword for symbols `%n` and `%m`.
512
513    ```mlir
514    %1 = affine.load %0[%i0 + symbol(%n), %i1 + symbol(%m)] : memref<100x100xf32>
515    ```
516  }];
517
518  let results = (outs AnyType:$result);
519
520  let builders = [
521    /// Builds an affine load op with the specified map and operands.
522    OpBuilder<(ins "AffineMap":$map, "ValueRange":$operands)>,
523    /// Builds an affine load op with an identity map and operands.
524    OpBuilder<(ins "Value":$memref, CArg<"ValueRange", "{}">:$indices)>,
525    /// Builds an affine load op with the specified map and its operands.
526    OpBuilder<(ins "Value":$memref, "AffineMap":$map,
527      "ValueRange":$mapOperands)>
528  ];
529
530  let extraClassDeclaration = extraClassDeclarationBase;
531
532  let hasCanonicalizer = 1;
533  let hasCustomAssemblyFormat = 1;
534  let hasFolder = 1;
535  let hasVerifier = 1;
536}
537
538class AffineMinMaxOpBase<string mnemonic, list<Trait> traits = []> :
539    Op<Affine_Dialect, mnemonic, traits> {
540  let arguments = (ins AffineMapAttr:$map, Variadic<Index>:$operands);
541  let results = (outs Index);
542
543  let extraClassDeclaration = [{
544    static StringRef getMapAttrStrName() { return "map"; }
545    AffineMap getAffineMap() { return getMap(); }
546    ValueRange getMapOperands() { return getOperands(); }
547    ValueRange getDimOperands() {
548      return OperandRange{getOperands().begin(),
549                          getOperands().begin() + getMap().getNumDims()};
550    }
551    ValueRange getSymbolOperands() {
552      return OperandRange{getOperands().begin() + getMap().getNumDims(),
553                          getOperands().end()};
554    }
555  }];
556  let hasCustomAssemblyFormat = 1;
557  let hasFolder = 1;
558  let hasCanonicalizer = 1;
559  let hasVerifier = 1;
560}
561
562def AffineMinOp : AffineMinMaxOpBase<"min", [Pure]> {
563  let summary = "min operation";
564  let description = [{
565    Syntax:
566
567    ```
568    operation ::= ssa-id `=` `affine.min` affine-map-attribute dim-and-symbol-use-list
569    ```
570
571    The `affine.min` operation applies an [affine mapping](#affine-expressions)
572    to a list of SSA values, and returns the minimum value of all result
573    expressions. The number of dimension and symbol arguments to `affine.min`
574    must be equal to the respective number of dimensional and symbolic inputs to
575    the affine mapping; the `affine.min` operation always returns one value. The
576    input operands and result must all have 'index' type.
577
578    Example:
579
580    ```mlir
581    %0 = affine.min affine_map<(d0)[s0] -> (1000, d0 + 512, s0)> (%arg0)[%arg1]
582    ```
583  }];
584}
585
586def AffineMaxOp : AffineMinMaxOpBase<"max", [Pure]> {
587  let summary = "max operation";
588  let description = [{
589    The `affine.max` operation computes the maximum value result from a multi-result
590    affine map.
591
592    Example:
593
594    ```mlir
595    %0 = affine.max (d0) -> (1000, d0 + 512) (%i0) : index
596    ```
597  }];
598}
599
600def AffineParallelOp : Affine_Op<"parallel",
601    [AutomaticAllocationScope, ImplicitAffineTerminator, RecursivelySpeculatable,
602     RecursiveMemoryEffects, DeclareOpInterfaceMethods<LoopLikeOpInterface>,
603     MemRefsNormalizable]> {
604  let summary = "multi-index parallel band operation";
605  let description = [{
606    The `affine.parallel` operation represents a hyper-rectangular affine
607    parallel band, defining zero or more SSA values for its induction variables.
608    It has one region capturing the parallel band body. The induction variables
609    are represented as arguments of this region. These SSA values always have
610    type index, which is the size of the machine word. The strides, represented
611    by steps, are positive constant integers which defaults to "1" if not
612    present. The lower and upper bounds specify a half-open range: the range
613    includes the lower bound but does not include the upper bound. The body
614    region must contain exactly one block that terminates with `affine.yield`.
615
616    The lower and upper bounds of a parallel operation are represented as an
617    application of an affine mapping to a list of SSA values passed to the map.
618    The same restrictions hold for these SSA values as for all bindings of SSA
619    values to dimensions and symbols. The list of expressions in each map is
620    interpreted according to the respective bounds group attribute. If a single
621    expression belongs to the group, then the result of this expression is taken
622    as a lower(upper) bound of the corresponding loop induction variable. If
623    multiple expressions belong to the group, then the lower(upper) bound is the
624    max(min) of these values obtained from these expressions. The loop band has
625    as many loops as elements in the group bounds attributes.
626
627    Each value yielded by `affine.yield` will be accumulated/reduced via one of
628    the reduction methods defined in the AtomicRMWKind enum.  The order of
629    reduction is unspecified, and lowering may produce any valid ordering.
630    Loops with a 0 trip count will produce as a result the identity value
631    associated with each reduction (i.e. 0.0 for addf, 1.0 for mulf).  Assign
632    reductions for loops with a trip count != 1 produces undefined results.
633
634    Note: Calling `AffineParallelOp::build` will create the required region and
635    block, and insert the required terminator if it is trivial (i.e. no values
636    are yielded).  Parsing will also create the required region, block, and
637    terminator, even when they are missing from the textual representation.
638
639    Example (3x3 valid convolution):
640
641    ```mlir
642    func.func @conv_2d(%D : memref<100x100xf32>, %K : memref<3x3xf32>) -> (memref<98x98xf32>) {
643      %O = memref.alloc() : memref<98x98xf32>
644      affine.parallel (%x, %y) = (0, 0) to (98, 98) {
645        %0 = affine.parallel (%kx, %ky) = (0, 0) to (2, 2) reduce ("addf") -> f32 {
646          %1 = affine.load %D[%x + %kx, %y + %ky] : memref<100x100xf32>
647          %2 = affine.load %K[%kx, %ky] : memref<3x3xf32>
648          %3 = arith.mulf %1, %2 : f32
649          affine.yield %3 : f32
650        }
651        affine.store %0, %O[%x, %y] : memref<98x98xf32>
652      }
653      return %O : memref<98x98xf32>
654    }
655    ```
656
657    Example (tiling by potentially imperfectly dividing sizes):
658
659    ```mlir
660    affine.parallel (%ii, %jj) = (0, 0) to (%N, %M) step (32, 32) {
661      affine.parallel (%i, %j) = (%ii, %jj)
662                              to (min(%ii + 32, %N), min(%jj + 32, %M)) {
663        call @f(%i, %j) : (index, index) -> ()
664      }
665    }
666    ```
667  }];
668
669  let arguments = (ins
670     TypedArrayAttrBase<AtomicRMWKindAttr, "Reduction ops">:$reductions,
671     AffineMapAttr:$lowerBoundsMap,
672     I32ElementsAttr:$lowerBoundsGroups,
673     AffineMapAttr:$upperBoundsMap,
674     I32ElementsAttr:$upperBoundsGroups,
675     I64SmallVectorArrayAttr:$steps,
676     Variadic<Index>:$mapOperands);
677  let results = (outs Variadic<AnyType>:$results);
678  let regions = (region SizedRegion<1>:$region);
679
680  let builders = [
681    OpBuilder<(ins "TypeRange":$resultTypes,
682      "ArrayRef<arith::AtomicRMWKind>":$reductions, "ArrayRef<int64_t>":$ranges)>,
683    OpBuilder<(ins "TypeRange":$resultTypes,
684      "ArrayRef<arith::AtomicRMWKind>":$reductions, "ArrayRef<AffineMap>":$lbMaps,
685      "ValueRange":$lbArgs, "ArrayRef<AffineMap>":$ubMaps, "ValueRange":$ubArgs,
686      "ArrayRef<int64_t>":$steps)>
687  ];
688
689  let extraClassDeclaration = [{
690    /// Get the number of dimensions.
691    unsigned getNumDims();
692
693    /// Get ranges as constants, may fail in dynamic case.
694    std::optional<SmallVector<int64_t, 8>> getConstantRanges();
695
696    Block *getBody();
697    OpBuilder getBodyBuilder();
698    MutableArrayRef<BlockArgument> getIVs() {
699      return getBody()->getArguments();
700    }
701
702    /// Returns elements of the loop lower bound.
703    AffineMap getLowerBoundMap(unsigned pos);
704    operand_range getLowerBoundsOperands();
705    AffineValueMap getLowerBoundsValueMap();
706
707    /// Sets elements of the loop lower bound.
708    void setLowerBounds(ValueRange operands, AffineMap map);
709
710    /// Returns elements of the loop upper bound.
711    AffineMap getUpperBoundMap(unsigned pos);
712    operand_range getUpperBoundsOperands();
713    AffineValueMap getUpperBoundsValueMap();
714
715    /// Sets elements fo the loop upper bound.
716    void setUpperBounds(ValueRange operands, AffineMap map);
717
718    void setSteps(ArrayRef<int64_t> newSteps);
719
720    /// Returns attribute names to use in op construction. Not expected to be
721    /// used directly.
722    static StringRef getReductionsAttrStrName() { return "reductions"; }
723    static StringRef getLowerBoundsMapAttrStrName() { return "lowerBoundsMap"; }
724    static StringRef getLowerBoundsGroupsAttrStrName() {
725      return "lowerBoundsGroups";
726    }
727    static StringRef getUpperBoundsMapAttrStrName() { return "upperBoundsMap"; }
728    static StringRef getUpperBoundsGroupsAttrStrName() {
729      return "upperBoundsGroups";
730    }
731    static StringRef getStepsAttrStrName() { return "steps"; }
732
733    /// Returns `true` if the loop bounds have min/max expressions.
734    bool hasMinMaxBounds() {
735      return getLowerBoundsMap().getNumResults() != getNumDims() ||
736             getUpperBoundsMap().getNumResults() != getNumDims();
737    }
738  }];
739
740  let hasCustomAssemblyFormat = 1;
741  let hasFolder = 1;
742  let hasVerifier = 1;
743}
744
745def AffinePrefetchOp : Affine_Op<"prefetch",
746  [DeclareOpInterfaceMethods<AffineMapAccessInterface>,
747   MemRefsNormalizable]> {
748  let summary = "affine prefetch operation";
749  let description = [{
750    The `affine.prefetch` op prefetches data from a memref location described
751    with an affine subscript similar to affine.load, and has three attributes:
752    a read/write specifier, a locality hint, and a cache type specifier as shown
753    below:
754
755    ```mlir
756    affine.prefetch %0[%i, %j + 5], read, locality<3>, data : memref<400x400xi32>
757    ```
758
759    The read/write specifier is either 'read' or 'write', the locality hint
760    specifier ranges from locality<0> (no locality) to locality<3> (extremely
761    local keep in cache). The cache type specifier is either 'data' or 'instr'
762    and specifies whether the prefetch is performed on data cache or on
763    instruction cache.
764  }];
765
766  let arguments = (ins AnyMemRef:$memref, Variadic<Index>:$indices,
767                   BoolAttr:$isWrite,
768                   ConfinedAttr<I32Attr, [IntMinValue<0>,
769                     IntMaxValue<3>]>:$localityHint,
770                   BoolAttr:$isDataCache,
771                   AffineMapAttr:$map);
772
773  let builders = [
774    OpBuilder<(ins "Value":$memref, "AffineMap":$map,
775      "ArrayRef<Value>":$mapOperands, "bool":$isWrite, "unsigned":$localityHint,
776      "bool":$isDataCache),
777    [{
778      assert(map.getNumInputs() == mapOperands.size()
779             && "inconsistent index info");
780      auto localityHintAttr = $_builder.getI32IntegerAttr(localityHint);
781      auto isWriteAttr = $_builder.getBoolAttr(isWrite);
782      auto isDataCacheAttr = $_builder.getBoolAttr(isDataCache);
783      $_state.addOperands(memref);
784      $_state.addOperands(mapOperands);
785      Properties &prop = $_state.getOrAddProperties<Properties>();
786      prop.map = AffineMapAttr::get(map);
787      prop.localityHint = localityHintAttr;
788      prop.isWrite = isWriteAttr;
789      prop.isDataCache = isDataCacheAttr;
790    }]>];
791
792  let extraClassDeclaration = [{
793    MemRefType getMemRefType() {
794      return ::llvm::cast<MemRefType>(getMemref().getType());
795    }
796
797    /// Returns the affine map used to index the memref for this operation.
798    AffineMap getAffineMap() { return getAffineMapAttr().getValue(); }
799    AffineMapAttr getAffineMapAttr() {
800      return getProperties().map;
801    }
802
803    /// Implements the AffineMapAccessInterface.
804    /// Returns the AffineMapAttr associated with 'memref'.
805    NamedAttribute getAffineMapAttrForMemRef(Value mref) {
806      assert(mref == getMemref() &&
807             "Expected mref argument to match memref operand");
808      return {StringAttr::get(getContext(), getMapAttrStrName()),
809        getAffineMapAttr()};
810    }
811
812    /// Get affine map operands.
813    operand_range getMapOperands() {
814      return {operand_begin() + 1, operand_end()};
815    }
816
817    static StringRef getMapAttrStrName() { return "map"; }
818    static StringRef getLocalityHintAttrStrName() { return "localityHint"; }
819    static StringRef getIsWriteAttrStrName() { return "isWrite"; }
820    static StringRef getIsDataCacheAttrStrName() { return "isDataCache"; }
821  }];
822
823  let hasCanonicalizer = 1;
824  let hasCustomAssemblyFormat = 1;
825  let hasFolder = 1;
826  let hasVerifier = 1;
827}
828
829class AffineStoreOpBase<string mnemonic, list<Trait> traits = []> :
830    Affine_Op<mnemonic, !listconcat(traits,
831    [DeclareOpInterfaceMethods<AffineWriteOpInterface>,
832    DeclareOpInterfaceMethods<AffineMapAccessInterface>,
833    MemRefsNormalizable])> {
834  code extraClassDeclarationBase = [{
835    /// Returns the operand index of the value to be stored.
836    unsigned getStoredValOperandIndex() { return 0; }
837
838    /// Returns the operand index of the memref.
839    unsigned getMemRefOperandIndex() { return 1; }
840
841    void setMemRef(Value value) { setOperand(getMemRefOperandIndex(), value); }
842
843    /// Returns the affine map used to index the memref for this operation.
844    AffineMapAttr getAffineMapAttr() {
845      return getProperties().map;
846    }
847
848    static StringRef getMapAttrStrName() { return "map"; }
849  }];
850}
851
852def AffineStoreOp : AffineStoreOpBase<"store"> {
853  let summary = "affine store operation";
854  let description = [{
855    Syntax:
856
857    ```
858    operation ::= `affine.store` ssa-use, ssa-use `[` multi-dim-affine-map-of-ssa-ids `]` `:` memref-type
859    ```
860
861    The `affine.store` op writes an element to a memref, where the index
862    for each memref dimension is an affine expression of loop induction
863    variables and symbols. The `affine.store` op stores a new value which is the
864    same type as the elements of the memref. An affine expression of loop IVs
865    and symbols must be specified for each dimension of the memref. The keyword
866    `symbol` can be used to indicate SSA identifiers which are symbolic.
867
868    Example 1:
869
870    ```mlir
871    affine.store %v0, %0[%i0 + 3, %i1 + 7] : memref<100x100xf32>
872    ```
873
874    Example 2: Uses `symbol` keyword for symbols `%n` and `%m`.
875
876    ```mlir
877    affine.store %v0, %0[%i0 + symbol(%n), %i1 + symbol(%m)] : memref<100x100xf32>
878    ```
879  }];
880  let arguments = (ins AnyType:$value,
881      Arg<AnyMemRef, "the reference to store to",
882      [MemWrite]>:$memref,
883      Variadic<Index>:$indices,
884      AffineMapAttr:$map);
885
886  let skipDefaultBuilders = 1;
887  let builders = [
888    OpBuilder<(ins "Value":$valueToStore, "Value":$memref,
889      "ValueRange":$indices)>,
890    OpBuilder<(ins "Value":$valueToStore, "Value":$memref, "AffineMap":$map,
891      "ValueRange":$mapOperands)>
892  ];
893
894  let extraClassDeclaration = extraClassDeclarationBase;
895
896  let hasCanonicalizer = 1;
897  let hasCustomAssemblyFormat = 1;
898  let hasFolder = 1;
899  let hasVerifier = 1;
900}
901
902def AffineYieldOp : Affine_Op<"yield", [Pure, Terminator, ReturnLike,
903    MemRefsNormalizable]> {
904  let summary = "Yield values to parent operation";
905  let description = [{
906    The `affine.yield` yields zero or more SSA values from an affine op region and
907    terminates the region. The semantics of how the values yielded are used
908    is defined by the parent operation.
909    If `affine.yield` has any operands, the operands must match the parent
910    operation's results.
911    If the parent operation defines no values, then the `affine.yield` may be
912    left out in the custom syntax and the builders will insert one implicitly.
913    Otherwise, it has to be present in the syntax to indicate which values are
914    yielded.
915  }];
916
917  let arguments = (ins Variadic<AnyType>:$operands);
918
919  let builders = [
920    OpBuilder<(ins), [{ build($_builder, $_state, std::nullopt); }]>
921  ];
922
923  let assemblyFormat = "attr-dict ($operands^ `:` type($operands))?";
924  let hasVerifier = 1;
925}
926
927def AffineVectorLoadOp : AffineLoadOpBase<"vector_load"> {
928  let summary = "affine vector load operation";
929  let description = [{
930    The `affine.vector_load` is the vector counterpart of
931    [affine.load](#affineload-mliraffineloadop). It reads a slice from a
932    [MemRef](Builtin.md/#memreftype), supplied as its first operand,
933    into a [vector](Builtin.md/#vectortype) of the same base elemental type.
934    The index for each memref dimension is an affine expression of loop induction
935    variables and symbols. These indices determine the start position of the read
936    within the memref. The shape of the return vector type determines the shape of
937    the slice read from the memref. This slice is contiguous along the respective
938    dimensions of the shape. Strided vector loads will be supported in the future.
939    An affine expression of loop IVs and symbols must be specified for each
940    dimension of the memref. The keyword `symbol` can be used to indicate SSA
941    identifiers which are symbolic.
942
943    Example 1: 8-wide f32 vector load.
944
945    ```mlir
946    %1 = affine.vector_load %0[%i0 + 3, %i1 + 7] : memref<100x100xf32>, vector<8xf32>
947    ```
948
949    Example 2: 4-wide f32 vector load. Uses `symbol` keyword for symbols `%n` and `%m`.
950
951    ```mlir
952    %1 = affine.vector_load %0[%i0 + symbol(%n), %i1 + symbol(%m)] : memref<100x100xf32>, vector<4xf32>
953    ```
954
955    Example 3: 2-dim f32 vector load.
956
957    ```mlir
958    %1 = affine.vector_load %0[%i0, %i1] : memref<100x100xf32>, vector<2x8xf32>
959    ```
960
961    TODOs:
962    * Add support for strided vector loads.
963    * Consider adding a permutation map to permute the slice that is read from memory
964    (see [vector.transfer_read](../Vector/#vectortransfer_read-mlirvectortransferreadop)).
965  }];
966
967  let results = (outs AnyVectorOfNonZeroRank:$result);
968
969  let builders = [
970    /// Builds an affine vector load op with the specified map and operands.
971    OpBuilder<(ins "VectorType":$resultType, "AffineMap":$map,
972      "ValueRange":$operands)>,
973    /// Builds an affine vector load op with an identity map and operands.
974    OpBuilder<(ins "VectorType":$resultType, "Value":$memref,
975      CArg<"ValueRange", "{}">:$indices)>,
976    /// Builds an affine vector load op with the specified map and its operands.
977    OpBuilder<(ins "VectorType":$resultType, "Value":$memref,
978      "AffineMap":$map, "ValueRange":$mapOperands)>
979  ];
980
981  let extraClassDeclaration = extraClassDeclarationBase # [{
982    VectorType getVectorType() {
983      return ::llvm::cast<VectorType>(getResult().getType());
984    }
985  }];
986
987  let hasCanonicalizer = 1;
988  let hasCustomAssemblyFormat = 1;
989  let hasVerifier = 1;
990}
991
992def AffineVectorStoreOp : AffineStoreOpBase<"vector_store"> {
993  let summary = "affine vector store operation";
994  let description = [{
995    The `affine.vector_store` is the vector counterpart of
996    [affine.store](#affinestore-mliraffinestoreop). It writes a
997    [vector](Builtin.md/#vectortype), supplied as its first operand,
998    into a slice within a [MemRef](Builtin.md/#memreftype) of the same base
999    elemental type, supplied as its second operand.
1000    The index for each memref dimension is an affine expression of loop
1001    induction variables and symbols. These indices determine the start position
1002    of the write within the memref. The shape of th input vector determines the
1003    shape of the slice written to the memref. This slice is contiguous along the
1004    respective dimensions of the shape. Strided vector stores will be supported
1005    in the future.
1006    An affine expression of loop IVs and symbols must be specified for each
1007    dimension of the memref. The keyword `symbol` can be used to indicate SSA
1008    identifiers which are symbolic.
1009
1010    Example 1: 8-wide f32 vector store.
1011
1012    ```mlir
1013    affine.vector_store %v0, %0[%i0 + 3, %i1 + 7] : memref<100x100xf32>, vector<8xf32>
1014    ```
1015
1016    Example 2: 4-wide f32 vector store. Uses `symbol` keyword for symbols `%n` and `%m`.
1017
1018    ```mlir
1019    affine.vector_store %v0, %0[%i0 + symbol(%n), %i1 + symbol(%m)] : memref<100x100xf32>, vector<4xf32>
1020    ```
1021
1022    Example 3: 2-dim f32 vector store.
1023
1024    ```mlir
1025    affine.vector_store %v0, %0[%i0, %i1] : memref<100x100xf32>, vector<2x8xf32>
1026    ```
1027
1028    TODOs:
1029    * Add support for strided vector stores.
1030    * Consider adding a permutation map to permute the slice that is written to memory
1031    (see [vector.transfer_write](../Vector/#vectortransfer_write-mlirvectortransferwriteop)).
1032  }];
1033
1034  let arguments = (ins AnyVectorOfNonZeroRank:$value,
1035      Arg<AnyMemRef, "the reference to store to",
1036      [MemWrite]>:$memref,
1037      Variadic<Index>:$indices,
1038      AffineMapAttr:$map);
1039
1040  let skipDefaultBuilders = 1;
1041  let builders = [
1042    OpBuilder<(ins "Value":$valueToStore, "Value":$memref,
1043      "ValueRange":$indices)>,
1044    OpBuilder<(ins "Value":$valueToStore, "Value":$memref, "AffineMap":$map,
1045      "ValueRange":$mapOperands)>
1046  ];
1047
1048  let extraClassDeclaration = extraClassDeclarationBase # [{
1049    VectorType getVectorType() {
1050      return ::llvm::cast<VectorType>(getValue().getType());
1051    }
1052  }];
1053
1054  let hasCanonicalizer = 1;
1055  let hasCustomAssemblyFormat = 1;
1056  let hasVerifier = 1;
1057}
1058
1059//===----------------------------------------------------------------------===//
1060// AffineDelinearizeIndexOp
1061//===----------------------------------------------------------------------===//
1062
1063def AffineDelinearizeIndexOp : Affine_Op<"delinearize_index", [Pure]> {
1064  let summary = "delinearize an index";
1065  let description = [{
1066    The `affine.delinearize_index` operation takes a single index value and
1067    calculates the multi-index according to the given basis.
1068
1069    Example:
1070
1071    ```
1072    %indices:3 = affine.delinearize_index %linear_index into (%c16, %c224, %c224) : index, index, index
1073    ```
1074
1075    In the above example, `%indices:3` conceptually holds the following:
1076
1077    ```
1078    #map0 = affine_map<()[s0] -> (s0 floordiv 50176)>
1079    #map1 = affine_map<()[s0] -> ((s0 mod 50176) floordiv 224)>
1080    #map2 = affine_map<()[s0] -> (s0 mod 224)>
1081    %indices_0 = affine.apply #map0()[%linear_index]
1082    %indices_1 = affine.apply #map1()[%linear_index]
1083    %indices_2 = affine.apply #map2()[%linear_index]
1084    ```
1085
1086    In other words, `%0:3 = affine.delinearize_index %x into (B, C)` produces
1087    `%0 = {%x / (B * C), (%x mod (B * C)) / C, %x mod C}`.
1088
1089    The basis may either contain `N` or `N-1` elements, where `N` is the number of results.
1090    If there are N basis elements, the first one will not be used during computations,
1091    but may be used during analysis and canonicalization to eliminate terms from
1092    the `affine.delinearize_index` or to enable conclusions about the total size of
1093    `%linear_index`.
1094
1095    If the basis is fully provided, the delinearize_index operation is said to "have
1096    an outer bound". The builders assume that an `affine.delinearize_index` has
1097    an outer bound by default, as this is how the operation was initially defined.
1098
1099    That is, the example above could also have been written
1100    ```mlir
1101    %0:3 = affine.delinearize_index %linear_index into (244, 244) : index, index
1102    ```
1103
1104    Note that, for symmetry with `getPaddedBasis()`, if `hasOuterBound` is `true`
1105    when one of the `OpFoldResult` builders is called but the first element of the
1106    basis is `nullptr`, that first element is ignored and the builder proceeds as if
1107    there was no outer bound.
1108
1109    Due to the constraints of affine maps, all the basis elements must
1110    be strictly positive. A dynamic basis element being 0 or negative causes
1111    undefined behavior.
1112  }];
1113
1114  let arguments = (ins Index:$linear_index,
1115    Variadic<Index>:$dynamic_basis,
1116    DenseI64ArrayAttr:$static_basis);
1117  let results = (outs Variadic<Index>:$multi_index);
1118
1119  let assemblyFormat = [{
1120    $linear_index `into`
1121    custom<DynamicIndexList>($dynamic_basis, $static_basis, "{}", "::mlir::AsmParser::Delimiter::Paren")
1122    attr-dict `:` type($multi_index)
1123  }];
1124
1125  let builders = [
1126    OpBuilder<(ins "Value":$linear_index, "ValueRange":$dynamic_basis, "ArrayRef<int64_t>":$static_basis, CArg<"bool", "true">:$hasOuterBound)>,
1127    OpBuilder<(ins "Value":$linear_index, "ValueRange":$basis, CArg<"bool", "true">:$hasOuterBound)>,
1128    OpBuilder<(ins "Value":$linear_index, "ArrayRef<OpFoldResult>":$basis, CArg<"bool", "true">:$hasOuterBound)>,
1129    OpBuilder<(ins "Value":$linear_index, "ArrayRef<int64_t>":$basis, CArg<"bool", "true">:$hasOuterBound)>
1130  ];
1131
1132  let extraClassDeclaration = [{
1133    /// Return true if the basis includes a bound on the first index input.
1134    bool hasOuterBound() {
1135      return getMultiIndex().size() == getStaticBasis().size();
1136    }
1137
1138    /// Returns a vector with all the static and dynamic basis values.
1139    SmallVector<OpFoldResult> getMixedBasis() {
1140      OpBuilder builder(getContext());
1141      return ::mlir::getMixedValues(getStaticBasis(), getDynamicBasis(), builder);
1142    }
1143
1144    /// Return a vector that contains the basis of the operation, removing
1145    /// the outer bound if one is present.
1146    SmallVector<OpFoldResult> getEffectiveBasis();
1147
1148    /// Return the vector with one basis element per result of the operation. If
1149    /// there is no outer bound specified, the leading entry of this result will be
1150    /// nullptr.
1151    SmallVector<OpFoldResult> getPaddedBasis();
1152  }];
1153
1154  let hasVerifier = 1;
1155  let hasFolder = 1;
1156  let hasCanonicalizer = 1;
1157}
1158
1159//===----------------------------------------------------------------------===//
1160// AffineLinearizeIndexOp
1161//===----------------------------------------------------------------------===//
1162def AffineLinearizeIndexOp : Affine_Op<"linearize_index",
1163    [Pure, AttrSizedOperandSegments]> {
1164  let summary = "linearize an index";
1165  let description = [{
1166    The `affine.linearize_index` operation takes a sequence of index values and a
1167    basis of the same length and linearizes the indices using that basis.
1168
1169    That is, for indices `%idx_0` to `%idx_{N-1}` and basis elements `b_0`
1170    (or `b_1`) up to `b_{N-1}` it computes
1171
1172    ```
1173    sum(i = 0 to N-1) %idx_i * product(j = i + 1 to N-1) B_j
1174    ```
1175
1176    In other words, `%0 = affine.linearize_index [%z, %y, %x] by (Z, Y, X)`
1177    gives `%0 = %x + %y * X + %z * X * Y`, or `%0 = %x + X * (%y + Y * (%z))`.
1178
1179    The basis may either have `N` or `N-1` elements, where `N` is the number of
1180    inputs to linearize_index. If `N` inputs are provided, the first one is not used
1181    in computation, but may be used during analysis or canonicalization as a bound
1182    on `%idx_0`.
1183
1184    If all `N` basis elements are provided, the linearize_index operation is said to
1185    "have an outer bound".
1186
1187    As a convenience, and for symmetry with `getPaddedBasis()`, ifg the first
1188    element of a set of `OpFoldResult`s passed to the builders of this operation is
1189    `nullptr`, that element is ignored.
1190
1191    If the `disjoint` property is present, this is an optimization hint that,
1192    for all `i`, `0 <= %idx_i < B_i` - that is, no index affects any other index,
1193    except that `%idx_0` may be negative to make the index as a whole negative.
1194
1195    Note that the outputs of `affine.delinearize_index` are, by definition, `disjoint`.
1196
1197    Example:
1198
1199    ```mlir
1200    %linear_index = affine.linearize_index [%index_0, %index_1, %index_2] by (2, 3, 5) : index
1201    // Same effect
1202    %linear_index = affine.linearize_index [%index_0, %index_1, %index_2] by (3, 5) : index
1203    ```
1204
1205    In the above example, `%linear_index` conceptually holds the following:
1206
1207    ```mlir
1208    #map = affine_map<()[s0, s1, s2] -> (s0 * 15 + s1 * 5 + s2)>
1209    %linear_index = affine.apply #map()[%index_0, %index_1, %index_2]
1210    ```
1211  }];
1212
1213  let arguments = (ins Variadic<Index>:$multi_index,
1214    Variadic<Index>:$dynamic_basis,
1215    DenseI64ArrayAttr:$static_basis,
1216    UnitProp:$disjoint);
1217  let results = (outs Index:$linear_index);
1218
1219  let assemblyFormat = [{
1220    (`disjoint` $disjoint^)? ` `
1221    `[` $multi_index `]` `by`
1222    custom<DynamicIndexList>($dynamic_basis, $static_basis, "{}", "::mlir::AsmParser::Delimiter::Paren")
1223    attr-dict `:` type($linear_index)
1224  }];
1225
1226  let builders = [
1227    OpBuilder<(ins "ValueRange":$multi_index, "ValueRange":$basis, CArg<"bool", "false">:$disjoint)>,
1228    OpBuilder<(ins "ValueRange":$multi_index, "ArrayRef<OpFoldResult>":$basis, CArg<"bool", "false">:$disjoint)>,
1229    OpBuilder<(ins "ValueRange":$multi_index, "ArrayRef<int64_t>":$basis, CArg<"bool", "false">:$disjoint)>
1230  ];
1231
1232  let extraClassDeclaration = [{
1233    /// Return true if the basis includes a bound on the first index input.
1234    bool hasOuterBound() {
1235      return getMultiIndex().size() == getStaticBasis().size();
1236    }
1237
1238    /// Return a vector with all the static and dynamic basis values.
1239    SmallVector<OpFoldResult> getMixedBasis() {
1240      OpBuilder builder(getContext());
1241      return ::mlir::getMixedValues(getStaticBasis(), getDynamicBasis(), builder);
1242    }
1243
1244    /// Return a vector that contains the basis of the operation, removing
1245    /// the outer bound if one is present.
1246    SmallVector<OpFoldResult> getEffectiveBasis();
1247
1248    /// Return the vector with one basis element per index operand of the operation.
1249    /// If there is no outer bound specified, the leading entry of this basis will be
1250    /// nullptr.
1251    SmallVector<OpFoldResult> getPaddedBasis();
1252  }];
1253
1254  let hasVerifier = 1;
1255  let hasFolder = 1;
1256  let hasCanonicalizer = 1;
1257}
1258
1259#endif // AFFINE_OPS
1260