1 //===- AffineOps.cpp - MLIR Affine Operations -----------------------------===//
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 #include "mlir/Dialect/Affine/IR/AffineOps.h"
10 #include "mlir/Dialect/Affine/IR/AffineValueMap.h"
11 #include "mlir/Dialect/MemRef/IR/MemRef.h"
12 #include "mlir/Dialect/StandardOps/IR/Ops.h"
13 #include "mlir/IR/BlockAndValueMapping.h"
14 #include "mlir/IR/BuiltinOps.h"
15 #include "mlir/IR/IntegerSet.h"
16 #include "mlir/IR/Matchers.h"
17 #include "mlir/IR/OpImplementation.h"
18 #include "mlir/IR/PatternMatch.h"
19 #include "mlir/Transforms/InliningUtils.h"
20 #include "llvm/ADT/SetVector.h"
21 #include "llvm/ADT/SmallBitVector.h"
22 #include "llvm/ADT/TypeSwitch.h"
23 #include "llvm/Support/Debug.h"
24 
25 using namespace mlir;
26 using llvm::dbgs;
27 
28 #define DEBUG_TYPE "affine-analysis"
29 
30 /// A utility function to check if a value is defined at the top level of
31 /// `region` or is an argument of `region`. A value of index type defined at the
32 /// top level of a `AffineScope` region is always a valid symbol for all
33 /// uses in that region.
34 static bool isTopLevelValue(Value value, Region *region) {
35   if (auto arg = value.dyn_cast<BlockArgument>())
36     return arg.getParentRegion() == region;
37   return value.getDefiningOp()->getParentRegion() == region;
38 }
39 
40 /// Checks if `value` known to be a legal affine dimension or symbol in `src`
41 /// region remains legal if the operation that uses it is inlined into `dest`
42 /// with the given value mapping. `legalityCheck` is either `isValidDim` or
43 /// `isValidSymbol`, depending on the value being required to remain a valid
44 /// dimension or symbol.
45 static bool
46 remainsLegalAfterInline(Value value, Region *src, Region *dest,
47                         const BlockAndValueMapping &mapping,
48                         function_ref<bool(Value, Region *)> legalityCheck) {
49   // If the value is a valid dimension for any other reason than being
50   // a top-level value, it will remain valid: constants get inlined
51   // with the function, transitive affine applies also get inlined and
52   // will be checked themselves, etc.
53   if (!isTopLevelValue(value, src))
54     return true;
55 
56   // If it's a top-level value because it's a block operand, i.e. a
57   // function argument, check whether the value replacing it after
58   // inlining is a valid dimension in the new region.
59   if (value.isa<BlockArgument>())
60     return legalityCheck(mapping.lookup(value), dest);
61 
62   // If it's a top-level value beacuse it's defined in the region,
63   // it can only be inlined if the defining op is a constant or a
64   // `dim`, which can appear anywhere and be valid, since the defining
65   // op won't be top-level anymore after inlining.
66   Attribute operandCst;
67   return matchPattern(value.getDefiningOp(), m_Constant(&operandCst)) ||
68          value.getDefiningOp<memref::DimOp>();
69 }
70 
71 /// Checks if all values known to be legal affine dimensions or symbols in `src`
72 /// remain so if their respective users are inlined into `dest`.
73 static bool
74 remainsLegalAfterInline(ValueRange values, Region *src, Region *dest,
75                         const BlockAndValueMapping &mapping,
76                         function_ref<bool(Value, Region *)> legalityCheck) {
77   return llvm::all_of(values, [&](Value v) {
78     return remainsLegalAfterInline(v, src, dest, mapping, legalityCheck);
79   });
80 }
81 
82 /// Checks if an affine read or write operation remains legal after inlining
83 /// from `src` to `dest`.
84 template <typename OpTy>
85 static bool remainsLegalAfterInline(OpTy op, Region *src, Region *dest,
86                                     const BlockAndValueMapping &mapping) {
87   static_assert(llvm::is_one_of<OpTy, AffineReadOpInterface,
88                                 AffineWriteOpInterface>::value,
89                 "only ops with affine read/write interface are supported");
90 
91   AffineMap map = op.getAffineMap();
92   ValueRange dimOperands = op.getMapOperands().take_front(map.getNumDims());
93   ValueRange symbolOperands =
94       op.getMapOperands().take_back(map.getNumSymbols());
95   if (!remainsLegalAfterInline(
96           dimOperands, src, dest, mapping,
97           static_cast<bool (*)(Value, Region *)>(isValidDim)))
98     return false;
99   if (!remainsLegalAfterInline(
100           symbolOperands, src, dest, mapping,
101           static_cast<bool (*)(Value, Region *)>(isValidSymbol)))
102     return false;
103   return true;
104 }
105 
106 /// Checks if an affine apply operation remains legal after inlining from `src`
107 /// to `dest`.
108 template <>
109 bool remainsLegalAfterInline(AffineApplyOp op, Region *src, Region *dest,
110                              const BlockAndValueMapping &mapping) {
111   // If it's a valid dimension, we need to check that it remains so.
112   if (isValidDim(op.getResult(), src))
113     return remainsLegalAfterInline(
114         op.getMapOperands(), src, dest, mapping,
115         static_cast<bool (*)(Value, Region *)>(isValidDim));
116 
117   // Otherwise it must be a valid symbol, check that it remains so.
118   return remainsLegalAfterInline(
119       op.getMapOperands(), src, dest, mapping,
120       static_cast<bool (*)(Value, Region *)>(isValidSymbol));
121 }
122 
123 //===----------------------------------------------------------------------===//
124 // AffineDialect Interfaces
125 //===----------------------------------------------------------------------===//
126 
127 namespace {
128 /// This class defines the interface for handling inlining with affine
129 /// operations.
130 struct AffineInlinerInterface : public DialectInlinerInterface {
131   using DialectInlinerInterface::DialectInlinerInterface;
132 
133   //===--------------------------------------------------------------------===//
134   // Analysis Hooks
135   //===--------------------------------------------------------------------===//
136 
137   /// Returns true if the given region 'src' can be inlined into the region
138   /// 'dest' that is attached to an operation registered to the current dialect.
139   /// 'wouldBeCloned' is set if the region is cloned into its new location
140   /// rather than moved, indicating there may be other users.
141   bool isLegalToInline(Region *dest, Region *src, bool wouldBeCloned,
142                        BlockAndValueMapping &valueMapping) const final {
143     // We can inline into affine loops and conditionals if this doesn't break
144     // affine value categorization rules.
145     Operation *destOp = dest->getParentOp();
146     if (!isa<AffineParallelOp, AffineForOp, AffineIfOp>(destOp))
147       return false;
148 
149     // Multi-block regions cannot be inlined into affine constructs, all of
150     // which require single-block regions.
151     if (!llvm::hasSingleElement(*src))
152       return false;
153 
154     // Side-effecting operations that the affine dialect cannot understand
155     // should not be inlined.
156     Block &srcBlock = src->front();
157     for (Operation &op : srcBlock) {
158       // Ops with no side effects are fine,
159       if (auto iface = dyn_cast<MemoryEffectOpInterface>(op)) {
160         if (iface.hasNoEffect())
161           continue;
162       }
163 
164       // Assuming the inlined region is valid, we only need to check if the
165       // inlining would change it.
166       bool remainsValid =
167           llvm::TypeSwitch<Operation *, bool>(&op)
168               .Case<AffineApplyOp, AffineReadOpInterface,
169                     AffineWriteOpInterface>([&](auto op) {
170                 return remainsLegalAfterInline(op, src, dest, valueMapping);
171               })
172               .Default([](Operation *) {
173                 // Conservatively disallow inlining ops we cannot reason about.
174                 return false;
175               });
176 
177       if (!remainsValid)
178         return false;
179     }
180 
181     return true;
182   }
183 
184   /// Returns true if the given operation 'op', that is registered to this
185   /// dialect, can be inlined into the given region, false otherwise.
186   bool isLegalToInline(Operation *op, Region *region, bool wouldBeCloned,
187                        BlockAndValueMapping &valueMapping) const final {
188     // Always allow inlining affine operations into a region that is marked as
189     // affine scope, or into affine loops and conditionals. There are some edge
190     // cases when inlining *into* affine structures, but that is handled in the
191     // other 'isLegalToInline' hook above.
192     Operation *parentOp = region->getParentOp();
193     return parentOp->hasTrait<OpTrait::AffineScope>() ||
194            isa<AffineForOp, AffineParallelOp, AffineIfOp>(parentOp);
195   }
196 
197   /// Affine regions should be analyzed recursively.
198   bool shouldAnalyzeRecursively(Operation *op) const final { return true; }
199 };
200 } // end anonymous namespace
201 
202 //===----------------------------------------------------------------------===//
203 // AffineDialect
204 //===----------------------------------------------------------------------===//
205 
206 void AffineDialect::initialize() {
207   addOperations<AffineDmaStartOp, AffineDmaWaitOp,
208 #define GET_OP_LIST
209 #include "mlir/Dialect/Affine/IR/AffineOps.cpp.inc"
210                 >();
211   addInterfaces<AffineInlinerInterface>();
212 }
213 
214 /// Materialize a single constant operation from a given attribute value with
215 /// the desired resultant type.
216 Operation *AffineDialect::materializeConstant(OpBuilder &builder,
217                                               Attribute value, Type type,
218                                               Location loc) {
219   return builder.create<ConstantOp>(loc, type, value);
220 }
221 
222 /// A utility function to check if a value is defined at the top level of an
223 /// op with trait `AffineScope`. If the value is defined in an unlinked region,
224 /// conservatively assume it is not top-level. A value of index type defined at
225 /// the top level is always a valid symbol.
226 bool mlir::isTopLevelValue(Value value) {
227   if (auto arg = value.dyn_cast<BlockArgument>()) {
228     // The block owning the argument may be unlinked, e.g. when the surrounding
229     // region has not yet been attached to an Op, at which point the parent Op
230     // is null.
231     Operation *parentOp = arg.getOwner()->getParentOp();
232     return parentOp && parentOp->hasTrait<OpTrait::AffineScope>();
233   }
234   // The defining Op may live in an unlinked block so its parent Op may be null.
235   Operation *parentOp = value.getDefiningOp()->getParentOp();
236   return parentOp && parentOp->hasTrait<OpTrait::AffineScope>();
237 }
238 
239 /// Returns the closest region enclosing `op` that is held by an operation with
240 /// trait `AffineScope`; `nullptr` if there is no such region.
241 //  TODO: getAffineScope should be publicly exposed for affine passes/utilities.
242 static Region *getAffineScope(Operation *op) {
243   auto *curOp = op;
244   while (auto *parentOp = curOp->getParentOp()) {
245     if (parentOp->hasTrait<OpTrait::AffineScope>())
246       return curOp->getParentRegion();
247     curOp = parentOp;
248   }
249   return nullptr;
250 }
251 
252 // A Value can be used as a dimension id iff it meets one of the following
253 // conditions:
254 // *) It is valid as a symbol.
255 // *) It is an induction variable.
256 // *) It is the result of affine apply operation with dimension id arguments.
257 bool mlir::isValidDim(Value value) {
258   // The value must be an index type.
259   if (!value.getType().isIndex())
260     return false;
261 
262   if (auto *defOp = value.getDefiningOp())
263     return isValidDim(value, getAffineScope(defOp));
264 
265   // This value has to be a block argument for an op that has the
266   // `AffineScope` trait or for an affine.for or affine.parallel.
267   auto *parentOp = value.cast<BlockArgument>().getOwner()->getParentOp();
268   return parentOp && (parentOp->hasTrait<OpTrait::AffineScope>() ||
269                       isa<AffineForOp, AffineParallelOp>(parentOp));
270 }
271 
272 // Value can be used as a dimension id iff it meets one of the following
273 // conditions:
274 // *) It is valid as a symbol.
275 // *) It is an induction variable.
276 // *) It is the result of an affine apply operation with dimension id operands.
277 bool mlir::isValidDim(Value value, Region *region) {
278   // The value must be an index type.
279   if (!value.getType().isIndex())
280     return false;
281 
282   // All valid symbols are okay.
283   if (isValidSymbol(value, region))
284     return true;
285 
286   auto *op = value.getDefiningOp();
287   if (!op) {
288     // This value has to be a block argument for an affine.for or an
289     // affine.parallel.
290     auto *parentOp = value.cast<BlockArgument>().getOwner()->getParentOp();
291     return isa<AffineForOp, AffineParallelOp>(parentOp);
292   }
293 
294   // Affine apply operation is ok if all of its operands are ok.
295   if (auto applyOp = dyn_cast<AffineApplyOp>(op))
296     return applyOp.isValidDim(region);
297   // The dim op is okay if its operand memref/tensor is defined at the top
298   // level.
299   if (auto dimOp = dyn_cast<memref::DimOp>(op))
300     return isTopLevelValue(dimOp.memrefOrTensor());
301   return false;
302 }
303 
304 /// Returns true if the 'index' dimension of the `memref` defined by
305 /// `memrefDefOp` is a statically  shaped one or defined using a valid symbol
306 /// for `region`.
307 template <typename AnyMemRefDefOp>
308 static bool isMemRefSizeValidSymbol(AnyMemRefDefOp memrefDefOp, unsigned index,
309                                     Region *region) {
310   auto memRefType = memrefDefOp.getType();
311   // Statically shaped.
312   if (!memRefType.isDynamicDim(index))
313     return true;
314   // Get the position of the dimension among dynamic dimensions;
315   unsigned dynamicDimPos = memRefType.getDynamicDimIndex(index);
316   return isValidSymbol(*(memrefDefOp.getDynamicSizes().begin() + dynamicDimPos),
317                        region);
318 }
319 
320 /// Returns true if the result of the dim op is a valid symbol for `region`.
321 static bool isDimOpValidSymbol(memref::DimOp dimOp, Region *region) {
322   // The dim op is okay if its operand memref is defined at the top level.
323   if (isTopLevelValue(dimOp.memrefOrTensor()))
324     return true;
325 
326   // Conservatively handle remaining BlockArguments as non-valid symbols.
327   // E.g. scf.for iterArgs.
328   if (dimOp.memrefOrTensor().isa<BlockArgument>())
329     return false;
330 
331   // The dim op is also okay if its operand memref is a view/subview whose
332   // corresponding size is a valid symbol.
333   Optional<int64_t> index = dimOp.getConstantIndex();
334   assert(index.hasValue() &&
335          "expect only `dim` operations with a constant index");
336   int64_t i = index.getValue();
337   return TypeSwitch<Operation *, bool>(dimOp.memrefOrTensor().getDefiningOp())
338       .Case<memref::ViewOp, memref::SubViewOp, memref::AllocOp>(
339           [&](auto op) { return isMemRefSizeValidSymbol(op, i, region); })
340       .Default([](Operation *) { return false; });
341 }
342 
343 // A value can be used as a symbol (at all its use sites) iff it meets one of
344 // the following conditions:
345 // *) It is a constant.
346 // *) Its defining op or block arg appearance is immediately enclosed by an op
347 //    with `AffineScope` trait.
348 // *) It is the result of an affine.apply operation with symbol operands.
349 // *) It is a result of the dim op on a memref whose corresponding size is a
350 //    valid symbol.
351 bool mlir::isValidSymbol(Value value) {
352   // The value must be an index type.
353   if (!value.getType().isIndex())
354     return false;
355 
356   // Check that the value is a top level value.
357   if (isTopLevelValue(value))
358     return true;
359 
360   if (auto *defOp = value.getDefiningOp())
361     return isValidSymbol(value, getAffineScope(defOp));
362 
363   return false;
364 }
365 
366 /// A value can be used as a symbol for `region` iff it meets onf of the the
367 /// following conditions:
368 /// *) It is a constant.
369 /// *) It is the result of an affine apply operation with symbol arguments.
370 /// *) It is a result of the dim op on a memref whose corresponding size is
371 ///    a valid symbol.
372 /// *) It is defined at the top level of 'region' or is its argument.
373 /// *) It dominates `region`'s parent op.
374 /// If `region` is null, conservatively assume the symbol definition scope does
375 /// not exist and only accept the values that would be symbols regardless of
376 /// the surrounding region structure, i.e. the first three cases above.
377 bool mlir::isValidSymbol(Value value, Region *region) {
378   // The value must be an index type.
379   if (!value.getType().isIndex())
380     return false;
381 
382   // A top-level value is a valid symbol.
383   if (region && ::isTopLevelValue(value, region))
384     return true;
385 
386   auto *defOp = value.getDefiningOp();
387   if (!defOp) {
388     // A block argument that is not a top-level value is a valid symbol if it
389     // dominates region's parent op.
390     Operation *regionOp = region ? region->getParentOp() : nullptr;
391     if (regionOp && !regionOp->hasTrait<OpTrait::IsIsolatedFromAbove>())
392       if (auto *parentOpRegion = region->getParentOp()->getParentRegion())
393         return isValidSymbol(value, parentOpRegion);
394     return false;
395   }
396 
397   // Constant operation is ok.
398   Attribute operandCst;
399   if (matchPattern(defOp, m_Constant(&operandCst)))
400     return true;
401 
402   // Affine apply operation is ok if all of its operands are ok.
403   if (auto applyOp = dyn_cast<AffineApplyOp>(defOp))
404     return applyOp.isValidSymbol(region);
405 
406   // Dim op results could be valid symbols at any level.
407   if (auto dimOp = dyn_cast<memref::DimOp>(defOp))
408     return isDimOpValidSymbol(dimOp, region);
409 
410   // Check for values dominating `region`'s parent op.
411   Operation *regionOp = region ? region->getParentOp() : nullptr;
412   if (regionOp && !regionOp->hasTrait<OpTrait::IsIsolatedFromAbove>())
413     if (auto *parentRegion = region->getParentOp()->getParentRegion())
414       return isValidSymbol(value, parentRegion);
415 
416   return false;
417 }
418 
419 // Returns true if 'value' is a valid index to an affine operation (e.g.
420 // affine.load, affine.store, affine.dma_start, affine.dma_wait) where
421 // `region` provides the polyhedral symbol scope. Returns false otherwise.
422 static bool isValidAffineIndexOperand(Value value, Region *region) {
423   return isValidDim(value, region) || isValidSymbol(value, region);
424 }
425 
426 /// Prints dimension and symbol list.
427 static void printDimAndSymbolList(Operation::operand_iterator begin,
428                                   Operation::operand_iterator end,
429                                   unsigned numDims, OpAsmPrinter &printer) {
430   OperandRange operands(begin, end);
431   printer << '(' << operands.take_front(numDims) << ')';
432   if (operands.size() > numDims)
433     printer << '[' << operands.drop_front(numDims) << ']';
434 }
435 
436 /// Parses dimension and symbol list and returns true if parsing failed.
437 ParseResult mlir::parseDimAndSymbolList(OpAsmParser &parser,
438                                         SmallVectorImpl<Value> &operands,
439                                         unsigned &numDims) {
440   SmallVector<OpAsmParser::OperandType, 8> opInfos;
441   if (parser.parseOperandList(opInfos, OpAsmParser::Delimiter::Paren))
442     return failure();
443   // Store number of dimensions for validation by caller.
444   numDims = opInfos.size();
445 
446   // Parse the optional symbol operands.
447   auto indexTy = parser.getBuilder().getIndexType();
448   return failure(parser.parseOperandList(
449                      opInfos, OpAsmParser::Delimiter::OptionalSquare) ||
450                  parser.resolveOperands(opInfos, indexTy, operands));
451 }
452 
453 /// Utility function to verify that a set of operands are valid dimension and
454 /// symbol identifiers. The operands should be laid out such that the dimension
455 /// operands are before the symbol operands. This function returns failure if
456 /// there was an invalid operand. An operation is provided to emit any necessary
457 /// errors.
458 template <typename OpTy>
459 static LogicalResult
460 verifyDimAndSymbolIdentifiers(OpTy &op, Operation::operand_range operands,
461                               unsigned numDims) {
462   unsigned opIt = 0;
463   for (auto operand : operands) {
464     if (opIt++ < numDims) {
465       if (!isValidDim(operand, getAffineScope(op)))
466         return op.emitOpError("operand cannot be used as a dimension id");
467     } else if (!isValidSymbol(operand, getAffineScope(op))) {
468       return op.emitOpError("operand cannot be used as a symbol");
469     }
470   }
471   return success();
472 }
473 
474 //===----------------------------------------------------------------------===//
475 // AffineApplyOp
476 //===----------------------------------------------------------------------===//
477 
478 AffineValueMap AffineApplyOp::getAffineValueMap() {
479   return AffineValueMap(getAffineMap(), getOperands(), getResult());
480 }
481 
482 static ParseResult parseAffineApplyOp(OpAsmParser &parser,
483                                       OperationState &result) {
484   auto &builder = parser.getBuilder();
485   auto indexTy = builder.getIndexType();
486 
487   AffineMapAttr mapAttr;
488   unsigned numDims;
489   if (parser.parseAttribute(mapAttr, "map", result.attributes) ||
490       parseDimAndSymbolList(parser, result.operands, numDims) ||
491       parser.parseOptionalAttrDict(result.attributes))
492     return failure();
493   auto map = mapAttr.getValue();
494 
495   if (map.getNumDims() != numDims ||
496       numDims + map.getNumSymbols() != result.operands.size()) {
497     return parser.emitError(parser.getNameLoc(),
498                             "dimension or symbol index mismatch");
499   }
500 
501   result.types.append(map.getNumResults(), indexTy);
502   return success();
503 }
504 
505 static void print(OpAsmPrinter &p, AffineApplyOp op) {
506   p << AffineApplyOp::getOperationName() << " " << op.mapAttr();
507   printDimAndSymbolList(op.operand_begin(), op.operand_end(),
508                         op.getAffineMap().getNumDims(), p);
509   p.printOptionalAttrDict(op->getAttrs(), /*elidedAttrs=*/{"map"});
510 }
511 
512 static LogicalResult verify(AffineApplyOp op) {
513   // Check input and output dimensions match.
514   auto map = op.map();
515 
516   // Verify that operand count matches affine map dimension and symbol count.
517   if (op.getNumOperands() != map.getNumDims() + map.getNumSymbols())
518     return op.emitOpError(
519         "operand count and affine map dimension and symbol count must match");
520 
521   // Verify that the map only produces one result.
522   if (map.getNumResults() != 1)
523     return op.emitOpError("mapping must produce one value");
524 
525   return success();
526 }
527 
528 // The result of the affine apply operation can be used as a dimension id if all
529 // its operands are valid dimension ids.
530 bool AffineApplyOp::isValidDim() {
531   return llvm::all_of(getOperands(),
532                       [](Value op) { return mlir::isValidDim(op); });
533 }
534 
535 // The result of the affine apply operation can be used as a dimension id if all
536 // its operands are valid dimension ids with the parent operation of `region`
537 // defining the polyhedral scope for symbols.
538 bool AffineApplyOp::isValidDim(Region *region) {
539   return llvm::all_of(getOperands(),
540                       [&](Value op) { return ::isValidDim(op, region); });
541 }
542 
543 // The result of the affine apply operation can be used as a symbol if all its
544 // operands are symbols.
545 bool AffineApplyOp::isValidSymbol() {
546   return llvm::all_of(getOperands(),
547                       [](Value op) { return mlir::isValidSymbol(op); });
548 }
549 
550 // The result of the affine apply operation can be used as a symbol in `region`
551 // if all its operands are symbols in `region`.
552 bool AffineApplyOp::isValidSymbol(Region *region) {
553   return llvm::all_of(getOperands(), [&](Value operand) {
554     return mlir::isValidSymbol(operand, region);
555   });
556 }
557 
558 OpFoldResult AffineApplyOp::fold(ArrayRef<Attribute> operands) {
559   auto map = getAffineMap();
560 
561   // Fold dims and symbols to existing values.
562   auto expr = map.getResult(0);
563   if (auto dim = expr.dyn_cast<AffineDimExpr>())
564     return getOperand(dim.getPosition());
565   if (auto sym = expr.dyn_cast<AffineSymbolExpr>())
566     return getOperand(map.getNumDims() + sym.getPosition());
567 
568   // Otherwise, default to folding the map.
569   SmallVector<Attribute, 1> result;
570   if (failed(map.constantFold(operands, result)))
571     return {};
572   return result[0];
573 }
574 
575 /// Replace all occurrences of AffineExpr at position `pos` in `map` by the
576 /// defining AffineApplyOp expression and operands.
577 /// When `dimOrSymbolPosition < dims.size()`, AffineDimExpr@[pos] is replaced.
578 /// When `dimOrSymbolPosition >= dims.size()`,
579 /// AffineSymbolExpr@[pos - dims.size()] is replaced.
580 /// Mutate `map`,`dims` and `syms` in place as follows:
581 ///   1. `dims` and `syms` are only appended to.
582 ///   2. `map` dim and symbols are gradually shifted to higer positions.
583 ///   3. Old `dim` and `sym` entries are replaced by nullptr
584 /// This avoids the need for any bookkeeping.
585 static LogicalResult replaceDimOrSym(AffineMap *map,
586                                      unsigned dimOrSymbolPosition,
587                                      SmallVectorImpl<Value> &dims,
588                                      SmallVectorImpl<Value> &syms) {
589   bool isDimReplacement = (dimOrSymbolPosition < dims.size());
590   unsigned pos = isDimReplacement ? dimOrSymbolPosition
591                                   : dimOrSymbolPosition - dims.size();
592   Value &v = isDimReplacement ? dims[pos] : syms[pos];
593   if (!v)
594     return failure();
595 
596   auto affineApply = v.getDefiningOp<AffineApplyOp>();
597   if (!affineApply)
598     return failure();
599 
600   // At this point we will perform a replacement of `v`, set the entry in `dim`
601   // or `sym` to nullptr immediately.
602   v = nullptr;
603 
604   // Compute the map, dims and symbols coming from the AffineApplyOp.
605   AffineMap composeMap = affineApply.getAffineMap();
606   assert(composeMap.getNumResults() == 1 && "affine.apply with >1 results");
607   AffineExpr composeExpr =
608       composeMap.shiftDims(dims.size()).shiftSymbols(syms.size()).getResult(0);
609   ValueRange composeDims =
610       affineApply.getMapOperands().take_front(composeMap.getNumDims());
611   ValueRange composeSyms =
612       affineApply.getMapOperands().take_back(composeMap.getNumSymbols());
613 
614   // Perform the replacement and append the dims and symbols where relevant.
615   MLIRContext *ctx = map->getContext();
616   AffineExpr toReplace = isDimReplacement ? getAffineDimExpr(pos, ctx)
617                                           : getAffineSymbolExpr(pos, ctx);
618   *map = map->replace(toReplace, composeExpr, dims.size(), syms.size());
619   dims.append(composeDims.begin(), composeDims.end());
620   syms.append(composeSyms.begin(), composeSyms.end());
621 
622   return success();
623 }
624 
625 /// Iterate over `operands` and fold away all those produced by an AffineApplyOp
626 /// iteratively. Perform canonicalization of map and operands as well as
627 /// AffineMap simplification. `map` and `operands` are mutated in place.
628 static void composeAffineMapAndOperands(AffineMap *map,
629                                         SmallVectorImpl<Value> *operands) {
630   if (map->getNumResults() == 0) {
631     canonicalizeMapAndOperands(map, operands);
632     *map = simplifyAffineMap(*map);
633     return;
634   }
635 
636   MLIRContext *ctx = map->getContext();
637   SmallVector<Value, 4> dims(operands->begin(),
638                              operands->begin() + map->getNumDims());
639   SmallVector<Value, 4> syms(operands->begin() + map->getNumDims(),
640                              operands->end());
641 
642   // Iterate over dims and symbols coming from AffineApplyOp and replace until
643   // exhaustion. This iteratively mutates `map`, `dims` and `syms`. Both `dims`
644   // and `syms` can only increase by construction.
645   // The implementation uses a `while` loop to support the case of symbols
646   // that may be constructed from dims ;this may be overkill.
647   while (true) {
648     bool changed = false;
649     for (unsigned pos = 0; pos != dims.size() + syms.size(); ++pos)
650       if ((changed |= succeeded(replaceDimOrSym(map, pos, dims, syms))))
651         break;
652     if (!changed)
653       break;
654   }
655 
656   // Clear operands so we can fill them anew.
657   operands->clear();
658 
659   // At this point we may have introduced null operands, prune them out before
660   // canonicalizing map and operands.
661   unsigned nDims = 0, nSyms = 0;
662   SmallVector<AffineExpr, 4> dimReplacements, symReplacements;
663   dimReplacements.reserve(dims.size());
664   symReplacements.reserve(syms.size());
665   for (auto *container : {&dims, &syms}) {
666     bool isDim = (container == &dims);
667     auto &repls = isDim ? dimReplacements : symReplacements;
668     for (auto en : llvm::enumerate(*container)) {
669       Value v = en.value();
670       if (!v) {
671         assert(isDim ? !map->isFunctionOfDim(en.index())
672                      : !map->isFunctionOfSymbol(en.index()) &&
673                            "map is function of unexpected expr@pos");
674         repls.push_back(getAffineConstantExpr(0, ctx));
675         continue;
676       }
677       repls.push_back(isDim ? getAffineDimExpr(nDims++, ctx)
678                             : getAffineSymbolExpr(nSyms++, ctx));
679       operands->push_back(v);
680     }
681   }
682   *map = map->replaceDimsAndSymbols(dimReplacements, symReplacements, nDims,
683                                     nSyms);
684 
685   // Canonicalize and simplify before returning.
686   canonicalizeMapAndOperands(map, operands);
687   *map = simplifyAffineMap(*map);
688 }
689 
690 void mlir::fullyComposeAffineMapAndOperands(AffineMap *map,
691                                             SmallVectorImpl<Value> *operands) {
692   while (llvm::any_of(*operands, [](Value v) {
693     return isa_and_nonnull<AffineApplyOp>(v.getDefiningOp());
694   })) {
695     composeAffineMapAndOperands(map, operands);
696   }
697 }
698 
699 AffineApplyOp mlir::makeComposedAffineApply(OpBuilder &b, Location loc,
700                                             AffineMap map,
701                                             ArrayRef<Value> operands) {
702   AffineMap normalizedMap = map;
703   SmallVector<Value, 8> normalizedOperands(operands.begin(), operands.end());
704   composeAffineMapAndOperands(&normalizedMap, &normalizedOperands);
705   assert(normalizedMap);
706   return b.create<AffineApplyOp>(loc, normalizedMap, normalizedOperands);
707 }
708 
709 // A symbol may appear as a dim in affine.apply operations. This function
710 // canonicalizes dims that are valid symbols into actual symbols.
711 template <class MapOrSet>
712 static void canonicalizePromotedSymbols(MapOrSet *mapOrSet,
713                                         SmallVectorImpl<Value> *operands) {
714   if (!mapOrSet || operands->empty())
715     return;
716 
717   assert(mapOrSet->getNumInputs() == operands->size() &&
718          "map/set inputs must match number of operands");
719 
720   auto *context = mapOrSet->getContext();
721   SmallVector<Value, 8> resultOperands;
722   resultOperands.reserve(operands->size());
723   SmallVector<Value, 8> remappedSymbols;
724   remappedSymbols.reserve(operands->size());
725   unsigned nextDim = 0;
726   unsigned nextSym = 0;
727   unsigned oldNumSyms = mapOrSet->getNumSymbols();
728   SmallVector<AffineExpr, 8> dimRemapping(mapOrSet->getNumDims());
729   for (unsigned i = 0, e = mapOrSet->getNumInputs(); i != e; ++i) {
730     if (i < mapOrSet->getNumDims()) {
731       if (isValidSymbol((*operands)[i])) {
732         // This is a valid symbol that appears as a dim, canonicalize it.
733         dimRemapping[i] = getAffineSymbolExpr(oldNumSyms + nextSym++, context);
734         remappedSymbols.push_back((*operands)[i]);
735       } else {
736         dimRemapping[i] = getAffineDimExpr(nextDim++, context);
737         resultOperands.push_back((*operands)[i]);
738       }
739     } else {
740       resultOperands.push_back((*operands)[i]);
741     }
742   }
743 
744   resultOperands.append(remappedSymbols.begin(), remappedSymbols.end());
745   *operands = resultOperands;
746   *mapOrSet = mapOrSet->replaceDimsAndSymbols(dimRemapping, {}, nextDim,
747                                               oldNumSyms + nextSym);
748 
749   assert(mapOrSet->getNumInputs() == operands->size() &&
750          "map/set inputs must match number of operands");
751 }
752 
753 // Works for either an affine map or an integer set.
754 template <class MapOrSet>
755 static void canonicalizeMapOrSetAndOperands(MapOrSet *mapOrSet,
756                                             SmallVectorImpl<Value> *operands) {
757   static_assert(llvm::is_one_of<MapOrSet, AffineMap, IntegerSet>::value,
758                 "Argument must be either of AffineMap or IntegerSet type");
759 
760   if (!mapOrSet || operands->empty())
761     return;
762 
763   assert(mapOrSet->getNumInputs() == operands->size() &&
764          "map/set inputs must match number of operands");
765 
766   canonicalizePromotedSymbols<MapOrSet>(mapOrSet, operands);
767 
768   // Check to see what dims are used.
769   llvm::SmallBitVector usedDims(mapOrSet->getNumDims());
770   llvm::SmallBitVector usedSyms(mapOrSet->getNumSymbols());
771   mapOrSet->walkExprs([&](AffineExpr expr) {
772     if (auto dimExpr = expr.dyn_cast<AffineDimExpr>())
773       usedDims[dimExpr.getPosition()] = true;
774     else if (auto symExpr = expr.dyn_cast<AffineSymbolExpr>())
775       usedSyms[symExpr.getPosition()] = true;
776   });
777 
778   auto *context = mapOrSet->getContext();
779 
780   SmallVector<Value, 8> resultOperands;
781   resultOperands.reserve(operands->size());
782 
783   llvm::SmallDenseMap<Value, AffineExpr, 8> seenDims;
784   SmallVector<AffineExpr, 8> dimRemapping(mapOrSet->getNumDims());
785   unsigned nextDim = 0;
786   for (unsigned i = 0, e = mapOrSet->getNumDims(); i != e; ++i) {
787     if (usedDims[i]) {
788       // Remap dim positions for duplicate operands.
789       auto it = seenDims.find((*operands)[i]);
790       if (it == seenDims.end()) {
791         dimRemapping[i] = getAffineDimExpr(nextDim++, context);
792         resultOperands.push_back((*operands)[i]);
793         seenDims.insert(std::make_pair((*operands)[i], dimRemapping[i]));
794       } else {
795         dimRemapping[i] = it->second;
796       }
797     }
798   }
799   llvm::SmallDenseMap<Value, AffineExpr, 8> seenSymbols;
800   SmallVector<AffineExpr, 8> symRemapping(mapOrSet->getNumSymbols());
801   unsigned nextSym = 0;
802   for (unsigned i = 0, e = mapOrSet->getNumSymbols(); i != e; ++i) {
803     if (!usedSyms[i])
804       continue;
805     // Handle constant operands (only needed for symbolic operands since
806     // constant operands in dimensional positions would have already been
807     // promoted to symbolic positions above).
808     IntegerAttr operandCst;
809     if (matchPattern((*operands)[i + mapOrSet->getNumDims()],
810                      m_Constant(&operandCst))) {
811       symRemapping[i] =
812           getAffineConstantExpr(operandCst.getValue().getSExtValue(), context);
813       continue;
814     }
815     // Remap symbol positions for duplicate operands.
816     auto it = seenSymbols.find((*operands)[i + mapOrSet->getNumDims()]);
817     if (it == seenSymbols.end()) {
818       symRemapping[i] = getAffineSymbolExpr(nextSym++, context);
819       resultOperands.push_back((*operands)[i + mapOrSet->getNumDims()]);
820       seenSymbols.insert(std::make_pair((*operands)[i + mapOrSet->getNumDims()],
821                                         symRemapping[i]));
822     } else {
823       symRemapping[i] = it->second;
824     }
825   }
826   *mapOrSet = mapOrSet->replaceDimsAndSymbols(dimRemapping, symRemapping,
827                                               nextDim, nextSym);
828   *operands = resultOperands;
829 }
830 
831 void mlir::canonicalizeMapAndOperands(AffineMap *map,
832                                       SmallVectorImpl<Value> *operands) {
833   canonicalizeMapOrSetAndOperands<AffineMap>(map, operands);
834 }
835 
836 void mlir::canonicalizeSetAndOperands(IntegerSet *set,
837                                       SmallVectorImpl<Value> *operands) {
838   canonicalizeMapOrSetAndOperands<IntegerSet>(set, operands);
839 }
840 
841 namespace {
842 /// Simplify AffineApply, AffineLoad, and AffineStore operations by composing
843 /// maps that supply results into them.
844 ///
845 template <typename AffineOpTy>
846 struct SimplifyAffineOp : public OpRewritePattern<AffineOpTy> {
847   using OpRewritePattern<AffineOpTy>::OpRewritePattern;
848 
849   /// Replace the affine op with another instance of it with the supplied
850   /// map and mapOperands.
851   void replaceAffineOp(PatternRewriter &rewriter, AffineOpTy affineOp,
852                        AffineMap map, ArrayRef<Value> mapOperands) const;
853 
854   LogicalResult matchAndRewrite(AffineOpTy affineOp,
855                                 PatternRewriter &rewriter) const override {
856     static_assert(llvm::is_one_of<AffineOpTy, AffineLoadOp, AffinePrefetchOp,
857                                   AffineStoreOp, AffineApplyOp, AffineMinOp,
858                                   AffineMaxOp>::value,
859                   "affine load/store/apply/prefetch/min/max op expected");
860     auto map = affineOp.getAffineMap();
861     AffineMap oldMap = map;
862     auto oldOperands = affineOp.getMapOperands();
863     SmallVector<Value, 8> resultOperands(oldOperands);
864     composeAffineMapAndOperands(&map, &resultOperands);
865     if (map == oldMap && std::equal(oldOperands.begin(), oldOperands.end(),
866                                     resultOperands.begin()))
867       return failure();
868 
869     replaceAffineOp(rewriter, affineOp, map, resultOperands);
870     return success();
871   }
872 };
873 
874 // Specialize the template to account for the different build signatures for
875 // affine load, store, and apply ops.
876 template <>
877 void SimplifyAffineOp<AffineLoadOp>::replaceAffineOp(
878     PatternRewriter &rewriter, AffineLoadOp load, AffineMap map,
879     ArrayRef<Value> mapOperands) const {
880   rewriter.replaceOpWithNewOp<AffineLoadOp>(load, load.getMemRef(), map,
881                                             mapOperands);
882 }
883 template <>
884 void SimplifyAffineOp<AffinePrefetchOp>::replaceAffineOp(
885     PatternRewriter &rewriter, AffinePrefetchOp prefetch, AffineMap map,
886     ArrayRef<Value> mapOperands) const {
887   rewriter.replaceOpWithNewOp<AffinePrefetchOp>(
888       prefetch, prefetch.memref(), map, mapOperands, prefetch.localityHint(),
889       prefetch.isWrite(), prefetch.isDataCache());
890 }
891 template <>
892 void SimplifyAffineOp<AffineStoreOp>::replaceAffineOp(
893     PatternRewriter &rewriter, AffineStoreOp store, AffineMap map,
894     ArrayRef<Value> mapOperands) const {
895   rewriter.replaceOpWithNewOp<AffineStoreOp>(
896       store, store.getValueToStore(), store.getMemRef(), map, mapOperands);
897 }
898 
899 // Generic version for ops that don't have extra operands.
900 template <typename AffineOpTy>
901 void SimplifyAffineOp<AffineOpTy>::replaceAffineOp(
902     PatternRewriter &rewriter, AffineOpTy op, AffineMap map,
903     ArrayRef<Value> mapOperands) const {
904   rewriter.replaceOpWithNewOp<AffineOpTy>(op, map, mapOperands);
905 }
906 } // end anonymous namespace.
907 
908 void AffineApplyOp::getCanonicalizationPatterns(RewritePatternSet &results,
909                                                 MLIRContext *context) {
910   results.add<SimplifyAffineOp<AffineApplyOp>>(context);
911 }
912 
913 //===----------------------------------------------------------------------===//
914 // Common canonicalization pattern support logic
915 //===----------------------------------------------------------------------===//
916 
917 /// This is a common class used for patterns of the form
918 /// "someop(memrefcast) -> someop".  It folds the source of any memref.cast
919 /// into the root operation directly.
920 static LogicalResult foldMemRefCast(Operation *op) {
921   bool folded = false;
922   for (OpOperand &operand : op->getOpOperands()) {
923     auto cast = operand.get().getDefiningOp<memref::CastOp>();
924     if (cast && !cast.getOperand().getType().isa<UnrankedMemRefType>()) {
925       operand.set(cast.getOperand());
926       folded = true;
927     }
928   }
929   return success(folded);
930 }
931 
932 //===----------------------------------------------------------------------===//
933 // AffineDmaStartOp
934 //===----------------------------------------------------------------------===//
935 
936 // TODO: Check that map operands are loop IVs or symbols.
937 void AffineDmaStartOp::build(OpBuilder &builder, OperationState &result,
938                              Value srcMemRef, AffineMap srcMap,
939                              ValueRange srcIndices, Value destMemRef,
940                              AffineMap dstMap, ValueRange destIndices,
941                              Value tagMemRef, AffineMap tagMap,
942                              ValueRange tagIndices, Value numElements,
943                              Value stride, Value elementsPerStride) {
944   result.addOperands(srcMemRef);
945   result.addAttribute(getSrcMapAttrName(), AffineMapAttr::get(srcMap));
946   result.addOperands(srcIndices);
947   result.addOperands(destMemRef);
948   result.addAttribute(getDstMapAttrName(), AffineMapAttr::get(dstMap));
949   result.addOperands(destIndices);
950   result.addOperands(tagMemRef);
951   result.addAttribute(getTagMapAttrName(), AffineMapAttr::get(tagMap));
952   result.addOperands(tagIndices);
953   result.addOperands(numElements);
954   if (stride) {
955     result.addOperands({stride, elementsPerStride});
956   }
957 }
958 
959 void AffineDmaStartOp::print(OpAsmPrinter &p) {
960   p << "affine.dma_start " << getSrcMemRef() << '[';
961   p.printAffineMapOfSSAIds(getSrcMapAttr(), getSrcIndices());
962   p << "], " << getDstMemRef() << '[';
963   p.printAffineMapOfSSAIds(getDstMapAttr(), getDstIndices());
964   p << "], " << getTagMemRef() << '[';
965   p.printAffineMapOfSSAIds(getTagMapAttr(), getTagIndices());
966   p << "], " << getNumElements();
967   if (isStrided()) {
968     p << ", " << getStride();
969     p << ", " << getNumElementsPerStride();
970   }
971   p << " : " << getSrcMemRefType() << ", " << getDstMemRefType() << ", "
972     << getTagMemRefType();
973 }
974 
975 // Parse AffineDmaStartOp.
976 // Ex:
977 //   affine.dma_start %src[%i, %j], %dst[%k, %l], %tag[%index], %size,
978 //     %stride, %num_elt_per_stride
979 //       : memref<3076 x f32, 0>, memref<1024 x f32, 2>, memref<1 x i32>
980 //
981 ParseResult AffineDmaStartOp::parse(OpAsmParser &parser,
982                                     OperationState &result) {
983   OpAsmParser::OperandType srcMemRefInfo;
984   AffineMapAttr srcMapAttr;
985   SmallVector<OpAsmParser::OperandType, 4> srcMapOperands;
986   OpAsmParser::OperandType dstMemRefInfo;
987   AffineMapAttr dstMapAttr;
988   SmallVector<OpAsmParser::OperandType, 4> dstMapOperands;
989   OpAsmParser::OperandType tagMemRefInfo;
990   AffineMapAttr tagMapAttr;
991   SmallVector<OpAsmParser::OperandType, 4> tagMapOperands;
992   OpAsmParser::OperandType numElementsInfo;
993   SmallVector<OpAsmParser::OperandType, 2> strideInfo;
994 
995   SmallVector<Type, 3> types;
996   auto indexType = parser.getBuilder().getIndexType();
997 
998   // Parse and resolve the following list of operands:
999   // *) dst memref followed by its affine maps operands (in square brackets).
1000   // *) src memref followed by its affine map operands (in square brackets).
1001   // *) tag memref followed by its affine map operands (in square brackets).
1002   // *) number of elements transferred by DMA operation.
1003   if (parser.parseOperand(srcMemRefInfo) ||
1004       parser.parseAffineMapOfSSAIds(srcMapOperands, srcMapAttr,
1005                                     getSrcMapAttrName(), result.attributes) ||
1006       parser.parseComma() || parser.parseOperand(dstMemRefInfo) ||
1007       parser.parseAffineMapOfSSAIds(dstMapOperands, dstMapAttr,
1008                                     getDstMapAttrName(), result.attributes) ||
1009       parser.parseComma() || parser.parseOperand(tagMemRefInfo) ||
1010       parser.parseAffineMapOfSSAIds(tagMapOperands, tagMapAttr,
1011                                     getTagMapAttrName(), result.attributes) ||
1012       parser.parseComma() || parser.parseOperand(numElementsInfo))
1013     return failure();
1014 
1015   // Parse optional stride and elements per stride.
1016   if (parser.parseTrailingOperandList(strideInfo)) {
1017     return failure();
1018   }
1019   if (!strideInfo.empty() && strideInfo.size() != 2) {
1020     return parser.emitError(parser.getNameLoc(),
1021                             "expected two stride related operands");
1022   }
1023   bool isStrided = strideInfo.size() == 2;
1024 
1025   if (parser.parseColonTypeList(types))
1026     return failure();
1027 
1028   if (types.size() != 3)
1029     return parser.emitError(parser.getNameLoc(), "expected three types");
1030 
1031   if (parser.resolveOperand(srcMemRefInfo, types[0], result.operands) ||
1032       parser.resolveOperands(srcMapOperands, indexType, result.operands) ||
1033       parser.resolveOperand(dstMemRefInfo, types[1], result.operands) ||
1034       parser.resolveOperands(dstMapOperands, indexType, result.operands) ||
1035       parser.resolveOperand(tagMemRefInfo, types[2], result.operands) ||
1036       parser.resolveOperands(tagMapOperands, indexType, result.operands) ||
1037       parser.resolveOperand(numElementsInfo, indexType, result.operands))
1038     return failure();
1039 
1040   if (isStrided) {
1041     if (parser.resolveOperands(strideInfo, indexType, result.operands))
1042       return failure();
1043   }
1044 
1045   // Check that src/dst/tag operand counts match their map.numInputs.
1046   if (srcMapOperands.size() != srcMapAttr.getValue().getNumInputs() ||
1047       dstMapOperands.size() != dstMapAttr.getValue().getNumInputs() ||
1048       tagMapOperands.size() != tagMapAttr.getValue().getNumInputs())
1049     return parser.emitError(parser.getNameLoc(),
1050                             "memref operand count not equal to map.numInputs");
1051   return success();
1052 }
1053 
1054 LogicalResult AffineDmaStartOp::verify() {
1055   if (!getOperand(getSrcMemRefOperandIndex()).getType().isa<MemRefType>())
1056     return emitOpError("expected DMA source to be of memref type");
1057   if (!getOperand(getDstMemRefOperandIndex()).getType().isa<MemRefType>())
1058     return emitOpError("expected DMA destination to be of memref type");
1059   if (!getOperand(getTagMemRefOperandIndex()).getType().isa<MemRefType>())
1060     return emitOpError("expected DMA tag to be of memref type");
1061 
1062   // DMAs from different memory spaces supported.
1063   if (getSrcMemorySpace() == getDstMemorySpace()) {
1064     return emitOpError("DMA should be between different memory spaces");
1065   }
1066   unsigned numInputsAllMaps = getSrcMap().getNumInputs() +
1067                               getDstMap().getNumInputs() +
1068                               getTagMap().getNumInputs();
1069   if (getNumOperands() != numInputsAllMaps + 3 + 1 &&
1070       getNumOperands() != numInputsAllMaps + 3 + 1 + 2) {
1071     return emitOpError("incorrect number of operands");
1072   }
1073 
1074   Region *scope = getAffineScope(*this);
1075   for (auto idx : getSrcIndices()) {
1076     if (!idx.getType().isIndex())
1077       return emitOpError("src index to dma_start must have 'index' type");
1078     if (!isValidAffineIndexOperand(idx, scope))
1079       return emitOpError("src index must be a dimension or symbol identifier");
1080   }
1081   for (auto idx : getDstIndices()) {
1082     if (!idx.getType().isIndex())
1083       return emitOpError("dst index to dma_start must have 'index' type");
1084     if (!isValidAffineIndexOperand(idx, scope))
1085       return emitOpError("dst index must be a dimension or symbol identifier");
1086   }
1087   for (auto idx : getTagIndices()) {
1088     if (!idx.getType().isIndex())
1089       return emitOpError("tag index to dma_start must have 'index' type");
1090     if (!isValidAffineIndexOperand(idx, scope))
1091       return emitOpError("tag index must be a dimension or symbol identifier");
1092   }
1093   return success();
1094 }
1095 
1096 LogicalResult AffineDmaStartOp::fold(ArrayRef<Attribute> cstOperands,
1097                                      SmallVectorImpl<OpFoldResult> &results) {
1098   /// dma_start(memrefcast) -> dma_start
1099   return foldMemRefCast(*this);
1100 }
1101 
1102 //===----------------------------------------------------------------------===//
1103 // AffineDmaWaitOp
1104 //===----------------------------------------------------------------------===//
1105 
1106 // TODO: Check that map operands are loop IVs or symbols.
1107 void AffineDmaWaitOp::build(OpBuilder &builder, OperationState &result,
1108                             Value tagMemRef, AffineMap tagMap,
1109                             ValueRange tagIndices, Value numElements) {
1110   result.addOperands(tagMemRef);
1111   result.addAttribute(getTagMapAttrName(), AffineMapAttr::get(tagMap));
1112   result.addOperands(tagIndices);
1113   result.addOperands(numElements);
1114 }
1115 
1116 void AffineDmaWaitOp::print(OpAsmPrinter &p) {
1117   p << "affine.dma_wait " << getTagMemRef() << '[';
1118   SmallVector<Value, 2> operands(getTagIndices());
1119   p.printAffineMapOfSSAIds(getTagMapAttr(), operands);
1120   p << "], ";
1121   p.printOperand(getNumElements());
1122   p << " : " << getTagMemRef().getType();
1123 }
1124 
1125 // Parse AffineDmaWaitOp.
1126 // Eg:
1127 //   affine.dma_wait %tag[%index], %num_elements
1128 //     : memref<1 x i32, (d0) -> (d0), 4>
1129 //
1130 ParseResult AffineDmaWaitOp::parse(OpAsmParser &parser,
1131                                    OperationState &result) {
1132   OpAsmParser::OperandType tagMemRefInfo;
1133   AffineMapAttr tagMapAttr;
1134   SmallVector<OpAsmParser::OperandType, 2> tagMapOperands;
1135   Type type;
1136   auto indexType = parser.getBuilder().getIndexType();
1137   OpAsmParser::OperandType numElementsInfo;
1138 
1139   // Parse tag memref, its map operands, and dma size.
1140   if (parser.parseOperand(tagMemRefInfo) ||
1141       parser.parseAffineMapOfSSAIds(tagMapOperands, tagMapAttr,
1142                                     getTagMapAttrName(), result.attributes) ||
1143       parser.parseComma() || parser.parseOperand(numElementsInfo) ||
1144       parser.parseColonType(type) ||
1145       parser.resolveOperand(tagMemRefInfo, type, result.operands) ||
1146       parser.resolveOperands(tagMapOperands, indexType, result.operands) ||
1147       parser.resolveOperand(numElementsInfo, indexType, result.operands))
1148     return failure();
1149 
1150   if (!type.isa<MemRefType>())
1151     return parser.emitError(parser.getNameLoc(),
1152                             "expected tag to be of memref type");
1153 
1154   if (tagMapOperands.size() != tagMapAttr.getValue().getNumInputs())
1155     return parser.emitError(parser.getNameLoc(),
1156                             "tag memref operand count != to map.numInputs");
1157   return success();
1158 }
1159 
1160 LogicalResult AffineDmaWaitOp::verify() {
1161   if (!getOperand(0).getType().isa<MemRefType>())
1162     return emitOpError("expected DMA tag to be of memref type");
1163   Region *scope = getAffineScope(*this);
1164   for (auto idx : getTagIndices()) {
1165     if (!idx.getType().isIndex())
1166       return emitOpError("index to dma_wait must have 'index' type");
1167     if (!isValidAffineIndexOperand(idx, scope))
1168       return emitOpError("index must be a dimension or symbol identifier");
1169   }
1170   return success();
1171 }
1172 
1173 LogicalResult AffineDmaWaitOp::fold(ArrayRef<Attribute> cstOperands,
1174                                     SmallVectorImpl<OpFoldResult> &results) {
1175   /// dma_wait(memrefcast) -> dma_wait
1176   return foldMemRefCast(*this);
1177 }
1178 
1179 //===----------------------------------------------------------------------===//
1180 // AffineForOp
1181 //===----------------------------------------------------------------------===//
1182 
1183 /// 'bodyBuilder' is used to build the body of affine.for. If iterArgs and
1184 /// bodyBuilder are empty/null, we include default terminator op.
1185 void AffineForOp::build(OpBuilder &builder, OperationState &result,
1186                         ValueRange lbOperands, AffineMap lbMap,
1187                         ValueRange ubOperands, AffineMap ubMap, int64_t step,
1188                         ValueRange iterArgs, BodyBuilderFn bodyBuilder) {
1189   assert(((!lbMap && lbOperands.empty()) ||
1190           lbOperands.size() == lbMap.getNumInputs()) &&
1191          "lower bound operand count does not match the affine map");
1192   assert(((!ubMap && ubOperands.empty()) ||
1193           ubOperands.size() == ubMap.getNumInputs()) &&
1194          "upper bound operand count does not match the affine map");
1195   assert(step > 0 && "step has to be a positive integer constant");
1196 
1197   for (Value val : iterArgs)
1198     result.addTypes(val.getType());
1199 
1200   // Add an attribute for the step.
1201   result.addAttribute(getStepAttrName(),
1202                       builder.getIntegerAttr(builder.getIndexType(), step));
1203 
1204   // Add the lower bound.
1205   result.addAttribute(getLowerBoundAttrName(), AffineMapAttr::get(lbMap));
1206   result.addOperands(lbOperands);
1207 
1208   // Add the upper bound.
1209   result.addAttribute(getUpperBoundAttrName(), AffineMapAttr::get(ubMap));
1210   result.addOperands(ubOperands);
1211 
1212   result.addOperands(iterArgs);
1213   // Create a region and a block for the body.  The argument of the region is
1214   // the loop induction variable.
1215   Region *bodyRegion = result.addRegion();
1216   bodyRegion->push_back(new Block);
1217   Block &bodyBlock = bodyRegion->front();
1218   Value inductionVar = bodyBlock.addArgument(builder.getIndexType());
1219   for (Value val : iterArgs)
1220     bodyBlock.addArgument(val.getType());
1221 
1222   // Create the default terminator if the builder is not provided and if the
1223   // iteration arguments are not provided. Otherwise, leave this to the caller
1224   // because we don't know which values to return from the loop.
1225   if (iterArgs.empty() && !bodyBuilder) {
1226     ensureTerminator(*bodyRegion, builder, result.location);
1227   } else if (bodyBuilder) {
1228     OpBuilder::InsertionGuard guard(builder);
1229     builder.setInsertionPointToStart(&bodyBlock);
1230     bodyBuilder(builder, result.location, inductionVar,
1231                 bodyBlock.getArguments().drop_front());
1232   }
1233 }
1234 
1235 void AffineForOp::build(OpBuilder &builder, OperationState &result, int64_t lb,
1236                         int64_t ub, int64_t step, ValueRange iterArgs,
1237                         BodyBuilderFn bodyBuilder) {
1238   auto lbMap = AffineMap::getConstantMap(lb, builder.getContext());
1239   auto ubMap = AffineMap::getConstantMap(ub, builder.getContext());
1240   return build(builder, result, {}, lbMap, {}, ubMap, step, iterArgs,
1241                bodyBuilder);
1242 }
1243 
1244 static LogicalResult verify(AffineForOp op) {
1245   // Check that the body defines as single block argument for the induction
1246   // variable.
1247   auto *body = op.getBody();
1248   if (body->getNumArguments() == 0 || !body->getArgument(0).getType().isIndex())
1249     return op.emitOpError(
1250         "expected body to have a single index argument for the "
1251         "induction variable");
1252 
1253   // Verify that the bound operands are valid dimension/symbols.
1254   /// Lower bound.
1255   if (op.getLowerBoundMap().getNumInputs() > 0)
1256     if (failed(
1257             verifyDimAndSymbolIdentifiers(op, op.getLowerBoundOperands(),
1258                                           op.getLowerBoundMap().getNumDims())))
1259       return failure();
1260   /// Upper bound.
1261   if (op.getUpperBoundMap().getNumInputs() > 0)
1262     if (failed(
1263             verifyDimAndSymbolIdentifiers(op, op.getUpperBoundOperands(),
1264                                           op.getUpperBoundMap().getNumDims())))
1265       return failure();
1266 
1267   unsigned opNumResults = op.getNumResults();
1268   if (opNumResults == 0)
1269     return success();
1270 
1271   // If ForOp defines values, check that the number and types of the defined
1272   // values match ForOp initial iter operands and backedge basic block
1273   // arguments.
1274   if (op.getNumIterOperands() != opNumResults)
1275     return op.emitOpError(
1276         "mismatch between the number of loop-carried values and results");
1277   if (op.getNumRegionIterArgs() != opNumResults)
1278     return op.emitOpError(
1279         "mismatch between the number of basic block args and results");
1280 
1281   return success();
1282 }
1283 
1284 /// Parse a for operation loop bounds.
1285 static ParseResult parseBound(bool isLower, OperationState &result,
1286                               OpAsmParser &p) {
1287   // 'min' / 'max' prefixes are generally syntactic sugar, but are required if
1288   // the map has multiple results.
1289   bool failedToParsedMinMax =
1290       failed(p.parseOptionalKeyword(isLower ? "max" : "min"));
1291 
1292   auto &builder = p.getBuilder();
1293   auto boundAttrName = isLower ? AffineForOp::getLowerBoundAttrName()
1294                                : AffineForOp::getUpperBoundAttrName();
1295 
1296   // Parse ssa-id as identity map.
1297   SmallVector<OpAsmParser::OperandType, 1> boundOpInfos;
1298   if (p.parseOperandList(boundOpInfos))
1299     return failure();
1300 
1301   if (!boundOpInfos.empty()) {
1302     // Check that only one operand was parsed.
1303     if (boundOpInfos.size() > 1)
1304       return p.emitError(p.getNameLoc(),
1305                          "expected only one loop bound operand");
1306 
1307     // TODO: improve error message when SSA value is not of index type.
1308     // Currently it is 'use of value ... expects different type than prior uses'
1309     if (p.resolveOperand(boundOpInfos.front(), builder.getIndexType(),
1310                          result.operands))
1311       return failure();
1312 
1313     // Create an identity map using symbol id. This representation is optimized
1314     // for storage. Analysis passes may expand it into a multi-dimensional map
1315     // if desired.
1316     AffineMap map = builder.getSymbolIdentityMap();
1317     result.addAttribute(boundAttrName, AffineMapAttr::get(map));
1318     return success();
1319   }
1320 
1321   // Get the attribute location.
1322   llvm::SMLoc attrLoc = p.getCurrentLocation();
1323 
1324   Attribute boundAttr;
1325   if (p.parseAttribute(boundAttr, builder.getIndexType(), boundAttrName,
1326                        result.attributes))
1327     return failure();
1328 
1329   // Parse full form - affine map followed by dim and symbol list.
1330   if (auto affineMapAttr = boundAttr.dyn_cast<AffineMapAttr>()) {
1331     unsigned currentNumOperands = result.operands.size();
1332     unsigned numDims;
1333     if (parseDimAndSymbolList(p, result.operands, numDims))
1334       return failure();
1335 
1336     auto map = affineMapAttr.getValue();
1337     if (map.getNumDims() != numDims)
1338       return p.emitError(
1339           p.getNameLoc(),
1340           "dim operand count and affine map dim count must match");
1341 
1342     unsigned numDimAndSymbolOperands =
1343         result.operands.size() - currentNumOperands;
1344     if (numDims + map.getNumSymbols() != numDimAndSymbolOperands)
1345       return p.emitError(
1346           p.getNameLoc(),
1347           "symbol operand count and affine map symbol count must match");
1348 
1349     // If the map has multiple results, make sure that we parsed the min/max
1350     // prefix.
1351     if (map.getNumResults() > 1 && failedToParsedMinMax) {
1352       if (isLower) {
1353         return p.emitError(attrLoc, "lower loop bound affine map with "
1354                                     "multiple results requires 'max' prefix");
1355       }
1356       return p.emitError(attrLoc, "upper loop bound affine map with multiple "
1357                                   "results requires 'min' prefix");
1358     }
1359     return success();
1360   }
1361 
1362   // Parse custom assembly form.
1363   if (auto integerAttr = boundAttr.dyn_cast<IntegerAttr>()) {
1364     result.attributes.pop_back();
1365     result.addAttribute(
1366         boundAttrName,
1367         AffineMapAttr::get(builder.getConstantAffineMap(integerAttr.getInt())));
1368     return success();
1369   }
1370 
1371   return p.emitError(
1372       p.getNameLoc(),
1373       "expected valid affine map representation for loop bounds");
1374 }
1375 
1376 static ParseResult parseAffineForOp(OpAsmParser &parser,
1377                                     OperationState &result) {
1378   auto &builder = parser.getBuilder();
1379   OpAsmParser::OperandType inductionVariable;
1380   // Parse the induction variable followed by '='.
1381   if (parser.parseRegionArgument(inductionVariable) || parser.parseEqual())
1382     return failure();
1383 
1384   // Parse loop bounds.
1385   if (parseBound(/*isLower=*/true, result, parser) ||
1386       parser.parseKeyword("to", " between bounds") ||
1387       parseBound(/*isLower=*/false, result, parser))
1388     return failure();
1389 
1390   // Parse the optional loop step, we default to 1 if one is not present.
1391   if (parser.parseOptionalKeyword("step")) {
1392     result.addAttribute(
1393         AffineForOp::getStepAttrName(),
1394         builder.getIntegerAttr(builder.getIndexType(), /*value=*/1));
1395   } else {
1396     llvm::SMLoc stepLoc = parser.getCurrentLocation();
1397     IntegerAttr stepAttr;
1398     if (parser.parseAttribute(stepAttr, builder.getIndexType(),
1399                               AffineForOp::getStepAttrName().data(),
1400                               result.attributes))
1401       return failure();
1402 
1403     if (stepAttr.getValue().getSExtValue() < 0)
1404       return parser.emitError(
1405           stepLoc,
1406           "expected step to be representable as a positive signed integer");
1407   }
1408 
1409   // Parse the optional initial iteration arguments.
1410   SmallVector<OpAsmParser::OperandType, 4> regionArgs, operands;
1411   SmallVector<Type, 4> argTypes;
1412   regionArgs.push_back(inductionVariable);
1413 
1414   if (succeeded(parser.parseOptionalKeyword("iter_args"))) {
1415     // Parse assignment list and results type list.
1416     if (parser.parseAssignmentList(regionArgs, operands) ||
1417         parser.parseArrowTypeList(result.types))
1418       return failure();
1419     // Resolve input operands.
1420     for (auto operandType : llvm::zip(operands, result.types))
1421       if (parser.resolveOperand(std::get<0>(operandType),
1422                                 std::get<1>(operandType), result.operands))
1423         return failure();
1424   }
1425   // Induction variable.
1426   Type indexType = builder.getIndexType();
1427   argTypes.push_back(indexType);
1428   // Loop carried variables.
1429   argTypes.append(result.types.begin(), result.types.end());
1430   // Parse the body region.
1431   Region *body = result.addRegion();
1432   if (regionArgs.size() != argTypes.size())
1433     return parser.emitError(
1434         parser.getNameLoc(),
1435         "mismatch between the number of loop-carried values and results");
1436   if (parser.parseRegion(*body, regionArgs, argTypes))
1437     return failure();
1438 
1439   AffineForOp::ensureTerminator(*body, builder, result.location);
1440 
1441   // Parse the optional attribute list.
1442   return parser.parseOptionalAttrDict(result.attributes);
1443 }
1444 
1445 static void printBound(AffineMapAttr boundMap,
1446                        Operation::operand_range boundOperands,
1447                        const char *prefix, OpAsmPrinter &p) {
1448   AffineMap map = boundMap.getValue();
1449 
1450   // Check if this bound should be printed using custom assembly form.
1451   // The decision to restrict printing custom assembly form to trivial cases
1452   // comes from the will to roundtrip MLIR binary -> text -> binary in a
1453   // lossless way.
1454   // Therefore, custom assembly form parsing and printing is only supported for
1455   // zero-operand constant maps and single symbol operand identity maps.
1456   if (map.getNumResults() == 1) {
1457     AffineExpr expr = map.getResult(0);
1458 
1459     // Print constant bound.
1460     if (map.getNumDims() == 0 && map.getNumSymbols() == 0) {
1461       if (auto constExpr = expr.dyn_cast<AffineConstantExpr>()) {
1462         p << constExpr.getValue();
1463         return;
1464       }
1465     }
1466 
1467     // Print bound that consists of a single SSA symbol if the map is over a
1468     // single symbol.
1469     if (map.getNumDims() == 0 && map.getNumSymbols() == 1) {
1470       if (auto symExpr = expr.dyn_cast<AffineSymbolExpr>()) {
1471         p.printOperand(*boundOperands.begin());
1472         return;
1473       }
1474     }
1475   } else {
1476     // Map has multiple results. Print 'min' or 'max' prefix.
1477     p << prefix << ' ';
1478   }
1479 
1480   // Print the map and its operands.
1481   p << boundMap;
1482   printDimAndSymbolList(boundOperands.begin(), boundOperands.end(),
1483                         map.getNumDims(), p);
1484 }
1485 
1486 unsigned AffineForOp::getNumIterOperands() {
1487   AffineMap lbMap = getLowerBoundMapAttr().getValue();
1488   AffineMap ubMap = getUpperBoundMapAttr().getValue();
1489 
1490   return getNumOperands() - lbMap.getNumInputs() - ubMap.getNumInputs();
1491 }
1492 
1493 static void print(OpAsmPrinter &p, AffineForOp op) {
1494   p << op.getOperationName() << ' ';
1495   p.printOperand(op.getBody()->getArgument(0));
1496   p << " = ";
1497   printBound(op.getLowerBoundMapAttr(), op.getLowerBoundOperands(), "max", p);
1498   p << " to ";
1499   printBound(op.getUpperBoundMapAttr(), op.getUpperBoundOperands(), "min", p);
1500 
1501   if (op.getStep() != 1)
1502     p << " step " << op.getStep();
1503 
1504   bool printBlockTerminators = false;
1505   if (op.getNumIterOperands() > 0) {
1506     p << " iter_args(";
1507     auto regionArgs = op.getRegionIterArgs();
1508     auto operands = op.getIterOperands();
1509 
1510     llvm::interleaveComma(llvm::zip(regionArgs, operands), p, [&](auto it) {
1511       p << std::get<0>(it) << " = " << std::get<1>(it);
1512     });
1513     p << ") -> (" << op.getResultTypes() << ")";
1514     printBlockTerminators = true;
1515   }
1516 
1517   p.printRegion(op.region(),
1518                 /*printEntryBlockArgs=*/false, printBlockTerminators);
1519   p.printOptionalAttrDict(op->getAttrs(),
1520                           /*elidedAttrs=*/{op.getLowerBoundAttrName(),
1521                                            op.getUpperBoundAttrName(),
1522                                            op.getStepAttrName()});
1523 }
1524 
1525 /// Fold the constant bounds of a loop.
1526 static LogicalResult foldLoopBounds(AffineForOp forOp) {
1527   auto foldLowerOrUpperBound = [&forOp](bool lower) {
1528     // Check to see if each of the operands is the result of a constant.  If
1529     // so, get the value.  If not, ignore it.
1530     SmallVector<Attribute, 8> operandConstants;
1531     auto boundOperands =
1532         lower ? forOp.getLowerBoundOperands() : forOp.getUpperBoundOperands();
1533     for (auto operand : boundOperands) {
1534       Attribute operandCst;
1535       matchPattern(operand, m_Constant(&operandCst));
1536       operandConstants.push_back(operandCst);
1537     }
1538 
1539     AffineMap boundMap =
1540         lower ? forOp.getLowerBoundMap() : forOp.getUpperBoundMap();
1541     assert(boundMap.getNumResults() >= 1 &&
1542            "bound maps should have at least one result");
1543     SmallVector<Attribute, 4> foldedResults;
1544     if (failed(boundMap.constantFold(operandConstants, foldedResults)))
1545       return failure();
1546 
1547     // Compute the max or min as applicable over the results.
1548     assert(!foldedResults.empty() && "bounds should have at least one result");
1549     auto maxOrMin = foldedResults[0].cast<IntegerAttr>().getValue();
1550     for (unsigned i = 1, e = foldedResults.size(); i < e; i++) {
1551       auto foldedResult = foldedResults[i].cast<IntegerAttr>().getValue();
1552       maxOrMin = lower ? llvm::APIntOps::smax(maxOrMin, foldedResult)
1553                        : llvm::APIntOps::smin(maxOrMin, foldedResult);
1554     }
1555     lower ? forOp.setConstantLowerBound(maxOrMin.getSExtValue())
1556           : forOp.setConstantUpperBound(maxOrMin.getSExtValue());
1557     return success();
1558   };
1559 
1560   // Try to fold the lower bound.
1561   bool folded = false;
1562   if (!forOp.hasConstantLowerBound())
1563     folded |= succeeded(foldLowerOrUpperBound(/*lower=*/true));
1564 
1565   // Try to fold the upper bound.
1566   if (!forOp.hasConstantUpperBound())
1567     folded |= succeeded(foldLowerOrUpperBound(/*lower=*/false));
1568   return success(folded);
1569 }
1570 
1571 /// Canonicalize the bounds of the given loop.
1572 static LogicalResult canonicalizeLoopBounds(AffineForOp forOp) {
1573   SmallVector<Value, 4> lbOperands(forOp.getLowerBoundOperands());
1574   SmallVector<Value, 4> ubOperands(forOp.getUpperBoundOperands());
1575 
1576   auto lbMap = forOp.getLowerBoundMap();
1577   auto ubMap = forOp.getUpperBoundMap();
1578   auto prevLbMap = lbMap;
1579   auto prevUbMap = ubMap;
1580 
1581   canonicalizeMapAndOperands(&lbMap, &lbOperands);
1582   lbMap = removeDuplicateExprs(lbMap);
1583 
1584   canonicalizeMapAndOperands(&ubMap, &ubOperands);
1585   ubMap = removeDuplicateExprs(ubMap);
1586 
1587   // Any canonicalization change always leads to updated map(s).
1588   if (lbMap == prevLbMap && ubMap == prevUbMap)
1589     return failure();
1590 
1591   if (lbMap != prevLbMap)
1592     forOp.setLowerBound(lbOperands, lbMap);
1593   if (ubMap != prevUbMap)
1594     forOp.setUpperBound(ubOperands, ubMap);
1595   return success();
1596 }
1597 
1598 namespace {
1599 /// This is a pattern to fold trivially empty loops.
1600 struct AffineForEmptyLoopFolder : public OpRewritePattern<AffineForOp> {
1601   using OpRewritePattern<AffineForOp>::OpRewritePattern;
1602 
1603   LogicalResult matchAndRewrite(AffineForOp forOp,
1604                                 PatternRewriter &rewriter) const override {
1605     // Check that the body only contains a yield.
1606     if (!llvm::hasSingleElement(*forOp.getBody()))
1607       return failure();
1608     rewriter.eraseOp(forOp);
1609     return success();
1610   }
1611 };
1612 } // end anonymous namespace
1613 
1614 void AffineForOp::getCanonicalizationPatterns(RewritePatternSet &results,
1615                                               MLIRContext *context) {
1616   results.add<AffineForEmptyLoopFolder>(context);
1617 }
1618 
1619 LogicalResult AffineForOp::fold(ArrayRef<Attribute> operands,
1620                                 SmallVectorImpl<OpFoldResult> &results) {
1621   bool folded = succeeded(foldLoopBounds(*this));
1622   folded |= succeeded(canonicalizeLoopBounds(*this));
1623   return success(folded);
1624 }
1625 
1626 AffineBound AffineForOp::getLowerBound() {
1627   auto lbMap = getLowerBoundMap();
1628   return AffineBound(AffineForOp(*this), 0, lbMap.getNumInputs(), lbMap);
1629 }
1630 
1631 AffineBound AffineForOp::getUpperBound() {
1632   auto lbMap = getLowerBoundMap();
1633   auto ubMap = getUpperBoundMap();
1634   return AffineBound(AffineForOp(*this), lbMap.getNumInputs(),
1635                      lbMap.getNumInputs() + ubMap.getNumInputs(), ubMap);
1636 }
1637 
1638 void AffineForOp::setLowerBound(ValueRange lbOperands, AffineMap map) {
1639   assert(lbOperands.size() == map.getNumInputs());
1640   assert(map.getNumResults() >= 1 && "bound map has at least one result");
1641 
1642   SmallVector<Value, 4> newOperands(lbOperands.begin(), lbOperands.end());
1643 
1644   auto ubOperands = getUpperBoundOperands();
1645   newOperands.append(ubOperands.begin(), ubOperands.end());
1646   auto iterOperands = getIterOperands();
1647   newOperands.append(iterOperands.begin(), iterOperands.end());
1648   (*this)->setOperands(newOperands);
1649 
1650   (*this)->setAttr(getLowerBoundAttrName(), AffineMapAttr::get(map));
1651 }
1652 
1653 void AffineForOp::setUpperBound(ValueRange ubOperands, AffineMap map) {
1654   assert(ubOperands.size() == map.getNumInputs());
1655   assert(map.getNumResults() >= 1 && "bound map has at least one result");
1656 
1657   SmallVector<Value, 4> newOperands(getLowerBoundOperands());
1658   newOperands.append(ubOperands.begin(), ubOperands.end());
1659   auto iterOperands = getIterOperands();
1660   newOperands.append(iterOperands.begin(), iterOperands.end());
1661   (*this)->setOperands(newOperands);
1662 
1663   (*this)->setAttr(getUpperBoundAttrName(), AffineMapAttr::get(map));
1664 }
1665 
1666 void AffineForOp::setLowerBoundMap(AffineMap map) {
1667   auto lbMap = getLowerBoundMap();
1668   assert(lbMap.getNumDims() == map.getNumDims() &&
1669          lbMap.getNumSymbols() == map.getNumSymbols());
1670   assert(map.getNumResults() >= 1 && "bound map has at least one result");
1671   (void)lbMap;
1672   (*this)->setAttr(getLowerBoundAttrName(), AffineMapAttr::get(map));
1673 }
1674 
1675 void AffineForOp::setUpperBoundMap(AffineMap map) {
1676   auto ubMap = getUpperBoundMap();
1677   assert(ubMap.getNumDims() == map.getNumDims() &&
1678          ubMap.getNumSymbols() == map.getNumSymbols());
1679   assert(map.getNumResults() >= 1 && "bound map has at least one result");
1680   (void)ubMap;
1681   (*this)->setAttr(getUpperBoundAttrName(), AffineMapAttr::get(map));
1682 }
1683 
1684 bool AffineForOp::hasConstantLowerBound() {
1685   return getLowerBoundMap().isSingleConstant();
1686 }
1687 
1688 bool AffineForOp::hasConstantUpperBound() {
1689   return getUpperBoundMap().isSingleConstant();
1690 }
1691 
1692 int64_t AffineForOp::getConstantLowerBound() {
1693   return getLowerBoundMap().getSingleConstantResult();
1694 }
1695 
1696 int64_t AffineForOp::getConstantUpperBound() {
1697   return getUpperBoundMap().getSingleConstantResult();
1698 }
1699 
1700 void AffineForOp::setConstantLowerBound(int64_t value) {
1701   setLowerBound({}, AffineMap::getConstantMap(value, getContext()));
1702 }
1703 
1704 void AffineForOp::setConstantUpperBound(int64_t value) {
1705   setUpperBound({}, AffineMap::getConstantMap(value, getContext()));
1706 }
1707 
1708 AffineForOp::operand_range AffineForOp::getLowerBoundOperands() {
1709   return {operand_begin(), operand_begin() + getLowerBoundMap().getNumInputs()};
1710 }
1711 
1712 AffineForOp::operand_range AffineForOp::getUpperBoundOperands() {
1713   return {operand_begin() + getLowerBoundMap().getNumInputs(),
1714           operand_begin() + getLowerBoundMap().getNumInputs() +
1715               getUpperBoundMap().getNumInputs()};
1716 }
1717 
1718 bool AffineForOp::matchingBoundOperandList() {
1719   auto lbMap = getLowerBoundMap();
1720   auto ubMap = getUpperBoundMap();
1721   if (lbMap.getNumDims() != ubMap.getNumDims() ||
1722       lbMap.getNumSymbols() != ubMap.getNumSymbols())
1723     return false;
1724 
1725   unsigned numOperands = lbMap.getNumInputs();
1726   for (unsigned i = 0, e = lbMap.getNumInputs(); i < e; i++) {
1727     // Compare Value 's.
1728     if (getOperand(i) != getOperand(numOperands + i))
1729       return false;
1730   }
1731   return true;
1732 }
1733 
1734 Region &AffineForOp::getLoopBody() { return region(); }
1735 
1736 bool AffineForOp::isDefinedOutsideOfLoop(Value value) {
1737   return !region().isAncestor(value.getParentRegion());
1738 }
1739 
1740 LogicalResult AffineForOp::moveOutOfLoop(ArrayRef<Operation *> ops) {
1741   for (auto *op : ops)
1742     op->moveBefore(*this);
1743   return success();
1744 }
1745 
1746 /// Returns true if the provided value is the induction variable of a
1747 /// AffineForOp.
1748 bool mlir::isForInductionVar(Value val) {
1749   return getForInductionVarOwner(val) != AffineForOp();
1750 }
1751 
1752 /// Returns the loop parent of an induction variable. If the provided value is
1753 /// not an induction variable, then return nullptr.
1754 AffineForOp mlir::getForInductionVarOwner(Value val) {
1755   auto ivArg = val.dyn_cast<BlockArgument>();
1756   if (!ivArg || !ivArg.getOwner())
1757     return AffineForOp();
1758   auto *containingInst = ivArg.getOwner()->getParent()->getParentOp();
1759   return dyn_cast<AffineForOp>(containingInst);
1760 }
1761 
1762 /// Extracts the induction variables from a list of AffineForOps and returns
1763 /// them.
1764 void mlir::extractForInductionVars(ArrayRef<AffineForOp> forInsts,
1765                                    SmallVectorImpl<Value> *ivs) {
1766   ivs->reserve(forInsts.size());
1767   for (auto forInst : forInsts)
1768     ivs->push_back(forInst.getInductionVar());
1769 }
1770 
1771 /// Builds an affine loop nest, using "loopCreatorFn" to create individual loop
1772 /// operations.
1773 template <typename BoundListTy, typename LoopCreatorTy>
1774 static void buildAffineLoopNestImpl(
1775     OpBuilder &builder, Location loc, BoundListTy lbs, BoundListTy ubs,
1776     ArrayRef<int64_t> steps,
1777     function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilderFn,
1778     LoopCreatorTy &&loopCreatorFn) {
1779   assert(lbs.size() == ubs.size() && "Mismatch in number of arguments");
1780   assert(lbs.size() == steps.size() && "Mismatch in number of arguments");
1781 
1782   // If there are no loops to be constructed, construct the body anyway.
1783   OpBuilder::InsertionGuard guard(builder);
1784   if (lbs.empty()) {
1785     if (bodyBuilderFn)
1786       bodyBuilderFn(builder, loc, ValueRange());
1787     return;
1788   }
1789 
1790   // Create the loops iteratively and store the induction variables.
1791   SmallVector<Value, 4> ivs;
1792   ivs.reserve(lbs.size());
1793   for (unsigned i = 0, e = lbs.size(); i < e; ++i) {
1794     // Callback for creating the loop body, always creates the terminator.
1795     auto loopBody = [&](OpBuilder &nestedBuilder, Location nestedLoc, Value iv,
1796                         ValueRange iterArgs) {
1797       ivs.push_back(iv);
1798       // In the innermost loop, call the body builder.
1799       if (i == e - 1 && bodyBuilderFn) {
1800         OpBuilder::InsertionGuard nestedGuard(nestedBuilder);
1801         bodyBuilderFn(nestedBuilder, nestedLoc, ivs);
1802       }
1803       nestedBuilder.create<AffineYieldOp>(nestedLoc);
1804     };
1805 
1806     // Delegate actual loop creation to the callback in order to dispatch
1807     // between constant- and variable-bound loops.
1808     auto loop = loopCreatorFn(builder, loc, lbs[i], ubs[i], steps[i], loopBody);
1809     builder.setInsertionPointToStart(loop.getBody());
1810   }
1811 }
1812 
1813 /// Creates an affine loop from the bounds known to be constants.
1814 static AffineForOp
1815 buildAffineLoopFromConstants(OpBuilder &builder, Location loc, int64_t lb,
1816                              int64_t ub, int64_t step,
1817                              AffineForOp::BodyBuilderFn bodyBuilderFn) {
1818   return builder.create<AffineForOp>(loc, lb, ub, step, /*iterArgs=*/llvm::None,
1819                                      bodyBuilderFn);
1820 }
1821 
1822 /// Creates an affine loop from the bounds that may or may not be constants.
1823 static AffineForOp
1824 buildAffineLoopFromValues(OpBuilder &builder, Location loc, Value lb, Value ub,
1825                           int64_t step,
1826                           AffineForOp::BodyBuilderFn bodyBuilderFn) {
1827   auto lbConst = lb.getDefiningOp<ConstantIndexOp>();
1828   auto ubConst = ub.getDefiningOp<ConstantIndexOp>();
1829   if (lbConst && ubConst)
1830     return buildAffineLoopFromConstants(builder, loc, lbConst.getValue(),
1831                                         ubConst.getValue(), step,
1832                                         bodyBuilderFn);
1833   return builder.create<AffineForOp>(loc, lb, builder.getDimIdentityMap(), ub,
1834                                      builder.getDimIdentityMap(), step,
1835                                      /*iterArgs=*/llvm::None, bodyBuilderFn);
1836 }
1837 
1838 void mlir::buildAffineLoopNest(
1839     OpBuilder &builder, Location loc, ArrayRef<int64_t> lbs,
1840     ArrayRef<int64_t> ubs, ArrayRef<int64_t> steps,
1841     function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilderFn) {
1842   buildAffineLoopNestImpl(builder, loc, lbs, ubs, steps, bodyBuilderFn,
1843                           buildAffineLoopFromConstants);
1844 }
1845 
1846 void mlir::buildAffineLoopNest(
1847     OpBuilder &builder, Location loc, ValueRange lbs, ValueRange ubs,
1848     ArrayRef<int64_t> steps,
1849     function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilderFn) {
1850   buildAffineLoopNestImpl(builder, loc, lbs, ubs, steps, bodyBuilderFn,
1851                           buildAffineLoopFromValues);
1852 }
1853 
1854 //===----------------------------------------------------------------------===//
1855 // AffineIfOp
1856 //===----------------------------------------------------------------------===//
1857 
1858 namespace {
1859 /// Remove else blocks that have nothing other than a zero value yield.
1860 struct SimplifyDeadElse : public OpRewritePattern<AffineIfOp> {
1861   using OpRewritePattern<AffineIfOp>::OpRewritePattern;
1862 
1863   LogicalResult matchAndRewrite(AffineIfOp ifOp,
1864                                 PatternRewriter &rewriter) const override {
1865     if (ifOp.elseRegion().empty() ||
1866         !llvm::hasSingleElement(*ifOp.getElseBlock()) || ifOp.getNumResults())
1867       return failure();
1868 
1869     rewriter.startRootUpdate(ifOp);
1870     rewriter.eraseBlock(ifOp.getElseBlock());
1871     rewriter.finalizeRootUpdate(ifOp);
1872     return success();
1873   }
1874 };
1875 } // end anonymous namespace.
1876 
1877 static LogicalResult verify(AffineIfOp op) {
1878   // Verify that we have a condition attribute.
1879   auto conditionAttr =
1880       op->getAttrOfType<IntegerSetAttr>(op.getConditionAttrName());
1881   if (!conditionAttr)
1882     return op.emitOpError(
1883         "requires an integer set attribute named 'condition'");
1884 
1885   // Verify that there are enough operands for the condition.
1886   IntegerSet condition = conditionAttr.getValue();
1887   if (op.getNumOperands() != condition.getNumInputs())
1888     return op.emitOpError(
1889         "operand count and condition integer set dimension and "
1890         "symbol count must match");
1891 
1892   // Verify that the operands are valid dimension/symbols.
1893   if (failed(verifyDimAndSymbolIdentifiers(op, op.getOperands(),
1894                                            condition.getNumDims())))
1895     return failure();
1896 
1897   return success();
1898 }
1899 
1900 static ParseResult parseAffineIfOp(OpAsmParser &parser,
1901                                    OperationState &result) {
1902   // Parse the condition attribute set.
1903   IntegerSetAttr conditionAttr;
1904   unsigned numDims;
1905   if (parser.parseAttribute(conditionAttr, AffineIfOp::getConditionAttrName(),
1906                             result.attributes) ||
1907       parseDimAndSymbolList(parser, result.operands, numDims))
1908     return failure();
1909 
1910   // Verify the condition operands.
1911   auto set = conditionAttr.getValue();
1912   if (set.getNumDims() != numDims)
1913     return parser.emitError(
1914         parser.getNameLoc(),
1915         "dim operand count and integer set dim count must match");
1916   if (numDims + set.getNumSymbols() != result.operands.size())
1917     return parser.emitError(
1918         parser.getNameLoc(),
1919         "symbol operand count and integer set symbol count must match");
1920 
1921   if (parser.parseOptionalArrowTypeList(result.types))
1922     return failure();
1923 
1924   // Create the regions for 'then' and 'else'.  The latter must be created even
1925   // if it remains empty for the validity of the operation.
1926   result.regions.reserve(2);
1927   Region *thenRegion = result.addRegion();
1928   Region *elseRegion = result.addRegion();
1929 
1930   // Parse the 'then' region.
1931   if (parser.parseRegion(*thenRegion, {}, {}))
1932     return failure();
1933   AffineIfOp::ensureTerminator(*thenRegion, parser.getBuilder(),
1934                                result.location);
1935 
1936   // If we find an 'else' keyword then parse the 'else' region.
1937   if (!parser.parseOptionalKeyword("else")) {
1938     if (parser.parseRegion(*elseRegion, {}, {}))
1939       return failure();
1940     AffineIfOp::ensureTerminator(*elseRegion, parser.getBuilder(),
1941                                  result.location);
1942   }
1943 
1944   // Parse the optional attribute list.
1945   if (parser.parseOptionalAttrDict(result.attributes))
1946     return failure();
1947 
1948   return success();
1949 }
1950 
1951 static void print(OpAsmPrinter &p, AffineIfOp op) {
1952   auto conditionAttr =
1953       op->getAttrOfType<IntegerSetAttr>(op.getConditionAttrName());
1954   p << "affine.if " << conditionAttr;
1955   printDimAndSymbolList(op.operand_begin(), op.operand_end(),
1956                         conditionAttr.getValue().getNumDims(), p);
1957   p.printOptionalArrowTypeList(op.getResultTypes());
1958   p.printRegion(op.thenRegion(),
1959                 /*printEntryBlockArgs=*/false,
1960                 /*printBlockTerminators=*/op.getNumResults());
1961 
1962   // Print the 'else' regions if it has any blocks.
1963   auto &elseRegion = op.elseRegion();
1964   if (!elseRegion.empty()) {
1965     p << " else";
1966     p.printRegion(elseRegion,
1967                   /*printEntryBlockArgs=*/false,
1968                   /*printBlockTerminators=*/op.getNumResults());
1969   }
1970 
1971   // Print the attribute list.
1972   p.printOptionalAttrDict(op->getAttrs(),
1973                           /*elidedAttrs=*/op.getConditionAttrName());
1974 }
1975 
1976 IntegerSet AffineIfOp::getIntegerSet() {
1977   return (*this)
1978       ->getAttrOfType<IntegerSetAttr>(getConditionAttrName())
1979       .getValue();
1980 }
1981 void AffineIfOp::setIntegerSet(IntegerSet newSet) {
1982   (*this)->setAttr(getConditionAttrName(), IntegerSetAttr::get(newSet));
1983 }
1984 
1985 void AffineIfOp::setConditional(IntegerSet set, ValueRange operands) {
1986   setIntegerSet(set);
1987   (*this)->setOperands(operands);
1988 }
1989 
1990 void AffineIfOp::build(OpBuilder &builder, OperationState &result,
1991                        TypeRange resultTypes, IntegerSet set, ValueRange args,
1992                        bool withElseRegion) {
1993   assert(resultTypes.empty() || withElseRegion);
1994   result.addTypes(resultTypes);
1995   result.addOperands(args);
1996   result.addAttribute(getConditionAttrName(), IntegerSetAttr::get(set));
1997 
1998   Region *thenRegion = result.addRegion();
1999   thenRegion->push_back(new Block());
2000   if (resultTypes.empty())
2001     AffineIfOp::ensureTerminator(*thenRegion, builder, result.location);
2002 
2003   Region *elseRegion = result.addRegion();
2004   if (withElseRegion) {
2005     elseRegion->push_back(new Block());
2006     if (resultTypes.empty())
2007       AffineIfOp::ensureTerminator(*elseRegion, builder, result.location);
2008   }
2009 }
2010 
2011 void AffineIfOp::build(OpBuilder &builder, OperationState &result,
2012                        IntegerSet set, ValueRange args, bool withElseRegion) {
2013   AffineIfOp::build(builder, result, /*resultTypes=*/{}, set, args,
2014                     withElseRegion);
2015 }
2016 
2017 /// Canonicalize an affine if op's conditional (integer set + operands).
2018 LogicalResult AffineIfOp::fold(ArrayRef<Attribute>,
2019                                SmallVectorImpl<OpFoldResult> &) {
2020   auto set = getIntegerSet();
2021   SmallVector<Value, 4> operands(getOperands());
2022   canonicalizeSetAndOperands(&set, &operands);
2023 
2024   // Any canonicalization change always leads to either a reduction in the
2025   // number of operands or a change in the number of symbolic operands
2026   // (promotion of dims to symbols).
2027   if (operands.size() < getIntegerSet().getNumInputs() ||
2028       set.getNumSymbols() > getIntegerSet().getNumSymbols()) {
2029     setConditional(set, operands);
2030     return success();
2031   }
2032 
2033   return failure();
2034 }
2035 
2036 void AffineIfOp::getCanonicalizationPatterns(RewritePatternSet &results,
2037                                              MLIRContext *context) {
2038   results.add<SimplifyDeadElse>(context);
2039 }
2040 
2041 //===----------------------------------------------------------------------===//
2042 // AffineLoadOp
2043 //===----------------------------------------------------------------------===//
2044 
2045 void AffineLoadOp::build(OpBuilder &builder, OperationState &result,
2046                          AffineMap map, ValueRange operands) {
2047   assert(operands.size() == 1 + map.getNumInputs() && "inconsistent operands");
2048   result.addOperands(operands);
2049   if (map)
2050     result.addAttribute(getMapAttrName(), AffineMapAttr::get(map));
2051   auto memrefType = operands[0].getType().cast<MemRefType>();
2052   result.types.push_back(memrefType.getElementType());
2053 }
2054 
2055 void AffineLoadOp::build(OpBuilder &builder, OperationState &result,
2056                          Value memref, AffineMap map, ValueRange mapOperands) {
2057   assert(map.getNumInputs() == mapOperands.size() && "inconsistent index info");
2058   result.addOperands(memref);
2059   result.addOperands(mapOperands);
2060   auto memrefType = memref.getType().cast<MemRefType>();
2061   result.addAttribute(getMapAttrName(), AffineMapAttr::get(map));
2062   result.types.push_back(memrefType.getElementType());
2063 }
2064 
2065 void AffineLoadOp::build(OpBuilder &builder, OperationState &result,
2066                          Value memref, ValueRange indices) {
2067   auto memrefType = memref.getType().cast<MemRefType>();
2068   int64_t rank = memrefType.getRank();
2069   // Create identity map for memrefs with at least one dimension or () -> ()
2070   // for zero-dimensional memrefs.
2071   auto map =
2072       rank ? builder.getMultiDimIdentityMap(rank) : builder.getEmptyAffineMap();
2073   build(builder, result, memref, map, indices);
2074 }
2075 
2076 static ParseResult parseAffineLoadOp(OpAsmParser &parser,
2077                                      OperationState &result) {
2078   auto &builder = parser.getBuilder();
2079   auto indexTy = builder.getIndexType();
2080 
2081   MemRefType type;
2082   OpAsmParser::OperandType memrefInfo;
2083   AffineMapAttr mapAttr;
2084   SmallVector<OpAsmParser::OperandType, 1> mapOperands;
2085   return failure(
2086       parser.parseOperand(memrefInfo) ||
2087       parser.parseAffineMapOfSSAIds(mapOperands, mapAttr,
2088                                     AffineLoadOp::getMapAttrName(),
2089                                     result.attributes) ||
2090       parser.parseOptionalAttrDict(result.attributes) ||
2091       parser.parseColonType(type) ||
2092       parser.resolveOperand(memrefInfo, type, result.operands) ||
2093       parser.resolveOperands(mapOperands, indexTy, result.operands) ||
2094       parser.addTypeToList(type.getElementType(), result.types));
2095 }
2096 
2097 static void print(OpAsmPrinter &p, AffineLoadOp op) {
2098   p << "affine.load " << op.getMemRef() << '[';
2099   if (AffineMapAttr mapAttr =
2100           op->getAttrOfType<AffineMapAttr>(op.getMapAttrName()))
2101     p.printAffineMapOfSSAIds(mapAttr, op.getMapOperands());
2102   p << ']';
2103   p.printOptionalAttrDict(op->getAttrs(),
2104                           /*elidedAttrs=*/{op.getMapAttrName()});
2105   p << " : " << op.getMemRefType();
2106 }
2107 
2108 /// Verify common indexing invariants of affine.load, affine.store,
2109 /// affine.vector_load and affine.vector_store.
2110 static LogicalResult
2111 verifyMemoryOpIndexing(Operation *op, AffineMapAttr mapAttr,
2112                        Operation::operand_range mapOperands,
2113                        MemRefType memrefType, unsigned numIndexOperands) {
2114   if (mapAttr) {
2115     AffineMap map = mapAttr.getValue();
2116     if (map.getNumResults() != memrefType.getRank())
2117       return op->emitOpError("affine map num results must equal memref rank");
2118     if (map.getNumInputs() != numIndexOperands)
2119       return op->emitOpError("expects as many subscripts as affine map inputs");
2120   } else {
2121     if (memrefType.getRank() != numIndexOperands)
2122       return op->emitOpError(
2123           "expects the number of subscripts to be equal to memref rank");
2124   }
2125 
2126   Region *scope = getAffineScope(op);
2127   for (auto idx : mapOperands) {
2128     if (!idx.getType().isIndex())
2129       return op->emitOpError("index to load must have 'index' type");
2130     if (!isValidAffineIndexOperand(idx, scope))
2131       return op->emitOpError("index must be a dimension or symbol identifier");
2132   }
2133 
2134   return success();
2135 }
2136 
2137 LogicalResult verify(AffineLoadOp op) {
2138   auto memrefType = op.getMemRefType();
2139   if (op.getType() != memrefType.getElementType())
2140     return op.emitOpError("result type must match element type of memref");
2141 
2142   if (failed(verifyMemoryOpIndexing(
2143           op.getOperation(),
2144           op->getAttrOfType<AffineMapAttr>(op.getMapAttrName()),
2145           op.getMapOperands(), memrefType,
2146           /*numIndexOperands=*/op.getNumOperands() - 1)))
2147     return failure();
2148 
2149   return success();
2150 }
2151 
2152 void AffineLoadOp::getCanonicalizationPatterns(RewritePatternSet &results,
2153                                                MLIRContext *context) {
2154   results.add<SimplifyAffineOp<AffineLoadOp>>(context);
2155 }
2156 
2157 OpFoldResult AffineLoadOp::fold(ArrayRef<Attribute> cstOperands) {
2158   /// load(memrefcast) -> load
2159   if (succeeded(foldMemRefCast(*this)))
2160     return getResult();
2161   return OpFoldResult();
2162 }
2163 
2164 //===----------------------------------------------------------------------===//
2165 // AffineStoreOp
2166 //===----------------------------------------------------------------------===//
2167 
2168 void AffineStoreOp::build(OpBuilder &builder, OperationState &result,
2169                           Value valueToStore, Value memref, AffineMap map,
2170                           ValueRange mapOperands) {
2171   assert(map.getNumInputs() == mapOperands.size() && "inconsistent index info");
2172   result.addOperands(valueToStore);
2173   result.addOperands(memref);
2174   result.addOperands(mapOperands);
2175   result.addAttribute(getMapAttrName(), AffineMapAttr::get(map));
2176 }
2177 
2178 // Use identity map.
2179 void AffineStoreOp::build(OpBuilder &builder, OperationState &result,
2180                           Value valueToStore, Value memref,
2181                           ValueRange indices) {
2182   auto memrefType = memref.getType().cast<MemRefType>();
2183   int64_t rank = memrefType.getRank();
2184   // Create identity map for memrefs with at least one dimension or () -> ()
2185   // for zero-dimensional memrefs.
2186   auto map =
2187       rank ? builder.getMultiDimIdentityMap(rank) : builder.getEmptyAffineMap();
2188   build(builder, result, valueToStore, memref, map, indices);
2189 }
2190 
2191 static ParseResult parseAffineStoreOp(OpAsmParser &parser,
2192                                       OperationState &result) {
2193   auto indexTy = parser.getBuilder().getIndexType();
2194 
2195   MemRefType type;
2196   OpAsmParser::OperandType storeValueInfo;
2197   OpAsmParser::OperandType memrefInfo;
2198   AffineMapAttr mapAttr;
2199   SmallVector<OpAsmParser::OperandType, 1> mapOperands;
2200   return failure(parser.parseOperand(storeValueInfo) || parser.parseComma() ||
2201                  parser.parseOperand(memrefInfo) ||
2202                  parser.parseAffineMapOfSSAIds(mapOperands, mapAttr,
2203                                                AffineStoreOp::getMapAttrName(),
2204                                                result.attributes) ||
2205                  parser.parseOptionalAttrDict(result.attributes) ||
2206                  parser.parseColonType(type) ||
2207                  parser.resolveOperand(storeValueInfo, type.getElementType(),
2208                                        result.operands) ||
2209                  parser.resolveOperand(memrefInfo, type, result.operands) ||
2210                  parser.resolveOperands(mapOperands, indexTy, result.operands));
2211 }
2212 
2213 static void print(OpAsmPrinter &p, AffineStoreOp op) {
2214   p << "affine.store " << op.getValueToStore();
2215   p << ", " << op.getMemRef() << '[';
2216   if (AffineMapAttr mapAttr =
2217           op->getAttrOfType<AffineMapAttr>(op.getMapAttrName()))
2218     p.printAffineMapOfSSAIds(mapAttr, op.getMapOperands());
2219   p << ']';
2220   p.printOptionalAttrDict(op->getAttrs(),
2221                           /*elidedAttrs=*/{op.getMapAttrName()});
2222   p << " : " << op.getMemRefType();
2223 }
2224 
2225 LogicalResult verify(AffineStoreOp op) {
2226   // First operand must have same type as memref element type.
2227   auto memrefType = op.getMemRefType();
2228   if (op.getValueToStore().getType() != memrefType.getElementType())
2229     return op.emitOpError(
2230         "first operand must have same type memref element type");
2231 
2232   if (failed(verifyMemoryOpIndexing(
2233           op.getOperation(),
2234           op->getAttrOfType<AffineMapAttr>(op.getMapAttrName()),
2235           op.getMapOperands(), memrefType,
2236           /*numIndexOperands=*/op.getNumOperands() - 2)))
2237     return failure();
2238 
2239   return success();
2240 }
2241 
2242 void AffineStoreOp::getCanonicalizationPatterns(RewritePatternSet &results,
2243                                                 MLIRContext *context) {
2244   results.add<SimplifyAffineOp<AffineStoreOp>>(context);
2245 }
2246 
2247 LogicalResult AffineStoreOp::fold(ArrayRef<Attribute> cstOperands,
2248                                   SmallVectorImpl<OpFoldResult> &results) {
2249   /// store(memrefcast) -> store
2250   return foldMemRefCast(*this);
2251 }
2252 
2253 //===----------------------------------------------------------------------===//
2254 // AffineMinMaxOpBase
2255 //===----------------------------------------------------------------------===//
2256 
2257 template <typename T>
2258 static LogicalResult verifyAffineMinMaxOp(T op) {
2259   // Verify that operand count matches affine map dimension and symbol count.
2260   if (op.getNumOperands() != op.map().getNumDims() + op.map().getNumSymbols())
2261     return op.emitOpError(
2262         "operand count and affine map dimension and symbol count must match");
2263   return success();
2264 }
2265 
2266 template <typename T>
2267 static void printAffineMinMaxOp(OpAsmPrinter &p, T op) {
2268   p << op.getOperationName() << ' ' << op->getAttr(T::getMapAttrName());
2269   auto operands = op.getOperands();
2270   unsigned numDims = op.map().getNumDims();
2271   p << '(' << operands.take_front(numDims) << ')';
2272 
2273   if (operands.size() != numDims)
2274     p << '[' << operands.drop_front(numDims) << ']';
2275   p.printOptionalAttrDict(op->getAttrs(),
2276                           /*elidedAttrs=*/{T::getMapAttrName()});
2277 }
2278 
2279 template <typename T>
2280 static ParseResult parseAffineMinMaxOp(OpAsmParser &parser,
2281                                        OperationState &result) {
2282   auto &builder = parser.getBuilder();
2283   auto indexType = builder.getIndexType();
2284   SmallVector<OpAsmParser::OperandType, 8> dim_infos;
2285   SmallVector<OpAsmParser::OperandType, 8> sym_infos;
2286   AffineMapAttr mapAttr;
2287   return failure(
2288       parser.parseAttribute(mapAttr, T::getMapAttrName(), result.attributes) ||
2289       parser.parseOperandList(dim_infos, OpAsmParser::Delimiter::Paren) ||
2290       parser.parseOperandList(sym_infos,
2291                               OpAsmParser::Delimiter::OptionalSquare) ||
2292       parser.parseOptionalAttrDict(result.attributes) ||
2293       parser.resolveOperands(dim_infos, indexType, result.operands) ||
2294       parser.resolveOperands(sym_infos, indexType, result.operands) ||
2295       parser.addTypeToList(indexType, result.types));
2296 }
2297 
2298 /// Fold an affine min or max operation with the given operands. The operand
2299 /// list may contain nulls, which are interpreted as the operand not being a
2300 /// constant.
2301 template <typename T>
2302 static OpFoldResult foldMinMaxOp(T op, ArrayRef<Attribute> operands) {
2303   static_assert(llvm::is_one_of<T, AffineMinOp, AffineMaxOp>::value,
2304                 "expected affine min or max op");
2305 
2306   // Fold the affine map.
2307   // TODO: Fold more cases:
2308   // min(some_affine, some_affine + constant, ...), etc.
2309   SmallVector<int64_t, 2> results;
2310   auto foldedMap = op.map().partialConstantFold(operands, &results);
2311 
2312   // If some of the map results are not constant, try changing the map in-place.
2313   if (results.empty()) {
2314     // If the map is the same, report that folding did not happen.
2315     if (foldedMap == op.map())
2316       return {};
2317     op->setAttr("map", AffineMapAttr::get(foldedMap));
2318     return op.getResult();
2319   }
2320 
2321   // Otherwise, completely fold the op into a constant.
2322   auto resultIt = std::is_same<T, AffineMinOp>::value
2323                       ? std::min_element(results.begin(), results.end())
2324                       : std::max_element(results.begin(), results.end());
2325   if (resultIt == results.end())
2326     return {};
2327   return IntegerAttr::get(IndexType::get(op.getContext()), *resultIt);
2328 }
2329 
2330 //===----------------------------------------------------------------------===//
2331 // AffineMinOp
2332 //===----------------------------------------------------------------------===//
2333 //
2334 //   %0 = affine.min (d0) -> (1000, d0 + 512) (%i0)
2335 //
2336 
2337 OpFoldResult AffineMinOp::fold(ArrayRef<Attribute> operands) {
2338   return foldMinMaxOp(*this, operands);
2339 }
2340 
2341 void AffineMinOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
2342                                               MLIRContext *context) {
2343   patterns.add<SimplifyAffineOp<AffineMinOp>>(context);
2344 }
2345 
2346 //===----------------------------------------------------------------------===//
2347 // AffineMaxOp
2348 //===----------------------------------------------------------------------===//
2349 //
2350 //   %0 = affine.max (d0) -> (1000, d0 + 512) (%i0)
2351 //
2352 
2353 OpFoldResult AffineMaxOp::fold(ArrayRef<Attribute> operands) {
2354   return foldMinMaxOp(*this, operands);
2355 }
2356 
2357 void AffineMaxOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
2358                                               MLIRContext *context) {
2359   patterns.add<SimplifyAffineOp<AffineMaxOp>>(context);
2360 }
2361 
2362 //===----------------------------------------------------------------------===//
2363 // AffinePrefetchOp
2364 //===----------------------------------------------------------------------===//
2365 
2366 //
2367 // affine.prefetch %0[%i, %j + 5], read, locality<3>, data : memref<400x400xi32>
2368 //
2369 static ParseResult parseAffinePrefetchOp(OpAsmParser &parser,
2370                                          OperationState &result) {
2371   auto &builder = parser.getBuilder();
2372   auto indexTy = builder.getIndexType();
2373 
2374   MemRefType type;
2375   OpAsmParser::OperandType memrefInfo;
2376   IntegerAttr hintInfo;
2377   auto i32Type = parser.getBuilder().getIntegerType(32);
2378   StringRef readOrWrite, cacheType;
2379 
2380   AffineMapAttr mapAttr;
2381   SmallVector<OpAsmParser::OperandType, 1> mapOperands;
2382   if (parser.parseOperand(memrefInfo) ||
2383       parser.parseAffineMapOfSSAIds(mapOperands, mapAttr,
2384                                     AffinePrefetchOp::getMapAttrName(),
2385                                     result.attributes) ||
2386       parser.parseComma() || parser.parseKeyword(&readOrWrite) ||
2387       parser.parseComma() || parser.parseKeyword("locality") ||
2388       parser.parseLess() ||
2389       parser.parseAttribute(hintInfo, i32Type,
2390                             AffinePrefetchOp::getLocalityHintAttrName(),
2391                             result.attributes) ||
2392       parser.parseGreater() || parser.parseComma() ||
2393       parser.parseKeyword(&cacheType) ||
2394       parser.parseOptionalAttrDict(result.attributes) ||
2395       parser.parseColonType(type) ||
2396       parser.resolveOperand(memrefInfo, type, result.operands) ||
2397       parser.resolveOperands(mapOperands, indexTy, result.operands))
2398     return failure();
2399 
2400   if (!readOrWrite.equals("read") && !readOrWrite.equals("write"))
2401     return parser.emitError(parser.getNameLoc(),
2402                             "rw specifier has to be 'read' or 'write'");
2403   result.addAttribute(
2404       AffinePrefetchOp::getIsWriteAttrName(),
2405       parser.getBuilder().getBoolAttr(readOrWrite.equals("write")));
2406 
2407   if (!cacheType.equals("data") && !cacheType.equals("instr"))
2408     return parser.emitError(parser.getNameLoc(),
2409                             "cache type has to be 'data' or 'instr'");
2410 
2411   result.addAttribute(
2412       AffinePrefetchOp::getIsDataCacheAttrName(),
2413       parser.getBuilder().getBoolAttr(cacheType.equals("data")));
2414 
2415   return success();
2416 }
2417 
2418 static void print(OpAsmPrinter &p, AffinePrefetchOp op) {
2419   p << AffinePrefetchOp::getOperationName() << " " << op.memref() << '[';
2420   AffineMapAttr mapAttr = op->getAttrOfType<AffineMapAttr>(op.getMapAttrName());
2421   if (mapAttr) {
2422     SmallVector<Value, 2> operands(op.getMapOperands());
2423     p.printAffineMapOfSSAIds(mapAttr, operands);
2424   }
2425   p << ']' << ", " << (op.isWrite() ? "write" : "read") << ", "
2426     << "locality<" << op.localityHint() << ">, "
2427     << (op.isDataCache() ? "data" : "instr");
2428   p.printOptionalAttrDict(
2429       op->getAttrs(),
2430       /*elidedAttrs=*/{op.getMapAttrName(), op.getLocalityHintAttrName(),
2431                        op.getIsDataCacheAttrName(), op.getIsWriteAttrName()});
2432   p << " : " << op.getMemRefType();
2433 }
2434 
2435 static LogicalResult verify(AffinePrefetchOp op) {
2436   auto mapAttr = op->getAttrOfType<AffineMapAttr>(op.getMapAttrName());
2437   if (mapAttr) {
2438     AffineMap map = mapAttr.getValue();
2439     if (map.getNumResults() != op.getMemRefType().getRank())
2440       return op.emitOpError("affine.prefetch affine map num results must equal"
2441                             " memref rank");
2442     if (map.getNumInputs() + 1 != op.getNumOperands())
2443       return op.emitOpError("too few operands");
2444   } else {
2445     if (op.getNumOperands() != 1)
2446       return op.emitOpError("too few operands");
2447   }
2448 
2449   Region *scope = getAffineScope(op);
2450   for (auto idx : op.getMapOperands()) {
2451     if (!isValidAffineIndexOperand(idx, scope))
2452       return op.emitOpError("index must be a dimension or symbol identifier");
2453   }
2454   return success();
2455 }
2456 
2457 void AffinePrefetchOp::getCanonicalizationPatterns(RewritePatternSet &results,
2458                                                    MLIRContext *context) {
2459   // prefetch(memrefcast) -> prefetch
2460   results.add<SimplifyAffineOp<AffinePrefetchOp>>(context);
2461 }
2462 
2463 LogicalResult AffinePrefetchOp::fold(ArrayRef<Attribute> cstOperands,
2464                                      SmallVectorImpl<OpFoldResult> &results) {
2465   /// prefetch(memrefcast) -> prefetch
2466   return foldMemRefCast(*this);
2467 }
2468 
2469 //===----------------------------------------------------------------------===//
2470 // AffineParallelOp
2471 //===----------------------------------------------------------------------===//
2472 
2473 void AffineParallelOp::build(OpBuilder &builder, OperationState &result,
2474                              TypeRange resultTypes,
2475                              ArrayRef<AtomicRMWKind> reductions,
2476                              ArrayRef<int64_t> ranges) {
2477   SmallVector<AffineExpr, 8> lbExprs(ranges.size(),
2478                                      builder.getAffineConstantExpr(0));
2479   auto lbMap = AffineMap::get(0, 0, lbExprs, builder.getContext());
2480   SmallVector<AffineExpr, 8> ubExprs;
2481   for (int64_t range : ranges)
2482     ubExprs.push_back(builder.getAffineConstantExpr(range));
2483   auto ubMap = AffineMap::get(0, 0, ubExprs, builder.getContext());
2484   build(builder, result, resultTypes, reductions, lbMap, /*lbArgs=*/{}, ubMap,
2485         /*ubArgs=*/{});
2486 }
2487 
2488 void AffineParallelOp::build(OpBuilder &builder, OperationState &result,
2489                              TypeRange resultTypes,
2490                              ArrayRef<AtomicRMWKind> reductions,
2491                              AffineMap lbMap, ValueRange lbArgs,
2492                              AffineMap ubMap, ValueRange ubArgs) {
2493   auto numDims = lbMap.getNumResults();
2494   // Verify that the dimensionality of both maps are the same.
2495   assert(numDims == ubMap.getNumResults() &&
2496          "num dims and num results mismatch");
2497   // Make default step sizes of 1.
2498   SmallVector<int64_t, 8> steps(numDims, 1);
2499   build(builder, result, resultTypes, reductions, lbMap, lbArgs, ubMap, ubArgs,
2500         steps);
2501 }
2502 
2503 void AffineParallelOp::build(OpBuilder &builder, OperationState &result,
2504                              TypeRange resultTypes,
2505                              ArrayRef<AtomicRMWKind> reductions,
2506                              AffineMap lbMap, ValueRange lbArgs,
2507                              AffineMap ubMap, ValueRange ubArgs,
2508                              ArrayRef<int64_t> steps) {
2509   auto numDims = lbMap.getNumResults();
2510   // Verify that the dimensionality of the maps matches the number of steps.
2511   assert(numDims == ubMap.getNumResults() &&
2512          "num dims and num results mismatch");
2513   assert(numDims == steps.size() && "num dims and num steps mismatch");
2514 
2515   result.addTypes(resultTypes);
2516   // Convert the reductions to integer attributes.
2517   SmallVector<Attribute, 4> reductionAttrs;
2518   for (AtomicRMWKind reduction : reductions)
2519     reductionAttrs.push_back(
2520         builder.getI64IntegerAttr(static_cast<int64_t>(reduction)));
2521   result.addAttribute(getReductionsAttrName(),
2522                       builder.getArrayAttr(reductionAttrs));
2523   result.addAttribute(getLowerBoundsMapAttrName(), AffineMapAttr::get(lbMap));
2524   result.addAttribute(getUpperBoundsMapAttrName(), AffineMapAttr::get(ubMap));
2525   result.addAttribute(getStepsAttrName(), builder.getI64ArrayAttr(steps));
2526   result.addOperands(lbArgs);
2527   result.addOperands(ubArgs);
2528   // Create a region and a block for the body.
2529   auto *bodyRegion = result.addRegion();
2530   auto *body = new Block();
2531   // Add all the block arguments.
2532   for (unsigned i = 0; i < numDims; ++i)
2533     body->addArgument(IndexType::get(builder.getContext()));
2534   bodyRegion->push_back(body);
2535   if (resultTypes.empty())
2536     ensureTerminator(*bodyRegion, builder, result.location);
2537 }
2538 
2539 Region &AffineParallelOp::getLoopBody() { return region(); }
2540 
2541 bool AffineParallelOp::isDefinedOutsideOfLoop(Value value) {
2542   return !region().isAncestor(value.getParentRegion());
2543 }
2544 
2545 LogicalResult AffineParallelOp::moveOutOfLoop(ArrayRef<Operation *> ops) {
2546   for (Operation *op : ops)
2547     op->moveBefore(*this);
2548   return success();
2549 }
2550 
2551 unsigned AffineParallelOp::getNumDims() { return steps().size(); }
2552 
2553 AffineParallelOp::operand_range AffineParallelOp::getLowerBoundsOperands() {
2554   return getOperands().take_front(lowerBoundsMap().getNumInputs());
2555 }
2556 
2557 AffineParallelOp::operand_range AffineParallelOp::getUpperBoundsOperands() {
2558   return getOperands().drop_front(lowerBoundsMap().getNumInputs());
2559 }
2560 
2561 AffineValueMap AffineParallelOp::getLowerBoundsValueMap() {
2562   return AffineValueMap(lowerBoundsMap(), getLowerBoundsOperands());
2563 }
2564 
2565 AffineValueMap AffineParallelOp::getUpperBoundsValueMap() {
2566   return AffineValueMap(upperBoundsMap(), getUpperBoundsOperands());
2567 }
2568 
2569 AffineValueMap AffineParallelOp::getRangesValueMap() {
2570   AffineValueMap out;
2571   AffineValueMap::difference(getUpperBoundsValueMap(), getLowerBoundsValueMap(),
2572                              &out);
2573   return out;
2574 }
2575 
2576 Optional<SmallVector<int64_t, 8>> AffineParallelOp::getConstantRanges() {
2577   // Try to convert all the ranges to constant expressions.
2578   SmallVector<int64_t, 8> out;
2579   AffineValueMap rangesValueMap = getRangesValueMap();
2580   out.reserve(rangesValueMap.getNumResults());
2581   for (unsigned i = 0, e = rangesValueMap.getNumResults(); i < e; ++i) {
2582     auto expr = rangesValueMap.getResult(i);
2583     auto cst = expr.dyn_cast<AffineConstantExpr>();
2584     if (!cst)
2585       return llvm::None;
2586     out.push_back(cst.getValue());
2587   }
2588   return out;
2589 }
2590 
2591 Block *AffineParallelOp::getBody() { return &region().front(); }
2592 
2593 OpBuilder AffineParallelOp::getBodyBuilder() {
2594   return OpBuilder(getBody(), std::prev(getBody()->end()));
2595 }
2596 
2597 void AffineParallelOp::setLowerBounds(ValueRange lbOperands, AffineMap map) {
2598   assert(lbOperands.size() == map.getNumInputs() &&
2599          "operands to map must match number of inputs");
2600   assert(map.getNumResults() >= 1 && "bounds map has at least one result");
2601 
2602   auto ubOperands = getUpperBoundsOperands();
2603 
2604   SmallVector<Value, 4> newOperands(lbOperands);
2605   newOperands.append(ubOperands.begin(), ubOperands.end());
2606   (*this)->setOperands(newOperands);
2607 
2608   lowerBoundsMapAttr(AffineMapAttr::get(map));
2609 }
2610 
2611 void AffineParallelOp::setUpperBounds(ValueRange ubOperands, AffineMap map) {
2612   assert(ubOperands.size() == map.getNumInputs() &&
2613          "operands to map must match number of inputs");
2614   assert(map.getNumResults() >= 1 && "bounds map has at least one result");
2615 
2616   SmallVector<Value, 4> newOperands(getLowerBoundsOperands());
2617   newOperands.append(ubOperands.begin(), ubOperands.end());
2618   (*this)->setOperands(newOperands);
2619 
2620   upperBoundsMapAttr(AffineMapAttr::get(map));
2621 }
2622 
2623 void AffineParallelOp::setLowerBoundsMap(AffineMap map) {
2624   AffineMap lbMap = lowerBoundsMap();
2625   assert(lbMap.getNumDims() == map.getNumDims() &&
2626          lbMap.getNumSymbols() == map.getNumSymbols());
2627   (void)lbMap;
2628   lowerBoundsMapAttr(AffineMapAttr::get(map));
2629 }
2630 
2631 void AffineParallelOp::setUpperBoundsMap(AffineMap map) {
2632   AffineMap ubMap = upperBoundsMap();
2633   assert(ubMap.getNumDims() == map.getNumDims() &&
2634          ubMap.getNumSymbols() == map.getNumSymbols());
2635   (void)ubMap;
2636   upperBoundsMapAttr(AffineMapAttr::get(map));
2637 }
2638 
2639 SmallVector<int64_t, 8> AffineParallelOp::getSteps() {
2640   SmallVector<int64_t, 8> result;
2641   for (Attribute attr : steps()) {
2642     result.push_back(attr.cast<IntegerAttr>().getInt());
2643   }
2644   return result;
2645 }
2646 
2647 void AffineParallelOp::setSteps(ArrayRef<int64_t> newSteps) {
2648   stepsAttr(getBodyBuilder().getI64ArrayAttr(newSteps));
2649 }
2650 
2651 static LogicalResult verify(AffineParallelOp op) {
2652   auto numDims = op.getNumDims();
2653   if (op.lowerBoundsMap().getNumResults() != numDims ||
2654       op.upperBoundsMap().getNumResults() != numDims ||
2655       op.steps().size() != numDims ||
2656       op.getBody()->getNumArguments() != numDims)
2657     return op.emitOpError("region argument count and num results of upper "
2658                           "bounds, lower bounds, and steps must all match");
2659 
2660   if (op.reductions().size() != op.getNumResults())
2661     return op.emitOpError("a reduction must be specified for each output");
2662 
2663   // Verify reduction  ops are all valid
2664   for (Attribute attr : op.reductions()) {
2665     auto intAttr = attr.dyn_cast<IntegerAttr>();
2666     if (!intAttr || !symbolizeAtomicRMWKind(intAttr.getInt()))
2667       return op.emitOpError("invalid reduction attribute");
2668   }
2669 
2670   // Verify that the bound operands are valid dimension/symbols.
2671   /// Lower bounds.
2672   if (failed(verifyDimAndSymbolIdentifiers(op, op.getLowerBoundsOperands(),
2673                                            op.lowerBoundsMap().getNumDims())))
2674     return failure();
2675   /// Upper bounds.
2676   if (failed(verifyDimAndSymbolIdentifiers(op, op.getUpperBoundsOperands(),
2677                                            op.upperBoundsMap().getNumDims())))
2678     return failure();
2679   return success();
2680 }
2681 
2682 LogicalResult AffineValueMap::canonicalize() {
2683   SmallVector<Value, 4> newOperands{operands};
2684   auto newMap = getAffineMap();
2685   composeAffineMapAndOperands(&newMap, &newOperands);
2686   if (newMap == getAffineMap() && newOperands == operands)
2687     return failure();
2688   reset(newMap, newOperands);
2689   return success();
2690 }
2691 
2692 /// Canonicalize the bounds of the given loop.
2693 static LogicalResult canonicalizeLoopBounds(AffineParallelOp op) {
2694   AffineValueMap lb = op.getLowerBoundsValueMap();
2695   bool lbCanonicalized = succeeded(lb.canonicalize());
2696 
2697   AffineValueMap ub = op.getUpperBoundsValueMap();
2698   bool ubCanonicalized = succeeded(ub.canonicalize());
2699 
2700   // Any canonicalization change always leads to updated map(s).
2701   if (!lbCanonicalized && !ubCanonicalized)
2702     return failure();
2703 
2704   if (lbCanonicalized)
2705     op.setLowerBounds(lb.getOperands(), lb.getAffineMap());
2706   if (ubCanonicalized)
2707     op.setUpperBounds(ub.getOperands(), ub.getAffineMap());
2708 
2709   return success();
2710 }
2711 
2712 LogicalResult AffineParallelOp::fold(ArrayRef<Attribute> operands,
2713                                      SmallVectorImpl<OpFoldResult> &results) {
2714   return canonicalizeLoopBounds(*this);
2715 }
2716 
2717 static void print(OpAsmPrinter &p, AffineParallelOp op) {
2718   p << op.getOperationName() << " (" << op.getBody()->getArguments() << ") = (";
2719   p.printAffineMapOfSSAIds(op.lowerBoundsMapAttr(),
2720                            op.getLowerBoundsOperands());
2721   p << ") to (";
2722   p.printAffineMapOfSSAIds(op.upperBoundsMapAttr(),
2723                            op.getUpperBoundsOperands());
2724   p << ')';
2725   SmallVector<int64_t, 8> steps = op.getSteps();
2726   bool elideSteps = llvm::all_of(steps, [](int64_t step) { return step == 1; });
2727   if (!elideSteps) {
2728     p << " step (";
2729     llvm::interleaveComma(steps, p);
2730     p << ')';
2731   }
2732   if (op.getNumResults()) {
2733     p << " reduce (";
2734     llvm::interleaveComma(op.reductions(), p, [&](auto &attr) {
2735       AtomicRMWKind sym =
2736           *symbolizeAtomicRMWKind(attr.template cast<IntegerAttr>().getInt());
2737       p << "\"" << stringifyAtomicRMWKind(sym) << "\"";
2738     });
2739     p << ") -> (" << op.getResultTypes() << ")";
2740   }
2741 
2742   p.printRegion(op.region(), /*printEntryBlockArgs=*/false,
2743                 /*printBlockTerminators=*/op.getNumResults());
2744   p.printOptionalAttrDict(
2745       op->getAttrs(),
2746       /*elidedAttrs=*/{AffineParallelOp::getReductionsAttrName(),
2747                        AffineParallelOp::getLowerBoundsMapAttrName(),
2748                        AffineParallelOp::getUpperBoundsMapAttrName(),
2749                        AffineParallelOp::getStepsAttrName()});
2750 }
2751 
2752 //
2753 // operation ::= `affine.parallel` `(` ssa-ids `)` `=` `(` map-of-ssa-ids `)`
2754 //               `to` `(` map-of-ssa-ids `)` steps? region attr-dict?
2755 // steps     ::= `steps` `(` integer-literals `)`
2756 //
2757 static ParseResult parseAffineParallelOp(OpAsmParser &parser,
2758                                          OperationState &result) {
2759   auto &builder = parser.getBuilder();
2760   auto indexType = builder.getIndexType();
2761   AffineMapAttr lowerBoundsAttr, upperBoundsAttr;
2762   SmallVector<OpAsmParser::OperandType, 4> ivs;
2763   SmallVector<OpAsmParser::OperandType, 4> lowerBoundsMapOperands;
2764   SmallVector<OpAsmParser::OperandType, 4> upperBoundsMapOperands;
2765   if (parser.parseRegionArgumentList(ivs, /*requiredOperandCount=*/-1,
2766                                      OpAsmParser::Delimiter::Paren) ||
2767       parser.parseEqual() ||
2768       parser.parseAffineMapOfSSAIds(
2769           lowerBoundsMapOperands, lowerBoundsAttr,
2770           AffineParallelOp::getLowerBoundsMapAttrName(), result.attributes,
2771           OpAsmParser::Delimiter::Paren) ||
2772       parser.resolveOperands(lowerBoundsMapOperands, indexType,
2773                              result.operands) ||
2774       parser.parseKeyword("to") ||
2775       parser.parseAffineMapOfSSAIds(
2776           upperBoundsMapOperands, upperBoundsAttr,
2777           AffineParallelOp::getUpperBoundsMapAttrName(), result.attributes,
2778           OpAsmParser::Delimiter::Paren) ||
2779       parser.resolveOperands(upperBoundsMapOperands, indexType,
2780                              result.operands))
2781     return failure();
2782 
2783   AffineMapAttr stepsMapAttr;
2784   NamedAttrList stepsAttrs;
2785   SmallVector<OpAsmParser::OperandType, 4> stepsMapOperands;
2786   if (failed(parser.parseOptionalKeyword("step"))) {
2787     SmallVector<int64_t, 4> steps(ivs.size(), 1);
2788     result.addAttribute(AffineParallelOp::getStepsAttrName(),
2789                         builder.getI64ArrayAttr(steps));
2790   } else {
2791     if (parser.parseAffineMapOfSSAIds(stepsMapOperands, stepsMapAttr,
2792                                       AffineParallelOp::getStepsAttrName(),
2793                                       stepsAttrs,
2794                                       OpAsmParser::Delimiter::Paren))
2795       return failure();
2796 
2797     // Convert steps from an AffineMap into an I64ArrayAttr.
2798     SmallVector<int64_t, 4> steps;
2799     auto stepsMap = stepsMapAttr.getValue();
2800     for (const auto &result : stepsMap.getResults()) {
2801       auto constExpr = result.dyn_cast<AffineConstantExpr>();
2802       if (!constExpr)
2803         return parser.emitError(parser.getNameLoc(),
2804                                 "steps must be constant integers");
2805       steps.push_back(constExpr.getValue());
2806     }
2807     result.addAttribute(AffineParallelOp::getStepsAttrName(),
2808                         builder.getI64ArrayAttr(steps));
2809   }
2810 
2811   // Parse optional clause of the form: `reduce ("addf", "maxf")`, where the
2812   // quoted strings are a member of the enum AtomicRMWKind.
2813   SmallVector<Attribute, 4> reductions;
2814   if (succeeded(parser.parseOptionalKeyword("reduce"))) {
2815     if (parser.parseLParen())
2816       return failure();
2817     do {
2818       // Parse a single quoted string via the attribute parsing, and then
2819       // verify it is a member of the enum and convert to it's integer
2820       // representation.
2821       StringAttr attrVal;
2822       NamedAttrList attrStorage;
2823       auto loc = parser.getCurrentLocation();
2824       if (parser.parseAttribute(attrVal, builder.getNoneType(), "reduce",
2825                                 attrStorage))
2826         return failure();
2827       llvm::Optional<AtomicRMWKind> reduction =
2828           symbolizeAtomicRMWKind(attrVal.getValue());
2829       if (!reduction)
2830         return parser.emitError(loc, "invalid reduction value: ") << attrVal;
2831       reductions.push_back(builder.getI64IntegerAttr(
2832           static_cast<int64_t>(reduction.getValue())));
2833       // While we keep getting commas, keep parsing.
2834     } while (succeeded(parser.parseOptionalComma()));
2835     if (parser.parseRParen())
2836       return failure();
2837   }
2838   result.addAttribute(AffineParallelOp::getReductionsAttrName(),
2839                       builder.getArrayAttr(reductions));
2840 
2841   // Parse return types of reductions (if any)
2842   if (parser.parseOptionalArrowTypeList(result.types))
2843     return failure();
2844 
2845   // Now parse the body.
2846   Region *body = result.addRegion();
2847   SmallVector<Type, 4> types(ivs.size(), indexType);
2848   if (parser.parseRegion(*body, ivs, types) ||
2849       parser.parseOptionalAttrDict(result.attributes))
2850     return failure();
2851 
2852   // Add a terminator if none was parsed.
2853   AffineParallelOp::ensureTerminator(*body, builder, result.location);
2854   return success();
2855 }
2856 
2857 //===----------------------------------------------------------------------===//
2858 // AffineYieldOp
2859 //===----------------------------------------------------------------------===//
2860 
2861 static LogicalResult verify(AffineYieldOp op) {
2862   auto *parentOp = op->getParentOp();
2863   auto results = parentOp->getResults();
2864   auto operands = op.getOperands();
2865 
2866   if (!isa<AffineParallelOp, AffineIfOp, AffineForOp>(parentOp))
2867     return op.emitOpError() << "only terminates affine.if/for/parallel regions";
2868   if (parentOp->getNumResults() != op.getNumOperands())
2869     return op.emitOpError() << "parent of yield must have same number of "
2870                                "results as the yield operands";
2871   for (auto it : llvm::zip(results, operands)) {
2872     if (std::get<0>(it).getType() != std::get<1>(it).getType())
2873       return op.emitOpError()
2874              << "types mismatch between yield op and its parent";
2875   }
2876 
2877   return success();
2878 }
2879 
2880 //===----------------------------------------------------------------------===//
2881 // AffineVectorLoadOp
2882 //===----------------------------------------------------------------------===//
2883 
2884 void AffineVectorLoadOp::build(OpBuilder &builder, OperationState &result,
2885                                VectorType resultType, AffineMap map,
2886                                ValueRange operands) {
2887   assert(operands.size() == 1 + map.getNumInputs() && "inconsistent operands");
2888   result.addOperands(operands);
2889   if (map)
2890     result.addAttribute(getMapAttrName(), AffineMapAttr::get(map));
2891   result.types.push_back(resultType);
2892 }
2893 
2894 void AffineVectorLoadOp::build(OpBuilder &builder, OperationState &result,
2895                                VectorType resultType, Value memref,
2896                                AffineMap map, ValueRange mapOperands) {
2897   assert(map.getNumInputs() == mapOperands.size() && "inconsistent index info");
2898   result.addOperands(memref);
2899   result.addOperands(mapOperands);
2900   result.addAttribute(getMapAttrName(), AffineMapAttr::get(map));
2901   result.types.push_back(resultType);
2902 }
2903 
2904 void AffineVectorLoadOp::build(OpBuilder &builder, OperationState &result,
2905                                VectorType resultType, Value memref,
2906                                ValueRange indices) {
2907   auto memrefType = memref.getType().cast<MemRefType>();
2908   int64_t rank = memrefType.getRank();
2909   // Create identity map for memrefs with at least one dimension or () -> ()
2910   // for zero-dimensional memrefs.
2911   auto map =
2912       rank ? builder.getMultiDimIdentityMap(rank) : builder.getEmptyAffineMap();
2913   build(builder, result, resultType, memref, map, indices);
2914 }
2915 
2916 static ParseResult parseAffineVectorLoadOp(OpAsmParser &parser,
2917                                            OperationState &result) {
2918   auto &builder = parser.getBuilder();
2919   auto indexTy = builder.getIndexType();
2920 
2921   MemRefType memrefType;
2922   VectorType resultType;
2923   OpAsmParser::OperandType memrefInfo;
2924   AffineMapAttr mapAttr;
2925   SmallVector<OpAsmParser::OperandType, 1> mapOperands;
2926   return failure(
2927       parser.parseOperand(memrefInfo) ||
2928       parser.parseAffineMapOfSSAIds(mapOperands, mapAttr,
2929                                     AffineVectorLoadOp::getMapAttrName(),
2930                                     result.attributes) ||
2931       parser.parseOptionalAttrDict(result.attributes) ||
2932       parser.parseColonType(memrefType) || parser.parseComma() ||
2933       parser.parseType(resultType) ||
2934       parser.resolveOperand(memrefInfo, memrefType, result.operands) ||
2935       parser.resolveOperands(mapOperands, indexTy, result.operands) ||
2936       parser.addTypeToList(resultType, result.types));
2937 }
2938 
2939 static void print(OpAsmPrinter &p, AffineVectorLoadOp op) {
2940   p << "affine.vector_load " << op.getMemRef() << '[';
2941   if (AffineMapAttr mapAttr =
2942           op->getAttrOfType<AffineMapAttr>(op.getMapAttrName()))
2943     p.printAffineMapOfSSAIds(mapAttr, op.getMapOperands());
2944   p << ']';
2945   p.printOptionalAttrDict(op->getAttrs(),
2946                           /*elidedAttrs=*/{op.getMapAttrName()});
2947   p << " : " << op.getMemRefType() << ", " << op.getType();
2948 }
2949 
2950 /// Verify common invariants of affine.vector_load and affine.vector_store.
2951 static LogicalResult verifyVectorMemoryOp(Operation *op, MemRefType memrefType,
2952                                           VectorType vectorType) {
2953   // Check that memref and vector element types match.
2954   if (memrefType.getElementType() != vectorType.getElementType())
2955     return op->emitOpError(
2956         "requires memref and vector types of the same elemental type");
2957   return success();
2958 }
2959 
2960 static LogicalResult verify(AffineVectorLoadOp op) {
2961   MemRefType memrefType = op.getMemRefType();
2962   if (failed(verifyMemoryOpIndexing(
2963           op.getOperation(),
2964           op->getAttrOfType<AffineMapAttr>(op.getMapAttrName()),
2965           op.getMapOperands(), memrefType,
2966           /*numIndexOperands=*/op.getNumOperands() - 1)))
2967     return failure();
2968 
2969   if (failed(verifyVectorMemoryOp(op.getOperation(), memrefType,
2970                                   op.getVectorType())))
2971     return failure();
2972 
2973   return success();
2974 }
2975 
2976 //===----------------------------------------------------------------------===//
2977 // AffineVectorStoreOp
2978 //===----------------------------------------------------------------------===//
2979 
2980 void AffineVectorStoreOp::build(OpBuilder &builder, OperationState &result,
2981                                 Value valueToStore, Value memref, AffineMap map,
2982                                 ValueRange mapOperands) {
2983   assert(map.getNumInputs() == mapOperands.size() && "inconsistent index info");
2984   result.addOperands(valueToStore);
2985   result.addOperands(memref);
2986   result.addOperands(mapOperands);
2987   result.addAttribute(getMapAttrName(), AffineMapAttr::get(map));
2988 }
2989 
2990 // Use identity map.
2991 void AffineVectorStoreOp::build(OpBuilder &builder, OperationState &result,
2992                                 Value valueToStore, Value memref,
2993                                 ValueRange indices) {
2994   auto memrefType = memref.getType().cast<MemRefType>();
2995   int64_t rank = memrefType.getRank();
2996   // Create identity map for memrefs with at least one dimension or () -> ()
2997   // for zero-dimensional memrefs.
2998   auto map =
2999       rank ? builder.getMultiDimIdentityMap(rank) : builder.getEmptyAffineMap();
3000   build(builder, result, valueToStore, memref, map, indices);
3001 }
3002 
3003 static ParseResult parseAffineVectorStoreOp(OpAsmParser &parser,
3004                                             OperationState &result) {
3005   auto indexTy = parser.getBuilder().getIndexType();
3006 
3007   MemRefType memrefType;
3008   VectorType resultType;
3009   OpAsmParser::OperandType storeValueInfo;
3010   OpAsmParser::OperandType memrefInfo;
3011   AffineMapAttr mapAttr;
3012   SmallVector<OpAsmParser::OperandType, 1> mapOperands;
3013   return failure(
3014       parser.parseOperand(storeValueInfo) || parser.parseComma() ||
3015       parser.parseOperand(memrefInfo) ||
3016       parser.parseAffineMapOfSSAIds(mapOperands, mapAttr,
3017                                     AffineVectorStoreOp::getMapAttrName(),
3018                                     result.attributes) ||
3019       parser.parseOptionalAttrDict(result.attributes) ||
3020       parser.parseColonType(memrefType) || parser.parseComma() ||
3021       parser.parseType(resultType) ||
3022       parser.resolveOperand(storeValueInfo, resultType, result.operands) ||
3023       parser.resolveOperand(memrefInfo, memrefType, result.operands) ||
3024       parser.resolveOperands(mapOperands, indexTy, result.operands));
3025 }
3026 
3027 static void print(OpAsmPrinter &p, AffineVectorStoreOp op) {
3028   p << "affine.vector_store " << op.getValueToStore();
3029   p << ", " << op.getMemRef() << '[';
3030   if (AffineMapAttr mapAttr =
3031           op->getAttrOfType<AffineMapAttr>(op.getMapAttrName()))
3032     p.printAffineMapOfSSAIds(mapAttr, op.getMapOperands());
3033   p << ']';
3034   p.printOptionalAttrDict(op->getAttrs(),
3035                           /*elidedAttrs=*/{op.getMapAttrName()});
3036   p << " : " << op.getMemRefType() << ", " << op.getValueToStore().getType();
3037 }
3038 
3039 static LogicalResult verify(AffineVectorStoreOp op) {
3040   MemRefType memrefType = op.getMemRefType();
3041   if (failed(verifyMemoryOpIndexing(
3042           op.getOperation(),
3043           op->getAttrOfType<AffineMapAttr>(op.getMapAttrName()),
3044           op.getMapOperands(), memrefType,
3045           /*numIndexOperands=*/op.getNumOperands() - 2)))
3046     return failure();
3047 
3048   if (failed(verifyVectorMemoryOp(op.getOperation(), memrefType,
3049                                   op.getVectorType())))
3050     return failure();
3051 
3052   return success();
3053 }
3054 
3055 //===----------------------------------------------------------------------===//
3056 // TableGen'd op method definitions
3057 //===----------------------------------------------------------------------===//
3058 
3059 #define GET_OP_CLASSES
3060 #include "mlir/Dialect/Affine/IR/AffineOps.cpp.inc"
3061