MLIR  21.0.0git
AffineOps.cpp
Go to the documentation of this file.
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 
15 #include "mlir/IR/IRMapping.h"
16 #include "mlir/IR/IntegerSet.h"
17 #include "mlir/IR/Matchers.h"
18 #include "mlir/IR/OpDefinition.h"
19 #include "mlir/IR/PatternMatch.h"
23 #include "llvm/ADT/STLExtras.h"
24 #include "llvm/ADT/ScopeExit.h"
25 #include "llvm/ADT/SmallBitVector.h"
26 #include "llvm/ADT/SmallVectorExtras.h"
27 #include "llvm/ADT/TypeSwitch.h"
28 #include "llvm/Support/Debug.h"
29 #include "llvm/Support/MathExtras.h"
30 #include <numeric>
31 #include <optional>
32 
33 using namespace mlir;
34 using namespace mlir::affine;
35 
36 using llvm::divideCeilSigned;
37 using llvm::divideFloorSigned;
38 using llvm::mod;
39 
40 #define DEBUG_TYPE "affine-ops"
41 
42 #include "mlir/Dialect/Affine/IR/AffineOpsDialect.cpp.inc"
43 
44 /// A utility function to check if a value is defined at the top level of
45 /// `region` or is an argument of `region`. A value of index type defined at the
46 /// top level of a `AffineScope` region is always a valid symbol for all
47 /// uses in that region.
49  if (auto arg = llvm::dyn_cast<BlockArgument>(value))
50  return arg.getParentRegion() == region;
51  return value.getDefiningOp()->getParentRegion() == region;
52 }
53 
54 /// Checks if `value` known to be a legal affine dimension or symbol in `src`
55 /// region remains legal if the operation that uses it is inlined into `dest`
56 /// with the given value mapping. `legalityCheck` is either `isValidDim` or
57 /// `isValidSymbol`, depending on the value being required to remain a valid
58 /// dimension or symbol.
59 static bool
61  const IRMapping &mapping,
62  function_ref<bool(Value, Region *)> legalityCheck) {
63  // If the value is a valid dimension for any other reason than being
64  // a top-level value, it will remain valid: constants get inlined
65  // with the function, transitive affine applies also get inlined and
66  // will be checked themselves, etc.
67  if (!isTopLevelValue(value, src))
68  return true;
69 
70  // If it's a top-level value because it's a block operand, i.e. a
71  // function argument, check whether the value replacing it after
72  // inlining is a valid dimension in the new region.
73  if (llvm::isa<BlockArgument>(value))
74  return legalityCheck(mapping.lookup(value), dest);
75 
76  // If it's a top-level value because it's defined in the region,
77  // it can only be inlined if the defining op is a constant or a
78  // `dim`, which can appear anywhere and be valid, since the defining
79  // op won't be top-level anymore after inlining.
80  Attribute operandCst;
81  bool isDimLikeOp = isa<ShapedDimOpInterface>(value.getDefiningOp());
82  return matchPattern(value.getDefiningOp(), m_Constant(&operandCst)) ||
83  isDimLikeOp;
84 }
85 
86 /// Checks if all values known to be legal affine dimensions or symbols in `src`
87 /// remain so if their respective users are inlined into `dest`.
88 static bool
90  const IRMapping &mapping,
91  function_ref<bool(Value, Region *)> legalityCheck) {
92  return llvm::all_of(values, [&](Value v) {
93  return remainsLegalAfterInline(v, src, dest, mapping, legalityCheck);
94  });
95 }
96 
97 /// Checks if an affine read or write operation remains legal after inlining
98 /// from `src` to `dest`.
99 template <typename OpTy>
100 static bool remainsLegalAfterInline(OpTy op, Region *src, Region *dest,
101  const IRMapping &mapping) {
102  static_assert(llvm::is_one_of<OpTy, AffineReadOpInterface,
103  AffineWriteOpInterface>::value,
104  "only ops with affine read/write interface are supported");
105 
106  AffineMap map = op.getAffineMap();
107  ValueRange dimOperands = op.getMapOperands().take_front(map.getNumDims());
108  ValueRange symbolOperands =
109  op.getMapOperands().take_back(map.getNumSymbols());
111  dimOperands, src, dest, mapping,
112  static_cast<bool (*)(Value, Region *)>(isValidDim)))
113  return false;
115  symbolOperands, src, dest, mapping,
116  static_cast<bool (*)(Value, Region *)>(isValidSymbol)))
117  return false;
118  return true;
119 }
120 
121 /// Checks if an affine apply operation remains legal after inlining from `src`
122 /// to `dest`.
123 // Use "unused attribute" marker to silence clang-tidy warning stemming from
124 // the inability to see through "llvm::TypeSwitch".
125 template <>
126 bool LLVM_ATTRIBUTE_UNUSED remainsLegalAfterInline(AffineApplyOp op,
127  Region *src, Region *dest,
128  const IRMapping &mapping) {
129  // If it's a valid dimension, we need to check that it remains so.
130  if (isValidDim(op.getResult(), src))
132  op.getMapOperands(), src, dest, mapping,
133  static_cast<bool (*)(Value, Region *)>(isValidDim));
134 
135  // Otherwise it must be a valid symbol, check that it remains so.
137  op.getMapOperands(), src, dest, mapping,
138  static_cast<bool (*)(Value, Region *)>(isValidSymbol));
139 }
140 
141 //===----------------------------------------------------------------------===//
142 // AffineDialect Interfaces
143 //===----------------------------------------------------------------------===//
144 
145 namespace {
146 /// This class defines the interface for handling inlining with affine
147 /// operations.
148 struct AffineInlinerInterface : public DialectInlinerInterface {
150 
151  //===--------------------------------------------------------------------===//
152  // Analysis Hooks
153  //===--------------------------------------------------------------------===//
154 
155  /// Returns true if the given region 'src' can be inlined into the region
156  /// 'dest' that is attached to an operation registered to the current dialect.
157  /// 'wouldBeCloned' is set if the region is cloned into its new location
158  /// rather than moved, indicating there may be other users.
159  bool isLegalToInline(Region *dest, Region *src, bool wouldBeCloned,
160  IRMapping &valueMapping) const final {
161  // We can inline into affine loops and conditionals if this doesn't break
162  // affine value categorization rules.
163  Operation *destOp = dest->getParentOp();
164  if (!isa<AffineParallelOp, AffineForOp, AffineIfOp>(destOp))
165  return false;
166 
167  // Multi-block regions cannot be inlined into affine constructs, all of
168  // which require single-block regions.
169  if (!llvm::hasSingleElement(*src))
170  return false;
171 
172  // Side-effecting operations that the affine dialect cannot understand
173  // should not be inlined.
174  Block &srcBlock = src->front();
175  for (Operation &op : srcBlock) {
176  // Ops with no side effects are fine,
177  if (auto iface = dyn_cast<MemoryEffectOpInterface>(op)) {
178  if (iface.hasNoEffect())
179  continue;
180  }
181 
182  // Assuming the inlined region is valid, we only need to check if the
183  // inlining would change it.
184  bool remainsValid =
186  .Case<AffineApplyOp, AffineReadOpInterface,
187  AffineWriteOpInterface>([&](auto op) {
188  return remainsLegalAfterInline(op, src, dest, valueMapping);
189  })
190  .Default([](Operation *) {
191  // Conservatively disallow inlining ops we cannot reason about.
192  return false;
193  });
194 
195  if (!remainsValid)
196  return false;
197  }
198 
199  return true;
200  }
201 
202  /// Returns true if the given operation 'op', that is registered to this
203  /// dialect, can be inlined into the given region, false otherwise.
204  bool isLegalToInline(Operation *op, Region *region, bool wouldBeCloned,
205  IRMapping &valueMapping) const final {
206  // Always allow inlining affine operations into a region that is marked as
207  // affine scope, or into affine loops and conditionals. There are some edge
208  // cases when inlining *into* affine structures, but that is handled in the
209  // other 'isLegalToInline' hook above.
210  Operation *parentOp = region->getParentOp();
211  return parentOp->hasTrait<OpTrait::AffineScope>() ||
212  isa<AffineForOp, AffineParallelOp, AffineIfOp>(parentOp);
213  }
214 
215  /// Affine regions should be analyzed recursively.
216  bool shouldAnalyzeRecursively(Operation *op) const final { return true; }
217 };
218 } // namespace
219 
220 //===----------------------------------------------------------------------===//
221 // AffineDialect
222 //===----------------------------------------------------------------------===//
223 
224 void AffineDialect::initialize() {
225  addOperations<AffineDmaStartOp, AffineDmaWaitOp,
226 #define GET_OP_LIST
227 #include "mlir/Dialect/Affine/IR/AffineOps.cpp.inc"
228  >();
229  addInterfaces<AffineInlinerInterface>();
230  declarePromisedInterfaces<ValueBoundsOpInterface, AffineApplyOp, AffineMaxOp,
231  AffineMinOp>();
232 }
233 
234 /// Materialize a single constant operation from a given attribute value with
235 /// the desired resultant type.
237  Attribute value, Type type,
238  Location loc) {
239  if (auto poison = dyn_cast<ub::PoisonAttr>(value))
240  return builder.create<ub::PoisonOp>(loc, type, poison);
241  return arith::ConstantOp::materialize(builder, value, type, loc);
242 }
243 
244 /// A utility function to check if a value is defined at the top level of an
245 /// op with trait `AffineScope`. If the value is defined in an unlinked region,
246 /// conservatively assume it is not top-level. A value of index type defined at
247 /// the top level is always a valid symbol.
249  if (auto arg = llvm::dyn_cast<BlockArgument>(value)) {
250  // The block owning the argument may be unlinked, e.g. when the surrounding
251  // region has not yet been attached to an Op, at which point the parent Op
252  // is null.
253  Operation *parentOp = arg.getOwner()->getParentOp();
254  return parentOp && parentOp->hasTrait<OpTrait::AffineScope>();
255  }
256  // The defining Op may live in an unlinked block so its parent Op may be null.
257  Operation *parentOp = value.getDefiningOp()->getParentOp();
258  return parentOp && parentOp->hasTrait<OpTrait::AffineScope>();
259 }
260 
261 /// Returns the closest region enclosing `op` that is held by an operation with
262 /// trait `AffineScope`; `nullptr` if there is no such region.
264  auto *curOp = op;
265  while (auto *parentOp = curOp->getParentOp()) {
266  if (parentOp->hasTrait<OpTrait::AffineScope>())
267  return curOp->getParentRegion();
268  curOp = parentOp;
269  }
270  return nullptr;
271 }
272 
274  Operation *curOp = op;
275  while (auto *parentOp = curOp->getParentOp()) {
276  if (!isa<AffineForOp, AffineIfOp, AffineParallelOp>(parentOp))
277  return curOp->getParentRegion();
278  curOp = parentOp;
279  }
280  return nullptr;
281 }
282 
283 // A Value can be used as a dimension id iff it meets one of the following
284 // conditions:
285 // *) It is valid as a symbol.
286 // *) It is an induction variable.
287 // *) It is the result of affine apply operation with dimension id arguments.
289  // The value must be an index type.
290  if (!value.getType().isIndex())
291  return false;
292 
293  if (auto *defOp = value.getDefiningOp())
294  return isValidDim(value, getAffineScope(defOp));
295 
296  // This value has to be a block argument for an op that has the
297  // `AffineScope` trait or an induction var of an affine.for or
298  // affine.parallel.
299  if (isAffineInductionVar(value))
300  return true;
301  auto *parentOp = llvm::cast<BlockArgument>(value).getOwner()->getParentOp();
302  return parentOp && parentOp->hasTrait<OpTrait::AffineScope>();
303 }
304 
305 // Value can be used as a dimension id iff it meets one of the following
306 // conditions:
307 // *) It is valid as a symbol.
308 // *) It is an induction variable.
309 // *) It is the result of an affine apply operation with dimension id operands.
310 bool mlir::affine::isValidDim(Value value, Region *region) {
311  // The value must be an index type.
312  if (!value.getType().isIndex())
313  return false;
314 
315  // All valid symbols are okay.
316  if (isValidSymbol(value, region))
317  return true;
318 
319  auto *op = value.getDefiningOp();
320  if (!op) {
321  // This value has to be an induction var for an affine.for or an
322  // affine.parallel.
323  return isAffineInductionVar(value);
324  }
325 
326  // Affine apply operation is ok if all of its operands are ok.
327  if (auto applyOp = dyn_cast<AffineApplyOp>(op))
328  return applyOp.isValidDim(region);
329  // The dim op is okay if its operand memref/tensor is defined at the top
330  // level.
331  if (auto dimOp = dyn_cast<ShapedDimOpInterface>(op))
332  return isTopLevelValue(dimOp.getShapedValue());
333  return false;
334 }
335 
336 /// Returns true if the 'index' dimension of the `memref` defined by
337 /// `memrefDefOp` is a statically shaped one or defined using a valid symbol
338 /// for `region`.
339 template <typename AnyMemRefDefOp>
340 static bool isMemRefSizeValidSymbol(AnyMemRefDefOp memrefDefOp, unsigned index,
341  Region *region) {
342  MemRefType memRefType = memrefDefOp.getType();
343 
344  // Dimension index is out of bounds.
345  if (index >= memRefType.getRank()) {
346  return false;
347  }
348 
349  // Statically shaped.
350  if (!memRefType.isDynamicDim(index))
351  return true;
352  // Get the position of the dimension among dynamic dimensions;
353  unsigned dynamicDimPos = memRefType.getDynamicDimIndex(index);
354  return isValidSymbol(*(memrefDefOp.getDynamicSizes().begin() + dynamicDimPos),
355  region);
356 }
357 
358 /// Returns true if the result of the dim op is a valid symbol for `region`.
359 static bool isDimOpValidSymbol(ShapedDimOpInterface dimOp, Region *region) {
360  // The dim op is okay if its source is defined at the top level.
361  if (isTopLevelValue(dimOp.getShapedValue()))
362  return true;
363 
364  // Conservatively handle remaining BlockArguments as non-valid symbols.
365  // E.g. scf.for iterArgs.
366  if (llvm::isa<BlockArgument>(dimOp.getShapedValue()))
367  return false;
368 
369  // The dim op is also okay if its operand memref is a view/subview whose
370  // corresponding size is a valid symbol.
371  std::optional<int64_t> index = getConstantIntValue(dimOp.getDimension());
372 
373  // Be conservative if we can't understand the dimension.
374  if (!index.has_value())
375  return false;
376 
377  // Skip over all memref.cast ops (if any).
378  Operation *op = dimOp.getShapedValue().getDefiningOp();
379  while (auto castOp = dyn_cast<memref::CastOp>(op)) {
380  // Bail on unranked memrefs.
381  if (isa<UnrankedMemRefType>(castOp.getSource().getType()))
382  return false;
383  op = castOp.getSource().getDefiningOp();
384  if (!op)
385  return false;
386  }
387 
388  int64_t i = index.value();
390  .Case<memref::ViewOp, memref::SubViewOp, memref::AllocOp>(
391  [&](auto op) { return isMemRefSizeValidSymbol(op, i, region); })
392  .Default([](Operation *) { return false; });
393 }
394 
395 // A value can be used as a symbol (at all its use sites) iff it meets one of
396 // the following conditions:
397 // *) It is a constant.
398 // *) Its defining op or block arg appearance is immediately enclosed by an op
399 // with `AffineScope` trait.
400 // *) It is the result of an affine.apply operation with symbol operands.
401 // *) It is a result of the dim op on a memref whose corresponding size is a
402 // valid symbol.
404  if (!value)
405  return false;
406 
407  // The value must be an index type.
408  if (!value.getType().isIndex())
409  return false;
410 
411  // Check that the value is a top level value.
412  if (isTopLevelValue(value))
413  return true;
414 
415  if (auto *defOp = value.getDefiningOp())
416  return isValidSymbol(value, getAffineScope(defOp));
417 
418  return false;
419 }
420 
421 /// A value can be used as a symbol for `region` iff it meets one of the
422 /// following conditions:
423 /// *) It is a constant.
424 /// *) It is a result of a `Pure` operation whose operands are valid symbolic
425 /// *) identifiers.
426 /// *) It is a result of the dim op on a memref whose corresponding size is
427 /// a valid symbol.
428 /// *) It is defined at the top level of 'region' or is its argument.
429 /// *) It dominates `region`'s parent op.
430 /// If `region` is null, conservatively assume the symbol definition scope does
431 /// not exist and only accept the values that would be symbols regardless of
432 /// the surrounding region structure, i.e. the first three cases above.
434  // The value must be an index type.
435  if (!value.getType().isIndex())
436  return false;
437 
438  // A top-level value is a valid symbol.
439  if (region && ::isTopLevelValue(value, region))
440  return true;
441 
442  auto *defOp = value.getDefiningOp();
443  if (!defOp) {
444  // A block argument that is not a top-level value is a valid symbol if it
445  // dominates region's parent op.
446  Operation *regionOp = region ? region->getParentOp() : nullptr;
447  if (regionOp && !regionOp->hasTrait<OpTrait::IsIsolatedFromAbove>())
448  if (auto *parentOpRegion = region->getParentOp()->getParentRegion())
449  return isValidSymbol(value, parentOpRegion);
450  return false;
451  }
452 
453  // Constant operation is ok.
454  Attribute operandCst;
455  if (matchPattern(defOp, m_Constant(&operandCst)))
456  return true;
457 
458  // `Pure` operation that whose operands are valid symbolic identifiers.
459  if (isPure(defOp) && llvm::all_of(defOp->getOperands(), [&](Value operand) {
460  return affine::isValidSymbol(operand, region);
461  })) {
462  return true;
463  }
464 
465  // Dim op results could be valid symbols at any level.
466  if (auto dimOp = dyn_cast<ShapedDimOpInterface>(defOp))
467  return isDimOpValidSymbol(dimOp, region);
468 
469  // Check for values dominating `region`'s parent op.
470  Operation *regionOp = region ? region->getParentOp() : nullptr;
471  if (regionOp && !regionOp->hasTrait<OpTrait::IsIsolatedFromAbove>())
472  if (auto *parentRegion = region->getParentOp()->getParentRegion())
473  return isValidSymbol(value, parentRegion);
474 
475  return false;
476 }
477 
478 // Returns true if 'value' is a valid index to an affine operation (e.g.
479 // affine.load, affine.store, affine.dma_start, affine.dma_wait) where
480 // `region` provides the polyhedral symbol scope. Returns false otherwise.
481 static bool isValidAffineIndexOperand(Value value, Region *region) {
482  return isValidDim(value, region) || isValidSymbol(value, region);
483 }
484 
485 /// Prints dimension and symbol list.
488  unsigned numDims, OpAsmPrinter &printer) {
489  OperandRange operands(begin, end);
490  printer << '(' << operands.take_front(numDims) << ')';
491  if (operands.size() > numDims)
492  printer << '[' << operands.drop_front(numDims) << ']';
493 }
494 
495 /// Parses dimension and symbol list and returns true if parsing failed.
497  OpAsmParser &parser, SmallVectorImpl<Value> &operands, unsigned &numDims) {
499  if (parser.parseOperandList(opInfos, OpAsmParser::Delimiter::Paren))
500  return failure();
501  // Store number of dimensions for validation by caller.
502  numDims = opInfos.size();
503 
504  // Parse the optional symbol operands.
505  auto indexTy = parser.getBuilder().getIndexType();
506  return failure(parser.parseOperandList(
508  parser.resolveOperands(opInfos, indexTy, operands));
509 }
510 
511 /// Utility function to verify that a set of operands are valid dimension and
512 /// symbol identifiers. The operands should be laid out such that the dimension
513 /// operands are before the symbol operands. This function returns failure if
514 /// there was an invalid operand. An operation is provided to emit any necessary
515 /// errors.
516 template <typename OpTy>
517 static LogicalResult
519  unsigned numDims) {
520  unsigned opIt = 0;
521  for (auto operand : operands) {
522  if (opIt++ < numDims) {
523  if (!isValidDim(operand, getAffineScope(op)))
524  return op.emitOpError("operand cannot be used as a dimension id");
525  } else if (!isValidSymbol(operand, getAffineScope(op))) {
526  return op.emitOpError("operand cannot be used as a symbol");
527  }
528  }
529  return success();
530 }
531 
532 //===----------------------------------------------------------------------===//
533 // AffineApplyOp
534 //===----------------------------------------------------------------------===//
535 
536 AffineValueMap AffineApplyOp::getAffineValueMap() {
537  return AffineValueMap(getAffineMap(), getOperands(), getResult());
538 }
539 
540 ParseResult AffineApplyOp::parse(OpAsmParser &parser, OperationState &result) {
541  auto &builder = parser.getBuilder();
542  auto indexTy = builder.getIndexType();
543 
544  AffineMapAttr mapAttr;
545  unsigned numDims;
546  if (parser.parseAttribute(mapAttr, "map", result.attributes) ||
547  parseDimAndSymbolList(parser, result.operands, numDims) ||
548  parser.parseOptionalAttrDict(result.attributes))
549  return failure();
550  auto map = mapAttr.getValue();
551 
552  if (map.getNumDims() != numDims ||
553  numDims + map.getNumSymbols() != result.operands.size()) {
554  return parser.emitError(parser.getNameLoc(),
555  "dimension or symbol index mismatch");
556  }
557 
558  result.types.append(map.getNumResults(), indexTy);
559  return success();
560 }
561 
563  p << " " << getMapAttr();
564  printDimAndSymbolList(operand_begin(), operand_end(),
565  getAffineMap().getNumDims(), p);
566  p.printOptionalAttrDict((*this)->getAttrs(), /*elidedAttrs=*/{"map"});
567 }
568 
569 LogicalResult AffineApplyOp::verify() {
570  // Check input and output dimensions match.
571  AffineMap affineMap = getMap();
572 
573  // Verify that operand count matches affine map dimension and symbol count.
574  if (getNumOperands() != affineMap.getNumDims() + affineMap.getNumSymbols())
575  return emitOpError(
576  "operand count and affine map dimension and symbol count must match");
577 
578  // Verify that the map only produces one result.
579  if (affineMap.getNumResults() != 1)
580  return emitOpError("mapping must produce one value");
581 
582  // Do not allow valid dims to be used in symbol positions. We do allow
583  // affine.apply to use operands for values that may neither qualify as affine
584  // dims or affine symbols due to usage outside of affine ops, analyses, etc.
585  Region *region = getAffineScope(*this);
586  for (Value operand : getMapOperands().drop_front(affineMap.getNumDims())) {
587  if (::isValidDim(operand, region) && !::isValidSymbol(operand, region))
588  return emitError("dimensional operand cannot be used as a symbol");
589  }
590 
591  return success();
592 }
593 
594 // The result of the affine apply operation can be used as a dimension id if all
595 // its operands are valid dimension ids.
597  return llvm::all_of(getOperands(),
598  [](Value op) { return affine::isValidDim(op); });
599 }
600 
601 // The result of the affine apply operation can be used as a dimension id if all
602 // its operands are valid dimension ids with the parent operation of `region`
603 // defining the polyhedral scope for symbols.
604 bool AffineApplyOp::isValidDim(Region *region) {
605  return llvm::all_of(getOperands(),
606  [&](Value op) { return ::isValidDim(op, region); });
607 }
608 
609 // The result of the affine apply operation can be used as a symbol if all its
610 // operands are symbols.
612  return llvm::all_of(getOperands(),
613  [](Value op) { return affine::isValidSymbol(op); });
614 }
615 
616 // The result of the affine apply operation can be used as a symbol in `region`
617 // if all its operands are symbols in `region`.
618 bool AffineApplyOp::isValidSymbol(Region *region) {
619  return llvm::all_of(getOperands(), [&](Value operand) {
620  return affine::isValidSymbol(operand, region);
621  });
622 }
623 
624 OpFoldResult AffineApplyOp::fold(FoldAdaptor adaptor) {
625  auto map = getAffineMap();
626 
627  // Fold dims and symbols to existing values.
628  auto expr = map.getResult(0);
629  if (auto dim = dyn_cast<AffineDimExpr>(expr))
630  return getOperand(dim.getPosition());
631  if (auto sym = dyn_cast<AffineSymbolExpr>(expr))
632  return getOperand(map.getNumDims() + sym.getPosition());
633 
634  // Otherwise, default to folding the map.
636  bool hasPoison = false;
637  auto foldResult =
638  map.constantFold(adaptor.getMapOperands(), result, &hasPoison);
639  if (hasPoison)
641  if (failed(foldResult))
642  return {};
643  return result[0];
644 }
645 
646 /// Returns the largest known divisor of `e`. Exploits information from the
647 /// values in `operands`.
648 static int64_t getLargestKnownDivisor(AffineExpr e, ArrayRef<Value> operands) {
649  // This method isn't aware of `operands`.
650  int64_t div = e.getLargestKnownDivisor();
651 
652  // We now make use of operands for the case `e` is a dim expression.
653  // TODO: More powerful simplification would have to modify
654  // getLargestKnownDivisor to take `operands` and exploit that information as
655  // well for dim/sym expressions, but in that case, getLargestKnownDivisor
656  // can't be part of the IR library but of the `Analysis` library. The IR
657  // library can only really depend on simple O(1) checks.
658  auto dimExpr = dyn_cast<AffineDimExpr>(e);
659  // If it's not a dim expr, `div` is the best we have.
660  if (!dimExpr)
661  return div;
662 
663  // We simply exploit information from loop IVs.
664  // We don't need to use mlir::getLargestKnownDivisorOfValue since the other
665  // desired simplifications are expected to be part of other
666  // canonicalizations. Also, mlir::getLargestKnownDivisorOfValue is part of the
667  // LoopAnalysis library.
668  Value operand = operands[dimExpr.getPosition()];
669  int64_t operandDivisor = 1;
670  // TODO: With the right accessors, this can be extended to
671  // LoopLikeOpInterface.
672  if (AffineForOp forOp = getForInductionVarOwner(operand)) {
673  if (forOp.hasConstantLowerBound() && forOp.getConstantLowerBound() == 0) {
674  operandDivisor = forOp.getStepAsInt();
675  } else {
676  uint64_t lbLargestKnownDivisor =
677  forOp.getLowerBoundMap().getLargestKnownDivisorOfMapExprs();
678  operandDivisor = std::gcd(lbLargestKnownDivisor, forOp.getStepAsInt());
679  }
680  }
681  return operandDivisor;
682 }
683 
684 /// Check if `e` is known to be: 0 <= `e` < `k`. Handles the simple cases of `e`
685 /// being an affine dim expression or a constant.
687  int64_t k) {
688  if (auto constExpr = dyn_cast<AffineConstantExpr>(e)) {
689  int64_t constVal = constExpr.getValue();
690  return constVal >= 0 && constVal < k;
691  }
692  auto dimExpr = dyn_cast<AffineDimExpr>(e);
693  if (!dimExpr)
694  return false;
695  Value operand = operands[dimExpr.getPosition()];
696  // TODO: With the right accessors, this can be extended to
697  // LoopLikeOpInterface.
698  if (AffineForOp forOp = getForInductionVarOwner(operand)) {
699  if (forOp.hasConstantLowerBound() && forOp.getConstantLowerBound() >= 0 &&
700  forOp.hasConstantUpperBound() && forOp.getConstantUpperBound() <= k) {
701  return true;
702  }
703  }
704 
705  // We don't consider other cases like `operand` being defined by a constant or
706  // an affine.apply op since such cases will already be handled by other
707  // patterns and propagation of loop IVs or constant would happen.
708  return false;
709 }
710 
711 /// Check if expression `e` is of the form d*e_1 + e_2 where 0 <= e_2 < d.
712 /// Set `div` to `d`, `quotientTimesDiv` to e_1 and `rem` to e_2 if the
713 /// expression is in that form.
714 static bool isQTimesDPlusR(AffineExpr e, ArrayRef<Value> operands, int64_t &div,
715  AffineExpr &quotientTimesDiv, AffineExpr &rem) {
716  auto bin = dyn_cast<AffineBinaryOpExpr>(e);
717  if (!bin || bin.getKind() != AffineExprKind::Add)
718  return false;
719 
720  AffineExpr llhs = bin.getLHS();
721  AffineExpr rlhs = bin.getRHS();
722  div = getLargestKnownDivisor(llhs, operands);
723  if (isNonNegativeBoundedBy(rlhs, operands, div)) {
724  quotientTimesDiv = llhs;
725  rem = rlhs;
726  return true;
727  }
728  div = getLargestKnownDivisor(rlhs, operands);
729  if (isNonNegativeBoundedBy(llhs, operands, div)) {
730  quotientTimesDiv = rlhs;
731  rem = llhs;
732  return true;
733  }
734  return false;
735 }
736 
737 /// Gets the constant lower bound on an `iv`.
738 static std::optional<int64_t> getLowerBound(Value iv) {
739  AffineForOp forOp = getForInductionVarOwner(iv);
740  if (forOp && forOp.hasConstantLowerBound())
741  return forOp.getConstantLowerBound();
742  return std::nullopt;
743 }
744 
745 /// Gets the constant upper bound on an affine.for `iv`.
746 static std::optional<int64_t> getUpperBound(Value iv) {
747  AffineForOp forOp = getForInductionVarOwner(iv);
748  if (!forOp || !forOp.hasConstantUpperBound())
749  return std::nullopt;
750 
751  // If its lower bound is also known, we can get a more precise bound
752  // whenever the step is not one.
753  if (forOp.hasConstantLowerBound()) {
754  return forOp.getConstantUpperBound() - 1 -
755  (forOp.getConstantUpperBound() - forOp.getConstantLowerBound() - 1) %
756  forOp.getStepAsInt();
757  }
758  return forOp.getConstantUpperBound() - 1;
759 }
760 
761 /// Determine a constant upper bound for `expr` if one exists while exploiting
762 /// values in `operands`. Note that the upper bound is an inclusive one. `expr`
763 /// is guaranteed to be less than or equal to it.
764 static std::optional<int64_t> getUpperBound(AffineExpr expr, unsigned numDims,
765  unsigned numSymbols,
766  ArrayRef<Value> operands) {
767  // Get the constant lower or upper bounds on the operands.
768  SmallVector<std::optional<int64_t>> constLowerBounds, constUpperBounds;
769  constLowerBounds.reserve(operands.size());
770  constUpperBounds.reserve(operands.size());
771  for (Value operand : operands) {
772  constLowerBounds.push_back(getLowerBound(operand));
773  constUpperBounds.push_back(getUpperBound(operand));
774  }
775 
776  if (auto constExpr = dyn_cast<AffineConstantExpr>(expr))
777  return constExpr.getValue();
778 
779  return getBoundForAffineExpr(expr, numDims, numSymbols, constLowerBounds,
780  constUpperBounds,
781  /*isUpper=*/true);
782 }
783 
784 /// Determine a constant lower bound for `expr` if one exists while exploiting
785 /// values in `operands`. Note that the upper bound is an inclusive one. `expr`
786 /// is guaranteed to be less than or equal to it.
787 static std::optional<int64_t> getLowerBound(AffineExpr expr, unsigned numDims,
788  unsigned numSymbols,
789  ArrayRef<Value> operands) {
790  // Get the constant lower or upper bounds on the operands.
791  SmallVector<std::optional<int64_t>> constLowerBounds, constUpperBounds;
792  constLowerBounds.reserve(operands.size());
793  constUpperBounds.reserve(operands.size());
794  for (Value operand : operands) {
795  constLowerBounds.push_back(getLowerBound(operand));
796  constUpperBounds.push_back(getUpperBound(operand));
797  }
798 
799  std::optional<int64_t> lowerBound;
800  if (auto constExpr = dyn_cast<AffineConstantExpr>(expr)) {
801  lowerBound = constExpr.getValue();
802  } else {
803  lowerBound = getBoundForAffineExpr(expr, numDims, numSymbols,
804  constLowerBounds, constUpperBounds,
805  /*isUpper=*/false);
806  }
807  return lowerBound;
808 }
809 
810 /// Simplify `expr` while exploiting information from the values in `operands`.
811 static void simplifyExprAndOperands(AffineExpr &expr, unsigned numDims,
812  unsigned numSymbols,
813  ArrayRef<Value> operands) {
814  // We do this only for certain floordiv/mod expressions.
815  auto binExpr = dyn_cast<AffineBinaryOpExpr>(expr);
816  if (!binExpr)
817  return;
818 
819  // Simplify the child expressions first.
820  AffineExpr lhs = binExpr.getLHS();
821  AffineExpr rhs = binExpr.getRHS();
822  simplifyExprAndOperands(lhs, numDims, numSymbols, operands);
823  simplifyExprAndOperands(rhs, numDims, numSymbols, operands);
824  expr = getAffineBinaryOpExpr(binExpr.getKind(), lhs, rhs);
825 
826  binExpr = dyn_cast<AffineBinaryOpExpr>(expr);
827  if (!binExpr || (expr.getKind() != AffineExprKind::FloorDiv &&
828  expr.getKind() != AffineExprKind::CeilDiv &&
829  expr.getKind() != AffineExprKind::Mod)) {
830  return;
831  }
832 
833  // The `lhs` and `rhs` may be different post construction of simplified expr.
834  lhs = binExpr.getLHS();
835  rhs = binExpr.getRHS();
836  auto rhsConst = dyn_cast<AffineConstantExpr>(rhs);
837  if (!rhsConst)
838  return;
839 
840  int64_t rhsConstVal = rhsConst.getValue();
841  // Undefined exprsessions aren't touched; IR can still be valid with them.
842  if (rhsConstVal <= 0)
843  return;
844 
845  // Exploit constant lower/upper bounds to simplify a floordiv or mod.
846  MLIRContext *context = expr.getContext();
847  std::optional<int64_t> lhsLbConst =
848  getLowerBound(lhs, numDims, numSymbols, operands);
849  std::optional<int64_t> lhsUbConst =
850  getUpperBound(lhs, numDims, numSymbols, operands);
851  if (lhsLbConst && lhsUbConst) {
852  int64_t lhsLbConstVal = *lhsLbConst;
853  int64_t lhsUbConstVal = *lhsUbConst;
854  // lhs floordiv c is a single value lhs is bounded in a range `c` that has
855  // the same quotient.
856  if (binExpr.getKind() == AffineExprKind::FloorDiv &&
857  divideFloorSigned(lhsLbConstVal, rhsConstVal) ==
858  divideFloorSigned(lhsUbConstVal, rhsConstVal)) {
859  expr = getAffineConstantExpr(
860  divideFloorSigned(lhsLbConstVal, rhsConstVal), context);
861  return;
862  }
863  // lhs ceildiv c is a single value if the entire range has the same ceil
864  // quotient.
865  if (binExpr.getKind() == AffineExprKind::CeilDiv &&
866  divideCeilSigned(lhsLbConstVal, rhsConstVal) ==
867  divideCeilSigned(lhsUbConstVal, rhsConstVal)) {
868  expr = getAffineConstantExpr(divideCeilSigned(lhsLbConstVal, rhsConstVal),
869  context);
870  return;
871  }
872  // lhs mod c is lhs if the entire range has quotient 0 w.r.t the rhs.
873  if (binExpr.getKind() == AffineExprKind::Mod && lhsLbConstVal >= 0 &&
874  lhsLbConstVal < rhsConstVal && lhsUbConstVal < rhsConstVal) {
875  expr = lhs;
876  return;
877  }
878  }
879 
880  // Simplify expressions of the form e = (e_1 + e_2) floordiv c or (e_1 + e_2)
881  // mod c, where e_1 is a multiple of `k` and 0 <= e_2 < k. In such cases, if
882  // `c` % `k` == 0, (e_1 + e_2) floordiv c can be simplified to e_1 floordiv c.
883  // And when k % c == 0, (e_1 + e_2) mod c can be simplified to e_2 mod c.
884  AffineExpr quotientTimesDiv, rem;
885  int64_t divisor;
886  if (isQTimesDPlusR(lhs, operands, divisor, quotientTimesDiv, rem)) {
887  if (rhsConstVal % divisor == 0 &&
888  binExpr.getKind() == AffineExprKind::FloorDiv) {
889  expr = quotientTimesDiv.floorDiv(rhsConst);
890  } else if (divisor % rhsConstVal == 0 &&
891  binExpr.getKind() == AffineExprKind::Mod) {
892  expr = rem % rhsConst;
893  }
894  return;
895  }
896 
897  // Handle the simple case when the LHS expression can be either upper
898  // bounded or is a known multiple of RHS constant.
899  // lhs floordiv c -> 0 if 0 <= lhs < c,
900  // lhs mod c -> 0 if lhs % c = 0.
901  if ((isNonNegativeBoundedBy(lhs, operands, rhsConstVal) &&
902  binExpr.getKind() == AffineExprKind::FloorDiv) ||
903  (getLargestKnownDivisor(lhs, operands) % rhsConstVal == 0 &&
904  binExpr.getKind() == AffineExprKind::Mod)) {
905  expr = getAffineConstantExpr(0, expr.getContext());
906  }
907 }
908 
909 /// Simplify the expressions in `map` while making use of lower or upper bounds
910 /// of its operands. If `isMax` is true, the map is to be treated as a max of
911 /// its result expressions, and min otherwise. Eg: min (d0, d1) -> (8, 4 * d0 +
912 /// d1) can be simplified to (8) if the operands are respectively lower bounded
913 /// by 2 and 0 (the second expression can't be lower than 8).
915  ArrayRef<Value> operands,
916  bool isMax) {
917  // Can't simplify.
918  if (operands.empty())
919  return;
920 
921  // Get the upper or lower bound on an affine.for op IV using its range.
922  // Get the constant lower or upper bounds on the operands.
923  SmallVector<std::optional<int64_t>> constLowerBounds, constUpperBounds;
924  constLowerBounds.reserve(operands.size());
925  constUpperBounds.reserve(operands.size());
926  for (Value operand : operands) {
927  constLowerBounds.push_back(getLowerBound(operand));
928  constUpperBounds.push_back(getUpperBound(operand));
929  }
930 
931  // We will compute the lower and upper bounds on each of the expressions
932  // Then, we will check (depending on max or min) as to whether a specific
933  // bound is redundant by checking if its highest (in case of max) and its
934  // lowest (in the case of min) value is already lower than (or higher than)
935  // the lower bound (or upper bound in the case of min) of another bound.
936  SmallVector<std::optional<int64_t>, 4> lowerBounds, upperBounds;
937  lowerBounds.reserve(map.getNumResults());
938  upperBounds.reserve(map.getNumResults());
939  for (AffineExpr e : map.getResults()) {
940  if (auto constExpr = dyn_cast<AffineConstantExpr>(e)) {
941  lowerBounds.push_back(constExpr.getValue());
942  upperBounds.push_back(constExpr.getValue());
943  } else {
944  lowerBounds.push_back(
946  constLowerBounds, constUpperBounds,
947  /*isUpper=*/false));
948  upperBounds.push_back(
950  constLowerBounds, constUpperBounds,
951  /*isUpper=*/true));
952  }
953  }
954 
955  // Collect expressions that are not redundant.
956  SmallVector<AffineExpr, 4> irredundantExprs;
957  for (auto exprEn : llvm::enumerate(map.getResults())) {
958  AffineExpr e = exprEn.value();
959  unsigned i = exprEn.index();
960  // Some expressions can be turned into constants.
961  if (lowerBounds[i] && upperBounds[i] && *lowerBounds[i] == *upperBounds[i])
962  e = getAffineConstantExpr(*lowerBounds[i], e.getContext());
963 
964  // Check if the expression is redundant.
965  if (isMax) {
966  if (!upperBounds[i]) {
967  irredundantExprs.push_back(e);
968  continue;
969  }
970  // If there exists another expression such that its lower bound is greater
971  // than this expression's upper bound, it's redundant.
972  if (!llvm::any_of(llvm::enumerate(lowerBounds), [&](const auto &en) {
973  auto otherLowerBound = en.value();
974  unsigned pos = en.index();
975  if (pos == i || !otherLowerBound)
976  return false;
977  if (*otherLowerBound > *upperBounds[i])
978  return true;
979  if (*otherLowerBound < *upperBounds[i])
980  return false;
981  // Equality case. When both expressions are considered redundant, we
982  // don't want to get both of them. We keep the one that appears
983  // first.
984  if (upperBounds[pos] && lowerBounds[i] &&
985  lowerBounds[i] == upperBounds[i] &&
986  otherLowerBound == *upperBounds[pos] && i < pos)
987  return false;
988  return true;
989  }))
990  irredundantExprs.push_back(e);
991  } else {
992  if (!lowerBounds[i]) {
993  irredundantExprs.push_back(e);
994  continue;
995  }
996  // Likewise for the `min` case. Use the complement of the condition above.
997  if (!llvm::any_of(llvm::enumerate(upperBounds), [&](const auto &en) {
998  auto otherUpperBound = en.value();
999  unsigned pos = en.index();
1000  if (pos == i || !otherUpperBound)
1001  return false;
1002  if (*otherUpperBound < *lowerBounds[i])
1003  return true;
1004  if (*otherUpperBound > *lowerBounds[i])
1005  return false;
1006  if (lowerBounds[pos] && upperBounds[i] &&
1007  lowerBounds[i] == upperBounds[i] &&
1008  otherUpperBound == lowerBounds[pos] && i < pos)
1009  return false;
1010  return true;
1011  }))
1012  irredundantExprs.push_back(e);
1013  }
1014  }
1015 
1016  // Create the map without the redundant expressions.
1017  map = AffineMap::get(map.getNumDims(), map.getNumSymbols(), irredundantExprs,
1018  map.getContext());
1019 }
1020 
1021 /// Simplify the map while exploiting information on the values in `operands`.
1022 // Use "unused attribute" marker to silence warning stemming from the inability
1023 // to see through the template expansion.
1024 static void LLVM_ATTRIBUTE_UNUSED
1026  assert(map.getNumInputs() == operands.size() && "invalid operands for map");
1027  SmallVector<AffineExpr> newResults;
1028  newResults.reserve(map.getNumResults());
1029  for (AffineExpr expr : map.getResults()) {
1031  operands);
1032  newResults.push_back(expr);
1033  }
1034  map = AffineMap::get(map.getNumDims(), map.getNumSymbols(), newResults,
1035  map.getContext());
1036 }
1037 
1038 /// Replace all occurrences of AffineExpr at position `pos` in `map` by the
1039 /// defining AffineApplyOp expression and operands.
1040 /// When `dimOrSymbolPosition < dims.size()`, AffineDimExpr@[pos] is replaced.
1041 /// When `dimOrSymbolPosition >= dims.size()`,
1042 /// AffineSymbolExpr@[pos - dims.size()] is replaced.
1043 /// Mutate `map`,`dims` and `syms` in place as follows:
1044 /// 1. `dims` and `syms` are only appended to.
1045 /// 2. `map` dim and symbols are gradually shifted to higher positions.
1046 /// 3. Old `dim` and `sym` entries are replaced by nullptr
1047 /// This avoids the need for any bookkeeping.
1048 static LogicalResult replaceDimOrSym(AffineMap *map,
1049  unsigned dimOrSymbolPosition,
1050  SmallVectorImpl<Value> &dims,
1051  SmallVectorImpl<Value> &syms) {
1052  MLIRContext *ctx = map->getContext();
1053  bool isDimReplacement = (dimOrSymbolPosition < dims.size());
1054  unsigned pos = isDimReplacement ? dimOrSymbolPosition
1055  : dimOrSymbolPosition - dims.size();
1056  Value &v = isDimReplacement ? dims[pos] : syms[pos];
1057  if (!v)
1058  return failure();
1059 
1060  auto affineApply = v.getDefiningOp<AffineApplyOp>();
1061  if (!affineApply)
1062  return failure();
1063 
1064  // At this point we will perform a replacement of `v`, set the entry in `dim`
1065  // or `sym` to nullptr immediately.
1066  v = nullptr;
1067 
1068  // Compute the map, dims and symbols coming from the AffineApplyOp.
1069  AffineMap composeMap = affineApply.getAffineMap();
1070  assert(composeMap.getNumResults() == 1 && "affine.apply with >1 results");
1071  SmallVector<Value> composeOperands(affineApply.getMapOperands().begin(),
1072  affineApply.getMapOperands().end());
1073  // Canonicalize the map to promote dims to symbols when possible. This is to
1074  // avoid generating invalid maps.
1075  canonicalizeMapAndOperands(&composeMap, &composeOperands);
1076  AffineExpr replacementExpr =
1077  composeMap.shiftDims(dims.size()).shiftSymbols(syms.size()).getResult(0);
1078  ValueRange composeDims =
1079  ArrayRef<Value>(composeOperands).take_front(composeMap.getNumDims());
1080  ValueRange composeSyms =
1081  ArrayRef<Value>(composeOperands).take_back(composeMap.getNumSymbols());
1082  AffineExpr toReplace = isDimReplacement ? getAffineDimExpr(pos, ctx)
1083  : getAffineSymbolExpr(pos, ctx);
1084 
1085  // Append the dims and symbols where relevant and perform the replacement.
1086  dims.append(composeDims.begin(), composeDims.end());
1087  syms.append(composeSyms.begin(), composeSyms.end());
1088  *map = map->replace(toReplace, replacementExpr, dims.size(), syms.size());
1089 
1090  return success();
1091 }
1092 
1093 /// Iterate over `operands` and fold away all those produced by an AffineApplyOp
1094 /// iteratively. Perform canonicalization of map and operands as well as
1095 /// AffineMap simplification. `map` and `operands` are mutated in place.
1097  SmallVectorImpl<Value> *operands) {
1098  if (map->getNumResults() == 0) {
1099  canonicalizeMapAndOperands(map, operands);
1100  *map = simplifyAffineMap(*map);
1101  return;
1102  }
1103 
1104  MLIRContext *ctx = map->getContext();
1105  SmallVector<Value, 4> dims(operands->begin(),
1106  operands->begin() + map->getNumDims());
1107  SmallVector<Value, 4> syms(operands->begin() + map->getNumDims(),
1108  operands->end());
1109 
1110  // Iterate over dims and symbols coming from AffineApplyOp and replace until
1111  // exhaustion. This iteratively mutates `map`, `dims` and `syms`. Both `dims`
1112  // and `syms` can only increase by construction.
1113  // The implementation uses a `while` loop to support the case of symbols
1114  // that may be constructed from dims ;this may be overkill.
1115  while (true) {
1116  bool changed = false;
1117  for (unsigned pos = 0; pos != dims.size() + syms.size(); ++pos)
1118  if ((changed |= succeeded(replaceDimOrSym(map, pos, dims, syms))))
1119  break;
1120  if (!changed)
1121  break;
1122  }
1123 
1124  // Clear operands so we can fill them anew.
1125  operands->clear();
1126 
1127  // At this point we may have introduced null operands, prune them out before
1128  // canonicalizing map and operands.
1129  unsigned nDims = 0, nSyms = 0;
1130  SmallVector<AffineExpr, 4> dimReplacements, symReplacements;
1131  dimReplacements.reserve(dims.size());
1132  symReplacements.reserve(syms.size());
1133  for (auto *container : {&dims, &syms}) {
1134  bool isDim = (container == &dims);
1135  auto &repls = isDim ? dimReplacements : symReplacements;
1136  for (const auto &en : llvm::enumerate(*container)) {
1137  Value v = en.value();
1138  if (!v) {
1139  assert(isDim ? !map->isFunctionOfDim(en.index())
1140  : !map->isFunctionOfSymbol(en.index()) &&
1141  "map is function of unexpected expr@pos");
1142  repls.push_back(getAffineConstantExpr(0, ctx));
1143  continue;
1144  }
1145  repls.push_back(isDim ? getAffineDimExpr(nDims++, ctx)
1146  : getAffineSymbolExpr(nSyms++, ctx));
1147  operands->push_back(v);
1148  }
1149  }
1150  *map = map->replaceDimsAndSymbols(dimReplacements, symReplacements, nDims,
1151  nSyms);
1152 
1153  // Canonicalize and simplify before returning.
1154  canonicalizeMapAndOperands(map, operands);
1155  *map = simplifyAffineMap(*map);
1156 }
1157 
1159  AffineMap *map, SmallVectorImpl<Value> *operands) {
1160  while (llvm::any_of(*operands, [](Value v) {
1161  return isa_and_nonnull<AffineApplyOp>(v.getDefiningOp());
1162  })) {
1163  composeAffineMapAndOperands(map, operands);
1164  }
1165 }
1166 
1167 AffineApplyOp
1169  ArrayRef<OpFoldResult> operands) {
1170  SmallVector<Value> valueOperands;
1171  map = foldAttributesIntoMap(b, map, operands, valueOperands);
1172  composeAffineMapAndOperands(&map, &valueOperands);
1173  assert(map);
1174  return b.create<AffineApplyOp>(loc, map, valueOperands);
1175 }
1176 
1177 AffineApplyOp
1179  ArrayRef<OpFoldResult> operands) {
1180  return makeComposedAffineApply(
1181  b, loc,
1183  .front(),
1184  operands);
1185 }
1186 
1187 /// Composes the given affine map with the given list of operands, pulling in
1188 /// the maps from any affine.apply operations that supply the operands.
1190  SmallVectorImpl<Value> &operands) {
1191  // Compose and canonicalize each expression in the map individually because
1192  // composition only applies to single-result maps, collecting potentially
1193  // duplicate operands in a single list with shifted dimensions and symbols.
1194  SmallVector<Value> dims, symbols;
1196  for (unsigned i : llvm::seq<unsigned>(0, map.getNumResults())) {
1197  SmallVector<Value> submapOperands(operands.begin(), operands.end());
1198  AffineMap submap = map.getSubMap({i});
1199  fullyComposeAffineMapAndOperands(&submap, &submapOperands);
1200  canonicalizeMapAndOperands(&submap, &submapOperands);
1201  unsigned numNewDims = submap.getNumDims();
1202  submap = submap.shiftDims(dims.size()).shiftSymbols(symbols.size());
1203  llvm::append_range(dims,
1204  ArrayRef<Value>(submapOperands).take_front(numNewDims));
1205  llvm::append_range(symbols,
1206  ArrayRef<Value>(submapOperands).drop_front(numNewDims));
1207  exprs.push_back(submap.getResult(0));
1208  }
1209 
1210  // Canonicalize the map created from composed expressions to deduplicate the
1211  // dimension and symbol operands.
1212  operands = llvm::to_vector(llvm::concat<Value>(dims, symbols));
1213  map = AffineMap::get(dims.size(), symbols.size(), exprs, map.getContext());
1214  canonicalizeMapAndOperands(&map, &operands);
1215 }
1216 
1219  AffineMap map,
1220  ArrayRef<OpFoldResult> operands) {
1221  assert(map.getNumResults() == 1 && "building affine.apply with !=1 result");
1222 
1223  // Create new builder without a listener, so that no notification is
1224  // triggered if the op is folded.
1225  // TODO: OpBuilder::createOrFold should return OpFoldResults, then this
1226  // workaround is no longer needed.
1227  OpBuilder newBuilder(b.getContext());
1229 
1230  // Create op.
1231  AffineApplyOp applyOp =
1232  makeComposedAffineApply(newBuilder, loc, map, operands);
1233 
1234  // Get constant operands.
1235  SmallVector<Attribute> constOperands(applyOp->getNumOperands());
1236  for (unsigned i = 0, e = constOperands.size(); i != e; ++i)
1237  matchPattern(applyOp->getOperand(i), m_Constant(&constOperands[i]));
1238 
1239  // Try to fold the operation.
1240  SmallVector<OpFoldResult> foldResults;
1241  if (failed(applyOp->fold(constOperands, foldResults)) ||
1242  foldResults.empty()) {
1243  if (OpBuilder::Listener *listener = b.getListener())
1244  listener->notifyOperationInserted(applyOp, /*previous=*/{});
1245  return applyOp.getResult();
1246  }
1247 
1248  applyOp->erase();
1249  return llvm::getSingleElement(foldResults);
1250 }
1251 
1254  AffineExpr expr,
1255  ArrayRef<OpFoldResult> operands) {
1257  b, loc,
1259  .front(),
1260  operands);
1261 }
1262 
1265  OpBuilder &b, Location loc, AffineMap map,
1266  ArrayRef<OpFoldResult> operands) {
1267  return llvm::map_to_vector(llvm::seq<unsigned>(0, map.getNumResults()),
1268  [&](unsigned i) {
1269  return makeComposedFoldedAffineApply(
1270  b, loc, map.getSubMap({i}), operands);
1271  });
1272 }
1273 
1274 template <typename OpTy>
1276  ArrayRef<OpFoldResult> operands) {
1277  SmallVector<Value> valueOperands;
1278  map = foldAttributesIntoMap(b, map, operands, valueOperands);
1279  composeMultiResultAffineMap(map, valueOperands);
1280  return b.create<OpTy>(loc, b.getIndexType(), map, valueOperands);
1281 }
1282 
1283 AffineMinOp
1285  ArrayRef<OpFoldResult> operands) {
1286  return makeComposedMinMax<AffineMinOp>(b, loc, map, operands);
1287 }
1288 
1289 template <typename OpTy>
1291  AffineMap map,
1292  ArrayRef<OpFoldResult> operands) {
1293  // Create new builder without a listener, so that no notification is
1294  // triggered if the op is folded.
1295  // TODO: OpBuilder::createOrFold should return OpFoldResults, then this
1296  // workaround is no longer needed.
1297  OpBuilder newBuilder(b.getContext());
1299 
1300  // Create op.
1301  auto minMaxOp = makeComposedMinMax<OpTy>(newBuilder, loc, map, operands);
1302 
1303  // Get constant operands.
1304  SmallVector<Attribute> constOperands(minMaxOp->getNumOperands());
1305  for (unsigned i = 0, e = constOperands.size(); i != e; ++i)
1306  matchPattern(minMaxOp->getOperand(i), m_Constant(&constOperands[i]));
1307 
1308  // Try to fold the operation.
1309  SmallVector<OpFoldResult> foldResults;
1310  if (failed(minMaxOp->fold(constOperands, foldResults)) ||
1311  foldResults.empty()) {
1312  if (OpBuilder::Listener *listener = b.getListener())
1313  listener->notifyOperationInserted(minMaxOp, /*previous=*/{});
1314  return minMaxOp.getResult();
1315  }
1316 
1317  minMaxOp->erase();
1318  return llvm::getSingleElement(foldResults);
1319 }
1320 
1323  AffineMap map,
1324  ArrayRef<OpFoldResult> operands) {
1325  return makeComposedFoldedMinMax<AffineMinOp>(b, loc, map, operands);
1326 }
1327 
1330  AffineMap map,
1331  ArrayRef<OpFoldResult> operands) {
1332  return makeComposedFoldedMinMax<AffineMaxOp>(b, loc, map, operands);
1333 }
1334 
1335 // A symbol may appear as a dim in affine.apply operations. This function
1336 // canonicalizes dims that are valid symbols into actual symbols.
1337 template <class MapOrSet>
1338 static void canonicalizePromotedSymbols(MapOrSet *mapOrSet,
1339  SmallVectorImpl<Value> *operands) {
1340  if (!mapOrSet || operands->empty())
1341  return;
1342 
1343  assert(mapOrSet->getNumInputs() == operands->size() &&
1344  "map/set inputs must match number of operands");
1345 
1346  auto *context = mapOrSet->getContext();
1347  SmallVector<Value, 8> resultOperands;
1348  resultOperands.reserve(operands->size());
1349  SmallVector<Value, 8> remappedSymbols;
1350  remappedSymbols.reserve(operands->size());
1351  unsigned nextDim = 0;
1352  unsigned nextSym = 0;
1353  unsigned oldNumSyms = mapOrSet->getNumSymbols();
1354  SmallVector<AffineExpr, 8> dimRemapping(mapOrSet->getNumDims());
1355  for (unsigned i = 0, e = mapOrSet->getNumInputs(); i != e; ++i) {
1356  if (i < mapOrSet->getNumDims()) {
1357  if (isValidSymbol((*operands)[i])) {
1358  // This is a valid symbol that appears as a dim, canonicalize it.
1359  dimRemapping[i] = getAffineSymbolExpr(oldNumSyms + nextSym++, context);
1360  remappedSymbols.push_back((*operands)[i]);
1361  } else {
1362  dimRemapping[i] = getAffineDimExpr(nextDim++, context);
1363  resultOperands.push_back((*operands)[i]);
1364  }
1365  } else {
1366  resultOperands.push_back((*operands)[i]);
1367  }
1368  }
1369 
1370  resultOperands.append(remappedSymbols.begin(), remappedSymbols.end());
1371  *operands = resultOperands;
1372  *mapOrSet = mapOrSet->replaceDimsAndSymbols(
1373  dimRemapping, /*symReplacements=*/{}, nextDim, oldNumSyms + nextSym);
1374 
1375  assert(mapOrSet->getNumInputs() == operands->size() &&
1376  "map/set inputs must match number of operands");
1377 }
1378 
1379 /// A valid affine dimension may appear as a symbol in affine.apply operations.
1380 /// Given an application of `operands` to an affine map or integer set
1381 /// `mapOrSet`, this function canonicalizes symbols of `mapOrSet` that are valid
1382 /// dims, but not valid symbols into actual dims. Without such a legalization,
1383 /// the affine.apply will be invalid. This method is the exact inverse of
1384 /// canonicalizePromotedSymbols.
1385 template <class MapOrSet>
1386 static void legalizeDemotedDims(MapOrSet &mapOrSet,
1387  SmallVectorImpl<Value> &operands) {
1388  if (!mapOrSet || operands.empty())
1389  return;
1390 
1391  unsigned numOperands = operands.size();
1392 
1393  assert(mapOrSet.getNumInputs() == numOperands &&
1394  "map/set inputs must match number of operands");
1395 
1396  auto *context = mapOrSet.getContext();
1397  SmallVector<Value, 8> resultOperands;
1398  resultOperands.reserve(numOperands);
1399  SmallVector<Value, 8> remappedDims;
1400  remappedDims.reserve(numOperands);
1401  SmallVector<Value, 8> symOperands;
1402  symOperands.reserve(mapOrSet.getNumSymbols());
1403  unsigned nextSym = 0;
1404  unsigned nextDim = 0;
1405  unsigned oldNumDims = mapOrSet.getNumDims();
1406  SmallVector<AffineExpr, 8> symRemapping(mapOrSet.getNumSymbols());
1407  resultOperands.assign(operands.begin(), operands.begin() + oldNumDims);
1408  for (unsigned i = oldNumDims, e = mapOrSet.getNumInputs(); i != e; ++i) {
1409  if (operands[i] && isValidDim(operands[i]) && !isValidSymbol(operands[i])) {
1410  // This is a valid dim that appears as a symbol, legalize it.
1411  symRemapping[i - oldNumDims] =
1412  getAffineDimExpr(oldNumDims + nextDim++, context);
1413  remappedDims.push_back(operands[i]);
1414  } else {
1415  symRemapping[i - oldNumDims] = getAffineSymbolExpr(nextSym++, context);
1416  symOperands.push_back(operands[i]);
1417  }
1418  }
1419 
1420  append_range(resultOperands, remappedDims);
1421  append_range(resultOperands, symOperands);
1422  operands = resultOperands;
1423  mapOrSet = mapOrSet.replaceDimsAndSymbols(
1424  /*dimReplacements=*/{}, symRemapping, oldNumDims + nextDim, nextSym);
1425 
1426  assert(mapOrSet.getNumInputs() == operands.size() &&
1427  "map/set inputs must match number of operands");
1428 }
1429 
1430 // Works for either an affine map or an integer set.
1431 template <class MapOrSet>
1432 static void canonicalizeMapOrSetAndOperands(MapOrSet *mapOrSet,
1433  SmallVectorImpl<Value> *operands) {
1434  static_assert(llvm::is_one_of<MapOrSet, AffineMap, IntegerSet>::value,
1435  "Argument must be either of AffineMap or IntegerSet type");
1436 
1437  if (!mapOrSet || operands->empty())
1438  return;
1439 
1440  assert(mapOrSet->getNumInputs() == operands->size() &&
1441  "map/set inputs must match number of operands");
1442 
1443  canonicalizePromotedSymbols<MapOrSet>(mapOrSet, operands);
1444  legalizeDemotedDims<MapOrSet>(*mapOrSet, *operands);
1445 
1446  // Check to see what dims are used.
1447  llvm::SmallBitVector usedDims(mapOrSet->getNumDims());
1448  llvm::SmallBitVector usedSyms(mapOrSet->getNumSymbols());
1449  mapOrSet->walkExprs([&](AffineExpr expr) {
1450  if (auto dimExpr = dyn_cast<AffineDimExpr>(expr))
1451  usedDims[dimExpr.getPosition()] = true;
1452  else if (auto symExpr = dyn_cast<AffineSymbolExpr>(expr))
1453  usedSyms[symExpr.getPosition()] = true;
1454  });
1455 
1456  auto *context = mapOrSet->getContext();
1457 
1458  SmallVector<Value, 8> resultOperands;
1459  resultOperands.reserve(operands->size());
1460 
1461  llvm::SmallDenseMap<Value, AffineExpr, 8> seenDims;
1462  SmallVector<AffineExpr, 8> dimRemapping(mapOrSet->getNumDims());
1463  unsigned nextDim = 0;
1464  for (unsigned i = 0, e = mapOrSet->getNumDims(); i != e; ++i) {
1465  if (usedDims[i]) {
1466  // Remap dim positions for duplicate operands.
1467  auto it = seenDims.find((*operands)[i]);
1468  if (it == seenDims.end()) {
1469  dimRemapping[i] = getAffineDimExpr(nextDim++, context);
1470  resultOperands.push_back((*operands)[i]);
1471  seenDims.insert(std::make_pair((*operands)[i], dimRemapping[i]));
1472  } else {
1473  dimRemapping[i] = it->second;
1474  }
1475  }
1476  }
1477  llvm::SmallDenseMap<Value, AffineExpr, 8> seenSymbols;
1478  SmallVector<AffineExpr, 8> symRemapping(mapOrSet->getNumSymbols());
1479  unsigned nextSym = 0;
1480  for (unsigned i = 0, e = mapOrSet->getNumSymbols(); i != e; ++i) {
1481  if (!usedSyms[i])
1482  continue;
1483  // Handle constant operands (only needed for symbolic operands since
1484  // constant operands in dimensional positions would have already been
1485  // promoted to symbolic positions above).
1486  IntegerAttr operandCst;
1487  if (matchPattern((*operands)[i + mapOrSet->getNumDims()],
1488  m_Constant(&operandCst))) {
1489  symRemapping[i] =
1490  getAffineConstantExpr(operandCst.getValue().getSExtValue(), context);
1491  continue;
1492  }
1493  // Remap symbol positions for duplicate operands.
1494  auto it = seenSymbols.find((*operands)[i + mapOrSet->getNumDims()]);
1495  if (it == seenSymbols.end()) {
1496  symRemapping[i] = getAffineSymbolExpr(nextSym++, context);
1497  resultOperands.push_back((*operands)[i + mapOrSet->getNumDims()]);
1498  seenSymbols.insert(std::make_pair((*operands)[i + mapOrSet->getNumDims()],
1499  symRemapping[i]));
1500  } else {
1501  symRemapping[i] = it->second;
1502  }
1503  }
1504  *mapOrSet = mapOrSet->replaceDimsAndSymbols(dimRemapping, symRemapping,
1505  nextDim, nextSym);
1506  *operands = resultOperands;
1507 }
1508 
1510  AffineMap *map, SmallVectorImpl<Value> *operands) {
1511  canonicalizeMapOrSetAndOperands<AffineMap>(map, operands);
1512 }
1513 
1515  IntegerSet *set, SmallVectorImpl<Value> *operands) {
1516  canonicalizeMapOrSetAndOperands<IntegerSet>(set, operands);
1517 }
1518 
1519 namespace {
1520 /// Simplify AffineApply, AffineLoad, and AffineStore operations by composing
1521 /// maps that supply results into them.
1522 ///
1523 template <typename AffineOpTy>
1524 struct SimplifyAffineOp : public OpRewritePattern<AffineOpTy> {
1526 
1527  /// Replace the affine op with another instance of it with the supplied
1528  /// map and mapOperands.
1529  void replaceAffineOp(PatternRewriter &rewriter, AffineOpTy affineOp,
1530  AffineMap map, ArrayRef<Value> mapOperands) const;
1531 
1532  LogicalResult matchAndRewrite(AffineOpTy affineOp,
1533  PatternRewriter &rewriter) const override {
1534  static_assert(
1535  llvm::is_one_of<AffineOpTy, AffineLoadOp, AffinePrefetchOp,
1536  AffineStoreOp, AffineApplyOp, AffineMinOp, AffineMaxOp,
1537  AffineVectorStoreOp, AffineVectorLoadOp>::value,
1538  "affine load/store/vectorstore/vectorload/apply/prefetch/min/max op "
1539  "expected");
1540  auto map = affineOp.getAffineMap();
1541  AffineMap oldMap = map;
1542  auto oldOperands = affineOp.getMapOperands();
1543  SmallVector<Value, 8> resultOperands(oldOperands);
1544  composeAffineMapAndOperands(&map, &resultOperands);
1545  canonicalizeMapAndOperands(&map, &resultOperands);
1546  simplifyMapWithOperands(map, resultOperands);
1547  if (map == oldMap && std::equal(oldOperands.begin(), oldOperands.end(),
1548  resultOperands.begin()))
1549  return failure();
1550 
1551  replaceAffineOp(rewriter, affineOp, map, resultOperands);
1552  return success();
1553  }
1554 };
1555 
1556 // Specialize the template to account for the different build signatures for
1557 // affine load, store, and apply ops.
1558 template <>
1559 void SimplifyAffineOp<AffineLoadOp>::replaceAffineOp(
1560  PatternRewriter &rewriter, AffineLoadOp load, AffineMap map,
1561  ArrayRef<Value> mapOperands) const {
1562  rewriter.replaceOpWithNewOp<AffineLoadOp>(load, load.getMemRef(), map,
1563  mapOperands);
1564 }
1565 template <>
1566 void SimplifyAffineOp<AffinePrefetchOp>::replaceAffineOp(
1567  PatternRewriter &rewriter, AffinePrefetchOp prefetch, AffineMap map,
1568  ArrayRef<Value> mapOperands) const {
1569  rewriter.replaceOpWithNewOp<AffinePrefetchOp>(
1570  prefetch, prefetch.getMemref(), map, mapOperands, prefetch.getIsWrite(),
1571  prefetch.getLocalityHint(), prefetch.getIsDataCache());
1572 }
1573 template <>
1574 void SimplifyAffineOp<AffineStoreOp>::replaceAffineOp(
1575  PatternRewriter &rewriter, AffineStoreOp store, AffineMap map,
1576  ArrayRef<Value> mapOperands) const {
1577  rewriter.replaceOpWithNewOp<AffineStoreOp>(
1578  store, store.getValueToStore(), store.getMemRef(), map, mapOperands);
1579 }
1580 template <>
1581 void SimplifyAffineOp<AffineVectorLoadOp>::replaceAffineOp(
1582  PatternRewriter &rewriter, AffineVectorLoadOp vectorload, AffineMap map,
1583  ArrayRef<Value> mapOperands) const {
1584  rewriter.replaceOpWithNewOp<AffineVectorLoadOp>(
1585  vectorload, vectorload.getVectorType(), vectorload.getMemRef(), map,
1586  mapOperands);
1587 }
1588 template <>
1589 void SimplifyAffineOp<AffineVectorStoreOp>::replaceAffineOp(
1590  PatternRewriter &rewriter, AffineVectorStoreOp vectorstore, AffineMap map,
1591  ArrayRef<Value> mapOperands) const {
1592  rewriter.replaceOpWithNewOp<AffineVectorStoreOp>(
1593  vectorstore, vectorstore.getValueToStore(), vectorstore.getMemRef(), map,
1594  mapOperands);
1595 }
1596 
1597 // Generic version for ops that don't have extra operands.
1598 template <typename AffineOpTy>
1599 void SimplifyAffineOp<AffineOpTy>::replaceAffineOp(
1600  PatternRewriter &rewriter, AffineOpTy op, AffineMap map,
1601  ArrayRef<Value> mapOperands) const {
1602  rewriter.replaceOpWithNewOp<AffineOpTy>(op, map, mapOperands);
1603 }
1604 } // namespace
1605 
1606 void AffineApplyOp::getCanonicalizationPatterns(RewritePatternSet &results,
1607  MLIRContext *context) {
1608  results.add<SimplifyAffineOp<AffineApplyOp>>(context);
1609 }
1610 
1611 //===----------------------------------------------------------------------===//
1612 // AffineDmaStartOp
1613 //===----------------------------------------------------------------------===//
1614 
1615 // TODO: Check that map operands are loop IVs or symbols.
1616 void AffineDmaStartOp::build(OpBuilder &builder, OperationState &result,
1617  Value srcMemRef, AffineMap srcMap,
1618  ValueRange srcIndices, Value destMemRef,
1619  AffineMap dstMap, ValueRange destIndices,
1620  Value tagMemRef, AffineMap tagMap,
1621  ValueRange tagIndices, Value numElements,
1622  Value stride, Value elementsPerStride) {
1623  result.addOperands(srcMemRef);
1624  result.addAttribute(getSrcMapAttrStrName(), AffineMapAttr::get(srcMap));
1625  result.addOperands(srcIndices);
1626  result.addOperands(destMemRef);
1627  result.addAttribute(getDstMapAttrStrName(), AffineMapAttr::get(dstMap));
1628  result.addOperands(destIndices);
1629  result.addOperands(tagMemRef);
1630  result.addAttribute(getTagMapAttrStrName(), AffineMapAttr::get(tagMap));
1631  result.addOperands(tagIndices);
1632  result.addOperands(numElements);
1633  if (stride) {
1634  result.addOperands({stride, elementsPerStride});
1635  }
1636 }
1637 
1639  p << " " << getSrcMemRef() << '[';
1640  p.printAffineMapOfSSAIds(getSrcMapAttr(), getSrcIndices());
1641  p << "], " << getDstMemRef() << '[';
1642  p.printAffineMapOfSSAIds(getDstMapAttr(), getDstIndices());
1643  p << "], " << getTagMemRef() << '[';
1644  p.printAffineMapOfSSAIds(getTagMapAttr(), getTagIndices());
1645  p << "], " << getNumElements();
1646  if (isStrided()) {
1647  p << ", " << getStride();
1648  p << ", " << getNumElementsPerStride();
1649  }
1650  p << " : " << getSrcMemRefType() << ", " << getDstMemRefType() << ", "
1651  << getTagMemRefType();
1652 }
1653 
1654 // Parse AffineDmaStartOp.
1655 // Ex:
1656 // affine.dma_start %src[%i, %j], %dst[%k, %l], %tag[%index], %size,
1657 // %stride, %num_elt_per_stride
1658 // : memref<3076 x f32, 0>, memref<1024 x f32, 2>, memref<1 x i32>
1659 //
1661  OperationState &result) {
1662  OpAsmParser::UnresolvedOperand srcMemRefInfo;
1663  AffineMapAttr srcMapAttr;
1665  OpAsmParser::UnresolvedOperand dstMemRefInfo;
1666  AffineMapAttr dstMapAttr;
1668  OpAsmParser::UnresolvedOperand tagMemRefInfo;
1669  AffineMapAttr tagMapAttr;
1671  OpAsmParser::UnresolvedOperand numElementsInfo;
1673 
1674  SmallVector<Type, 3> types;
1675  auto indexType = parser.getBuilder().getIndexType();
1676 
1677  // Parse and resolve the following list of operands:
1678  // *) dst memref followed by its affine maps operands (in square brackets).
1679  // *) src memref followed by its affine map operands (in square brackets).
1680  // *) tag memref followed by its affine map operands (in square brackets).
1681  // *) number of elements transferred by DMA operation.
1682  if (parser.parseOperand(srcMemRefInfo) ||
1683  parser.parseAffineMapOfSSAIds(srcMapOperands, srcMapAttr,
1684  getSrcMapAttrStrName(),
1685  result.attributes) ||
1686  parser.parseComma() || parser.parseOperand(dstMemRefInfo) ||
1687  parser.parseAffineMapOfSSAIds(dstMapOperands, dstMapAttr,
1688  getDstMapAttrStrName(),
1689  result.attributes) ||
1690  parser.parseComma() || parser.parseOperand(tagMemRefInfo) ||
1691  parser.parseAffineMapOfSSAIds(tagMapOperands, tagMapAttr,
1692  getTagMapAttrStrName(),
1693  result.attributes) ||
1694  parser.parseComma() || parser.parseOperand(numElementsInfo))
1695  return failure();
1696 
1697  // Parse optional stride and elements per stride.
1698  if (parser.parseTrailingOperandList(strideInfo))
1699  return failure();
1700 
1701  if (!strideInfo.empty() && strideInfo.size() != 2) {
1702  return parser.emitError(parser.getNameLoc(),
1703  "expected two stride related operands");
1704  }
1705  bool isStrided = strideInfo.size() == 2;
1706 
1707  if (parser.parseColonTypeList(types))
1708  return failure();
1709 
1710  if (types.size() != 3)
1711  return parser.emitError(parser.getNameLoc(), "expected three types");
1712 
1713  if (parser.resolveOperand(srcMemRefInfo, types[0], result.operands) ||
1714  parser.resolveOperands(srcMapOperands, indexType, result.operands) ||
1715  parser.resolveOperand(dstMemRefInfo, types[1], result.operands) ||
1716  parser.resolveOperands(dstMapOperands, indexType, result.operands) ||
1717  parser.resolveOperand(tagMemRefInfo, types[2], result.operands) ||
1718  parser.resolveOperands(tagMapOperands, indexType, result.operands) ||
1719  parser.resolveOperand(numElementsInfo, indexType, result.operands))
1720  return failure();
1721 
1722  if (isStrided) {
1723  if (parser.resolveOperands(strideInfo, indexType, result.operands))
1724  return failure();
1725  }
1726 
1727  // Check that src/dst/tag operand counts match their map.numInputs.
1728  if (srcMapOperands.size() != srcMapAttr.getValue().getNumInputs() ||
1729  dstMapOperands.size() != dstMapAttr.getValue().getNumInputs() ||
1730  tagMapOperands.size() != tagMapAttr.getValue().getNumInputs())
1731  return parser.emitError(parser.getNameLoc(),
1732  "memref operand count not equal to map.numInputs");
1733  return success();
1734 }
1735 
1736 LogicalResult AffineDmaStartOp::verifyInvariantsImpl() {
1737  if (!llvm::isa<MemRefType>(getOperand(getSrcMemRefOperandIndex()).getType()))
1738  return emitOpError("expected DMA source to be of memref type");
1739  if (!llvm::isa<MemRefType>(getOperand(getDstMemRefOperandIndex()).getType()))
1740  return emitOpError("expected DMA destination to be of memref type");
1741  if (!llvm::isa<MemRefType>(getOperand(getTagMemRefOperandIndex()).getType()))
1742  return emitOpError("expected DMA tag to be of memref type");
1743 
1744  unsigned numInputsAllMaps = getSrcMap().getNumInputs() +
1745  getDstMap().getNumInputs() +
1746  getTagMap().getNumInputs();
1747  if (getNumOperands() != numInputsAllMaps + 3 + 1 &&
1748  getNumOperands() != numInputsAllMaps + 3 + 1 + 2) {
1749  return emitOpError("incorrect number of operands");
1750  }
1751 
1752  Region *scope = getAffineScope(*this);
1753  for (auto idx : getSrcIndices()) {
1754  if (!idx.getType().isIndex())
1755  return emitOpError("src index to dma_start must have 'index' type");
1756  if (!isValidAffineIndexOperand(idx, scope))
1757  return emitOpError(
1758  "src index must be a valid dimension or symbol identifier");
1759  }
1760  for (auto idx : getDstIndices()) {
1761  if (!idx.getType().isIndex())
1762  return emitOpError("dst index to dma_start must have 'index' type");
1763  if (!isValidAffineIndexOperand(idx, scope))
1764  return emitOpError(
1765  "dst index must be a valid dimension or symbol identifier");
1766  }
1767  for (auto idx : getTagIndices()) {
1768  if (!idx.getType().isIndex())
1769  return emitOpError("tag index to dma_start must have 'index' type");
1770  if (!isValidAffineIndexOperand(idx, scope))
1771  return emitOpError(
1772  "tag index must be a valid dimension or symbol identifier");
1773  }
1774  return success();
1775 }
1776 
1777 LogicalResult AffineDmaStartOp::fold(ArrayRef<Attribute> cstOperands,
1778  SmallVectorImpl<OpFoldResult> &results) {
1779  /// dma_start(memrefcast) -> dma_start
1780  return memref::foldMemRefCast(*this);
1781 }
1782 
1783 void AffineDmaStartOp::getEffects(
1785  &effects) {
1786  effects.emplace_back(MemoryEffects::Read::get(), &getSrcMemRefMutable(),
1788  effects.emplace_back(MemoryEffects::Write::get(), &getDstMemRefMutable(),
1790  effects.emplace_back(MemoryEffects::Read::get(), &getTagMemRefMutable(),
1792 }
1793 
1794 //===----------------------------------------------------------------------===//
1795 // AffineDmaWaitOp
1796 //===----------------------------------------------------------------------===//
1797 
1798 // TODO: Check that map operands are loop IVs or symbols.
1799 void AffineDmaWaitOp::build(OpBuilder &builder, OperationState &result,
1800  Value tagMemRef, AffineMap tagMap,
1801  ValueRange tagIndices, Value numElements) {
1802  result.addOperands(tagMemRef);
1803  result.addAttribute(getTagMapAttrStrName(), AffineMapAttr::get(tagMap));
1804  result.addOperands(tagIndices);
1805  result.addOperands(numElements);
1806 }
1807 
1809  p << " " << getTagMemRef() << '[';
1810  SmallVector<Value, 2> operands(getTagIndices());
1811  p.printAffineMapOfSSAIds(getTagMapAttr(), operands);
1812  p << "], ";
1814  p << " : " << getTagMemRef().getType();
1815 }
1816 
1817 // Parse AffineDmaWaitOp.
1818 // Eg:
1819 // affine.dma_wait %tag[%index], %num_elements
1820 // : memref<1 x i32, (d0) -> (d0), 4>
1821 //
1823  OperationState &result) {
1824  OpAsmParser::UnresolvedOperand tagMemRefInfo;
1825  AffineMapAttr tagMapAttr;
1827  Type type;
1828  auto indexType = parser.getBuilder().getIndexType();
1829  OpAsmParser::UnresolvedOperand numElementsInfo;
1830 
1831  // Parse tag memref, its map operands, and dma size.
1832  if (parser.parseOperand(tagMemRefInfo) ||
1833  parser.parseAffineMapOfSSAIds(tagMapOperands, tagMapAttr,
1834  getTagMapAttrStrName(),
1835  result.attributes) ||
1836  parser.parseComma() || parser.parseOperand(numElementsInfo) ||
1837  parser.parseColonType(type) ||
1838  parser.resolveOperand(tagMemRefInfo, type, result.operands) ||
1839  parser.resolveOperands(tagMapOperands, indexType, result.operands) ||
1840  parser.resolveOperand(numElementsInfo, indexType, result.operands))
1841  return failure();
1842 
1843  if (!llvm::isa<MemRefType>(type))
1844  return parser.emitError(parser.getNameLoc(),
1845  "expected tag to be of memref type");
1846 
1847  if (tagMapOperands.size() != tagMapAttr.getValue().getNumInputs())
1848  return parser.emitError(parser.getNameLoc(),
1849  "tag memref operand count != to map.numInputs");
1850  return success();
1851 }
1852 
1853 LogicalResult AffineDmaWaitOp::verifyInvariantsImpl() {
1854  if (!llvm::isa<MemRefType>(getOperand(0).getType()))
1855  return emitOpError("expected DMA tag to be of memref type");
1856  Region *scope = getAffineScope(*this);
1857  for (auto idx : getTagIndices()) {
1858  if (!idx.getType().isIndex())
1859  return emitOpError("index to dma_wait must have 'index' type");
1860  if (!isValidAffineIndexOperand(idx, scope))
1861  return emitOpError(
1862  "index must be a valid dimension or symbol identifier");
1863  }
1864  return success();
1865 }
1866 
1867 LogicalResult AffineDmaWaitOp::fold(ArrayRef<Attribute> cstOperands,
1868  SmallVectorImpl<OpFoldResult> &results) {
1869  /// dma_wait(memrefcast) -> dma_wait
1870  return memref::foldMemRefCast(*this);
1871 }
1872 
1873 void AffineDmaWaitOp::getEffects(
1875  &effects) {
1876  effects.emplace_back(MemoryEffects::Read::get(), &getTagMemRefMutable(),
1878 }
1879 
1880 //===----------------------------------------------------------------------===//
1881 // AffineForOp
1882 //===----------------------------------------------------------------------===//
1883 
1884 /// 'bodyBuilder' is used to build the body of affine.for. If iterArgs and
1885 /// bodyBuilder are empty/null, we include default terminator op.
1886 void AffineForOp::build(OpBuilder &builder, OperationState &result,
1887  ValueRange lbOperands, AffineMap lbMap,
1888  ValueRange ubOperands, AffineMap ubMap, int64_t step,
1889  ValueRange iterArgs, BodyBuilderFn bodyBuilder) {
1890  assert(((!lbMap && lbOperands.empty()) ||
1891  lbOperands.size() == lbMap.getNumInputs()) &&
1892  "lower bound operand count does not match the affine map");
1893  assert(((!ubMap && ubOperands.empty()) ||
1894  ubOperands.size() == ubMap.getNumInputs()) &&
1895  "upper bound operand count does not match the affine map");
1896  assert(step > 0 && "step has to be a positive integer constant");
1897 
1898  OpBuilder::InsertionGuard guard(builder);
1899 
1900  // Set variadic segment sizes.
1901  result.addAttribute(
1902  getOperandSegmentSizeAttr(),
1903  builder.getDenseI32ArrayAttr({static_cast<int32_t>(lbOperands.size()),
1904  static_cast<int32_t>(ubOperands.size()),
1905  static_cast<int32_t>(iterArgs.size())}));
1906 
1907  for (Value val : iterArgs)
1908  result.addTypes(val.getType());
1909 
1910  // Add an attribute for the step.
1911  result.addAttribute(getStepAttrName(result.name),
1912  builder.getIntegerAttr(builder.getIndexType(), step));
1913 
1914  // Add the lower bound.
1915  result.addAttribute(getLowerBoundMapAttrName(result.name),
1916  AffineMapAttr::get(lbMap));
1917  result.addOperands(lbOperands);
1918 
1919  // Add the upper bound.
1920  result.addAttribute(getUpperBoundMapAttrName(result.name),
1921  AffineMapAttr::get(ubMap));
1922  result.addOperands(ubOperands);
1923 
1924  result.addOperands(iterArgs);
1925  // Create a region and a block for the body. The argument of the region is
1926  // the loop induction variable.
1927  Region *bodyRegion = result.addRegion();
1928  Block *bodyBlock = builder.createBlock(bodyRegion);
1929  Value inductionVar =
1930  bodyBlock->addArgument(builder.getIndexType(), result.location);
1931  for (Value val : iterArgs)
1932  bodyBlock->addArgument(val.getType(), val.getLoc());
1933 
1934  // Create the default terminator if the builder is not provided and if the
1935  // iteration arguments are not provided. Otherwise, leave this to the caller
1936  // because we don't know which values to return from the loop.
1937  if (iterArgs.empty() && !bodyBuilder) {
1938  ensureTerminator(*bodyRegion, builder, result.location);
1939  } else if (bodyBuilder) {
1940  OpBuilder::InsertionGuard guard(builder);
1941  builder.setInsertionPointToStart(bodyBlock);
1942  bodyBuilder(builder, result.location, inductionVar,
1943  bodyBlock->getArguments().drop_front());
1944  }
1945 }
1946 
1947 void AffineForOp::build(OpBuilder &builder, OperationState &result, int64_t lb,
1948  int64_t ub, int64_t step, ValueRange iterArgs,
1949  BodyBuilderFn bodyBuilder) {
1950  auto lbMap = AffineMap::getConstantMap(lb, builder.getContext());
1951  auto ubMap = AffineMap::getConstantMap(ub, builder.getContext());
1952  return build(builder, result, {}, lbMap, {}, ubMap, step, iterArgs,
1953  bodyBuilder);
1954 }
1955 
1956 LogicalResult AffineForOp::verifyRegions() {
1957  // Check that the body defines as single block argument for the induction
1958  // variable.
1959  auto *body = getBody();
1960  if (body->getNumArguments() == 0 || !body->getArgument(0).getType().isIndex())
1961  return emitOpError("expected body to have a single index argument for the "
1962  "induction variable");
1963 
1964  // Verify that the bound operands are valid dimension/symbols.
1965  /// Lower bound.
1966  if (getLowerBoundMap().getNumInputs() > 0)
1968  getLowerBoundMap().getNumDims())))
1969  return failure();
1970  /// Upper bound.
1971  if (getUpperBoundMap().getNumInputs() > 0)
1973  getUpperBoundMap().getNumDims())))
1974  return failure();
1975  if (getLowerBoundMap().getNumResults() < 1)
1976  return emitOpError("expected lower bound map to have at least one result");
1977  if (getUpperBoundMap().getNumResults() < 1)
1978  return emitOpError("expected upper bound map to have at least one result");
1979 
1980  unsigned opNumResults = getNumResults();
1981  if (opNumResults == 0)
1982  return success();
1983 
1984  // If ForOp defines values, check that the number and types of the defined
1985  // values match ForOp initial iter operands and backedge basic block
1986  // arguments.
1987  if (getNumIterOperands() != opNumResults)
1988  return emitOpError(
1989  "mismatch between the number of loop-carried values and results");
1990  if (getNumRegionIterArgs() != opNumResults)
1991  return emitOpError(
1992  "mismatch between the number of basic block args and results");
1993 
1994  return success();
1995 }
1996 
1997 /// Parse a for operation loop bounds.
1998 static ParseResult parseBound(bool isLower, OperationState &result,
1999  OpAsmParser &p) {
2000  // 'min' / 'max' prefixes are generally syntactic sugar, but are required if
2001  // the map has multiple results.
2002  bool failedToParsedMinMax =
2003  failed(p.parseOptionalKeyword(isLower ? "max" : "min"));
2004 
2005  auto &builder = p.getBuilder();
2006  auto boundAttrStrName =
2007  isLower ? AffineForOp::getLowerBoundMapAttrName(result.name)
2008  : AffineForOp::getUpperBoundMapAttrName(result.name);
2009 
2010  // Parse ssa-id as identity map.
2012  if (p.parseOperandList(boundOpInfos))
2013  return failure();
2014 
2015  if (!boundOpInfos.empty()) {
2016  // Check that only one operand was parsed.
2017  if (boundOpInfos.size() > 1)
2018  return p.emitError(p.getNameLoc(),
2019  "expected only one loop bound operand");
2020 
2021  // TODO: improve error message when SSA value is not of index type.
2022  // Currently it is 'use of value ... expects different type than prior uses'
2023  if (p.resolveOperand(boundOpInfos.front(), builder.getIndexType(),
2024  result.operands))
2025  return failure();
2026 
2027  // Create an identity map using symbol id. This representation is optimized
2028  // for storage. Analysis passes may expand it into a multi-dimensional map
2029  // if desired.
2030  AffineMap map = builder.getSymbolIdentityMap();
2031  result.addAttribute(boundAttrStrName, AffineMapAttr::get(map));
2032  return success();
2033  }
2034 
2035  // Get the attribute location.
2036  SMLoc attrLoc = p.getCurrentLocation();
2037 
2038  Attribute boundAttr;
2039  if (p.parseAttribute(boundAttr, builder.getIndexType(), boundAttrStrName,
2040  result.attributes))
2041  return failure();
2042 
2043  // Parse full form - affine map followed by dim and symbol list.
2044  if (auto affineMapAttr = llvm::dyn_cast<AffineMapAttr>(boundAttr)) {
2045  unsigned currentNumOperands = result.operands.size();
2046  unsigned numDims;
2047  if (parseDimAndSymbolList(p, result.operands, numDims))
2048  return failure();
2049 
2050  auto map = affineMapAttr.getValue();
2051  if (map.getNumDims() != numDims)
2052  return p.emitError(
2053  p.getNameLoc(),
2054  "dim operand count and affine map dim count must match");
2055 
2056  unsigned numDimAndSymbolOperands =
2057  result.operands.size() - currentNumOperands;
2058  if (numDims + map.getNumSymbols() != numDimAndSymbolOperands)
2059  return p.emitError(
2060  p.getNameLoc(),
2061  "symbol operand count and affine map symbol count must match");
2062 
2063  // If the map has multiple results, make sure that we parsed the min/max
2064  // prefix.
2065  if (map.getNumResults() > 1 && failedToParsedMinMax) {
2066  if (isLower) {
2067  return p.emitError(attrLoc, "lower loop bound affine map with "
2068  "multiple results requires 'max' prefix");
2069  }
2070  return p.emitError(attrLoc, "upper loop bound affine map with multiple "
2071  "results requires 'min' prefix");
2072  }
2073  return success();
2074  }
2075 
2076  // Parse custom assembly form.
2077  if (auto integerAttr = llvm::dyn_cast<IntegerAttr>(boundAttr)) {
2078  result.attributes.pop_back();
2079  result.addAttribute(
2080  boundAttrStrName,
2081  AffineMapAttr::get(builder.getConstantAffineMap(integerAttr.getInt())));
2082  return success();
2083  }
2084 
2085  return p.emitError(
2086  p.getNameLoc(),
2087  "expected valid affine map representation for loop bounds");
2088 }
2089 
2090 ParseResult AffineForOp::parse(OpAsmParser &parser, OperationState &result) {
2091  auto &builder = parser.getBuilder();
2092  OpAsmParser::Argument inductionVariable;
2093  inductionVariable.type = builder.getIndexType();
2094  // Parse the induction variable followed by '='.
2095  if (parser.parseArgument(inductionVariable) || parser.parseEqual())
2096  return failure();
2097 
2098  // Parse loop bounds.
2099  int64_t numOperands = result.operands.size();
2100  if (parseBound(/*isLower=*/true, result, parser))
2101  return failure();
2102  int64_t numLbOperands = result.operands.size() - numOperands;
2103  if (parser.parseKeyword("to", " between bounds"))
2104  return failure();
2105  numOperands = result.operands.size();
2106  if (parseBound(/*isLower=*/false, result, parser))
2107  return failure();
2108  int64_t numUbOperands = result.operands.size() - numOperands;
2109 
2110  // Parse the optional loop step, we default to 1 if one is not present.
2111  if (parser.parseOptionalKeyword("step")) {
2112  result.addAttribute(
2113  getStepAttrName(result.name),
2114  builder.getIntegerAttr(builder.getIndexType(), /*value=*/1));
2115  } else {
2116  SMLoc stepLoc = parser.getCurrentLocation();
2117  IntegerAttr stepAttr;
2118  if (parser.parseAttribute(stepAttr, builder.getIndexType(),
2119  getStepAttrName(result.name).data(),
2120  result.attributes))
2121  return failure();
2122 
2123  if (stepAttr.getValue().isNegative())
2124  return parser.emitError(
2125  stepLoc,
2126  "expected step to be representable as a positive signed integer");
2127  }
2128 
2129  // Parse the optional initial iteration arguments.
2132 
2133  // Induction variable.
2134  regionArgs.push_back(inductionVariable);
2135 
2136  if (succeeded(parser.parseOptionalKeyword("iter_args"))) {
2137  // Parse assignment list and results type list.
2138  if (parser.parseAssignmentList(regionArgs, operands) ||
2139  parser.parseArrowTypeList(result.types))
2140  return failure();
2141  // Resolve input operands.
2142  for (auto argOperandType :
2143  llvm::zip(llvm::drop_begin(regionArgs), operands, result.types)) {
2144  Type type = std::get<2>(argOperandType);
2145  std::get<0>(argOperandType).type = type;
2146  if (parser.resolveOperand(std::get<1>(argOperandType), type,
2147  result.operands))
2148  return failure();
2149  }
2150  }
2151 
2152  result.addAttribute(
2153  getOperandSegmentSizeAttr(),
2154  builder.getDenseI32ArrayAttr({static_cast<int32_t>(numLbOperands),
2155  static_cast<int32_t>(numUbOperands),
2156  static_cast<int32_t>(operands.size())}));
2157 
2158  // Parse the body region.
2159  Region *body = result.addRegion();
2160  if (regionArgs.size() != result.types.size() + 1)
2161  return parser.emitError(
2162  parser.getNameLoc(),
2163  "mismatch between the number of loop-carried values and results");
2164  if (parser.parseRegion(*body, regionArgs))
2165  return failure();
2166 
2167  AffineForOp::ensureTerminator(*body, builder, result.location);
2168 
2169  // Parse the optional attribute list.
2170  return parser.parseOptionalAttrDict(result.attributes);
2171 }
2172 
2173 static void printBound(AffineMapAttr boundMap,
2174  Operation::operand_range boundOperands,
2175  const char *prefix, OpAsmPrinter &p) {
2176  AffineMap map = boundMap.getValue();
2177 
2178  // Check if this bound should be printed using custom assembly form.
2179  // The decision to restrict printing custom assembly form to trivial cases
2180  // comes from the will to roundtrip MLIR binary -> text -> binary in a
2181  // lossless way.
2182  // Therefore, custom assembly form parsing and printing is only supported for
2183  // zero-operand constant maps and single symbol operand identity maps.
2184  if (map.getNumResults() == 1) {
2185  AffineExpr expr = map.getResult(0);
2186 
2187  // Print constant bound.
2188  if (map.getNumDims() == 0 && map.getNumSymbols() == 0) {
2189  if (auto constExpr = dyn_cast<AffineConstantExpr>(expr)) {
2190  p << constExpr.getValue();
2191  return;
2192  }
2193  }
2194 
2195  // Print bound that consists of a single SSA symbol if the map is over a
2196  // single symbol.
2197  if (map.getNumDims() == 0 && map.getNumSymbols() == 1) {
2198  if (dyn_cast<AffineSymbolExpr>(expr)) {
2199  p.printOperand(*boundOperands.begin());
2200  return;
2201  }
2202  }
2203  } else {
2204  // Map has multiple results. Print 'min' or 'max' prefix.
2205  p << prefix << ' ';
2206  }
2207 
2208  // Print the map and its operands.
2209  p << boundMap;
2210  printDimAndSymbolList(boundOperands.begin(), boundOperands.end(),
2211  map.getNumDims(), p);
2212 }
2213 
2214 unsigned AffineForOp::getNumIterOperands() {
2215  AffineMap lbMap = getLowerBoundMapAttr().getValue();
2216  AffineMap ubMap = getUpperBoundMapAttr().getValue();
2217 
2218  return getNumOperands() - lbMap.getNumInputs() - ubMap.getNumInputs();
2219 }
2220 
2221 std::optional<MutableArrayRef<OpOperand>>
2222 AffineForOp::getYieldedValuesMutable() {
2223  return cast<AffineYieldOp>(getBody()->getTerminator()).getOperandsMutable();
2224 }
2225 
2227  p << ' ';
2228  p.printRegionArgument(getBody()->getArgument(0), /*argAttrs=*/{},
2229  /*omitType=*/true);
2230  p << " = ";
2231  printBound(getLowerBoundMapAttr(), getLowerBoundOperands(), "max", p);
2232  p << " to ";
2233  printBound(getUpperBoundMapAttr(), getUpperBoundOperands(), "min", p);
2234 
2235  if (getStepAsInt() != 1)
2236  p << " step " << getStepAsInt();
2237 
2238  bool printBlockTerminators = false;
2239  if (getNumIterOperands() > 0) {
2240  p << " iter_args(";
2241  auto regionArgs = getRegionIterArgs();
2242  auto operands = getInits();
2243 
2244  llvm::interleaveComma(llvm::zip(regionArgs, operands), p, [&](auto it) {
2245  p << std::get<0>(it) << " = " << std::get<1>(it);
2246  });
2247  p << ") -> (" << getResultTypes() << ")";
2248  printBlockTerminators = true;
2249  }
2250 
2251  p << ' ';
2252  p.printRegion(getRegion(), /*printEntryBlockArgs=*/false,
2253  printBlockTerminators);
2255  (*this)->getAttrs(),
2256  /*elidedAttrs=*/{getLowerBoundMapAttrName(getOperation()->getName()),
2257  getUpperBoundMapAttrName(getOperation()->getName()),
2258  getStepAttrName(getOperation()->getName()),
2259  getOperandSegmentSizeAttr()});
2260 }
2261 
2262 /// Fold the constant bounds of a loop.
2263 static LogicalResult foldLoopBounds(AffineForOp forOp) {
2264  auto foldLowerOrUpperBound = [&forOp](bool lower) {
2265  // Check to see if each of the operands is the result of a constant. If
2266  // so, get the value. If not, ignore it.
2267  SmallVector<Attribute, 8> operandConstants;
2268  auto boundOperands =
2269  lower ? forOp.getLowerBoundOperands() : forOp.getUpperBoundOperands();
2270  for (auto operand : boundOperands) {
2271  Attribute operandCst;
2272  matchPattern(operand, m_Constant(&operandCst));
2273  operandConstants.push_back(operandCst);
2274  }
2275 
2276  AffineMap boundMap =
2277  lower ? forOp.getLowerBoundMap() : forOp.getUpperBoundMap();
2278  assert(boundMap.getNumResults() >= 1 &&
2279  "bound maps should have at least one result");
2280  SmallVector<Attribute, 4> foldedResults;
2281  if (failed(boundMap.constantFold(operandConstants, foldedResults)))
2282  return failure();
2283 
2284  // Compute the max or min as applicable over the results.
2285  assert(!foldedResults.empty() && "bounds should have at least one result");
2286  auto maxOrMin = llvm::cast<IntegerAttr>(foldedResults[0]).getValue();
2287  for (unsigned i = 1, e = foldedResults.size(); i < e; i++) {
2288  auto foldedResult = llvm::cast<IntegerAttr>(foldedResults[i]).getValue();
2289  maxOrMin = lower ? llvm::APIntOps::smax(maxOrMin, foldedResult)
2290  : llvm::APIntOps::smin(maxOrMin, foldedResult);
2291  }
2292  lower ? forOp.setConstantLowerBound(maxOrMin.getSExtValue())
2293  : forOp.setConstantUpperBound(maxOrMin.getSExtValue());
2294  return success();
2295  };
2296 
2297  // Try to fold the lower bound.
2298  bool folded = false;
2299  if (!forOp.hasConstantLowerBound())
2300  folded |= succeeded(foldLowerOrUpperBound(/*lower=*/true));
2301 
2302  // Try to fold the upper bound.
2303  if (!forOp.hasConstantUpperBound())
2304  folded |= succeeded(foldLowerOrUpperBound(/*lower=*/false));
2305  return success(folded);
2306 }
2307 
2308 /// Canonicalize the bounds of the given loop.
2309 static LogicalResult canonicalizeLoopBounds(AffineForOp forOp) {
2310  SmallVector<Value, 4> lbOperands(forOp.getLowerBoundOperands());
2311  SmallVector<Value, 4> ubOperands(forOp.getUpperBoundOperands());
2312 
2313  auto lbMap = forOp.getLowerBoundMap();
2314  auto ubMap = forOp.getUpperBoundMap();
2315  auto prevLbMap = lbMap;
2316  auto prevUbMap = ubMap;
2317 
2318  composeAffineMapAndOperands(&lbMap, &lbOperands);
2319  canonicalizeMapAndOperands(&lbMap, &lbOperands);
2320  simplifyMinOrMaxExprWithOperands(lbMap, lbOperands, /*isMax=*/true);
2321  simplifyMinOrMaxExprWithOperands(ubMap, ubOperands, /*isMax=*/false);
2322  lbMap = removeDuplicateExprs(lbMap);
2323 
2324  composeAffineMapAndOperands(&ubMap, &ubOperands);
2325  canonicalizeMapAndOperands(&ubMap, &ubOperands);
2326  ubMap = removeDuplicateExprs(ubMap);
2327 
2328  // Any canonicalization change always leads to updated map(s).
2329  if (lbMap == prevLbMap && ubMap == prevUbMap)
2330  return failure();
2331 
2332  if (lbMap != prevLbMap)
2333  forOp.setLowerBound(lbOperands, lbMap);
2334  if (ubMap != prevUbMap)
2335  forOp.setUpperBound(ubOperands, ubMap);
2336  return success();
2337 }
2338 
2339 namespace {
2340 /// Returns constant trip count in trivial cases.
2341 static std::optional<uint64_t> getTrivialConstantTripCount(AffineForOp forOp) {
2342  int64_t step = forOp.getStepAsInt();
2343  if (!forOp.hasConstantBounds() || step <= 0)
2344  return std::nullopt;
2345  int64_t lb = forOp.getConstantLowerBound();
2346  int64_t ub = forOp.getConstantUpperBound();
2347  return ub - lb <= 0 ? 0 : (ub - lb + step - 1) / step;
2348 }
2349 
2350 /// This is a pattern to fold trivially empty loop bodies.
2351 /// TODO: This should be moved into the folding hook.
2352 struct AffineForEmptyLoopFolder : public OpRewritePattern<AffineForOp> {
2354 
2355  LogicalResult matchAndRewrite(AffineForOp forOp,
2356  PatternRewriter &rewriter) const override {
2357  // Check that the body only contains a yield.
2358  if (!llvm::hasSingleElement(*forOp.getBody()))
2359  return failure();
2360  if (forOp.getNumResults() == 0)
2361  return success();
2362  std::optional<uint64_t> tripCount = getTrivialConstantTripCount(forOp);
2363  if (tripCount && *tripCount == 0) {
2364  // The initial values of the iteration arguments would be the op's
2365  // results.
2366  rewriter.replaceOp(forOp, forOp.getInits());
2367  return success();
2368  }
2369  SmallVector<Value, 4> replacements;
2370  auto yieldOp = cast<AffineYieldOp>(forOp.getBody()->getTerminator());
2371  auto iterArgs = forOp.getRegionIterArgs();
2372  bool hasValDefinedOutsideLoop = false;
2373  bool iterArgsNotInOrder = false;
2374  for (unsigned i = 0, e = yieldOp->getNumOperands(); i < e; ++i) {
2375  Value val = yieldOp.getOperand(i);
2376  auto *iterArgIt = llvm::find(iterArgs, val);
2377  // TODO: It should be possible to perform a replacement by computing the
2378  // last value of the IV based on the bounds and the step.
2379  if (val == forOp.getInductionVar())
2380  return failure();
2381  if (iterArgIt == iterArgs.end()) {
2382  // `val` is defined outside of the loop.
2383  assert(forOp.isDefinedOutsideOfLoop(val) &&
2384  "must be defined outside of the loop");
2385  hasValDefinedOutsideLoop = true;
2386  replacements.push_back(val);
2387  } else {
2388  unsigned pos = std::distance(iterArgs.begin(), iterArgIt);
2389  if (pos != i)
2390  iterArgsNotInOrder = true;
2391  replacements.push_back(forOp.getInits()[pos]);
2392  }
2393  }
2394  // Bail out when the trip count is unknown and the loop returns any value
2395  // defined outside of the loop or any iterArg out of order.
2396  if (!tripCount.has_value() &&
2397  (hasValDefinedOutsideLoop || iterArgsNotInOrder))
2398  return failure();
2399  // Bail out when the loop iterates more than once and it returns any iterArg
2400  // out of order.
2401  if (tripCount.has_value() && tripCount.value() >= 2 && iterArgsNotInOrder)
2402  return failure();
2403  rewriter.replaceOp(forOp, replacements);
2404  return success();
2405  }
2406 };
2407 } // namespace
2408 
2409 void AffineForOp::getCanonicalizationPatterns(RewritePatternSet &results,
2410  MLIRContext *context) {
2411  results.add<AffineForEmptyLoopFolder>(context);
2412 }
2413 
2414 OperandRange AffineForOp::getEntrySuccessorOperands(RegionBranchPoint point) {
2415  assert((point.isParent() || point == getRegion()) && "invalid region point");
2416 
2417  // The initial operands map to the loop arguments after the induction
2418  // variable or are forwarded to the results when the trip count is zero.
2419  return getInits();
2420 }
2421 
2422 void AffineForOp::getSuccessorRegions(
2424  assert((point.isParent() || point == getRegion()) && "expected loop region");
2425  // The loop may typically branch back to its body or to the parent operation.
2426  // If the predecessor is the parent op and the trip count is known to be at
2427  // least one, branch into the body using the iterator arguments. And in cases
2428  // we know the trip count is zero, it can only branch back to its parent.
2429  std::optional<uint64_t> tripCount = getTrivialConstantTripCount(*this);
2430  if (point.isParent() && tripCount.has_value()) {
2431  if (tripCount.value() > 0) {
2432  regions.push_back(RegionSuccessor(&getRegion(), getRegionIterArgs()));
2433  return;
2434  }
2435  if (tripCount.value() == 0) {
2436  regions.push_back(RegionSuccessor(getResults()));
2437  return;
2438  }
2439  }
2440 
2441  // From the loop body, if the trip count is one, we can only branch back to
2442  // the parent.
2443  if (!point.isParent() && tripCount && *tripCount == 1) {
2444  regions.push_back(RegionSuccessor(getResults()));
2445  return;
2446  }
2447 
2448  // In all other cases, the loop may branch back to itself or the parent
2449  // operation.
2450  regions.push_back(RegionSuccessor(&getRegion(), getRegionIterArgs()));
2451  regions.push_back(RegionSuccessor(getResults()));
2452 }
2453 
2454 /// Returns true if the affine.for has zero iterations in trivial cases.
2455 static bool hasTrivialZeroTripCount(AffineForOp op) {
2456  std::optional<uint64_t> tripCount = getTrivialConstantTripCount(op);
2457  return tripCount && *tripCount == 0;
2458 }
2459 
2460 LogicalResult AffineForOp::fold(FoldAdaptor adaptor,
2461  SmallVectorImpl<OpFoldResult> &results) {
2462  bool folded = succeeded(foldLoopBounds(*this));
2463  folded |= succeeded(canonicalizeLoopBounds(*this));
2464  if (hasTrivialZeroTripCount(*this) && getNumResults() != 0) {
2465  // The initial values of the loop-carried variables (iter_args) are the
2466  // results of the op. But this must be avoided for an affine.for op that
2467  // does not return any results. Since ops that do not return results cannot
2468  // be folded away, we would enter an infinite loop of folds on the same
2469  // affine.for op.
2470  results.assign(getInits().begin(), getInits().end());
2471  folded = true;
2472  }
2473  return success(folded);
2474 }
2475 
2477  return AffineBound(*this, getLowerBoundOperands(), getLowerBoundMap());
2478 }
2479 
2481  return AffineBound(*this, getUpperBoundOperands(), getUpperBoundMap());
2482 }
2483 
2484 void AffineForOp::setLowerBound(ValueRange lbOperands, AffineMap map) {
2485  assert(lbOperands.size() == map.getNumInputs());
2486  assert(map.getNumResults() >= 1 && "bound map has at least one result");
2487  getLowerBoundOperandsMutable().assign(lbOperands);
2488  setLowerBoundMap(map);
2489 }
2490 
2491 void AffineForOp::setUpperBound(ValueRange ubOperands, AffineMap map) {
2492  assert(ubOperands.size() == map.getNumInputs());
2493  assert(map.getNumResults() >= 1 && "bound map has at least one result");
2494  getUpperBoundOperandsMutable().assign(ubOperands);
2495  setUpperBoundMap(map);
2496 }
2497 
2498 bool AffineForOp::hasConstantLowerBound() {
2499  return getLowerBoundMap().isSingleConstant();
2500 }
2501 
2502 bool AffineForOp::hasConstantUpperBound() {
2503  return getUpperBoundMap().isSingleConstant();
2504 }
2505 
2506 int64_t AffineForOp::getConstantLowerBound() {
2507  return getLowerBoundMap().getSingleConstantResult();
2508 }
2509 
2510 int64_t AffineForOp::getConstantUpperBound() {
2511  return getUpperBoundMap().getSingleConstantResult();
2512 }
2513 
2514 void AffineForOp::setConstantLowerBound(int64_t value) {
2515  setLowerBound({}, AffineMap::getConstantMap(value, getContext()));
2516 }
2517 
2518 void AffineForOp::setConstantUpperBound(int64_t value) {
2519  setUpperBound({}, AffineMap::getConstantMap(value, getContext()));
2520 }
2521 
2522 AffineForOp::operand_range AffineForOp::getControlOperands() {
2523  return {operand_begin(), operand_begin() + getLowerBoundOperands().size() +
2524  getUpperBoundOperands().size()};
2525 }
2526 
2527 bool AffineForOp::matchingBoundOperandList() {
2528  auto lbMap = getLowerBoundMap();
2529  auto ubMap = getUpperBoundMap();
2530  if (lbMap.getNumDims() != ubMap.getNumDims() ||
2531  lbMap.getNumSymbols() != ubMap.getNumSymbols())
2532  return false;
2533 
2534  unsigned numOperands = lbMap.getNumInputs();
2535  for (unsigned i = 0, e = lbMap.getNumInputs(); i < e; i++) {
2536  // Compare Value 's.
2537  if (getOperand(i) != getOperand(numOperands + i))
2538  return false;
2539  }
2540  return true;
2541 }
2542 
2543 SmallVector<Region *> AffineForOp::getLoopRegions() { return {&getRegion()}; }
2544 
2545 std::optional<SmallVector<Value>> AffineForOp::getLoopInductionVars() {
2546  return SmallVector<Value>{getInductionVar()};
2547 }
2548 
2549 std::optional<SmallVector<OpFoldResult>> AffineForOp::getLoopLowerBounds() {
2550  if (!hasConstantLowerBound())
2551  return std::nullopt;
2552  OpBuilder b(getContext());
2554  OpFoldResult(b.getI64IntegerAttr(getConstantLowerBound()))};
2555 }
2556 
2557 std::optional<SmallVector<OpFoldResult>> AffineForOp::getLoopSteps() {
2558  OpBuilder b(getContext());
2560  OpFoldResult(b.getI64IntegerAttr(getStepAsInt()))};
2561 }
2562 
2563 std::optional<SmallVector<OpFoldResult>> AffineForOp::getLoopUpperBounds() {
2564  if (!hasConstantUpperBound())
2565  return {};
2566  OpBuilder b(getContext());
2568  OpFoldResult(b.getI64IntegerAttr(getConstantUpperBound()))};
2569 }
2570 
2571 FailureOr<LoopLikeOpInterface> AffineForOp::replaceWithAdditionalYields(
2572  RewriterBase &rewriter, ValueRange newInitOperands,
2573  bool replaceInitOperandUsesInLoop,
2574  const NewYieldValuesFn &newYieldValuesFn) {
2575  // Create a new loop before the existing one, with the extra operands.
2576  OpBuilder::InsertionGuard g(rewriter);
2577  rewriter.setInsertionPoint(getOperation());
2578  auto inits = llvm::to_vector(getInits());
2579  inits.append(newInitOperands.begin(), newInitOperands.end());
2580  AffineForOp newLoop = rewriter.create<AffineForOp>(
2581  getLoc(), getLowerBoundOperands(), getLowerBoundMap(),
2582  getUpperBoundOperands(), getUpperBoundMap(), getStepAsInt(), inits);
2583 
2584  // Generate the new yield values and append them to the scf.yield operation.
2585  auto yieldOp = cast<AffineYieldOp>(getBody()->getTerminator());
2586  ArrayRef<BlockArgument> newIterArgs =
2587  newLoop.getBody()->getArguments().take_back(newInitOperands.size());
2588  {
2589  OpBuilder::InsertionGuard g(rewriter);
2590  rewriter.setInsertionPoint(yieldOp);
2591  SmallVector<Value> newYieldedValues =
2592  newYieldValuesFn(rewriter, getLoc(), newIterArgs);
2593  assert(newInitOperands.size() == newYieldedValues.size() &&
2594  "expected as many new yield values as new iter operands");
2595  rewriter.modifyOpInPlace(yieldOp, [&]() {
2596  yieldOp.getOperandsMutable().append(newYieldedValues);
2597  });
2598  }
2599 
2600  // Move the loop body to the new op.
2601  rewriter.mergeBlocks(getBody(), newLoop.getBody(),
2602  newLoop.getBody()->getArguments().take_front(
2603  getBody()->getNumArguments()));
2604 
2605  if (replaceInitOperandUsesInLoop) {
2606  // Replace all uses of `newInitOperands` with the corresponding basic block
2607  // arguments.
2608  for (auto it : llvm::zip(newInitOperands, newIterArgs)) {
2609  rewriter.replaceUsesWithIf(std::get<0>(it), std::get<1>(it),
2610  [&](OpOperand &use) {
2611  Operation *user = use.getOwner();
2612  return newLoop->isProperAncestor(user);
2613  });
2614  }
2615  }
2616 
2617  // Replace the old loop.
2618  rewriter.replaceOp(getOperation(),
2619  newLoop->getResults().take_front(getNumResults()));
2620  return cast<LoopLikeOpInterface>(newLoop.getOperation());
2621 }
2622 
2623 Speculation::Speculatability AffineForOp::getSpeculatability() {
2624  // `affine.for (I = Start; I < End; I += 1)` terminates for all values of
2625  // Start and End.
2626  //
2627  // For Step != 1, the loop may not terminate. We can add more smarts here if
2628  // needed.
2629  return getStepAsInt() == 1 ? Speculation::RecursivelySpeculatable
2631 }
2632 
2633 /// Returns true if the provided value is the induction variable of a
2634 /// AffineForOp.
2636  return getForInductionVarOwner(val) != AffineForOp();
2637 }
2638 
2640  return getAffineParallelInductionVarOwner(val) != nullptr;
2641 }
2642 
2645 }
2646 
2648  auto ivArg = llvm::dyn_cast<BlockArgument>(val);
2649  if (!ivArg || !ivArg.getOwner() || !ivArg.getOwner()->getParent())
2650  return AffineForOp();
2651  if (auto forOp =
2652  ivArg.getOwner()->getParent()->getParentOfType<AffineForOp>())
2653  // Check to make sure `val` is the induction variable, not an iter_arg.
2654  return forOp.getInductionVar() == val ? forOp : AffineForOp();
2655  return AffineForOp();
2656 }
2657 
2659  auto ivArg = llvm::dyn_cast<BlockArgument>(val);
2660  if (!ivArg || !ivArg.getOwner())
2661  return nullptr;
2662  Operation *containingOp = ivArg.getOwner()->getParentOp();
2663  auto parallelOp = dyn_cast<AffineParallelOp>(containingOp);
2664  if (parallelOp && llvm::is_contained(parallelOp.getIVs(), val))
2665  return parallelOp;
2666  return nullptr;
2667 }
2668 
2669 /// Extracts the induction variables from a list of AffineForOps and returns
2670 /// them.
2672  SmallVectorImpl<Value> *ivs) {
2673  ivs->reserve(forInsts.size());
2674  for (auto forInst : forInsts)
2675  ivs->push_back(forInst.getInductionVar());
2676 }
2677 
2680  ivs.reserve(affineOps.size());
2681  for (Operation *op : affineOps) {
2682  // Add constraints from forOp's bounds.
2683  if (auto forOp = dyn_cast<AffineForOp>(op))
2684  ivs.push_back(forOp.getInductionVar());
2685  else if (auto parallelOp = dyn_cast<AffineParallelOp>(op))
2686  for (size_t i = 0; i < parallelOp.getBody()->getNumArguments(); i++)
2687  ivs.push_back(parallelOp.getBody()->getArgument(i));
2688  }
2689 }
2690 
2691 /// Builds an affine loop nest, using "loopCreatorFn" to create individual loop
2692 /// operations.
2693 template <typename BoundListTy, typename LoopCreatorTy>
2695  OpBuilder &builder, Location loc, BoundListTy lbs, BoundListTy ubs,
2696  ArrayRef<int64_t> steps,
2697  function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilderFn,
2698  LoopCreatorTy &&loopCreatorFn) {
2699  assert(lbs.size() == ubs.size() && "Mismatch in number of arguments");
2700  assert(lbs.size() == steps.size() && "Mismatch in number of arguments");
2701 
2702  // If there are no loops to be constructed, construct the body anyway.
2703  OpBuilder::InsertionGuard guard(builder);
2704  if (lbs.empty()) {
2705  if (bodyBuilderFn)
2706  bodyBuilderFn(builder, loc, ValueRange());
2707  return;
2708  }
2709 
2710  // Create the loops iteratively and store the induction variables.
2712  ivs.reserve(lbs.size());
2713  for (unsigned i = 0, e = lbs.size(); i < e; ++i) {
2714  // Callback for creating the loop body, always creates the terminator.
2715  auto loopBody = [&](OpBuilder &nestedBuilder, Location nestedLoc, Value iv,
2716  ValueRange iterArgs) {
2717  ivs.push_back(iv);
2718  // In the innermost loop, call the body builder.
2719  if (i == e - 1 && bodyBuilderFn) {
2720  OpBuilder::InsertionGuard nestedGuard(nestedBuilder);
2721  bodyBuilderFn(nestedBuilder, nestedLoc, ivs);
2722  }
2723  nestedBuilder.create<AffineYieldOp>(nestedLoc);
2724  };
2725 
2726  // Delegate actual loop creation to the callback in order to dispatch
2727  // between constant- and variable-bound loops.
2728  auto loop = loopCreatorFn(builder, loc, lbs[i], ubs[i], steps[i], loopBody);
2729  builder.setInsertionPointToStart(loop.getBody());
2730  }
2731 }
2732 
2733 /// Creates an affine loop from the bounds known to be constants.
2734 static AffineForOp
2736  int64_t ub, int64_t step,
2737  AffineForOp::BodyBuilderFn bodyBuilderFn) {
2738  return builder.create<AffineForOp>(loc, lb, ub, step,
2739  /*iterArgs=*/std::nullopt, bodyBuilderFn);
2740 }
2741 
2742 /// Creates an affine loop from the bounds that may or may not be constants.
2743 static AffineForOp
2745  int64_t step,
2746  AffineForOp::BodyBuilderFn bodyBuilderFn) {
2747  std::optional<int64_t> lbConst = getConstantIntValue(lb);
2748  std::optional<int64_t> ubConst = getConstantIntValue(ub);
2749  if (lbConst && ubConst)
2750  return buildAffineLoopFromConstants(builder, loc, lbConst.value(),
2751  ubConst.value(), step, bodyBuilderFn);
2752  return builder.create<AffineForOp>(loc, lb, builder.getDimIdentityMap(), ub,
2753  builder.getDimIdentityMap(), step,
2754  /*iterArgs=*/std::nullopt, bodyBuilderFn);
2755 }
2756 
2758  OpBuilder &builder, Location loc, ArrayRef<int64_t> lbs,
2760  function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilderFn) {
2761  buildAffineLoopNestImpl(builder, loc, lbs, ubs, steps, bodyBuilderFn,
2763 }
2764 
2766  OpBuilder &builder, Location loc, ValueRange lbs, ValueRange ubs,
2767  ArrayRef<int64_t> steps,
2768  function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilderFn) {
2769  buildAffineLoopNestImpl(builder, loc, lbs, ubs, steps, bodyBuilderFn,
2771 }
2772 
2773 //===----------------------------------------------------------------------===//
2774 // AffineIfOp
2775 //===----------------------------------------------------------------------===//
2776 
2777 namespace {
2778 /// Remove else blocks that have nothing other than a zero value yield.
2779 struct SimplifyDeadElse : public OpRewritePattern<AffineIfOp> {
2781 
2782  LogicalResult matchAndRewrite(AffineIfOp ifOp,
2783  PatternRewriter &rewriter) const override {
2784  if (ifOp.getElseRegion().empty() ||
2785  !llvm::hasSingleElement(*ifOp.getElseBlock()) || ifOp.getNumResults())
2786  return failure();
2787 
2788  rewriter.startOpModification(ifOp);
2789  rewriter.eraseBlock(ifOp.getElseBlock());
2790  rewriter.finalizeOpModification(ifOp);
2791  return success();
2792  }
2793 };
2794 
2795 /// Removes affine.if cond if the condition is always true or false in certain
2796 /// trivial cases. Promotes the then/else block in the parent operation block.
2797 struct AlwaysTrueOrFalseIf : public OpRewritePattern<AffineIfOp> {
2799 
2800  LogicalResult matchAndRewrite(AffineIfOp op,
2801  PatternRewriter &rewriter) const override {
2802 
2803  auto isTriviallyFalse = [](IntegerSet iSet) {
2804  return iSet.isEmptyIntegerSet();
2805  };
2806 
2807  auto isTriviallyTrue = [](IntegerSet iSet) {
2808  return (iSet.getNumEqualities() == 1 && iSet.getNumInequalities() == 0 &&
2809  iSet.getConstraint(0) == 0);
2810  };
2811 
2812  IntegerSet affineIfConditions = op.getIntegerSet();
2813  Block *blockToMove;
2814  if (isTriviallyFalse(affineIfConditions)) {
2815  // The absence, or equivalently, the emptiness of the else region need not
2816  // be checked when affine.if is returning results because if an affine.if
2817  // operation is returning results, it always has a non-empty else region.
2818  if (op.getNumResults() == 0 && !op.hasElse()) {
2819  // If the else region is absent, or equivalently, empty, remove the
2820  // affine.if operation (which is not returning any results).
2821  rewriter.eraseOp(op);
2822  return success();
2823  }
2824  blockToMove = op.getElseBlock();
2825  } else if (isTriviallyTrue(affineIfConditions)) {
2826  blockToMove = op.getThenBlock();
2827  } else {
2828  return failure();
2829  }
2830  Operation *blockToMoveTerminator = blockToMove->getTerminator();
2831  // Promote the "blockToMove" block to the parent operation block between the
2832  // prologue and epilogue of "op".
2833  rewriter.inlineBlockBefore(blockToMove, op);
2834  // Replace the "op" operation with the operands of the
2835  // "blockToMoveTerminator" operation. Note that "blockToMoveTerminator" is
2836  // the affine.yield operation present in the "blockToMove" block. It has no
2837  // operands when affine.if is not returning results and therefore, in that
2838  // case, replaceOp just erases "op". When affine.if is not returning
2839  // results, the affine.yield operation can be omitted. It gets inserted
2840  // implicitly.
2841  rewriter.replaceOp(op, blockToMoveTerminator->getOperands());
2842  // Erase the "blockToMoveTerminator" operation since it is now in the parent
2843  // operation block, which already has its own terminator.
2844  rewriter.eraseOp(blockToMoveTerminator);
2845  return success();
2846  }
2847 };
2848 } // namespace
2849 
2850 /// AffineIfOp has two regions -- `then` and `else`. The flow of data should be
2851 /// as follows: AffineIfOp -> `then`/`else` -> AffineIfOp
2852 void AffineIfOp::getSuccessorRegions(
2854  // If the predecessor is an AffineIfOp, then branching into both `then` and
2855  // `else` region is valid.
2856  if (point.isParent()) {
2857  regions.reserve(2);
2858  regions.push_back(
2859  RegionSuccessor(&getThenRegion(), getThenRegion().getArguments()));
2860  // If the "else" region is empty, branch bach into parent.
2861  if (getElseRegion().empty()) {
2862  regions.push_back(getResults());
2863  } else {
2864  regions.push_back(
2865  RegionSuccessor(&getElseRegion(), getElseRegion().getArguments()));
2866  }
2867  return;
2868  }
2869 
2870  // If the predecessor is the `else`/`then` region, then branching into parent
2871  // op is valid.
2872  regions.push_back(RegionSuccessor(getResults()));
2873 }
2874 
2875 LogicalResult AffineIfOp::verify() {
2876  // Verify that we have a condition attribute.
2877  // FIXME: This should be specified in the arguments list in ODS.
2878  auto conditionAttr =
2879  (*this)->getAttrOfType<IntegerSetAttr>(getConditionAttrStrName());
2880  if (!conditionAttr)
2881  return emitOpError("requires an integer set attribute named 'condition'");
2882 
2883  // Verify that there are enough operands for the condition.
2884  IntegerSet condition = conditionAttr.getValue();
2885  if (getNumOperands() != condition.getNumInputs())
2886  return emitOpError("operand count and condition integer set dimension and "
2887  "symbol count must match");
2888 
2889  // Verify that the operands are valid dimension/symbols.
2890  if (failed(verifyDimAndSymbolIdentifiers(*this, getOperands(),
2891  condition.getNumDims())))
2892  return failure();
2893 
2894  return success();
2895 }
2896 
2897 ParseResult AffineIfOp::parse(OpAsmParser &parser, OperationState &result) {
2898  // Parse the condition attribute set.
2899  IntegerSetAttr conditionAttr;
2900  unsigned numDims;
2901  if (parser.parseAttribute(conditionAttr,
2902  AffineIfOp::getConditionAttrStrName(),
2903  result.attributes) ||
2904  parseDimAndSymbolList(parser, result.operands, numDims))
2905  return failure();
2906 
2907  // Verify the condition operands.
2908  auto set = conditionAttr.getValue();
2909  if (set.getNumDims() != numDims)
2910  return parser.emitError(
2911  parser.getNameLoc(),
2912  "dim operand count and integer set dim count must match");
2913  if (numDims + set.getNumSymbols() != result.operands.size())
2914  return parser.emitError(
2915  parser.getNameLoc(),
2916  "symbol operand count and integer set symbol count must match");
2917 
2918  if (parser.parseOptionalArrowTypeList(result.types))
2919  return failure();
2920 
2921  // Create the regions for 'then' and 'else'. The latter must be created even
2922  // if it remains empty for the validity of the operation.
2923  result.regions.reserve(2);
2924  Region *thenRegion = result.addRegion();
2925  Region *elseRegion = result.addRegion();
2926 
2927  // Parse the 'then' region.
2928  if (parser.parseRegion(*thenRegion, {}, {}))
2929  return failure();
2930  AffineIfOp::ensureTerminator(*thenRegion, parser.getBuilder(),
2931  result.location);
2932 
2933  // If we find an 'else' keyword then parse the 'else' region.
2934  if (!parser.parseOptionalKeyword("else")) {
2935  if (parser.parseRegion(*elseRegion, {}, {}))
2936  return failure();
2937  AffineIfOp::ensureTerminator(*elseRegion, parser.getBuilder(),
2938  result.location);
2939  }
2940 
2941  // Parse the optional attribute list.
2942  if (parser.parseOptionalAttrDict(result.attributes))
2943  return failure();
2944 
2945  return success();
2946 }
2947 
2949  auto conditionAttr =
2950  (*this)->getAttrOfType<IntegerSetAttr>(getConditionAttrStrName());
2951  p << " " << conditionAttr;
2952  printDimAndSymbolList(operand_begin(), operand_end(),
2953  conditionAttr.getValue().getNumDims(), p);
2954  p.printOptionalArrowTypeList(getResultTypes());
2955  p << ' ';
2956  p.printRegion(getThenRegion(), /*printEntryBlockArgs=*/false,
2957  /*printBlockTerminators=*/getNumResults());
2958 
2959  // Print the 'else' regions if it has any blocks.
2960  auto &elseRegion = this->getElseRegion();
2961  if (!elseRegion.empty()) {
2962  p << " else ";
2963  p.printRegion(elseRegion,
2964  /*printEntryBlockArgs=*/false,
2965  /*printBlockTerminators=*/getNumResults());
2966  }
2967 
2968  // Print the attribute list.
2969  p.printOptionalAttrDict((*this)->getAttrs(),
2970  /*elidedAttrs=*/getConditionAttrStrName());
2971 }
2972 
2973 IntegerSet AffineIfOp::getIntegerSet() {
2974  return (*this)
2975  ->getAttrOfType<IntegerSetAttr>(getConditionAttrStrName())
2976  .getValue();
2977 }
2978 
2979 void AffineIfOp::setIntegerSet(IntegerSet newSet) {
2980  (*this)->setAttr(getConditionAttrStrName(), IntegerSetAttr::get(newSet));
2981 }
2982 
2983 void AffineIfOp::setConditional(IntegerSet set, ValueRange operands) {
2984  setIntegerSet(set);
2985  (*this)->setOperands(operands);
2986 }
2987 
2988 void AffineIfOp::build(OpBuilder &builder, OperationState &result,
2989  TypeRange resultTypes, IntegerSet set, ValueRange args,
2990  bool withElseRegion) {
2991  assert(resultTypes.empty() || withElseRegion);
2992  OpBuilder::InsertionGuard guard(builder);
2993 
2994  result.addTypes(resultTypes);
2995  result.addOperands(args);
2996  result.addAttribute(getConditionAttrStrName(), IntegerSetAttr::get(set));
2997 
2998  Region *thenRegion = result.addRegion();
2999  builder.createBlock(thenRegion);
3000  if (resultTypes.empty())
3001  AffineIfOp::ensureTerminator(*thenRegion, builder, result.location);
3002 
3003  Region *elseRegion = result.addRegion();
3004  if (withElseRegion) {
3005  builder.createBlock(elseRegion);
3006  if (resultTypes.empty())
3007  AffineIfOp::ensureTerminator(*elseRegion, builder, result.location);
3008  }
3009 }
3010 
3011 void AffineIfOp::build(OpBuilder &builder, OperationState &result,
3012  IntegerSet set, ValueRange args, bool withElseRegion) {
3013  AffineIfOp::build(builder, result, /*resultTypes=*/{}, set, args,
3014  withElseRegion);
3015 }
3016 
3017 /// Compose any affine.apply ops feeding into `operands` of the integer set
3018 /// `set` by composing the maps of such affine.apply ops with the integer
3019 /// set constraints.
3021  SmallVectorImpl<Value> &operands) {
3022  // We will simply reuse the API of the map composition by viewing the LHSs of
3023  // the equalities and inequalities of `set` as the affine exprs of an affine
3024  // map. Convert to equivalent map, compose, and convert back to set.
3025  auto map = AffineMap::get(set.getNumDims(), set.getNumSymbols(),
3026  set.getConstraints(), set.getContext());
3027  // Check if any composition is possible.
3028  if (llvm::none_of(operands,
3029  [](Value v) { return v.getDefiningOp<AffineApplyOp>(); }))
3030  return;
3031 
3032  composeAffineMapAndOperands(&map, &operands);
3033  set = IntegerSet::get(map.getNumDims(), map.getNumSymbols(), map.getResults(),
3034  set.getEqFlags());
3035 }
3036 
3037 /// Canonicalize an affine if op's conditional (integer set + operands).
3038 LogicalResult AffineIfOp::fold(FoldAdaptor, SmallVectorImpl<OpFoldResult> &) {
3039  auto set = getIntegerSet();
3040  SmallVector<Value, 4> operands(getOperands());
3041  composeSetAndOperands(set, operands);
3042  canonicalizeSetAndOperands(&set, &operands);
3043 
3044  // Check if the canonicalization or composition led to any change.
3045  if (getIntegerSet() == set && llvm::equal(operands, getOperands()))
3046  return failure();
3047 
3048  setConditional(set, operands);
3049  return success();
3050 }
3051 
3052 void AffineIfOp::getCanonicalizationPatterns(RewritePatternSet &results,
3053  MLIRContext *context) {
3054  results.add<SimplifyDeadElse, AlwaysTrueOrFalseIf>(context);
3055 }
3056 
3057 //===----------------------------------------------------------------------===//
3058 // AffineLoadOp
3059 //===----------------------------------------------------------------------===//
3060 
3061 void AffineLoadOp::build(OpBuilder &builder, OperationState &result,
3062  AffineMap map, ValueRange operands) {
3063  assert(operands.size() == 1 + map.getNumInputs() && "inconsistent operands");
3064  result.addOperands(operands);
3065  if (map)
3066  result.addAttribute(getMapAttrStrName(), AffineMapAttr::get(map));
3067  auto memrefType = llvm::cast<MemRefType>(operands[0].getType());
3068  result.types.push_back(memrefType.getElementType());
3069 }
3070 
3071 void AffineLoadOp::build(OpBuilder &builder, OperationState &result,
3072  Value memref, AffineMap map, ValueRange mapOperands) {
3073  assert(map.getNumInputs() == mapOperands.size() && "inconsistent index info");
3074  result.addOperands(memref);
3075  result.addOperands(mapOperands);
3076  auto memrefType = llvm::cast<MemRefType>(memref.getType());
3077  result.addAttribute(getMapAttrStrName(), AffineMapAttr::get(map));
3078  result.types.push_back(memrefType.getElementType());
3079 }
3080 
3081 void AffineLoadOp::build(OpBuilder &builder, OperationState &result,
3082  Value memref, ValueRange indices) {
3083  auto memrefType = llvm::cast<MemRefType>(memref.getType());
3084  int64_t rank = memrefType.getRank();
3085  // Create identity map for memrefs with at least one dimension or () -> ()
3086  // for zero-dimensional memrefs.
3087  auto map =
3088  rank ? builder.getMultiDimIdentityMap(rank) : builder.getEmptyAffineMap();
3089  build(builder, result, memref, map, indices);
3090 }
3091 
3092 ParseResult AffineLoadOp::parse(OpAsmParser &parser, OperationState &result) {
3093  auto &builder = parser.getBuilder();
3094  auto indexTy = builder.getIndexType();
3095 
3096  MemRefType type;
3097  OpAsmParser::UnresolvedOperand memrefInfo;
3098  AffineMapAttr mapAttr;
3100  return failure(
3101  parser.parseOperand(memrefInfo) ||
3102  parser.parseAffineMapOfSSAIds(mapOperands, mapAttr,
3103  AffineLoadOp::getMapAttrStrName(),
3104  result.attributes) ||
3105  parser.parseOptionalAttrDict(result.attributes) ||
3106  parser.parseColonType(type) ||
3107  parser.resolveOperand(memrefInfo, type, result.operands) ||
3108  parser.resolveOperands(mapOperands, indexTy, result.operands) ||
3109  parser.addTypeToList(type.getElementType(), result.types));
3110 }
3111 
3113  p << " " << getMemRef() << '[';
3114  if (AffineMapAttr mapAttr =
3115  (*this)->getAttrOfType<AffineMapAttr>(getMapAttrStrName()))
3116  p.printAffineMapOfSSAIds(mapAttr, getMapOperands());
3117  p << ']';
3118  p.printOptionalAttrDict((*this)->getAttrs(),
3119  /*elidedAttrs=*/{getMapAttrStrName()});
3120  p << " : " << getMemRefType();
3121 }
3122 
3123 /// Verify common indexing invariants of affine.load, affine.store,
3124 /// affine.vector_load and affine.vector_store.
3125 template <typename AffineMemOpTy>
3126 static LogicalResult
3127 verifyMemoryOpIndexing(AffineMemOpTy op, AffineMapAttr mapAttr,
3128  Operation::operand_range mapOperands,
3129  MemRefType memrefType, unsigned numIndexOperands) {
3130  AffineMap map = mapAttr.getValue();
3131  if (map.getNumResults() != memrefType.getRank())
3132  return op->emitOpError("affine map num results must equal memref rank");
3133  if (map.getNumInputs() != numIndexOperands)
3134  return op->emitOpError("expects as many subscripts as affine map inputs");
3135 
3136  for (auto idx : mapOperands) {
3137  if (!idx.getType().isIndex())
3138  return op->emitOpError("index to load must have 'index' type");
3139  }
3140  if (failed(verifyDimAndSymbolIdentifiers(op, mapOperands, map.getNumDims())))
3141  return failure();
3142 
3143  return success();
3144 }
3145 
3146 LogicalResult AffineLoadOp::verify() {
3147  auto memrefType = getMemRefType();
3148  if (getType() != memrefType.getElementType())
3149  return emitOpError("result type must match element type of memref");
3150 
3151  if (failed(verifyMemoryOpIndexing(
3152  *this, (*this)->getAttrOfType<AffineMapAttr>(getMapAttrStrName()),
3153  getMapOperands(), memrefType,
3154  /*numIndexOperands=*/getNumOperands() - 1)))
3155  return failure();
3156 
3157  return success();
3158 }
3159 
3160 void AffineLoadOp::getCanonicalizationPatterns(RewritePatternSet &results,
3161  MLIRContext *context) {
3162  results.add<SimplifyAffineOp<AffineLoadOp>>(context);
3163 }
3164 
3165 OpFoldResult AffineLoadOp::fold(FoldAdaptor adaptor) {
3166  /// load(memrefcast) -> load
3167  if (succeeded(memref::foldMemRefCast(*this)))
3168  return getResult();
3169 
3170  // Fold load from a global constant memref.
3171  auto getGlobalOp = getMemref().getDefiningOp<memref::GetGlobalOp>();
3172  if (!getGlobalOp)
3173  return {};
3174  // Get to the memref.global defining the symbol.
3175  auto *symbolTableOp = getGlobalOp->getParentWithTrait<OpTrait::SymbolTable>();
3176  if (!symbolTableOp)
3177  return {};
3178  auto global = dyn_cast_or_null<memref::GlobalOp>(
3179  SymbolTable::lookupSymbolIn(symbolTableOp, getGlobalOp.getNameAttr()));
3180  if (!global)
3181  return {};
3182 
3183  // Check if the global memref is a constant.
3184  auto cstAttr =
3185  llvm::dyn_cast_or_null<DenseElementsAttr>(global.getConstantInitValue());
3186  if (!cstAttr)
3187  return {};
3188  // If it's a splat constant, we can fold irrespective of indices.
3189  if (auto splatAttr = llvm::dyn_cast<SplatElementsAttr>(cstAttr))
3190  return splatAttr.getSplatValue<Attribute>();
3191  // Otherwise, we can fold only if we know the indices.
3192  if (!getAffineMap().isConstant())
3193  return {};
3194  auto indices = llvm::to_vector<4>(
3195  llvm::map_range(getAffineMap().getConstantResults(),
3196  [](int64_t v) -> uint64_t { return v; }));
3197  return cstAttr.getValues<Attribute>()[indices];
3198 }
3199 
3200 //===----------------------------------------------------------------------===//
3201 // AffineStoreOp
3202 //===----------------------------------------------------------------------===//
3203 
3204 void AffineStoreOp::build(OpBuilder &builder, OperationState &result,
3205  Value valueToStore, Value memref, AffineMap map,
3206  ValueRange mapOperands) {
3207  assert(map.getNumInputs() == mapOperands.size() && "inconsistent index info");
3208  result.addOperands(valueToStore);
3209  result.addOperands(memref);
3210  result.addOperands(mapOperands);
3211  result.getOrAddProperties<Properties>().map = AffineMapAttr::get(map);
3212 }
3213 
3214 // Use identity map.
3215 void AffineStoreOp::build(OpBuilder &builder, OperationState &result,
3216  Value valueToStore, Value memref,
3217  ValueRange indices) {
3218  auto memrefType = llvm::cast<MemRefType>(memref.getType());
3219  int64_t rank = memrefType.getRank();
3220  // Create identity map for memrefs with at least one dimension or () -> ()
3221  // for zero-dimensional memrefs.
3222  auto map =
3223  rank ? builder.getMultiDimIdentityMap(rank) : builder.getEmptyAffineMap();
3224  build(builder, result, valueToStore, memref, map, indices);
3225 }
3226 
3227 ParseResult AffineStoreOp::parse(OpAsmParser &parser, OperationState &result) {
3228  auto indexTy = parser.getBuilder().getIndexType();
3229 
3230  MemRefType type;
3231  OpAsmParser::UnresolvedOperand storeValueInfo;
3232  OpAsmParser::UnresolvedOperand memrefInfo;
3233  AffineMapAttr mapAttr;
3235  return failure(parser.parseOperand(storeValueInfo) || parser.parseComma() ||
3236  parser.parseOperand(memrefInfo) ||
3237  parser.parseAffineMapOfSSAIds(
3238  mapOperands, mapAttr, AffineStoreOp::getMapAttrStrName(),
3239  result.attributes) ||
3240  parser.parseOptionalAttrDict(result.attributes) ||
3241  parser.parseColonType(type) ||
3242  parser.resolveOperand(storeValueInfo, type.getElementType(),
3243  result.operands) ||
3244  parser.resolveOperand(memrefInfo, type, result.operands) ||
3245  parser.resolveOperands(mapOperands, indexTy, result.operands));
3246 }
3247 
3249  p << " " << getValueToStore();
3250  p << ", " << getMemRef() << '[';
3251  if (AffineMapAttr mapAttr =
3252  (*this)->getAttrOfType<AffineMapAttr>(getMapAttrStrName()))
3253  p.printAffineMapOfSSAIds(mapAttr, getMapOperands());
3254  p << ']';
3255  p.printOptionalAttrDict((*this)->getAttrs(),
3256  /*elidedAttrs=*/{getMapAttrStrName()});
3257  p << " : " << getMemRefType();
3258 }
3259 
3260 LogicalResult AffineStoreOp::verify() {
3261  // The value to store must have the same type as memref element type.
3262  auto memrefType = getMemRefType();
3263  if (getValueToStore().getType() != memrefType.getElementType())
3264  return emitOpError(
3265  "value to store must have the same type as memref element type");
3266 
3267  if (failed(verifyMemoryOpIndexing(
3268  *this, (*this)->getAttrOfType<AffineMapAttr>(getMapAttrStrName()),
3269  getMapOperands(), memrefType,
3270  /*numIndexOperands=*/getNumOperands() - 2)))
3271  return failure();
3272 
3273  return success();
3274 }
3275 
3276 void AffineStoreOp::getCanonicalizationPatterns(RewritePatternSet &results,
3277  MLIRContext *context) {
3278  results.add<SimplifyAffineOp<AffineStoreOp>>(context);
3279 }
3280 
3281 LogicalResult AffineStoreOp::fold(FoldAdaptor adaptor,
3282  SmallVectorImpl<OpFoldResult> &results) {
3283  /// store(memrefcast) -> store
3284  return memref::foldMemRefCast(*this, getValueToStore());
3285 }
3286 
3287 //===----------------------------------------------------------------------===//
3288 // AffineMinMaxOpBase
3289 //===----------------------------------------------------------------------===//
3290 
3291 template <typename T>
3292 static LogicalResult verifyAffineMinMaxOp(T op) {
3293  // Verify that operand count matches affine map dimension and symbol count.
3294  if (op.getNumOperands() !=
3295  op.getMap().getNumDims() + op.getMap().getNumSymbols())
3296  return op.emitOpError(
3297  "operand count and affine map dimension and symbol count must match");
3298 
3299  if (op.getMap().getNumResults() == 0)
3300  return op.emitOpError("affine map expect at least one result");
3301  return success();
3302 }
3303 
3304 template <typename T>
3305 static void printAffineMinMaxOp(OpAsmPrinter &p, T op) {
3306  p << ' ' << op->getAttr(T::getMapAttrStrName());
3307  auto operands = op.getOperands();
3308  unsigned numDims = op.getMap().getNumDims();
3309  p << '(' << operands.take_front(numDims) << ')';
3310 
3311  if (operands.size() != numDims)
3312  p << '[' << operands.drop_front(numDims) << ']';
3313  p.printOptionalAttrDict(op->getAttrs(),
3314  /*elidedAttrs=*/{T::getMapAttrStrName()});
3315 }
3316 
3317 template <typename T>
3318 static ParseResult parseAffineMinMaxOp(OpAsmParser &parser,
3319  OperationState &result) {
3320  auto &builder = parser.getBuilder();
3321  auto indexType = builder.getIndexType();
3324  AffineMapAttr mapAttr;
3325  return failure(
3326  parser.parseAttribute(mapAttr, T::getMapAttrStrName(),
3327  result.attributes) ||
3328  parser.parseOperandList(dimInfos, OpAsmParser::Delimiter::Paren) ||
3329  parser.parseOperandList(symInfos,
3331  parser.parseOptionalAttrDict(result.attributes) ||
3332  parser.resolveOperands(dimInfos, indexType, result.operands) ||
3333  parser.resolveOperands(symInfos, indexType, result.operands) ||
3334  parser.addTypeToList(indexType, result.types));
3335 }
3336 
3337 /// Fold an affine min or max operation with the given operands. The operand
3338 /// list may contain nulls, which are interpreted as the operand not being a
3339 /// constant.
3340 template <typename T>
3342  static_assert(llvm::is_one_of<T, AffineMinOp, AffineMaxOp>::value,
3343  "expected affine min or max op");
3344 
3345  // Fold the affine map.
3346  // TODO: Fold more cases:
3347  // min(some_affine, some_affine + constant, ...), etc.
3348  SmallVector<int64_t, 2> results;
3349  auto foldedMap = op.getMap().partialConstantFold(operands, &results);
3350 
3351  if (foldedMap.getNumSymbols() == 1 && foldedMap.isSymbolIdentity())
3352  return op.getOperand(0);
3353 
3354  // If some of the map results are not constant, try changing the map in-place.
3355  if (results.empty()) {
3356  // If the map is the same, report that folding did not happen.
3357  if (foldedMap == op.getMap())
3358  return {};
3359  op->setAttr("map", AffineMapAttr::get(foldedMap));
3360  return op.getResult();
3361  }
3362 
3363  // Otherwise, completely fold the op into a constant.
3364  auto resultIt = std::is_same<T, AffineMinOp>::value
3365  ? llvm::min_element(results)
3366  : llvm::max_element(results);
3367  if (resultIt == results.end())
3368  return {};
3369  return IntegerAttr::get(IndexType::get(op.getContext()), *resultIt);
3370 }
3371 
3372 /// Remove duplicated expressions in affine min/max ops.
3373 template <typename T>
3376 
3377  LogicalResult matchAndRewrite(T affineOp,
3378  PatternRewriter &rewriter) const override {
3379  AffineMap oldMap = affineOp.getAffineMap();
3380 
3381  SmallVector<AffineExpr, 4> newExprs;
3382  for (AffineExpr expr : oldMap.getResults()) {
3383  // This is a linear scan over newExprs, but it should be fine given that
3384  // we typically just have a few expressions per op.
3385  if (!llvm::is_contained(newExprs, expr))
3386  newExprs.push_back(expr);
3387  }
3388 
3389  if (newExprs.size() == oldMap.getNumResults())
3390  return failure();
3391 
3392  auto newMap = AffineMap::get(oldMap.getNumDims(), oldMap.getNumSymbols(),
3393  newExprs, rewriter.getContext());
3394  rewriter.replaceOpWithNewOp<T>(affineOp, newMap, affineOp.getMapOperands());
3395 
3396  return success();
3397  }
3398 };
3399 
3400 /// Merge an affine min/max op to its consumers if its consumer is also an
3401 /// affine min/max op.
3402 ///
3403 /// This pattern requires the producer affine min/max op is bound to a
3404 /// dimension/symbol that is used as a standalone expression in the consumer
3405 /// affine op's map.
3406 ///
3407 /// For example, a pattern like the following:
3408 ///
3409 /// %0 = affine.min affine_map<()[s0] -> (s0 + 16, s0 * 8)> ()[%sym1]
3410 /// %1 = affine.min affine_map<(d0)[s0] -> (s0 + 4, d0)> (%0)[%sym2]
3411 ///
3412 /// Can be turned into:
3413 ///
3414 /// %1 = affine.min affine_map<
3415 /// ()[s0, s1] -> (s0 + 4, s1 + 16, s1 * 8)> ()[%sym2, %sym1]
3416 template <typename T>
3419 
3420  LogicalResult matchAndRewrite(T affineOp,
3421  PatternRewriter &rewriter) const override {
3422  AffineMap oldMap = affineOp.getAffineMap();
3423  ValueRange dimOperands =
3424  affineOp.getMapOperands().take_front(oldMap.getNumDims());
3425  ValueRange symOperands =
3426  affineOp.getMapOperands().take_back(oldMap.getNumSymbols());
3427 
3428  auto newDimOperands = llvm::to_vector<8>(dimOperands);
3429  auto newSymOperands = llvm::to_vector<8>(symOperands);
3430  SmallVector<AffineExpr, 4> newExprs;
3431  SmallVector<T, 4> producerOps;
3432 
3433  // Go over each expression to see whether it's a single dimension/symbol
3434  // with the corresponding operand which is the result of another affine
3435  // min/max op. If So it can be merged into this affine op.
3436  for (AffineExpr expr : oldMap.getResults()) {
3437  if (auto symExpr = dyn_cast<AffineSymbolExpr>(expr)) {
3438  Value symValue = symOperands[symExpr.getPosition()];
3439  if (auto producerOp = symValue.getDefiningOp<T>()) {
3440  producerOps.push_back(producerOp);
3441  continue;
3442  }
3443  } else if (auto dimExpr = dyn_cast<AffineDimExpr>(expr)) {
3444  Value dimValue = dimOperands[dimExpr.getPosition()];
3445  if (auto producerOp = dimValue.getDefiningOp<T>()) {
3446  producerOps.push_back(producerOp);
3447  continue;
3448  }
3449  }
3450  // For the above cases we will remove the expression by merging the
3451  // producer affine min/max's affine expressions. Otherwise we need to
3452  // keep the existing expression.
3453  newExprs.push_back(expr);
3454  }
3455 
3456  if (producerOps.empty())
3457  return failure();
3458 
3459  unsigned numUsedDims = oldMap.getNumDims();
3460  unsigned numUsedSyms = oldMap.getNumSymbols();
3461 
3462  // Now go over all producer affine ops and merge their expressions.
3463  for (T producerOp : producerOps) {
3464  AffineMap producerMap = producerOp.getAffineMap();
3465  unsigned numProducerDims = producerMap.getNumDims();
3466  unsigned numProducerSyms = producerMap.getNumSymbols();
3467 
3468  // Collect all dimension/symbol values.
3469  ValueRange dimValues =
3470  producerOp.getMapOperands().take_front(numProducerDims);
3471  ValueRange symValues =
3472  producerOp.getMapOperands().take_back(numProducerSyms);
3473  newDimOperands.append(dimValues.begin(), dimValues.end());
3474  newSymOperands.append(symValues.begin(), symValues.end());
3475 
3476  // For expressions we need to shift to avoid overlap.
3477  for (AffineExpr expr : producerMap.getResults()) {
3478  newExprs.push_back(expr.shiftDims(numProducerDims, numUsedDims)
3479  .shiftSymbols(numProducerSyms, numUsedSyms));
3480  }
3481 
3482  numUsedDims += numProducerDims;
3483  numUsedSyms += numProducerSyms;
3484  }
3485 
3486  auto newMap = AffineMap::get(numUsedDims, numUsedSyms, newExprs,
3487  rewriter.getContext());
3488  auto newOperands =
3489  llvm::to_vector<8>(llvm::concat<Value>(newDimOperands, newSymOperands));
3490  rewriter.replaceOpWithNewOp<T>(affineOp, newMap, newOperands);
3491 
3492  return success();
3493  }
3494 };
3495 
3496 /// Canonicalize the result expression order of an affine map and return success
3497 /// if the order changed.
3498 ///
3499 /// The function flattens the map's affine expressions to coefficient arrays and
3500 /// sorts them in lexicographic order. A coefficient array contains a multiplier
3501 /// for every dimension/symbol and a constant term. The canonicalization fails
3502 /// if a result expression is not pure or if the flattening requires local
3503 /// variables that, unlike dimensions and symbols, have no global order.
3504 static LogicalResult canonicalizeMapExprAndTermOrder(AffineMap &map) {
3505  SmallVector<SmallVector<int64_t>> flattenedExprs;
3506  for (const AffineExpr &resultExpr : map.getResults()) {
3507  // Fail if the expression is not pure.
3508  if (!resultExpr.isPureAffine())
3509  return failure();
3510 
3511  SimpleAffineExprFlattener flattener(map.getNumDims(), map.getNumSymbols());
3512  auto flattenResult = flattener.walkPostOrder(resultExpr);
3513  if (failed(flattenResult))
3514  return failure();
3515 
3516  // Fail if the flattened expression has local variables.
3517  if (flattener.operandExprStack.back().size() !=
3518  map.getNumDims() + map.getNumSymbols() + 1)
3519  return failure();
3520 
3521  flattenedExprs.emplace_back(flattener.operandExprStack.back().begin(),
3522  flattener.operandExprStack.back().end());
3523  }
3524 
3525  // Fail if sorting is not necessary.
3526  if (llvm::is_sorted(flattenedExprs))
3527  return failure();
3528 
3529  // Reorder the result expressions according to their flattened form.
3530  SmallVector<unsigned> resultPermutation =
3531  llvm::to_vector(llvm::seq<unsigned>(0, map.getNumResults()));
3532  llvm::sort(resultPermutation, [&](unsigned lhs, unsigned rhs) {
3533  return flattenedExprs[lhs] < flattenedExprs[rhs];
3534  });
3535  SmallVector<AffineExpr> newExprs;
3536  for (unsigned idx : resultPermutation)
3537  newExprs.push_back(map.getResult(idx));
3538 
3539  map = AffineMap::get(map.getNumDims(), map.getNumSymbols(), newExprs,
3540  map.getContext());
3541  return success();
3542 }
3543 
3544 /// Canonicalize the affine map result expression order of an affine min/max
3545 /// operation.
3546 ///
3547 /// The pattern calls `canonicalizeMapExprAndTermOrder` to order the result
3548 /// expressions and replaces the operation if the order changed.
3549 ///
3550 /// For example, the following operation:
3551 ///
3552 /// %0 = affine.min affine_map<(d0, d1) -> (d0 + d1, d1 + 16, 32)> (%i0, %i1)
3553 ///
3554 /// Turns into:
3555 ///
3556 /// %0 = affine.min affine_map<(d0, d1) -> (32, d1 + 16, d0 + d1)> (%i0, %i1)
3557 template <typename T>
3560 
3561  LogicalResult matchAndRewrite(T affineOp,
3562  PatternRewriter &rewriter) const override {
3563  AffineMap map = affineOp.getAffineMap();
3564  if (failed(canonicalizeMapExprAndTermOrder(map)))
3565  return failure();
3566  rewriter.replaceOpWithNewOp<T>(affineOp, map, affineOp.getMapOperands());
3567  return success();
3568  }
3569 };
3570 
3571 template <typename T>
3574 
3575  LogicalResult matchAndRewrite(T affineOp,
3576  PatternRewriter &rewriter) const override {
3577  if (affineOp.getMap().getNumResults() != 1)
3578  return failure();
3579  rewriter.replaceOpWithNewOp<AffineApplyOp>(affineOp, affineOp.getMap(),
3580  affineOp.getOperands());
3581  return success();
3582  }
3583 };
3584 
3585 //===----------------------------------------------------------------------===//
3586 // AffineMinOp
3587 //===----------------------------------------------------------------------===//
3588 //
3589 // %0 = affine.min (d0) -> (1000, d0 + 512) (%i0)
3590 //
3591 
3592 OpFoldResult AffineMinOp::fold(FoldAdaptor adaptor) {
3593  return foldMinMaxOp(*this, adaptor.getOperands());
3594 }
3595 
3596 void AffineMinOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
3597  MLIRContext *context) {
3600  MergeAffineMinMaxOp<AffineMinOp>, SimplifyAffineOp<AffineMinOp>,
3602  context);
3603 }
3604 
3605 LogicalResult AffineMinOp::verify() { return verifyAffineMinMaxOp(*this); }
3606 
3607 ParseResult AffineMinOp::parse(OpAsmParser &parser, OperationState &result) {
3608  return parseAffineMinMaxOp<AffineMinOp>(parser, result);
3609 }
3610 
3611 void AffineMinOp::print(OpAsmPrinter &p) { printAffineMinMaxOp(p, *this); }
3612 
3613 //===----------------------------------------------------------------------===//
3614 // AffineMaxOp
3615 //===----------------------------------------------------------------------===//
3616 //
3617 // %0 = affine.max (d0) -> (1000, d0 + 512) (%i0)
3618 //
3619 
3620 OpFoldResult AffineMaxOp::fold(FoldAdaptor adaptor) {
3621  return foldMinMaxOp(*this, adaptor.getOperands());
3622 }
3623 
3624 void AffineMaxOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
3625  MLIRContext *context) {
3628  MergeAffineMinMaxOp<AffineMaxOp>, SimplifyAffineOp<AffineMaxOp>,
3630  context);
3631 }
3632 
3633 LogicalResult AffineMaxOp::verify() { return verifyAffineMinMaxOp(*this); }
3634 
3635 ParseResult AffineMaxOp::parse(OpAsmParser &parser, OperationState &result) {
3636  return parseAffineMinMaxOp<AffineMaxOp>(parser, result);
3637 }
3638 
3639 void AffineMaxOp::print(OpAsmPrinter &p) { printAffineMinMaxOp(p, *this); }
3640 
3641 //===----------------------------------------------------------------------===//
3642 // AffinePrefetchOp
3643 //===----------------------------------------------------------------------===//
3644 
3645 //
3646 // affine.prefetch %0[%i, %j + 5], read, locality<3>, data : memref<400x400xi32>
3647 //
3648 ParseResult AffinePrefetchOp::parse(OpAsmParser &parser,
3649  OperationState &result) {
3650  auto &builder = parser.getBuilder();
3651  auto indexTy = builder.getIndexType();
3652 
3653  MemRefType type;
3654  OpAsmParser::UnresolvedOperand memrefInfo;
3655  IntegerAttr hintInfo;
3656  auto i32Type = parser.getBuilder().getIntegerType(32);
3657  StringRef readOrWrite, cacheType;
3658 
3659  AffineMapAttr mapAttr;
3661  if (parser.parseOperand(memrefInfo) ||
3662  parser.parseAffineMapOfSSAIds(mapOperands, mapAttr,
3663  AffinePrefetchOp::getMapAttrStrName(),
3664  result.attributes) ||
3665  parser.parseComma() || parser.parseKeyword(&readOrWrite) ||
3666  parser.parseComma() || parser.parseKeyword("locality") ||
3667  parser.parseLess() ||
3668  parser.parseAttribute(hintInfo, i32Type,
3669  AffinePrefetchOp::getLocalityHintAttrStrName(),
3670  result.attributes) ||
3671  parser.parseGreater() || parser.parseComma() ||
3672  parser.parseKeyword(&cacheType) ||
3673  parser.parseOptionalAttrDict(result.attributes) ||
3674  parser.parseColonType(type) ||
3675  parser.resolveOperand(memrefInfo, type, result.operands) ||
3676  parser.resolveOperands(mapOperands, indexTy, result.operands))
3677  return failure();
3678 
3679  if (readOrWrite != "read" && readOrWrite != "write")
3680  return parser.emitError(parser.getNameLoc(),
3681  "rw specifier has to be 'read' or 'write'");
3682  result.addAttribute(AffinePrefetchOp::getIsWriteAttrStrName(),
3683  parser.getBuilder().getBoolAttr(readOrWrite == "write"));
3684 
3685  if (cacheType != "data" && cacheType != "instr")
3686  return parser.emitError(parser.getNameLoc(),
3687  "cache type has to be 'data' or 'instr'");
3688 
3689  result.addAttribute(AffinePrefetchOp::getIsDataCacheAttrStrName(),
3690  parser.getBuilder().getBoolAttr(cacheType == "data"));
3691 
3692  return success();
3693 }
3694 
3696  p << " " << getMemref() << '[';
3697  AffineMapAttr mapAttr =
3698  (*this)->getAttrOfType<AffineMapAttr>(getMapAttrStrName());
3699  if (mapAttr)
3700  p.printAffineMapOfSSAIds(mapAttr, getMapOperands());
3701  p << ']' << ", " << (getIsWrite() ? "write" : "read") << ", "
3702  << "locality<" << getLocalityHint() << ">, "
3703  << (getIsDataCache() ? "data" : "instr");
3705  (*this)->getAttrs(),
3706  /*elidedAttrs=*/{getMapAttrStrName(), getLocalityHintAttrStrName(),
3707  getIsDataCacheAttrStrName(), getIsWriteAttrStrName()});
3708  p << " : " << getMemRefType();
3709 }
3710 
3711 LogicalResult AffinePrefetchOp::verify() {
3712  auto mapAttr = (*this)->getAttrOfType<AffineMapAttr>(getMapAttrStrName());
3713  if (mapAttr) {
3714  AffineMap map = mapAttr.getValue();
3715  if (map.getNumResults() != getMemRefType().getRank())
3716  return emitOpError("affine.prefetch affine map num results must equal"
3717  " memref rank");
3718  if (map.getNumInputs() + 1 != getNumOperands())
3719  return emitOpError("too few operands");
3720  } else {
3721  if (getNumOperands() != 1)
3722  return emitOpError("too few operands");
3723  }
3724 
3725  Region *scope = getAffineScope(*this);
3726  for (auto idx : getMapOperands()) {
3727  if (!isValidAffineIndexOperand(idx, scope))
3728  return emitOpError(
3729  "index must be a valid dimension or symbol identifier");
3730  }
3731  return success();
3732 }
3733 
3734 void AffinePrefetchOp::getCanonicalizationPatterns(RewritePatternSet &results,
3735  MLIRContext *context) {
3736  // prefetch(memrefcast) -> prefetch
3737  results.add<SimplifyAffineOp<AffinePrefetchOp>>(context);
3738 }
3739 
3740 LogicalResult AffinePrefetchOp::fold(FoldAdaptor adaptor,
3741  SmallVectorImpl<OpFoldResult> &results) {
3742  /// prefetch(memrefcast) -> prefetch
3743  return memref::foldMemRefCast(*this);
3744 }
3745 
3746 //===----------------------------------------------------------------------===//
3747 // AffineParallelOp
3748 //===----------------------------------------------------------------------===//
3749 
3750 void AffineParallelOp::build(OpBuilder &builder, OperationState &result,
3751  TypeRange resultTypes,
3752  ArrayRef<arith::AtomicRMWKind> reductions,
3753  ArrayRef<int64_t> ranges) {
3754  SmallVector<AffineMap> lbs(ranges.size(), builder.getConstantAffineMap(0));
3755  auto ubs = llvm::to_vector<4>(llvm::map_range(ranges, [&](int64_t value) {
3756  return builder.getConstantAffineMap(value);
3757  }));
3758  SmallVector<int64_t> steps(ranges.size(), 1);
3759  build(builder, result, resultTypes, reductions, lbs, /*lbArgs=*/{}, ubs,
3760  /*ubArgs=*/{}, steps);
3761 }
3762 
3763 void AffineParallelOp::build(OpBuilder &builder, OperationState &result,
3764  TypeRange resultTypes,
3765  ArrayRef<arith::AtomicRMWKind> reductions,
3766  ArrayRef<AffineMap> lbMaps, ValueRange lbArgs,
3767  ArrayRef<AffineMap> ubMaps, ValueRange ubArgs,
3768  ArrayRef<int64_t> steps) {
3769  assert(llvm::all_of(lbMaps,
3770  [lbMaps](AffineMap m) {
3771  return m.getNumDims() == lbMaps[0].getNumDims() &&
3772  m.getNumSymbols() == lbMaps[0].getNumSymbols();
3773  }) &&
3774  "expected all lower bounds maps to have the same number of dimensions "
3775  "and symbols");
3776  assert(llvm::all_of(ubMaps,
3777  [ubMaps](AffineMap m) {
3778  return m.getNumDims() == ubMaps[0].getNumDims() &&
3779  m.getNumSymbols() == ubMaps[0].getNumSymbols();
3780  }) &&
3781  "expected all upper bounds maps to have the same number of dimensions "
3782  "and symbols");
3783  assert((lbMaps.empty() || lbMaps[0].getNumInputs() == lbArgs.size()) &&
3784  "expected lower bound maps to have as many inputs as lower bound "
3785  "operands");
3786  assert((ubMaps.empty() || ubMaps[0].getNumInputs() == ubArgs.size()) &&
3787  "expected upper bound maps to have as many inputs as upper bound "
3788  "operands");
3789 
3790  OpBuilder::InsertionGuard guard(builder);
3791  result.addTypes(resultTypes);
3792 
3793  // Convert the reductions to integer attributes.
3794  SmallVector<Attribute, 4> reductionAttrs;
3795  for (arith::AtomicRMWKind reduction : reductions)
3796  reductionAttrs.push_back(
3797  builder.getI64IntegerAttr(static_cast<int64_t>(reduction)));
3798  result.addAttribute(getReductionsAttrStrName(),
3799  builder.getArrayAttr(reductionAttrs));
3800 
3801  // Concatenates maps defined in the same input space (same dimensions and
3802  // symbols), assumes there is at least one map.
3803  auto concatMapsSameInput = [&builder](ArrayRef<AffineMap> maps,
3804  SmallVectorImpl<int32_t> &groups) {
3805  if (maps.empty())
3806  return AffineMap::get(builder.getContext());
3808  groups.reserve(groups.size() + maps.size());
3809  exprs.reserve(maps.size());
3810  for (AffineMap m : maps) {
3811  llvm::append_range(exprs, m.getResults());
3812  groups.push_back(m.getNumResults());
3813  }
3814  return AffineMap::get(maps[0].getNumDims(), maps[0].getNumSymbols(), exprs,
3815  maps[0].getContext());
3816  };
3817 
3818  // Set up the bounds.
3819  SmallVector<int32_t> lbGroups, ubGroups;
3820  AffineMap lbMap = concatMapsSameInput(lbMaps, lbGroups);
3821  AffineMap ubMap = concatMapsSameInput(ubMaps, ubGroups);
3822  result.addAttribute(getLowerBoundsMapAttrStrName(),
3823  AffineMapAttr::get(lbMap));
3824  result.addAttribute(getLowerBoundsGroupsAttrStrName(),
3825  builder.getI32TensorAttr(lbGroups));
3826  result.addAttribute(getUpperBoundsMapAttrStrName(),
3827  AffineMapAttr::get(ubMap));
3828  result.addAttribute(getUpperBoundsGroupsAttrStrName(),
3829  builder.getI32TensorAttr(ubGroups));
3830  result.addAttribute(getStepsAttrStrName(), builder.getI64ArrayAttr(steps));
3831  result.addOperands(lbArgs);
3832  result.addOperands(ubArgs);
3833 
3834  // Create a region and a block for the body.
3835  auto *bodyRegion = result.addRegion();
3836  Block *body = builder.createBlock(bodyRegion);
3837 
3838  // Add all the block arguments.
3839  for (unsigned i = 0, e = steps.size(); i < e; ++i)
3840  body->addArgument(IndexType::get(builder.getContext()), result.location);
3841  if (resultTypes.empty())
3842  ensureTerminator(*bodyRegion, builder, result.location);
3843 }
3844 
3845 SmallVector<Region *> AffineParallelOp::getLoopRegions() {
3846  return {&getRegion()};
3847 }
3848 
3849 unsigned AffineParallelOp::getNumDims() { return getSteps().size(); }
3850 
3851 AffineParallelOp::operand_range AffineParallelOp::getLowerBoundsOperands() {
3852  return getOperands().take_front(getLowerBoundsMap().getNumInputs());
3853 }
3854 
3855 AffineParallelOp::operand_range AffineParallelOp::getUpperBoundsOperands() {
3856  return getOperands().drop_front(getLowerBoundsMap().getNumInputs());
3857 }
3858 
3859 AffineMap AffineParallelOp::getLowerBoundMap(unsigned pos) {
3860  auto values = getLowerBoundsGroups().getValues<int32_t>();
3861  unsigned start = 0;
3862  for (unsigned i = 0; i < pos; ++i)
3863  start += values[i];
3864  return getLowerBoundsMap().getSliceMap(start, values[pos]);
3865 }
3866 
3867 AffineMap AffineParallelOp::getUpperBoundMap(unsigned pos) {
3868  auto values = getUpperBoundsGroups().getValues<int32_t>();
3869  unsigned start = 0;
3870  for (unsigned i = 0; i < pos; ++i)
3871  start += values[i];
3872  return getUpperBoundsMap().getSliceMap(start, values[pos]);
3873 }
3874 
3875 AffineValueMap AffineParallelOp::getLowerBoundsValueMap() {
3876  return AffineValueMap(getLowerBoundsMap(), getLowerBoundsOperands());
3877 }
3878 
3879 AffineValueMap AffineParallelOp::getUpperBoundsValueMap() {
3880  return AffineValueMap(getUpperBoundsMap(), getUpperBoundsOperands());
3881 }
3882 
3883 std::optional<SmallVector<int64_t, 8>> AffineParallelOp::getConstantRanges() {
3884  if (hasMinMaxBounds())
3885  return std::nullopt;
3886 
3887  // Try to convert all the ranges to constant expressions.
3889  AffineValueMap rangesValueMap;
3890  AffineValueMap::difference(getUpperBoundsValueMap(), getLowerBoundsValueMap(),
3891  &rangesValueMap);
3892  out.reserve(rangesValueMap.getNumResults());
3893  for (unsigned i = 0, e = rangesValueMap.getNumResults(); i < e; ++i) {
3894  auto expr = rangesValueMap.getResult(i);
3895  auto cst = dyn_cast<AffineConstantExpr>(expr);
3896  if (!cst)
3897  return std::nullopt;
3898  out.push_back(cst.getValue());
3899  }
3900  return out;
3901 }
3902 
3903 Block *AffineParallelOp::getBody() { return &getRegion().front(); }
3904 
3905 OpBuilder AffineParallelOp::getBodyBuilder() {
3906  return OpBuilder(getBody(), std::prev(getBody()->end()));
3907 }
3908 
3909 void AffineParallelOp::setLowerBounds(ValueRange lbOperands, AffineMap map) {
3910  assert(lbOperands.size() == map.getNumInputs() &&
3911  "operands to map must match number of inputs");
3912 
3913  auto ubOperands = getUpperBoundsOperands();
3914 
3915  SmallVector<Value, 4> newOperands(lbOperands);
3916  newOperands.append(ubOperands.begin(), ubOperands.end());
3917  (*this)->setOperands(newOperands);
3918 
3919  setLowerBoundsMapAttr(AffineMapAttr::get(map));
3920 }
3921 
3922 void AffineParallelOp::setUpperBounds(ValueRange ubOperands, AffineMap map) {
3923  assert(ubOperands.size() == map.getNumInputs() &&
3924  "operands to map must match number of inputs");
3925 
3926  SmallVector<Value, 4> newOperands(getLowerBoundsOperands());
3927  newOperands.append(ubOperands.begin(), ubOperands.end());
3928  (*this)->setOperands(newOperands);
3929 
3930  setUpperBoundsMapAttr(AffineMapAttr::get(map));
3931 }
3932 
3933 void AffineParallelOp::setSteps(ArrayRef<int64_t> newSteps) {
3934  setStepsAttr(getBodyBuilder().getI64ArrayAttr(newSteps));
3935 }
3936 
3937 // check whether resultType match op or not in affine.parallel
3938 static bool isResultTypeMatchAtomicRMWKind(Type resultType,
3939  arith::AtomicRMWKind op) {
3940  switch (op) {
3941  case arith::AtomicRMWKind::addf:
3942  return isa<FloatType>(resultType);
3943  case arith::AtomicRMWKind::addi:
3944  return isa<IntegerType>(resultType);
3945  case arith::AtomicRMWKind::assign:
3946  return true;
3947  case arith::AtomicRMWKind::mulf:
3948  return isa<FloatType>(resultType);
3949  case arith::AtomicRMWKind::muli:
3950  return isa<IntegerType>(resultType);
3951  case arith::AtomicRMWKind::maximumf:
3952  return isa<FloatType>(resultType);
3953  case arith::AtomicRMWKind::minimumf:
3954  return isa<FloatType>(resultType);
3955  case arith::AtomicRMWKind::maxs: {
3956  auto intType = llvm::dyn_cast<IntegerType>(resultType);
3957  return intType && intType.isSigned();
3958  }
3959  case arith::AtomicRMWKind::mins: {
3960  auto intType = llvm::dyn_cast<IntegerType>(resultType);
3961  return intType && intType.isSigned();
3962  }
3963  case arith::AtomicRMWKind::maxu: {
3964  auto intType = llvm::dyn_cast<IntegerType>(resultType);
3965  return intType && intType.isUnsigned();
3966  }
3967  case arith::AtomicRMWKind::minu: {
3968  auto intType = llvm::dyn_cast<IntegerType>(resultType);
3969  return intType && intType.isUnsigned();
3970  }
3971  case arith::AtomicRMWKind::ori:
3972  return isa<IntegerType>(resultType);
3973  case arith::AtomicRMWKind::andi:
3974  return isa<IntegerType>(resultType);
3975  default:
3976  return false;
3977  }
3978 }
3979 
3980 LogicalResult AffineParallelOp::verify() {
3981  auto numDims = getNumDims();
3982  if (getLowerBoundsGroups().getNumElements() != numDims ||
3983  getUpperBoundsGroups().getNumElements() != numDims ||
3984  getSteps().size() != numDims || getBody()->getNumArguments() != numDims) {
3985  return emitOpError() << "the number of region arguments ("
3986  << getBody()->getNumArguments()
3987  << ") and the number of map groups for lower ("
3988  << getLowerBoundsGroups().getNumElements()
3989  << ") and upper bound ("
3990  << getUpperBoundsGroups().getNumElements()
3991  << "), and the number of steps (" << getSteps().size()
3992  << ") must all match";
3993  }
3994 
3995  unsigned expectedNumLBResults = 0;
3996  for (APInt v : getLowerBoundsGroups()) {
3997  unsigned results = v.getZExtValue();
3998  if (results == 0)
3999  return emitOpError()
4000  << "expected lower bound map to have at least one result";
4001  expectedNumLBResults += results;
4002  }
4003  if (expectedNumLBResults != getLowerBoundsMap().getNumResults())
4004  return emitOpError() << "expected lower bounds map to have "
4005  << expectedNumLBResults << " results";
4006  unsigned expectedNumUBResults = 0;
4007  for (APInt v : getUpperBoundsGroups()) {
4008  unsigned results = v.getZExtValue();
4009  if (results == 0)
4010  return emitOpError()
4011  << "expected upper bound map to have at least one result";
4012  expectedNumUBResults += results;
4013  }
4014  if (expectedNumUBResults != getUpperBoundsMap().getNumResults())
4015  return emitOpError() << "expected upper bounds map to have "
4016  << expectedNumUBResults << " results";
4017 
4018  if (getReductions().size() != getNumResults())
4019  return emitOpError("a reduction must be specified for each output");
4020 
4021  // Verify reduction ops are all valid and each result type matches reduction
4022  // ops
4023  for (auto it : llvm::enumerate((getReductions()))) {
4024  Attribute attr = it.value();
4025  auto intAttr = llvm::dyn_cast<IntegerAttr>(attr);
4026  if (!intAttr || !arith::symbolizeAtomicRMWKind(intAttr.getInt()))
4027  return emitOpError("invalid reduction attribute");
4028  auto kind = arith::symbolizeAtomicRMWKind(intAttr.getInt()).value();
4029  if (!isResultTypeMatchAtomicRMWKind(getResult(it.index()).getType(), kind))
4030  return emitOpError("result type cannot match reduction attribute");
4031  }
4032 
4033  // Verify that the bound operands are valid dimension/symbols.
4034  /// Lower bounds.
4035  if (failed(verifyDimAndSymbolIdentifiers(*this, getLowerBoundsOperands(),
4036  getLowerBoundsMap().getNumDims())))
4037  return failure();
4038  /// Upper bounds.
4039  if (failed(verifyDimAndSymbolIdentifiers(*this, getUpperBoundsOperands(),
4040  getUpperBoundsMap().getNumDims())))
4041  return failure();
4042  return success();
4043 }
4044 
4045 LogicalResult AffineValueMap::canonicalize() {
4046  SmallVector<Value, 4> newOperands{operands};
4047  auto newMap = getAffineMap();
4048  composeAffineMapAndOperands(&newMap, &newOperands);
4049  if (newMap == getAffineMap() && newOperands == operands)
4050  return failure();
4051  reset(newMap, newOperands);
4052  return success();
4053 }
4054 
4055 /// Canonicalize the bounds of the given loop.
4056 static LogicalResult canonicalizeLoopBounds(AffineParallelOp op) {
4057  AffineValueMap lb = op.getLowerBoundsValueMap();
4058  bool lbCanonicalized = succeeded(lb.canonicalize());
4059 
4060  AffineValueMap ub = op.getUpperBoundsValueMap();
4061  bool ubCanonicalized = succeeded(ub.canonicalize());
4062 
4063  // Any canonicalization change always leads to updated map(s).
4064  if (!lbCanonicalized && !ubCanonicalized)
4065  return failure();
4066 
4067  if (lbCanonicalized)
4068  op.setLowerBounds(lb.getOperands(), lb.getAffineMap());
4069  if (ubCanonicalized)
4070  op.setUpperBounds(ub.getOperands(), ub.getAffineMap());
4071 
4072  return success();
4073 }
4074 
4075 LogicalResult AffineParallelOp::fold(FoldAdaptor adaptor,
4076  SmallVectorImpl<OpFoldResult> &results) {
4077  return canonicalizeLoopBounds(*this);
4078 }
4079 
4080 /// Prints a lower(upper) bound of an affine parallel loop with max(min)
4081 /// conditions in it. `mapAttr` is a flat list of affine expressions and `group`
4082 /// identifies which of the those expressions form max/min groups. `operands`
4083 /// are the SSA values of dimensions and symbols and `keyword` is either "min"
4084 /// or "max".
4085 static void printMinMaxBound(OpAsmPrinter &p, AffineMapAttr mapAttr,
4086  DenseIntElementsAttr group, ValueRange operands,
4087  StringRef keyword) {
4088  AffineMap map = mapAttr.getValue();
4089  unsigned numDims = map.getNumDims();
4090  ValueRange dimOperands = operands.take_front(numDims);
4091  ValueRange symOperands = operands.drop_front(numDims);
4092  unsigned start = 0;
4093  for (llvm::APInt groupSize : group) {
4094  if (start != 0)
4095  p << ", ";
4096 
4097  unsigned size = groupSize.getZExtValue();
4098  if (size == 1) {
4099  p.printAffineExprOfSSAIds(map.getResult(start), dimOperands, symOperands);
4100  ++start;
4101  } else {
4102  p << keyword << '(';
4103  AffineMap submap = map.getSliceMap(start, size);
4104  p.printAffineMapOfSSAIds(AffineMapAttr::get(submap), operands);
4105  p << ')';
4106  start += size;
4107  }
4108  }
4109 }
4110 
4112  p << " (" << getBody()->getArguments() << ") = (";
4113  printMinMaxBound(p, getLowerBoundsMapAttr(), getLowerBoundsGroupsAttr(),
4114  getLowerBoundsOperands(), "max");
4115  p << ") to (";
4116  printMinMaxBound(p, getUpperBoundsMapAttr(), getUpperBoundsGroupsAttr(),
4117  getUpperBoundsOperands(), "min");
4118  p << ')';
4119  SmallVector<int64_t, 8> steps = getSteps();
4120  bool elideSteps = llvm::all_of(steps, [](int64_t step) { return step == 1; });
4121  if (!elideSteps) {
4122  p << " step (";
4123  llvm::interleaveComma(steps, p);
4124  p << ')';
4125  }
4126  if (getNumResults()) {
4127  p << " reduce (";
4128  llvm::interleaveComma(getReductions(), p, [&](auto &attr) {
4129  arith::AtomicRMWKind sym = *arith::symbolizeAtomicRMWKind(
4130  llvm::cast<IntegerAttr>(attr).getInt());
4131  p << "\"" << arith::stringifyAtomicRMWKind(sym) << "\"";
4132  });
4133  p << ") -> (" << getResultTypes() << ")";
4134  }
4135 
4136  p << ' ';
4137  p.printRegion(getRegion(), /*printEntryBlockArgs=*/false,
4138  /*printBlockTerminators=*/getNumResults());
4140  (*this)->getAttrs(),
4141  /*elidedAttrs=*/{AffineParallelOp::getReductionsAttrStrName(),
4142  AffineParallelOp::getLowerBoundsMapAttrStrName(),
4143  AffineParallelOp::getLowerBoundsGroupsAttrStrName(),
4144  AffineParallelOp::getUpperBoundsMapAttrStrName(),
4145  AffineParallelOp::getUpperBoundsGroupsAttrStrName(),
4146  AffineParallelOp::getStepsAttrStrName()});
4147 }
4148 
4149 /// Given a list of lists of parsed operands, populates `uniqueOperands` with
4150 /// unique operands. Also populates `replacements with affine expressions of
4151 /// `kind` that can be used to update affine maps previously accepting a
4152 /// `operands` to accept `uniqueOperands` instead.
4154  OpAsmParser &parser,
4156  SmallVectorImpl<Value> &uniqueOperands,
4159  "expected operands to be dim or symbol expression");
4160 
4161  Type indexType = parser.getBuilder().getIndexType();
4162  for (const auto &list : operands) {
4163  SmallVector<Value> valueOperands;
4164  if (parser.resolveOperands(list, indexType, valueOperands))
4165  return failure();
4166  for (Value operand : valueOperands) {
4167  unsigned pos = std::distance(uniqueOperands.begin(),
4168  llvm::find(uniqueOperands, operand));
4169  if (pos == uniqueOperands.size())
4170  uniqueOperands.push_back(operand);
4171  replacements.push_back(
4173  ? getAffineDimExpr(pos, parser.getContext())
4174  : getAffineSymbolExpr(pos, parser.getContext()));
4175  }
4176  }
4177  return success();
4178 }
4179 
4180 namespace {
4181 enum class MinMaxKind { Min, Max };
4182 } // namespace
4183 
4184 /// Parses an affine map that can contain a min/max for groups of its results,
4185 /// e.g., max(expr-1, expr-2), expr-3, max(expr-4, expr-5, expr-6). Populates
4186 /// `result` attributes with the map (flat list of expressions) and the grouping
4187 /// (list of integers that specify how many expressions to put into each
4188 /// min/max) attributes. Deduplicates repeated operands.
4189 ///
4190 /// parallel-bound ::= `(` parallel-group-list `)`
4191 /// parallel-group-list ::= parallel-group (`,` parallel-group-list)?
4192 /// parallel-group ::= simple-group | min-max-group
4193 /// simple-group ::= expr-of-ssa-ids
4194 /// min-max-group ::= ( `min` | `max` ) `(` expr-of-ssa-ids-list `)`
4195 /// expr-of-ssa-ids-list ::= expr-of-ssa-ids (`,` expr-of-ssa-id-list)?
4196 ///
4197 /// Examples:
4198 /// (%0, min(%1 + %2, %3), %4, min(%5 floordiv 32, %6))
4199 /// (%0, max(%1 - 2 * %2))
4200 static ParseResult parseAffineMapWithMinMax(OpAsmParser &parser,
4201  OperationState &result,
4202  MinMaxKind kind) {
4203  // Using `const` not `constexpr` below to workaround a MSVC optimizer bug,
4204  // see: https://reviews.llvm.org/D134227#3821753
4205  const llvm::StringLiteral tmpAttrStrName = "__pseudo_bound_map";
4206 
4207  StringRef mapName = kind == MinMaxKind::Min
4208  ? AffineParallelOp::getUpperBoundsMapAttrStrName()
4209  : AffineParallelOp::getLowerBoundsMapAttrStrName();
4210  StringRef groupsName =
4211  kind == MinMaxKind::Min
4212  ? AffineParallelOp::getUpperBoundsGroupsAttrStrName()
4213  : AffineParallelOp::getLowerBoundsGroupsAttrStrName();
4214 
4215  if (failed(parser.parseLParen()))
4216  return failure();
4217 
4218  if (succeeded(parser.parseOptionalRParen())) {
4219  result.addAttribute(
4220  mapName, AffineMapAttr::get(parser.getBuilder().getEmptyAffineMap()));
4221  result.addAttribute(groupsName, parser.getBuilder().getI32TensorAttr({}));
4222  return success();
4223  }
4224 
4225  SmallVector<AffineExpr> flatExprs;
4228  SmallVector<int32_t> numMapsPerGroup;
4230  auto parseOperands = [&]() {
4231  if (succeeded(parser.parseOptionalKeyword(
4232  kind == MinMaxKind::Min ? "min" : "max"))) {
4233  mapOperands.clear();
4234  AffineMapAttr map;
4235  if (failed(parser.parseAffineMapOfSSAIds(mapOperands, map, tmpAttrStrName,
4236  result.attributes,
4238  return failure();
4239  result.attributes.erase(tmpAttrStrName);
4240  llvm::append_range(flatExprs, map.getValue().getResults());
4241  auto operandsRef = llvm::ArrayRef(mapOperands);
4242  auto dimsRef = operandsRef.take_front(map.getValue().getNumDims());
4244  auto symsRef = operandsRef.drop_front(map.getValue().getNumDims());
4246  flatDimOperands.append(map.getValue().getNumResults(), dims);
4247  flatSymOperands.append(map.getValue().getNumResults(), syms);
4248  numMapsPerGroup.push_back(map.getValue().getNumResults());
4249  } else {
4250  if (failed(parser.parseAffineExprOfSSAIds(flatDimOperands.emplace_back(),
4251  flatSymOperands.emplace_back(),
4252  flatExprs.emplace_back())))
4253  return failure();
4254  numMapsPerGroup.push_back(1);
4255  }
4256  return success();
4257  };
4258  if (parser.parseCommaSeparatedList(parseOperands) || parser.parseRParen())
4259  return failure();
4260 
4261  unsigned totalNumDims = 0;
4262  unsigned totalNumSyms = 0;
4263  for (unsigned i = 0, e = flatExprs.size(); i < e; ++i) {
4264  unsigned numDims = flatDimOperands[i].size();
4265  unsigned numSyms = flatSymOperands[i].size();
4266  flatExprs[i] = flatExprs[i]
4267  .shiftDims(numDims, totalNumDims)
4268  .shiftSymbols(numSyms, totalNumSyms);
4269  totalNumDims += numDims;
4270  totalNumSyms += numSyms;
4271  }
4272 
4273  // Deduplicate map operands.
4274  SmallVector<Value> dimOperands, symOperands;
4275  SmallVector<AffineExpr> dimRplacements, symRepacements;
4276  if (deduplicateAndResolveOperands(parser, flatDimOperands, dimOperands,
4277  dimRplacements, AffineExprKind::DimId) ||
4278  deduplicateAndResolveOperands(parser, flatSymOperands, symOperands,
4279  symRepacements, AffineExprKind::SymbolId))
4280  return failure();
4281 
4282  result.operands.append(dimOperands.begin(), dimOperands.end());
4283  result.operands.append(symOperands.begin(), symOperands.end());
4284 
4285  Builder &builder = parser.getBuilder();
4286  auto flatMap = AffineMap::get(totalNumDims, totalNumSyms, flatExprs,
4287  parser.getContext());
4288  flatMap = flatMap.replaceDimsAndSymbols(
4289  dimRplacements, symRepacements, dimOperands.size(), symOperands.size());
4290 
4291  result.addAttribute(mapName, AffineMapAttr::get(flatMap));
4292  result.addAttribute(groupsName, builder.getI32TensorAttr(numMapsPerGroup));
4293  return success();
4294 }
4295 
4296 //
4297 // operation ::= `affine.parallel` `(` ssa-ids `)` `=` parallel-bound
4298 // `to` parallel-bound steps? region attr-dict?
4299 // steps ::= `steps` `(` integer-literals `)`
4300 //
4301 ParseResult AffineParallelOp::parse(OpAsmParser &parser,
4302  OperationState &result) {
4303  auto &builder = parser.getBuilder();
4304  auto indexType = builder.getIndexType();
4307  parser.parseEqual() ||
4308  parseAffineMapWithMinMax(parser, result, MinMaxKind::Max) ||
4309  parser.parseKeyword("to") ||
4310  parseAffineMapWithMinMax(parser, result, MinMaxKind::Min))
4311  return failure();
4312 
4313  AffineMapAttr stepsMapAttr;
4314  NamedAttrList stepsAttrs;
4316  if (failed(parser.parseOptionalKeyword("step"))) {
4317  SmallVector<int64_t, 4> steps(ivs.size(), 1);
4318  result.addAttribute(AffineParallelOp::getStepsAttrStrName(),
4319  builder.getI64ArrayAttr(steps));
4320  } else {
4321  if (parser.parseAffineMapOfSSAIds(stepsMapOperands, stepsMapAttr,
4322  AffineParallelOp::getStepsAttrStrName(),
4323  stepsAttrs,
4325  return failure();
4326 
4327  // Convert steps from an AffineMap into an I64ArrayAttr.
4329  auto stepsMap = stepsMapAttr.getValue();
4330  for (const auto &result : stepsMap.getResults()) {
4331  auto constExpr = dyn_cast<AffineConstantExpr>(result);
4332  if (!constExpr)
4333  return parser.emitError(parser.getNameLoc(),
4334  "steps must be constant integers");
4335  steps.push_back(constExpr.getValue());
4336  }
4337  result.addAttribute(AffineParallelOp::getStepsAttrStrName(),
4338  builder.getI64ArrayAttr(steps));
4339  }
4340 
4341  // Parse optional clause of the form: `reduce ("addf", "maxf")`, where the
4342  // quoted strings are a member of the enum AtomicRMWKind.
4343  SmallVector<Attribute, 4> reductions;
4344  if (succeeded(parser.parseOptionalKeyword("reduce"))) {
4345  if (parser.parseLParen())
4346  return failure();
4347  auto parseAttributes = [&]() -> ParseResult {
4348  // Parse a single quoted string via the attribute parsing, and then
4349  // verify it is a member of the enum and convert to it's integer
4350  // representation.
4351  StringAttr attrVal;
4352  NamedAttrList attrStorage;
4353  auto loc = parser.getCurrentLocation();
4354  if (parser.parseAttribute(attrVal, builder.getNoneType(), "reduce",
4355  attrStorage))
4356  return failure();
4357  std::optional<arith::AtomicRMWKind> reduction =
4358  arith::symbolizeAtomicRMWKind(attrVal.getValue());
4359  if (!reduction)
4360  return parser.emitError(loc, "invalid reduction value: ") << attrVal;
4361  reductions.push_back(
4362  builder.getI64IntegerAttr(static_cast<int64_t>(reduction.value())));
4363  // While we keep getting commas, keep parsing.
4364  return success();
4365  };
4366  if (parser.parseCommaSeparatedList(parseAttributes) || parser.parseRParen())
4367  return failure();
4368  }
4369  result.addAttribute(AffineParallelOp::getReductionsAttrStrName(),
4370  builder.getArrayAttr(reductions));
4371 
4372  // Parse return types of reductions (if any)
4373  if (parser.parseOptionalArrowTypeList(result.types))
4374  return failure();
4375 
4376  // Now parse the body.
4377  Region *body = result.addRegion();
4378  for (auto &iv : ivs)
4379  iv.type = indexType;
4380  if (parser.parseRegion(*body, ivs) ||
4381  parser.parseOptionalAttrDict(result.attributes))
4382  return failure();
4383 
4384  // Add a terminator if none was parsed.
4385  AffineParallelOp::ensureTerminator(*body, builder, result.location);
4386  return success();
4387 }
4388 
4389 //===----------------------------------------------------------------------===//
4390 // AffineYieldOp
4391 //===----------------------------------------------------------------------===//
4392 
4393 LogicalResult AffineYieldOp::verify() {
4394  auto *parentOp = (*this)->getParentOp();
4395  auto results = parentOp->getResults();
4396  auto operands = getOperands();
4397 
4398  if (!isa<AffineParallelOp, AffineIfOp, AffineForOp>(parentOp))
4399  return emitOpError() << "only terminates affine.if/for/parallel regions";
4400  if (parentOp->getNumResults() != getNumOperands())
4401  return emitOpError() << "parent of yield must have same number of "
4402  "results as the yield operands";
4403  for (auto it : llvm::zip(results, operands)) {
4404  if (std::get<0>(it).getType() != std::get<1>(it).getType())
4405  return emitOpError() << "types mismatch between yield op and its parent";
4406  }
4407 
4408  return success();
4409 }
4410 
4411 //===----------------------------------------------------------------------===//
4412 // AffineVectorLoadOp
4413 //===----------------------------------------------------------------------===//
4414 
4415 void AffineVectorLoadOp::build(OpBuilder &builder, OperationState &result,
4416  VectorType resultType, AffineMap map,
4417  ValueRange operands) {
4418  assert(operands.size() == 1 + map.getNumInputs() && "inconsistent operands");
4419  result.addOperands(operands);
4420  if (map)
4421  result.addAttribute(getMapAttrStrName(), AffineMapAttr::get(map));
4422  result.types.push_back(resultType);
4423 }
4424 
4425 void AffineVectorLoadOp::build(OpBuilder &builder, OperationState &result,
4426  VectorType resultType, Value memref,
4427  AffineMap map, ValueRange mapOperands) {
4428  assert(map.getNumInputs() == mapOperands.size() && "inconsistent index info");
4429  result.addOperands(memref);
4430  result.addOperands(mapOperands);
4431  result.addAttribute(getMapAttrStrName(), AffineMapAttr::get(map));
4432  result.types.push_back(resultType);
4433 }
4434 
4435 void AffineVectorLoadOp::build(OpBuilder &builder, OperationState &result,
4436  VectorType resultType, Value memref,
4437  ValueRange indices) {
4438  auto memrefType = llvm::cast<MemRefType>(memref.getType());
4439  int64_t rank = memrefType.getRank();
4440  // Create identity map for memrefs with at least one dimension or () -> ()
4441  // for zero-dimensional memrefs.
4442  auto map =
4443  rank ? builder.getMultiDimIdentityMap(rank) : builder.getEmptyAffineMap();
4444  build(builder, result, resultType, memref, map, indices);
4445 }
4446 
4447 void AffineVectorLoadOp::getCanonicalizationPatterns(RewritePatternSet &results,
4448  MLIRContext *context) {
4449  results.add<SimplifyAffineOp<AffineVectorLoadOp>>(context);
4450 }
4451 
4452 ParseResult AffineVectorLoadOp::parse(OpAsmParser &parser,
4453  OperationState &result) {
4454  auto &builder = parser.getBuilder();
4455  auto indexTy = builder.getIndexType();
4456 
4457  MemRefType memrefType;
4458  VectorType resultType;
4459  OpAsmParser::UnresolvedOperand memrefInfo;
4460  AffineMapAttr mapAttr;
4462  return failure(
4463  parser.parseOperand(memrefInfo) ||
4464  parser.parseAffineMapOfSSAIds(mapOperands, mapAttr,
4465  AffineVectorLoadOp::getMapAttrStrName(),
4466  result.attributes) ||
4467  parser.parseOptionalAttrDict(result.attributes) ||
4468  parser.parseColonType(memrefType) || parser.parseComma() ||
4469  parser.parseType(resultType) ||
4470  parser.resolveOperand(memrefInfo, memrefType, result.operands) ||
4471  parser.resolveOperands(mapOperands, indexTy, result.operands) ||
4472  parser.addTypeToList(resultType, result.types));
4473 }
4474 
4476  p << " " << getMemRef() << '[';
4477  if (AffineMapAttr mapAttr =
4478  (*this)->getAttrOfType<AffineMapAttr>(getMapAttrStrName()))
4479  p.printAffineMapOfSSAIds(mapAttr, getMapOperands());
4480  p << ']';
4481  p.printOptionalAttrDict((*this)->getAttrs(),
4482  /*elidedAttrs=*/{getMapAttrStrName()});
4483  p << " : " << getMemRefType() << ", " << getType();
4484 }
4485 
4486 /// Verify common invariants of affine.vector_load and affine.vector_store.
4487 static LogicalResult verifyVectorMemoryOp(Operation *op, MemRefType memrefType,
4488  VectorType vectorType) {
4489  // Check that memref and vector element types match.
4490  if (memrefType.getElementType() != vectorType.getElementType())
4491  return op->emitOpError(
4492  "requires memref and vector types of the same elemental type");
4493  return success();
4494 }
4495 
4496 LogicalResult AffineVectorLoadOp::verify() {
4497  MemRefType memrefType = getMemRefType();
4498  if (failed(verifyMemoryOpIndexing(
4499  *this, (*this)->getAttrOfType<AffineMapAttr>(getMapAttrStrName()),
4500  getMapOperands(), memrefType,
4501  /*numIndexOperands=*/getNumOperands() - 1)))
4502  return failure();
4503 
4504  if (failed(verifyVectorMemoryOp(getOperation(), memrefType, getVectorType())))
4505  return failure();
4506 
4507  return success();
4508 }
4509 
4510 //===----------------------------------------------------------------------===//
4511 // AffineVectorStoreOp
4512 //===----------------------------------------------------------------------===//
4513 
4514 void AffineVectorStoreOp::build(OpBuilder &builder, OperationState &result,
4515  Value valueToStore, Value memref, AffineMap map,
4516  ValueRange mapOperands) {
4517  assert(map.getNumInputs() == mapOperands.size() && "inconsistent index info");
4518  result.addOperands(valueToStore);
4519  result.addOperands(memref);
4520  result.addOperands(mapOperands);
4521  result.addAttribute(getMapAttrStrName(), AffineMapAttr::get(map));
4522 }
4523 
4524 // Use identity map.
4525 void AffineVectorStoreOp::build(OpBuilder &builder, OperationState &result,
4526  Value valueToStore, Value memref,
4527  ValueRange indices) {
4528  auto memrefType = llvm::cast<MemRefType>(memref.getType());
4529  int64_t rank = memrefType.getRank();
4530  // Create identity map for memrefs with at least one dimension or () -> ()
4531  // for zero-dimensional memrefs.
4532  auto map =
4533  rank ? builder.getMultiDimIdentityMap(rank) : builder.getEmptyAffineMap();
4534  build(builder, result, valueToStore, memref, map, indices);
4535 }
4536 void AffineVectorStoreOp::getCanonicalizationPatterns(
4537  RewritePatternSet &results, MLIRContext *context) {
4538  results.add<SimplifyAffineOp<AffineVectorStoreOp>>(context);
4539 }
4540 
4541 ParseResult AffineVectorStoreOp::parse(OpAsmParser &parser,
4542  OperationState &result) {
4543  auto indexTy = parser.getBuilder().getIndexType();
4544 
4545  MemRefType memrefType;
4546  VectorType resultType;
4547  OpAsmParser::UnresolvedOperand storeValueInfo;
4548  OpAsmParser::UnresolvedOperand memrefInfo;
4549  AffineMapAttr mapAttr;
4551  return failure(
4552  parser.parseOperand(storeValueInfo) || parser.parseComma() ||
4553  parser.parseOperand(memrefInfo) ||
4554  parser.parseAffineMapOfSSAIds(mapOperands, mapAttr,
4555  AffineVectorStoreOp::getMapAttrStrName(),
4556  result.attributes) ||
4557  parser.parseOptionalAttrDict(result.attributes) ||
4558  parser.parseColonType(memrefType) || parser.parseComma() ||
4559  parser.parseType(resultType) ||
4560  parser.resolveOperand(storeValueInfo, resultType, result.operands) ||
4561  parser.resolveOperand(memrefInfo, memrefType, result.operands) ||
4562  parser.resolveOperands(mapOperands, indexTy, result.operands));
4563 }
4564 
4566  p << " " << getValueToStore();
4567  p << ", " << getMemRef() << '[';
4568  if (AffineMapAttr mapAttr =
4569  (*this)->getAttrOfType<AffineMapAttr>(getMapAttrStrName()))
4570  p.printAffineMapOfSSAIds(mapAttr, getMapOperands());
4571  p << ']';
4572  p.printOptionalAttrDict((*this)->getAttrs(),
4573  /*elidedAttrs=*/{getMapAttrStrName()});
4574  p << " : " << getMemRefType() << ", " << getValueToStore().getType();
4575 }
4576 
4577 LogicalResult AffineVectorStoreOp::verify() {
4578  MemRefType memrefType = getMemRefType();
4579  if (failed(verifyMemoryOpIndexing(
4580  *this, (*this)->getAttrOfType<AffineMapAttr>(getMapAttrStrName()),
4581  getMapOperands(), memrefType,
4582  /*numIndexOperands=*/getNumOperands() - 2)))
4583  return failure();
4584 
4585  if (failed(verifyVectorMemoryOp(*this, memrefType, getVectorType())))
4586  return failure();
4587 
4588  return success();
4589 }
4590 
4591 //===----------------------------------------------------------------------===//
4592 // DelinearizeIndexOp
4593 //===----------------------------------------------------------------------===//
4594 
4595 void AffineDelinearizeIndexOp::build(OpBuilder &odsBuilder,
4596  OperationState &odsState,
4597  Value linearIndex, ValueRange dynamicBasis,
4598  ArrayRef<int64_t> staticBasis,
4599  bool hasOuterBound) {
4600  SmallVector<Type> returnTypes(hasOuterBound ? staticBasis.size()
4601  : staticBasis.size() + 1,
4602  linearIndex.getType());
4603  build(odsBuilder, odsState, returnTypes, linearIndex, dynamicBasis,
4604  staticBasis);
4605 }
4606 
4607 void AffineDelinearizeIndexOp::build(OpBuilder &odsBuilder,
4608  OperationState &odsState,
4609  Value linearIndex, ValueRange basis,
4610  bool hasOuterBound) {
4611  if (hasOuterBound && !basis.empty() && basis.front() == nullptr) {
4612  hasOuterBound = false;
4613  basis = basis.drop_front();
4614  }
4615  SmallVector<Value> dynamicBasis;
4616  SmallVector<int64_t> staticBasis;
4617  dispatchIndexOpFoldResults(getAsOpFoldResult(basis), dynamicBasis,
4618  staticBasis);
4619  build(odsBuilder, odsState, linearIndex, dynamicBasis, staticBasis,
4620  hasOuterBound);
4621 }
4622 
4623 void AffineDelinearizeIndexOp::build(OpBuilder &odsBuilder,
4624  OperationState &odsState,
4625  Value linearIndex,
4626  ArrayRef<OpFoldResult> basis,
4627  bool hasOuterBound) {
4628  if (hasOuterBound && !basis.empty() && basis.front() == OpFoldResult()) {
4629  hasOuterBound = false;
4630  basis = basis.drop_front();
4631  }
4632  SmallVector<Value> dynamicBasis;
4633  SmallVector<int64_t> staticBasis;
4634  dispatchIndexOpFoldResults(basis, dynamicBasis, staticBasis);
4635  build(odsBuilder, odsState, linearIndex, dynamicBasis, staticBasis,
4636  hasOuterBound);
4637 }
4638 
4639 void AffineDelinearizeIndexOp::build(OpBuilder &odsBuilder,
4640  OperationState &odsState,
4641  Value linearIndex, ArrayRef<int64_t> basis,
4642  bool hasOuterBound) {
4643  build(odsBuilder, odsState, linearIndex, ValueRange{}, basis, hasOuterBound);
4644 }
4645 
4646 LogicalResult AffineDelinearizeIndexOp::verify() {
4647  ArrayRef<int64_t> staticBasis = getStaticBasis();
4648  if (getNumResults() != staticBasis.size() &&
4649  getNumResults() != staticBasis.size() + 1)
4650  return emitOpError("should return an index for each basis element and up "
4651  "to one extra index");
4652 
4653  auto dynamicMarkersCount = llvm::count_if(staticBasis, ShapedType::isDynamic);
4654  if (static_cast<size_t>(dynamicMarkersCount) != getDynamicBasis().size())
4655  return emitOpError(
4656  "mismatch between dynamic and static basis (kDynamic marker but no "
4657  "corresponding dynamic basis entry) -- this can only happen due to an "
4658  "incorrect fold/rewrite");
4659 
4660  if (!llvm::all_of(staticBasis, [](int64_t v) {
4661  return v > 0 || ShapedType::isDynamic(v);
4662  }))
4663  return emitOpError("no basis element may be statically non-positive");
4664 
4665  return success();
4666 }
4667 
4668 /// Given mixed basis of affine.delinearize_index/linearize_index replace
4669 /// constant SSA values with the constant integer value and return the new
4670 /// static basis. In case no such candidate for replacement exists, this utility
4671 /// returns std::nullopt.
4672 static std::optional<SmallVector<int64_t>>
4674  MutableOperandRange mutableDynamicBasis,
4675  ArrayRef<Attribute> dynamicBasis) {
4676  uint64_t dynamicBasisIndex = 0;
4677  for (OpFoldResult basis : dynamicBasis) {
4678  if (basis) {
4679  mutableDynamicBasis.erase(dynamicBasisIndex);
4680  } else {
4681  ++dynamicBasisIndex;
4682  }
4683  }
4684 
4685  // No constant SSA value exists.
4686  if (dynamicBasisIndex == dynamicBasis.size())
4687  return std::nullopt;
4688 
4689  SmallVector<int64_t> staticBasis;
4690  for (OpFoldResult basis : mixedBasis) {
4691  std::optional<int64_t> basisVal = getConstantIntValue(basis);
4692  if (!basisVal)
4693  staticBasis.push_back(ShapedType::kDynamic);
4694  else
4695  staticBasis.push_back(*basisVal);
4696  }
4697 
4698  return staticBasis;
4699 }
4700 
4701 LogicalResult
4702 AffineDelinearizeIndexOp::fold(FoldAdaptor adaptor,
4704  std::optional<SmallVector<int64_t>> maybeStaticBasis =
4705  foldCstValueToCstAttrBasis(getMixedBasis(), getDynamicBasisMutable(),
4706  adaptor.getDynamicBasis());
4707  if (maybeStaticBasis) {
4708  setStaticBasis(*maybeStaticBasis);
4709  return success();
4710  }
4711  // If we won't be doing any division or modulo (no basis or the one basis
4712  // element is purely advisory), simply return the input value.
4713  if (getNumResults() == 1) {
4714  result.push_back(getLinearIndex());
4715  return success();
4716  }
4717 
4718  if (adaptor.getLinearIndex() == nullptr)
4719  return failure();
4720 
4721  if (!adaptor.getDynamicBasis().empty())
4722  return failure();
4723 
4724  int64_t highPart = cast<IntegerAttr>(adaptor.getLinearIndex()).getInt();
4725  Type attrType = getLinearIndex().getType();
4726 
4727  ArrayRef<int64_t> staticBasis = getStaticBasis();
4728  if (hasOuterBound())
4729  staticBasis = staticBasis.drop_front();
4730  for (int64_t modulus : llvm::reverse(staticBasis)) {
4731  result.push_back(IntegerAttr::get(attrType, llvm::mod(highPart, modulus)));
4732  highPart = llvm::divideFloorSigned(highPart, modulus);
4733  }
4734  result.push_back(IntegerAttr::get(attrType, highPart));
4735  std::reverse(result.begin(), result.end());
4736  return success();
4737 }
4738 
4739 SmallVector<OpFoldResult> AffineDelinearizeIndexOp::getEffectiveBasis() {
4740  OpBuilder builder(getContext());
4741  if (hasOuterBound()) {
4742  if (getStaticBasis().front() == ::mlir::ShapedType::kDynamic)
4743  return getMixedValues(getStaticBasis().drop_front(),
4744  getDynamicBasis().drop_front(), builder);
4745 
4746  return getMixedValues(getStaticBasis().drop_front(), getDynamicBasis(),
4747  builder);
4748  }
4749 
4750  return getMixedValues(getStaticBasis(), getDynamicBasis(), builder);
4751 }
4752 
4753 SmallVector<OpFoldResult> AffineDelinearizeIndexOp::getPaddedBasis() {
4754  SmallVector<OpFoldResult> ret = getMixedBasis();
4755  if (!hasOuterBound())
4756  ret.insert(ret.begin(), OpFoldResult());
4757  return ret;
4758 }
4759 
4760 namespace {
4761 
4762 // Drops delinearization indices that correspond to unit-extent basis
4763 struct DropUnitExtentBasis
4764  : public OpRewritePattern<affine::AffineDelinearizeIndexOp> {
4766 
4767  LogicalResult matchAndRewrite(affine::AffineDelinearizeIndexOp delinearizeOp,
4768  PatternRewriter &rewriter) const override {
4769  SmallVector<Value> replacements(delinearizeOp->getNumResults(), nullptr);
4770  std::optional<Value> zero = std::nullopt;
4771  Location loc = delinearizeOp->getLoc();
4772  auto getZero = [&]() -> Value {
4773  if (!zero)
4774  zero = rewriter.create<arith::ConstantIndexOp>(loc, 0);
4775  return zero.value();
4776  };
4777 
4778  // Replace all indices corresponding to unit-extent basis with 0.
4779  // Remaining basis can be used to get a new `affine.delinearize_index` op.
4780  SmallVector<OpFoldResult> newBasis;
4781  for (auto [index, basis] :
4782  llvm::enumerate(delinearizeOp.getPaddedBasis())) {
4783  std::optional<int64_t> basisVal =
4784  basis ? getConstantIntValue(basis) : std::nullopt;
4785  if (basisVal && *basisVal == 1)
4786  replacements[index] = getZero();
4787  else
4788  newBasis.push_back(basis);
4789  }
4790 
4791  if (newBasis.size() == delinearizeOp.getNumResults())
4792  return rewriter.notifyMatchFailure(delinearizeOp,
4793  "no unit basis elements");
4794 
4795  if (!newBasis.empty()) {
4796  // Will drop the leading nullptr from `basis` if there was no outer bound.
4797  auto newDelinearizeOp = rewriter.create<affine::AffineDelinearizeIndexOp>(
4798  loc, delinearizeOp.getLinearIndex(), newBasis);
4799  int newIndex = 0;
4800  // Map back the new delinearized indices to the values they replace.
4801  for (auto &replacement : replacements) {
4802  if (replacement)
4803  continue;
4804  replacement = newDelinearizeOp->getResult(newIndex++);
4805  }
4806  }
4807 
4808  rewriter.replaceOp(delinearizeOp, replacements);
4809  return success();
4810  }
4811 };
4812 
4813 /// If a `affine.delinearize_index`'s input is a `affine.linearize_index
4814 /// disjoint` and the two operations end with the same basis elements,
4815 /// cancel those parts of the operations out because they are inverses
4816 /// of each other.
4817 ///
4818 /// If the operations have the same basis, cancel them entirely.
4819 ///
4820 /// The `disjoint` flag is needed on the `affine.linearize_index` because
4821 /// otherwise, there is no guarantee that the inputs to the linearization are
4822 /// in-bounds the way the outputs of the delinearization would be.
4823 struct CancelDelinearizeOfLinearizeDisjointExactTail
4824  : public OpRewritePattern<affine::AffineDelinearizeIndexOp> {
4826 
4827  LogicalResult matchAndRewrite(affine::AffineDelinearizeIndexOp delinearizeOp,
4828  PatternRewriter &rewriter) const override {
4829  auto linearizeOp = delinearizeOp.getLinearIndex()
4830  .getDefiningOp<affine::AffineLinearizeIndexOp>();
4831  if (!linearizeOp)
4832  return rewriter.notifyMatchFailure(delinearizeOp,
4833  "index doesn't come from linearize");
4834 
4835  if (!linearizeOp.getDisjoint())
4836  return rewriter.notifyMatchFailure(linearizeOp, "not disjoint");
4837 
4838  ValueRange linearizeIns = linearizeOp.getMultiIndex();
4839  // Note: we use the full basis so we don't lose outer bounds later.
4840  SmallVector<OpFoldResult> linearizeBasis = linearizeOp.getMixedBasis();
4841  SmallVector<OpFoldResult> delinearizeBasis = delinearizeOp.getMixedBasis();
4842  size_t numMatches = 0;
4843  for (auto [linSize, delinSize] : llvm::zip(
4844  llvm::reverse(linearizeBasis), llvm::reverse(delinearizeBasis))) {
4845  if (linSize != delinSize)
4846  break;
4847  ++numMatches;
4848  }
4849 
4850  if (numMatches == 0)
4851  return rewriter.notifyMatchFailure(
4852  delinearizeOp, "final basis element doesn't match linearize");
4853 
4854  // The easy case: everything lines up and the basis match sup completely.
4855  if (numMatches == linearizeBasis.size() &&
4856  numMatches == delinearizeBasis.size() &&
4857  linearizeIns.size() == delinearizeOp.getNumResults()) {
4858  rewriter.replaceOp(delinearizeOp, linearizeOp.getMultiIndex());
4859  return success();
4860  }
4861 
4862  Value newLinearize = rewriter.create<affine::AffineLinearizeIndexOp>(
4863  linearizeOp.getLoc(), linearizeIns.drop_back(numMatches),
4864  ArrayRef<OpFoldResult>{linearizeBasis}.drop_back(numMatches),
4865  linearizeOp.getDisjoint());
4866  auto newDelinearize = rewriter.create<affine::AffineDelinearizeIndexOp>(
4867  delinearizeOp.getLoc(), newLinearize,
4868  ArrayRef<OpFoldResult>{delinearizeBasis}.drop_back(numMatches),
4869  delinearizeOp.hasOuterBound());
4870  SmallVector<Value> mergedResults(newDelinearize.getResults());
4871  mergedResults.append(linearizeIns.take_back(numMatches).begin(),
4872  linearizeIns.take_back(numMatches).end());
4873  rewriter.replaceOp(delinearizeOp, mergedResults);
4874  return success();
4875  }
4876 };
4877 
4878 /// If the input to a delinearization is a disjoint linearization, and the
4879 /// last k > 1 components of the delinearization basis multiply to the
4880 /// last component of the linearization basis, break the linearization and
4881 /// delinearization into two parts, peeling off the last input to linearization.
4882 ///
4883 /// For example:
4884 /// %0 = affine.linearize_index [%z, %y, %x] by (3, 2, 32) : index
4885 /// %1:4 = affine.delinearize_index %0 by (2, 3, 8, 4) : index, ...
4886 /// becomes
4887 /// %0 = affine.linearize_index [%z, %y] by (3, 2) : index
4888 /// %1:2 = affine.delinearize_index %0 by (2, 3) : index
4889 /// %2:2 = affine.delinearize_index %x by (8, 4) : index
4890 /// where the original %1:4 is replaced by %1:2 ++ %2:2
4891 struct SplitDelinearizeSpanningLastLinearizeArg final
4892  : OpRewritePattern<affine::AffineDelinearizeIndexOp> {
4894 
4895  LogicalResult matchAndRewrite(affine::AffineDelinearizeIndexOp delinearizeOp,
4896  PatternRewriter &rewriter) const override {
4897  auto linearizeOp = delinearizeOp.getLinearIndex()
4898  .getDefiningOp<affine::AffineLinearizeIndexOp>();
4899  if (!linearizeOp)
4900  return rewriter.notifyMatchFailure(delinearizeOp,
4901  "index doesn't come from linearize");
4902 
4903  if (!linearizeOp.getDisjoint())
4904  return rewriter.notifyMatchFailure(linearizeOp,
4905  "linearize isn't disjoint");
4906 
4907  int64_t target = linearizeOp.getStaticBasis().back();
4908  if (ShapedType::isDynamic(target))
4909  return rewriter.notifyMatchFailure(
4910  linearizeOp, "linearize ends with dynamic basis value");
4911 
4912  int64_t sizeToSplit = 1;
4913  size_t elemsToSplit = 0;
4914  ArrayRef<int64_t> basis = delinearizeOp.getStaticBasis();
4915  for (int64_t basisElem : llvm::reverse(basis)) {
4916  if (ShapedType::isDynamic(basisElem))
4917  return rewriter.notifyMatchFailure(
4918  delinearizeOp, "dynamic basis element while scanning for split");
4919  sizeToSplit *= basisElem;
4920  elemsToSplit += 1;
4921 
4922  if (sizeToSplit > target)
4923  return rewriter.notifyMatchFailure(delinearizeOp,
4924  "overshot last argument size");
4925  if (sizeToSplit == target)
4926  break;
4927  }
4928 
4929  if (sizeToSplit < target)
4930  return rewriter.notifyMatchFailure(
4931  delinearizeOp, "product of known basis elements doesn't exceed last "
4932  "linearize argument");
4933 
4934  if (elemsToSplit < 2)
4935  return rewriter.notifyMatchFailure(
4936  delinearizeOp,
4937  "need at least two elements to form the basis product");
4938 
4939  Value linearizeWithoutBack =
4940  rewriter.create<affine::AffineLinearizeIndexOp>(
4941  linearizeOp.getLoc(), linearizeOp.getMultiIndex().drop_back(),
4942  linearizeOp.getDynamicBasis(),
4943  linearizeOp.getStaticBasis().drop_back(),
4944  linearizeOp.getDisjoint());
4945  auto delinearizeWithoutSplitPart =
4946  rewriter.create<affine::AffineDelinearizeIndexOp>(
4947  delinearizeOp.getLoc(), linearizeWithoutBack,
4948  delinearizeOp.getDynamicBasis(), basis.drop_back(elemsToSplit),
4949  delinearizeOp.hasOuterBound());
4950  auto delinearizeBack = rewriter.create<affine::AffineDelinearizeIndexOp>(
4951  delinearizeOp.getLoc(), linearizeOp.getMultiIndex().back(),
4952  basis.take_back(elemsToSplit), /*hasOuterBound=*/true);
4953  SmallVector<Value> results = llvm::to_vector(
4954  llvm::concat<Value>(delinearizeWithoutSplitPart.getResults(),
4955  delinearizeBack.getResults()));
4956  rewriter.replaceOp(delinearizeOp, results);
4957 
4958  return success();
4959  }
4960 };
4961 } // namespace
4962 
4963 void affine::AffineDelinearizeIndexOp::getCanonicalizationPatterns(
4964  RewritePatternSet &patterns, MLIRContext *context) {
4965  patterns
4966  .insert<CancelDelinearizeOfLinearizeDisjointExactTail,
4967  DropUnitExtentBasis, SplitDelinearizeSpanningLastLinearizeArg>(
4968  context);
4969 }
4970 
4971 //===----------------------------------------------------------------------===//
4972 // LinearizeIndexOp
4973 //===----------------------------------------------------------------------===//
4974 
4975 void AffineLinearizeIndexOp::build(OpBuilder &odsBuilder,
4976  OperationState &odsState,
4977  ValueRange multiIndex, ValueRange basis,
4978  bool disjoint) {
4979  if (!basis.empty() && basis.front() == Value())
4980  basis = basis.drop_front();
4981  SmallVector<Value> dynamicBasis;
4982  SmallVector<int64_t> staticBasis;
4983  dispatchIndexOpFoldResults(getAsOpFoldResult(basis), dynamicBasis,
4984  staticBasis);
4985  build(odsBuilder, odsState, multiIndex, dynamicBasis, staticBasis, disjoint);
4986 }
4987 
4988 void AffineLinearizeIndexOp::build(OpBuilder &odsBuilder,
4989  OperationState &odsState,
4990  ValueRange multiIndex,
4991  ArrayRef<OpFoldResult> basis,
4992  bool disjoint) {
4993  if (!basis.empty() && basis.front() == OpFoldResult())
4994  basis = basis.drop_front();
4995  SmallVector<Value> dynamicBasis;
4996  SmallVector<int64_t> staticBasis;
4997  dispatchIndexOpFoldResults(basis, dynamicBasis, staticBasis);
4998  build(odsBuilder, odsState, multiIndex, dynamicBasis, staticBasis, disjoint);
4999 }
5000 
5001 void AffineLinearizeIndexOp::build(OpBuilder &odsBuilder,
5002  OperationState &odsState,
5003  ValueRange multiIndex,
5004  ArrayRef<int64_t> basis, bool disjoint) {
5005  build(odsBuilder, odsState, multiIndex, ValueRange{}, basis, disjoint);
5006 }
5007 
5008 LogicalResult AffineLinearizeIndexOp::verify() {
5009  size_t numIndexes = getMultiIndex().size();
5010  size_t numBasisElems = getStaticBasis().size();
5011  if (numIndexes != numBasisElems && numIndexes != numBasisElems + 1)
5012  return emitOpError("should be passed a basis element for each index except "
5013  "possibly the first");
5014 
5015  auto dynamicMarkersCount =
5016  llvm::count_if(getStaticBasis(), ShapedType::isDynamic);
5017  if (static_cast<size_t>(dynamicMarkersCount) != getDynamicBasis().size())
5018  return emitOpError(
5019  "mismatch between dynamic and static basis (kDynamic marker but no "
5020  "corresponding dynamic basis entry) -- this can only happen due to an "
5021  "incorrect fold/rewrite");
5022 
5023  return success();
5024 }
5025 
5026 OpFoldResult AffineLinearizeIndexOp::fold(FoldAdaptor adaptor) {
5027  std::optional<SmallVector<int64_t>> maybeStaticBasis =
5028  foldCstValueToCstAttrBasis(getMixedBasis(), getDynamicBasisMutable(),
5029  adaptor.getDynamicBasis());
5030  if (maybeStaticBasis) {
5031  setStaticBasis(*maybeStaticBasis);
5032  return getResult();
5033  }
5034  // No indices linearizes to zero.
5035  if (getMultiIndex().empty())
5036  return IntegerAttr::get(getResult().getType(), 0);
5037 
5038  // One single index linearizes to itself.
5039  if (getMultiIndex().size() == 1)
5040  return getMultiIndex().front();
5041 
5042  if (llvm::any_of(adaptor.getMultiIndex(),
5043  [](Attribute a) { return a == nullptr; }))
5044  return nullptr;
5045 
5046  if (!adaptor.getDynamicBasis().empty())
5047  return nullptr;
5048 
5049  int64_t result = 0;
5050  int64_t stride = 1;
5051  for (auto [length, indexAttr] :
5052  llvm::zip_first(llvm::reverse(getStaticBasis()),
5053  llvm::reverse(adaptor.getMultiIndex()))) {
5054  result = result + cast<IntegerAttr>(indexAttr).getInt() * stride;
5055  stride = stride * length;
5056  }
5057  // Handle the index element with no basis element.
5058  if (!hasOuterBound())
5059  result =
5060  result +
5061  cast<IntegerAttr>(adaptor.getMultiIndex().front()).getInt() * stride;
5062 
5063  return IntegerAttr::get(getResult().getType(), result);
5064 }
5065 
5066 SmallVector<OpFoldResult> AffineLinearizeIndexOp::getEffectiveBasis() {
5067  OpBuilder builder(getContext());
5068  if (hasOuterBound()) {
5069  if (getStaticBasis().front() == ::mlir::ShapedType::kDynamic)
5070  return getMixedValues(getStaticBasis().drop_front(),
5071  getDynamicBasis().drop_front(), builder);
5072 
5073  return getMixedValues(getStaticBasis().drop_front(), getDynamicBasis(),
5074  builder);
5075  }
5076 
5077  return getMixedValues(getStaticBasis(), getDynamicBasis(), builder);
5078 }
5079 
5080 SmallVector<OpFoldResult> AffineLinearizeIndexOp::getPaddedBasis() {
5081  SmallVector<OpFoldResult> ret = getMixedBasis();
5082  if (!hasOuterBound())
5083  ret.insert(ret.begin(), OpFoldResult());
5084  return ret;
5085 }
5086 
5087 namespace {
5088 /// Rewrite `affine.linearize_index disjoint [%...a, %x, %...b] by (%...c, 1,
5089 /// %...d)` to `affine.linearize_index disjoint [%...a, %...b] by (%...c,
5090 /// %...d)`.
5091 
5092 /// Note that `disjoint` is required here, because, without it, we could have
5093 /// `affine.linearize_index [%...a, %c64, %...b] by (%...c, 1, %...d)`
5094 /// is a valid operation where the `%c64` cannot be trivially dropped.
5095 ///
5096 /// Alternatively, if `%x` in the above is a known constant 0, remove it even if
5097 /// the operation isn't asserted to be `disjoint`.
5098 struct DropLinearizeUnitComponentsIfDisjointOrZero final
5099  : OpRewritePattern<affine::AffineLinearizeIndexOp> {
5101 
5102  LogicalResult matchAndRewrite(affine::AffineLinearizeIndexOp op,
5103  PatternRewriter &rewriter) const override {
5104  ValueRange multiIndex = op.getMultiIndex();
5105  size_t numIndices = multiIndex.size();
5106  SmallVector<Value> newIndices;
5107  newIndices.reserve(numIndices);
5108  SmallVector<OpFoldResult> newBasis;
5109  newBasis.reserve(numIndices);
5110 
5111  if (!op.hasOuterBound()) {
5112  newIndices.push_back(multiIndex.front());
5113  multiIndex = multiIndex.drop_front();
5114  }
5115 
5116  SmallVector<OpFoldResult> basis = op.getMixedBasis();
5117  for (auto [index, basisElem] : llvm::zip_equal(multiIndex, basis)) {
5118  std::optional<int64_t> basisEntry = getConstantIntValue(basisElem);
5119  if (!basisEntry || *basisEntry != 1) {
5120  newIndices.push_back(index);
5121  newBasis.push_back(basisElem);
5122  continue;
5123  }
5124 
5125  std::optional<int64_t> indexValue = getConstantIntValue(index);
5126  if (!op.getDisjoint() && (!indexValue || *indexValue != 0)) {
5127  newIndices.push_back(index);
5128  newBasis.push_back(basisElem);
5129  continue;
5130  }
5131  }
5132  if (newIndices.size() == numIndices)
5133  return rewriter.notifyMatchFailure(op,
5134  "no unit basis entries to replace");
5135 
5136  if (newIndices.size() == 0) {
5137  rewriter.replaceOpWithNewOp<arith::ConstantIndexOp>(op, 0);
5138  return success();
5139  }
5140  rewriter.replaceOpWithNewOp<affine::AffineLinearizeIndexOp>(
5141  op, newIndices, newBasis, op.getDisjoint());
5142  return success();
5143  }
5144 };
5145 
5147  ArrayRef<OpFoldResult> terms) {
5148  int64_t nDynamic = 0;
5149  SmallVector<Value> dynamicPart;
5150  AffineExpr result = builder.getAffineConstantExpr(1);
5151  for (OpFoldResult term : terms) {
5152  if (!term)
5153  return term;
5154  std::optional<int64_t> maybeConst = getConstantIntValue(term);
5155  if (maybeConst) {
5156  result = result * builder.getAffineConstantExpr(*maybeConst);
5157  } else {
5158  dynamicPart.push_back(cast<Value>(term));
5159  result = result * builder.getAffineSymbolExpr(nDynamic++);
5160  }
5161  }
5162  if (auto constant = dyn_cast<AffineConstantExpr>(result))
5163  return getAsIndexOpFoldResult(builder.getContext(), constant.getValue());
5164  return builder.create<AffineApplyOp>(loc, result, dynamicPart).getResult();
5165 }
5166 
5167 /// If conseceutive outputs of a delinearize_index are linearized with the same
5168 /// bounds, canonicalize away the redundant arithmetic.
5169 ///
5170 /// That is, if we have
5171 /// ```
5172 /// %s:N = affine.delinearize_index %x into (...a, B1, B2, ... BK, ...b)
5173 /// %t = affine.linearize_index [...c, %s#I, %s#(I + 1), ... %s#(I+K-1), ...d]
5174 /// by (...e, B1, B2, ..., BK, ...f)
5175 /// ```
5176 ///
5177 /// We can rewrite this to
5178 /// ```
5179 /// B = B1 * B2 ... BK
5180 /// %sMerged:(N-K+1) affine.delinearize_index %x into (...a, B, ...b)
5181 /// %t = affine.linearize_index [...c, %s#I, ...d] by (...e, B, ...f)
5182 /// ```
5183 /// where we replace all results of %s unaffected by the change with results
5184 /// from %sMerged.
5185 ///
5186 /// As a special case, if all results of the delinearize are merged in this way
5187 /// we can replace those usages with %x, thus cancelling the delinearization
5188 /// entirely, as in
5189 /// ```
5190 /// %s:3 = affine.delinearize_index %x into (2, 4, 8)
5191 /// %t = affine.linearize_index [%s#0, %s#1, %s#2, %c0] by (2, 4, 8, 16)
5192 /// ```
5193 /// becoming `%t = affine.linearize_index [%x, %c0] by (64, 16)`
5194 struct CancelLinearizeOfDelinearizePortion final
5195  : OpRewritePattern<affine::AffineLinearizeIndexOp> {
5197 
5198 private:
5199  // Struct representing a case where the cancellation pattern
5200  // applies. A `Match` means that `length` inputs to the linearize operation
5201  // starting at `linStart` can be cancelled with `length` outputs of
5202  // `delinearize`, starting from `delinStart`.
5203  struct Match {
5204  AffineDelinearizeIndexOp delinearize;
5205  unsigned linStart = 0;
5206  unsigned delinStart = 0;
5207  unsigned length = 0;
5208  };
5209 
5210 public:
5211  LogicalResult matchAndRewrite(affine::AffineLinearizeIndexOp linearizeOp,
5212  PatternRewriter &rewriter) const override {
5213  SmallVector<Match> matches;
5214 
5215  const SmallVector<OpFoldResult> linBasis = linearizeOp.getPaddedBasis();
5216  ArrayRef<OpFoldResult> linBasisRef = linBasis;
5217 
5218  ValueRange multiIndex = linearizeOp.getMultiIndex();
5219  unsigned numLinArgs = multiIndex.size();
5220  unsigned linArgIdx = 0;
5221  // We only want to replace one run from the same delinearize op per
5222  // pattern invocation lest we run into invalidation issues.
5223  llvm::SmallPtrSet<Operation *, 2> alreadyMatchedDelinearize;
5224  while (linArgIdx < numLinArgs) {
5225  auto asResult = dyn_cast<OpResult>(multiIndex[linArgIdx]);
5226  if (!asResult) {
5227  linArgIdx++;
5228  continue;
5229  }
5230 
5231  auto delinearizeOp =
5232  dyn_cast<AffineDelinearizeIndexOp>(asResult.getOwner());
5233  if (!delinearizeOp) {
5234  linArgIdx++;
5235  continue;
5236  }
5237 
5238  /// Result 0 of the delinearize and argument 0 of the linearize can
5239  /// leave their maximum value unspecified. However, even if this happens
5240  /// we can still sometimes start the match process. Specifically, if
5241  /// - The argument we're matching is result 0 and argument 0 (so the
5242  /// bounds don't matter). For example,
5243  ///
5244  /// %0:2 = affine.delinearize_index %x into (8) : index, index
5245  /// %1 = affine.linearize_index [%s#0, %s#1, ...] (8, ...)
5246  /// allows cancellation
5247  /// - The delinearization doesn't specify a bound, but the linearization
5248  /// is `disjoint`, which asserts that the bound on the linearization is
5249  /// correct.
5250  unsigned delinArgIdx = asResult.getResultNumber();
5251  SmallVector<OpFoldResult> delinBasis = delinearizeOp.getPaddedBasis();
5252  OpFoldResult firstDelinBound = delinBasis[delinArgIdx];
5253  OpFoldResult firstLinBound = linBasis[linArgIdx];
5254  bool boundsMatch = firstDelinBound == firstLinBound;
5255  bool bothAtFront = linArgIdx == 0 && delinArgIdx == 0;
5256  bool knownByDisjoint =
5257  linearizeOp.getDisjoint() && delinArgIdx == 0 && !firstDelinBound;
5258  if (!boundsMatch && !bothAtFront && !knownByDisjoint) {
5259  linArgIdx++;
5260  continue;
5261  }
5262 
5263  unsigned j = 1;
5264  unsigned numDelinOuts = delinearizeOp.getNumResults();
5265  for (; j + linArgIdx < numLinArgs && j + delinArgIdx < numDelinOuts;
5266  ++j) {
5267  if (multiIndex[linArgIdx + j] !=
5268  delinearizeOp.getResult(delinArgIdx + j))
5269  break;
5270  if (linBasis[linArgIdx + j] != delinBasis[delinArgIdx + j])
5271  break;
5272  }
5273  // If there're multiple matches against the same delinearize_index,
5274  // only rewrite the first one we find to prevent invalidations. The next
5275  // ones will be taken care of by subsequent pattern invocations.
5276  if (j <= 1 || !alreadyMatchedDelinearize.insert(delinearizeOp).second) {
5277  linArgIdx++;
5278  continue;
5279  }
5280  matches.push_back(Match{delinearizeOp, linArgIdx, delinArgIdx, j});
5281  linArgIdx += j;
5282  }
5283 
5284  if (matches.empty())
5285  return rewriter.notifyMatchFailure(
5286  linearizeOp, "no run of delinearize outputs to deal with");
5287 
5288  // Record all the delinearize replacements so we can do them after creating
5289  // the new linearization operation, since the new operation might use
5290  // outputs of something we're replacing.
5291  SmallVector<SmallVector<Value>> delinearizeReplacements;
5292 
5293  SmallVector<Value> newIndex;
5294  newIndex.reserve(numLinArgs);
5295  SmallVector<OpFoldResult> newBasis;
5296  newBasis.reserve(numLinArgs);
5297  unsigned prevMatchEnd = 0;
5298  for (Match m : matches) {
5299  unsigned gap = m.linStart - prevMatchEnd;
5300  llvm::append_range(newIndex, multiIndex.slice(prevMatchEnd, gap));
5301  llvm::append_range(newBasis, linBasisRef.slice(prevMatchEnd, gap));
5302  // Update here so we don't forget this during early continues
5303  prevMatchEnd = m.linStart + m.length;
5304 
5305  PatternRewriter::InsertionGuard g(rewriter);
5306  rewriter.setInsertionPoint(m.delinearize);
5307 
5308  ArrayRef<OpFoldResult> basisToMerge =
5309  linBasisRef.slice(m.linStart, m.length);
5310  // We use the slice from the linearize's basis above because of the
5311  // "bounds inferred from `disjoint`" case above.
5312  OpFoldResult newSize =
5313  computeProduct(linearizeOp.getLoc(), rewriter, basisToMerge);
5314 
5315  // Trivial case where we can just skip past the delinearize all together
5316  if (m.length == m.delinearize.getNumResults()) {
5317  newIndex.push_back(m.delinearize.getLinearIndex());
5318  newBasis.push_back(newSize);
5319  // Pad out set of replacements so we don't do anything with this one.
5320  delinearizeReplacements.push_back(SmallVector<Value>());
5321  continue;
5322  }
5323 
5324  SmallVector<Value> newDelinResults;
5325  SmallVector<OpFoldResult> newDelinBasis = m.delinearize.getPaddedBasis();
5326  newDelinBasis.erase(newDelinBasis.begin() + m.delinStart,
5327  newDelinBasis.begin() + m.delinStart + m.length);
5328  newDelinBasis.insert(newDelinBasis.begin() + m.delinStart, newSize);
5329  auto newDelinearize = rewriter.create<AffineDelinearizeIndexOp>(
5330  m.delinearize.getLoc(), m.delinearize.getLinearIndex(),
5331  newDelinBasis);
5332 
5333  // Since there may be other uses of the indices we just merged together,
5334  // create a residual affine.delinearize_index that delinearizes the
5335  // merged output into its component parts.
5336  Value combinedElem = newDelinearize.getResult(m.delinStart);
5337  auto residualDelinearize = rewriter.create<AffineDelinearizeIndexOp>(
5338  m.delinearize.getLoc(), combinedElem, basisToMerge);
5339 
5340  // Swap all the uses of the unaffected delinearize outputs to the new
5341  // delinearization so that the old code can be removed if this
5342  // linearize_index is the only user of the merged results.
5343  llvm::append_range(newDelinResults,
5344  newDelinearize.getResults().take_front(m.delinStart));
5345  llvm::append_range(newDelinResults, residualDelinearize.getResults());
5346  llvm::append_range(
5347  newDelinResults,
5348  newDelinearize.getResults().drop_front(m.delinStart + 1));
5349 
5350  delinearizeReplacements.push_back(newDelinResults);
5351  newIndex.push_back(combinedElem);
5352  newBasis.push_back(newSize);
5353  }
5354  llvm::append_range(newIndex, multiIndex.drop_front(prevMatchEnd));
5355  llvm::append_range(newBasis, linBasisRef.drop_front(prevMatchEnd));
5356  rewriter.replaceOpWithNewOp<AffineLinearizeIndexOp>(
5357  linearizeOp, newIndex, newBasis, linearizeOp.getDisjoint());
5358 
5359  for (auto [m, newResults] :
5360  llvm::zip_equal(matches, delinearizeReplacements)) {
5361  if (newResults.empty())
5362  continue;
5363  rewriter.replaceOp(m.delinearize, newResults);
5364  }
5365 
5366  return success();
5367  }
5368 };
5369 
5370 /// Strip leading zero from affine.linearize_index.
5371 ///
5372 /// `affine.linearize_index [%c0, ...a] by (%x, ...b)` can be rewritten
5373 /// to `affine.linearize_index [...a] by (...b)` in all cases.
5374 struct DropLinearizeLeadingZero final
5375  : OpRewritePattern<affine::AffineLinearizeIndexOp> {
5377 
5378  LogicalResult matchAndRewrite(affine::AffineLinearizeIndexOp op,
5379  PatternRewriter &rewriter) const override {
5380  Value leadingIdx = op.getMultiIndex().front();
5381  if (!matchPattern(leadingIdx, m_Zero()))
5382  return failure();
5383 
5384  if (op.getMultiIndex().size() == 1) {
5385  rewriter.replaceOp(op, leadingIdx);
5386  return success();
5387  }
5388 
5389  SmallVector<OpFoldResult> mixedBasis = op.getMixedBasis();
5390  ArrayRef<OpFoldResult> newMixedBasis = mixedBasis;
5391  if (op.hasOuterBound())
5392  newMixedBasis = newMixedBasis.drop_front();
5393 
5394  rewriter.replaceOpWithNewOp<affine::AffineLinearizeIndexOp>(
5395  op, op.getMultiIndex().drop_front(), newMixedBasis, op.getDisjoint());
5396  return success();
5397  }
5398 };
5399 } // namespace
5400 
5401 void affine::AffineLinearizeIndexOp::getCanonicalizationPatterns(
5402  RewritePatternSet &patterns, MLIRContext *context) {
5403  patterns.add<CancelLinearizeOfDelinearizePortion, DropLinearizeLeadingZero,
5404  DropLinearizeUnitComponentsIfDisjointOrZero>(context);
5405 }
5406 
5407 //===----------------------------------------------------------------------===//
5408 // TableGen'd op method definitions
5409 //===----------------------------------------------------------------------===//
5410 
5411 #define GET_OP_CLASSES
5412 #include "mlir/Dialect/Affine/IR/AffineOps.cpp.inc"
static AffineForOp buildAffineLoopFromConstants(OpBuilder &builder, Location loc, int64_t lb, int64_t ub, int64_t step, AffineForOp::BodyBuilderFn bodyBuilderFn)
Creates an affine loop from the bounds known to be constants.
Definition: AffineOps.cpp:2735
static bool hasTrivialZeroTripCount(AffineForOp op)
Returns true if the affine.for has zero iterations in trivial cases.
Definition: AffineOps.cpp:2455
static void composeMultiResultAffineMap(AffineMap &map, SmallVectorImpl< Value > &operands)
Composes the given affine map with the given list of operands, pulling in the maps from any affine....
Definition: AffineOps.cpp:1189
static LogicalResult verifyMemoryOpIndexing(AffineMemOpTy op, AffineMapAttr mapAttr, Operation::operand_range mapOperands, MemRefType memrefType, unsigned numIndexOperands)
Verify common indexing invariants of affine.load, affine.store, affine.vector_load and affine....
Definition: AffineOps.cpp:3127
static void printAffineMinMaxOp(OpAsmPrinter &p, T op)
Definition: AffineOps.cpp:3305
static bool isResultTypeMatchAtomicRMWKind(Type resultType, arith::AtomicRMWKind op)
Definition: AffineOps.cpp:3938
static bool remainsLegalAfterInline(Value value, Region *src, Region *dest, const IRMapping &mapping, function_ref< bool(Value, Region *)> legalityCheck)
Checks if value known to be a legal affine dimension or symbol in src region remains legal if the ope...
Definition: AffineOps.cpp:60
static void printMinMaxBound(OpAsmPrinter &p, AffineMapAttr mapAttr, DenseIntElementsAttr group, ValueRange operands, StringRef keyword)
Prints a lower(upper) bound of an affine parallel loop with max(min) conditions in it.
Definition: AffineOps.cpp:4085
static void LLVM_ATTRIBUTE_UNUSED simplifyMapWithOperands(AffineMap &map, ArrayRef< Value > operands)
Simplify the map while exploiting information on the values in operands.
Definition: AffineOps.cpp:1025
static OpFoldResult foldMinMaxOp(T op, ArrayRef< Attribute > operands)
Fold an affine min or max operation with the given operands.
Definition: AffineOps.cpp:3341
static LogicalResult canonicalizeLoopBounds(AffineForOp forOp)
Canonicalize the bounds of the given loop.
Definition: AffineOps.cpp:2309
static void simplifyExprAndOperands(AffineExpr &expr, unsigned numDims, unsigned numSymbols, ArrayRef< Value > operands)
Simplify expr while exploiting information from the values in operands.
Definition: AffineOps.cpp:811
static bool isValidAffineIndexOperand(Value value, Region *region)
Definition: AffineOps.cpp:481
static void canonicalizeMapOrSetAndOperands(MapOrSet *mapOrSet, SmallVectorImpl< Value > *operands)
Definition: AffineOps.cpp:1432
static void composeAffineMapAndOperands(AffineMap *map, SmallVectorImpl< Value > *operands)
Iterate over operands and fold away all those produced by an AffineApplyOp iteratively.
Definition: AffineOps.cpp:1096
static std::optional< int64_t > getUpperBound(Value iv)
Gets the constant upper bound on an affine.for iv.
Definition: AffineOps.cpp:746
static ParseResult parseBound(bool isLower, OperationState &result, OpAsmParser &p)
Parse a for operation loop bounds.
Definition: AffineOps.cpp:1998
static std::optional< int64_t > getLowerBound(Value iv)
Gets the constant lower bound on an iv.
Definition: AffineOps.cpp:738
static void composeSetAndOperands(IntegerSet &set, SmallVectorImpl< Value > &operands)
Compose any affine.apply ops feeding into operands of the integer set set by composing the maps of su...
Definition: AffineOps.cpp:3020
static LogicalResult replaceDimOrSym(AffineMap *map, unsigned dimOrSymbolPosition, SmallVectorImpl< Value > &dims, SmallVectorImpl< Value > &syms)
Replace all occurrences of AffineExpr at position pos in map by the defining AffineApplyOp expression...
Definition: AffineOps.cpp:1048
static void canonicalizePromotedSymbols(MapOrSet *mapOrSet, SmallVectorImpl< Value > *operands)
Definition: AffineOps.cpp:1338
static LogicalResult verifyVectorMemoryOp(Operation *op, MemRefType memrefType, VectorType vectorType)
Verify common invariants of affine.vector_load and affine.vector_store.
Definition: AffineOps.cpp:4487
static void simplifyMinOrMaxExprWithOperands(AffineMap &map, ArrayRef< Value > operands, bool isMax)
Simplify the expressions in map while making use of lower or upper bounds of its operands.
Definition: AffineOps.cpp:914
static ParseResult parseAffineMinMaxOp(OpAsmParser &parser, OperationState &result)
Definition: AffineOps.cpp:3318
static bool isMemRefSizeValidSymbol(AnyMemRefDefOp memrefDefOp, unsigned index, Region *region)
Returns true if the 'index' dimension of the memref defined by memrefDefOp is a statically shaped one...
Definition: AffineOps.cpp:340
static bool isNonNegativeBoundedBy(AffineExpr e, ArrayRef< Value > operands, int64_t k)
Check if e is known to be: 0 <= e < k.
Definition: AffineOps.cpp:686
static ParseResult parseAffineMapWithMinMax(OpAsmParser &parser, OperationState &result, MinMaxKind kind)
Parses an affine map that can contain a min/max for groups of its results, e.g., max(expr-1,...
Definition: AffineOps.cpp:4200
static AffineForOp buildAffineLoopFromValues(OpBuilder &builder, Location loc, Value lb, Value ub, int64_t step, AffineForOp::BodyBuilderFn bodyBuilderFn)
Creates an affine loop from the bounds that may or may not be constants.
Definition: AffineOps.cpp:2744
static void printDimAndSymbolList(Operation::operand_iterator begin, Operation::operand_iterator end, unsigned numDims, OpAsmPrinter &printer)
Prints dimension and symbol list.
Definition: AffineOps.cpp:486
static int64_t getLargestKnownDivisor(AffineExpr e, ArrayRef< Value > operands)
Returns the largest known divisor of e.
Definition: AffineOps.cpp:648
static void legalizeDemotedDims(MapOrSet &mapOrSet, SmallVectorImpl< Value > &operands)
A valid affine dimension may appear as a symbol in affine.apply operations.
Definition: AffineOps.cpp:1386
static OpTy makeComposedMinMax(OpBuilder &b, Location loc, AffineMap map, ArrayRef< OpFoldResult > operands)
Definition: AffineOps.cpp:1275
static void buildAffineLoopNestImpl(OpBuilder &builder, Location loc, BoundListTy lbs, BoundListTy ubs, ArrayRef< int64_t > steps, function_ref< void(OpBuilder &, Location, ValueRange)> bodyBuilderFn, LoopCreatorTy &&loopCreatorFn)
Builds an affine loop nest, using "loopCreatorFn" to create individual loop operations.
Definition: AffineOps.cpp:2694
static LogicalResult foldLoopBounds(AffineForOp forOp)
Fold the constant bounds of a loop.
Definition: AffineOps.cpp:2263
static LogicalResult verifyDimAndSymbolIdentifiers(OpTy &op, Operation::operand_range operands, unsigned numDims)
Utility function to verify that a set of operands are valid dimension and symbol identifiers.
Definition: AffineOps.cpp:518
static OpFoldResult makeComposedFoldedMinMax(OpBuilder &b, Location loc, AffineMap map, ArrayRef< OpFoldResult > operands)
Definition: AffineOps.cpp:1290
static bool isDimOpValidSymbol(ShapedDimOpInterface dimOp, Region *region)
Returns true if the result of the dim op is a valid symbol for region.
Definition: AffineOps.cpp:359
static bool isQTimesDPlusR(AffineExpr e, ArrayRef< Value > operands, int64_t &div, AffineExpr &quotientTimesDiv, AffineExpr &rem)
Check if expression e is of the form d*e_1 + e_2 where 0 <= e_2 < d.
Definition: AffineOps.cpp:714
static ParseResult deduplicateAndResolveOperands(OpAsmParser &parser, ArrayRef< SmallVector< OpAsmParser::UnresolvedOperand >> operands, SmallVectorImpl< Value > &uniqueOperands, SmallVectorImpl< AffineExpr > &replacements, AffineExprKind kind)
Given a list of lists of parsed operands, populates uniqueOperands with unique operands.
Definition: AffineOps.cpp:4153
static LogicalResult verifyAffineMinMaxOp(T op)
Definition: AffineOps.cpp:3292
static void printBound(AffineMapAttr boundMap, Operation::operand_range boundOperands, const char *prefix, OpAsmPrinter &p)
Definition: AffineOps.cpp:2173
static std::optional< SmallVector< int64_t > > foldCstValueToCstAttrBasis(ArrayRef< OpFoldResult > mixedBasis, MutableOperandRange mutableDynamicBasis, ArrayRef< Attribute > dynamicBasis)
Given mixed basis of affine.delinearize_index/linearize_index replace constant SSA values with the co...
Definition: AffineOps.cpp:4673
static LogicalResult canonicalizeMapExprAndTermOrder(AffineMap &map)
Canonicalize the result expression order of an affine map and return success if the order changed.
Definition: AffineOps.cpp:3504
static Value getZero(OpBuilder &b, Location loc, Type elementType)
Get zero value for an element type.
static Operation * materializeConstant(Dialect *dialect, OpBuilder &builder, Attribute value, Type type, Location loc)
A utility function used to materialize a constant for a given attribute and type.
Definition: FoldUtils.cpp:50
static MLIRContext * getContext(OpFoldResult val)
static bool isLegalToInline(InlinerInterface &interface, Region *src, Region *insertRegion, bool shouldCloneInlinedRegion, IRMapping &valueMapping)
Utility to check that all of the operations within 'src' can be inlined.
static int64_t getNumElements(Type t)
Compute the total number of elements in the given type, also taking into account nested types.
union mlir::linalg::@1197::ArityGroupAndKind::Kind kind
static Operation::operand_range getLowerBoundOperands(AffineForOp forOp)
Definition: SCFToGPU.cpp:76
static Operation::operand_range getUpperBoundOperands(AffineForOp forOp)
Definition: SCFToGPU.cpp:81
static void print(spirv::VerCapExtAttr triple, DialectAsmPrinter &printer)
static VectorType getVectorType(Type scalarTy, const VectorizationStrategy *strategy)
Returns the vector type resulting from applying the provided vectorization strategy on the scalar typ...
RetTy walkPostOrder(AffineExpr expr)
Base type for affine expression.
Definition: AffineExpr.h:68
AffineExpr floorDiv(uint64_t v) const
Definition: AffineExpr.cpp:921
AffineExprKind getKind() const
Return the classification for this type.
Definition: AffineExpr.cpp:35
int64_t getLargestKnownDivisor() const
Returns the greatest known integral divisor of this affine expression.
Definition: AffineExpr.cpp:243
MLIRContext * getContext() const
Definition: AffineExpr.cpp:33
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
Definition: AffineMap.h:46
AffineMap getSliceMap(unsigned start, unsigned length) const
Returns the map consisting of length expressions starting from start.
Definition: AffineMap.cpp:662
MLIRContext * getContext() const
Definition: AffineMap.cpp:343
bool isFunctionOfDim(unsigned position) const
Return true if any affine expression involves AffineDimExpr position.
Definition: AffineMap.h:221
static AffineMap get(MLIRContext *context)
Returns a zero result affine map with no dimensions or symbols: () -> ().
AffineMap shiftDims(unsigned shift, unsigned offset=0) const
Replace dims[offset ...
Definition: AffineMap.h:267
unsigned getNumSymbols() const
Definition: AffineMap.cpp:398
unsigned getNumDims() const
Definition: AffineMap.cpp:394
ArrayRef< AffineExpr > getResults() const
Definition: AffineMap.cpp:407
bool isFunctionOfSymbol(unsigned position) const
Return true if any affine expression involves AffineSymbolExpr position.
Definition: AffineMap.h:228
unsigned getNumResults() const
Definition: AffineMap.cpp:402
AffineMap replaceDimsAndSymbols(ArrayRef< AffineExpr > dimReplacements, ArrayRef< AffineExpr > symReplacements, unsigned numResultDims, unsigned numResultSyms) const
This method substitutes any uses of dimensions and symbols (e.g.
Definition: AffineMap.cpp:500
unsigned getNumInputs() const
Definition: AffineMap.cpp:403
AffineMap shiftSymbols(unsigned shift, unsigned offset=0) const
Replace symbols[offset ...
Definition: AffineMap.h:280
AffineExpr getResult(unsigned idx) const
Definition: AffineMap.cpp:411
AffineMap replace(AffineExpr expr, AffineExpr replacement, unsigned numResultDims, unsigned numResultSyms) const
Sparse replace method.
Definition: AffineMap.cpp:515
static AffineMap getConstantMap(int64_t val, MLIRContext *context)
Returns a single constant result affine map.
Definition: AffineMap.cpp:128
AffineMap getSubMap(ArrayRef< unsigned > resultPos) const
Returns the map consisting of the resultPos subset.
Definition: AffineMap.cpp:654
LogicalResult constantFold(ArrayRef< Attribute > operandConstants, SmallVectorImpl< Attribute > &results, bool *hasPoison=nullptr) const
Folds the results of the application of an affine map on the provided operands to a constant if possi...
Definition: AffineMap.cpp:434
static SmallVector< AffineMap, 4 > inferFromExprList(ArrayRef< ArrayRef< AffineExpr >> exprsList, MLIRContext *context)
Returns a vector of AffineMaps; each with as many results as exprs.size(), as many dims as the larges...
Definition: AffineMap.cpp:312
@ Paren
Parens surrounding zero or more operands.
@ OptionalSquare
Square brackets supporting zero or more ops, or nothing.
virtual ParseResult parseColonTypeList(SmallVectorImpl< Type > &result)=0
Parse a colon followed by a type list, which must have at least one type.
virtual ParseResult parseCommaSeparatedList(Delimiter delimiter, function_ref< ParseResult()> parseElementFn, StringRef contextMessage=StringRef())=0
Parse a list of comma-separated items with an optional delimiter.
virtual Builder & getBuilder() const =0
Return a builder which provides useful access to MLIRContext, global objects like types and attribute...
virtual ParseResult parseOptionalAttrDict(NamedAttrList &result)=0
Parse a named dictionary into 'result' if it is present.
virtual ParseResult parseOptionalKeyword(StringRef keyword)=0
Parse the given keyword if present.
MLIRContext * getContext() const
Definition: AsmPrinter.cpp:73
virtual ParseResult parseRParen()=0
Parse a ) token.
virtual InFlightDiagnostic emitError(SMLoc loc, const Twine &message={})=0
Emit a diagnostic at the specified location and return failure.
ParseResult addTypeToList(Type type, SmallVectorImpl< Type > &result)
Add the specified type to the end of the specified type list and return success.
virtual ParseResult parseOptionalRParen()=0
Parse a ) token if present.
virtual ParseResult parseLess()=0
Parse a '<' token.
virtual ParseResult parseEqual()=0
Parse a = token.
virtual ParseResult parseColonType(Type &result)=0
Parse a colon followed by a type.
virtual SMLoc getCurrentLocation()=0
Get the location of the next token and store it into the argument.
virtual SMLoc getNameLoc() const =0
Return the location of the original name token.
virtual ParseResult parseGreater()=0
Parse a '>' token.
virtual ParseResult parseLParen()=0
Parse a ( token.
virtual ParseResult parseType(Type &result)=0
Parse a type.
virtual ParseResult parseComma()=0
Parse a , token.
virtual ParseResult parseOptionalArrowTypeList(SmallVectorImpl< Type > &result)=0
Parse an optional arrow followed by a type list.
virtual ParseResult parseArrowTypeList(SmallVectorImpl< Type > &result)=0
Parse an arrow followed by a type list.
ParseResult parseKeyword(StringRef keyword)
Parse a given keyword.
virtual ParseResult parseAttribute(Attribute &result, Type type={})=0
Parse an arbitrary attribute of a given type and return it in result.
void printOptionalArrowTypeList(TypeRange &&types)
Print an optional arrow followed by a type list.
Attributes are known-constant values of operations.
Definition: Attributes.h:25
Block represents an ordered list of Operations.
Definition: Block.h:33
Operation * getTerminator()
Get the terminator operation of this block.
Definition: Block.cpp:246
BlockArgument addArgument(Type type, Location loc)
Add one value to the argument list.
Definition: Block.cpp:155
BlockArgListType getArguments()
Definition: Block.h:87
Operation & front()
Definition: Block.h:153
This class is a general helper class for creating context-global objects like types,...
Definition: Builders.h:51
DenseI32ArrayAttr getDenseI32ArrayAttr(ArrayRef< int32_t > values)
Definition: Builders.cpp:159
IntegerAttr getIntegerAttr(Type type, int64_t value)
Definition: Builders.cpp:224
AffineMap getDimIdentityMap()
Definition: Builders.cpp:379
AffineMap getMultiDimIdentityMap(unsigned rank)
Definition: Builders.cpp:383
AffineExpr getAffineSymbolExpr(unsigned position)
Definition: Builders.cpp:364
AffineExpr getAffineConstantExpr(int64_t constant)
Definition: Builders.cpp:368
DenseIntElementsAttr getI32TensorAttr(ArrayRef< int32_t > values)
Tensor-typed DenseIntElementsAttr getters.
Definition: Builders.cpp:175
IntegerAttr getI64IntegerAttr(int64_t value)
Definition: Builders.cpp:108
IntegerType getIntegerType(unsigned width)
Definition: Builders.cpp:67
NoneType getNoneType()
Definition: Builders.cpp:84
BoolAttr getBoolAttr(bool value)
Definition: Builders.cpp:96
AffineMap getEmptyAffineMap()
Returns a zero result affine map with no dimensions or symbols: () -> ().
Definition: Builders.cpp:372
AffineMap getConstantAffineMap(int64_t val)
Returns a single constant result affine map with 0 dimensions and 0 symbols.
Definition: Builders.cpp:374
MLIRContext * getContext() const
Definition: Builders.h:56
AffineMap getSymbolIdentityMap()
Definition: Builders.cpp:392
ArrayAttr getArrayAttr(ArrayRef< Attribute > value)
Definition: Builders.cpp:262
ArrayAttr getI64ArrayAttr(ArrayRef< int64_t > values)
Definition: Builders.cpp:277
IndexType getIndexType()
Definition: Builders.cpp:51
An attribute that represents a reference to a dense integer vector or tensor object.
This is the interface that must be implemented by the dialects of operations to be inlined.
Definition: InliningUtils.h:44
DialectInlinerInterface(Dialect *dialect)
Definition: InliningUtils.h:46
This is a utility class for mapping one set of IR entities to another.
Definition: IRMapping.h:26
auto lookup(T from) const
Lookup a mapped value within the map.
Definition: IRMapping.h:72
An integer set representing a conjunction of one or more affine equalities and inequalities.
Definition: IntegerSet.h:44
unsigned getNumDims() const
Definition: IntegerSet.cpp:15
static IntegerSet get(unsigned dimCount, unsigned symbolCount, ArrayRef< AffineExpr > constraints, ArrayRef< bool > eqFlags)
MLIRContext * getContext() const
Definition: IntegerSet.cpp:57
unsigned getNumInputs() const
Definition: IntegerSet.cpp:17
ArrayRef< AffineExpr > getConstraints() const
Definition: IntegerSet.cpp:41
ArrayRef< bool > getEqFlags() const
Returns the equality bits, which specify whether each of the constraints is an equality or inequality...
Definition: IntegerSet.cpp:51
unsigned getNumSymbols() const
Definition: IntegerSet.cpp:16
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:66
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:60
This class provides a mutable adaptor for a range of operands.
Definition: ValueRange.h:118
void erase(unsigned subStart, unsigned subLen=1)
Erase the operands within the given sub-range.
NamedAttrList is array of NamedAttributes that tracks whether it is sorted and does some basic work t...
void pop_back()
Pop last element from list.
Attribute erase(StringAttr name)
Erase the attribute with the given name from the list.
The OpAsmParser has methods for interacting with the asm parser: parsing things from it,...
virtual ParseResult parseRegion(Region &region, ArrayRef< Argument > arguments={}, bool enableNameShadowing=false)=0
Parses a region.
virtual ParseResult parseArgument(Argument &result, bool allowType=false, bool allowAttrs=false)=0
Parse a single argument with the following syntax:
ParseResult parseTrailingOperandList(SmallVectorImpl< UnresolvedOperand > &result, Delimiter delimiter=Delimiter::None)
Parse zero or more trailing SSA comma-separated trailing operand references with a specified surround...
virtual ParseResult parseArgumentList(SmallVectorImpl< Argument > &result, Delimiter delimiter=Delimiter::None, bool allowType=false, bool allowAttrs=false)=0
Parse zero or more arguments with a specified surrounding delimiter.
virtual ParseResult parseAffineMapOfSSAIds(SmallVectorImpl< UnresolvedOperand > &operands, Attribute &map, StringRef attrName, NamedAttrList &attrs, Delimiter delimiter=Delimiter::Square)=0
Parses an affine map attribute where dims and symbols are SSA operands.
ParseResult parseAssignmentList(SmallVectorImpl< Argument > &lhs, SmallVectorImpl< UnresolvedOperand > &rhs)
Parse a list of assignments of the form (x1 = y1, x2 = y2, ...)
virtual ParseResult resolveOperand(const UnresolvedOperand &operand, Type type, SmallVectorImpl< Value > &result)=0
Resolve an operand to an SSA value, emitting an error on failure.
ParseResult resolveOperands(Operands &&operands, Type type, SmallVectorImpl< Value > &result)
Resolve a list of operands to SSA values, emitting an error on failure, or appending the results to t...
virtual ParseResult parseOperand(UnresolvedOperand &result, bool allowResultNumber=true)=0
Parse a single SSA value operand name along with a result number if allowResultNumber is true.
virtual ParseResult parseAffineExprOfSSAIds(SmallVectorImpl< UnresolvedOperand > &dimOperands, SmallVectorImpl< UnresolvedOperand > &symbOperands, AffineExpr &expr)=0
Parses an affine expression where dims and symbols are SSA operands.
virtual ParseResult parseOperandList(SmallVectorImpl< UnresolvedOperand > &result, Delimiter delimiter=Delimiter::None, bool allowResultNumber=true, int requiredOperandCount=-1)=0
Parse zero or more SSA comma-separated operand references with a specified surrounding delimiter,...
This is a pure-virtual base class that exposes the asmprinter hooks necessary to implement a custom p...
virtual void printOptionalAttrDict(ArrayRef< NamedAttribute > attrs, ArrayRef< StringRef > elidedAttrs={})=0
If the specified operation has attributes, print out an attribute dictionary with their values.
virtual void printAffineExprOfSSAIds(AffineExpr expr, ValueRange dimOperands, ValueRange symOperands)=0
Prints an affine expression of SSA ids with SSA id names used instead of dims and symbols.
virtual void printAffineMapOfSSAIds(AffineMapAttr mapAttr, ValueRange operands)=0
Prints an affine map of SSA ids, where SSA id names are used in place of dims/symbols.
virtual void printRegion(Region &blocks, bool printEntryBlockArgs=true, bool printBlockTerminators=true, bool printEmptyBlock=false)=0
Prints a region.
virtual void printRegionArgument(BlockArgument arg, ArrayRef< NamedAttribute > argAttrs={}, bool omitType=false)=0
Print a block argument in the usual format of: ssaName : type {attr1=42} loc("here") where location p...
virtual void printOperand(Value value)=0
Print implementations for various things an operation contains.
RAII guard to reset the insertion point of the builder when destroyed.
Definition: Builders.h:346
This class helps build Operations.
Definition: Builders.h:205
Block::iterator getInsertionPoint() const
Returns the current insertion point of the builder.
Definition: Builders.h:443
void setInsertionPointToStart(Block *block)
Sets the insertion point to the start of the specified block.
Definition: Builders.h:429
void setInsertionPoint(Block *block, Block::iterator insertPoint)
Set the insertion point to the specified location.
Definition: Builders.h:396
Listener * getListener() const
Returns the current listener of this builder, or nullptr if this builder doesn't have a listener.
Definition: Builders.h:318
Block * createBlock(Region *parent, Region::iterator insertPt={}, TypeRange argTypes=std::nullopt, ArrayRef< Location > locs=std::nullopt)
Add new block with 'argTypes' arguments and set the insertion point to the end of it.
Definition: Builders.cpp:426
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
Definition: Builders.cpp:453
Block * getInsertionBlock() const
Return the block the current insertion point belongs to.
Definition: Builders.h:440
This class represents a single result from folding an operation.
Definition: OpDefinition.h:271
This class represents an operand of an operation.
Definition: Value.h:243
A trait of region holding operations that defines a new scope for polyhedral optimization purposes.
This class provides the API for ops that are known to be isolated from above.
A trait used to provide symbol table functionalities to a region operation.
Definition: SymbolTable.h:442
This class implements the operand iterators for the Operation class.
Definition: ValueRange.h:43
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
bool hasTrait()
Returns true if the operation was registered with a particular trait, e.g.
Definition: Operation.h:750
OpResult getResult(unsigned idx)
Get the 'idx'th result of this operation.
Definition: Operation.h:407
Operation * getParentOp()
Returns the closest surrounding operation that contains this operation or nullptr if this is a top-le...
Definition: Operation.h:234
operand_range getOperands()
Returns an iterator on the underlying Value's.
Definition: Operation.h:378
Region * getParentRegion()
Returns the region to which the instruction belongs.
Definition: Operation.h:230
bool isProperAncestor(Operation *other)
Return true if this operation is a proper ancestor of the other operation.
Definition: Operation.cpp:219
operand_range::iterator operand_iterator
Definition: Operation.h:372
InFlightDiagnostic emitOpError(const Twine &message={})
Emit an error with the op name prefixed, like "'dim' op " which is convenient for verifiers.
Definition: Operation.cpp:673
A special type of RewriterBase that coordinates the application of a rewrite pattern on the current I...
Definition: PatternMatch.h:749
This class represents a point being branched from in the methods of the RegionBranchOpInterface.
bool isParent() const
Returns true if branching from the parent op.
This class represents a successor of a region.
This class contains a list of basic blocks and a link to the parent operation it is attached to.
Definition: Region.h:26
Operation * getParentOp()
Return the parent operation this region is attached to.
Definition: Region.h:200
bool empty()
Definition: Region.h:60
Block & front()
Definition: Region.h:65
RewritePatternSet & add(ConstructorArg &&arg, ConstructorArgs &&...args)
Add an instance of each of the pattern types 'Ts' to the pattern list with the given arguments.
Definition: PatternMatch.h:811
This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...
Definition: PatternMatch.h:358
std::enable_if_t<!std::is_convertible< CallbackT, Twine >::value, LogicalResult > notifyMatchFailure(Location loc, CallbackT &&reasonCallback)
Used to notify the listener that the IR failed to be rewritten because of a match failure,...
Definition: PatternMatch.h:682
virtual void eraseBlock(Block *block)
This method erases all operations in a block.
virtual void replaceOp(Operation *op, ValueRange newValues)
Replace the results of the given (original) operation with the specified list of values (replacements...
void mergeBlocks(Block *source, Block *dest, ValueRange argValues=std::nullopt)
Inline the operations of block 'source' into the end of block 'dest'.
virtual void finalizeOpModification(Operation *op)
This method is used to signal the end of an in-place modification of the given operation.
virtual void eraseOp(Operation *op)
This method erases an operation that is known to have no uses.
void replaceUsesWithIf(Value from, Value to, function_ref< bool(OpOperand &)> functor, bool *allUsesReplaced=nullptr)
Find uses of from and replace them with to if the functor returns true.
void modifyOpInPlace(Operation *root, CallableT &&callable)
This method is a utility wrapper around an in-place modification of an operation.
Definition: PatternMatch.h:594
virtual void inlineBlockBefore(Block *source, Block *dest, Block::iterator before, ValueRange argValues=std::nullopt)
Inline the operations of block 'source' into block 'dest' before the given position.
virtual void startOpModification(Operation *op)
This method is used to notify the rewriter that an in-place operation modification is about to happen...
Definition: PatternMatch.h:578
OpTy replaceOpWithNewOp(Operation *op, Args &&...args)
Replace the results of the given (original) op with a new op that is created without verification (re...
Definition: PatternMatch.h:500
This class represents a specific instance of an effect.
static DerivedEffect * get()
Returns a unique instance for the derived effect class.
static DefaultResource * get()
Returns a unique instance for the given effect class.
std::vector< SmallVector< int64_t, 8 > > operandExprStack
static Operation * lookupSymbolIn(Operation *op, StringAttr symbol)
Returns the operation registered with the given symbol name with the regions of 'symbolTableOp'.
This class provides an abstraction over the various different ranges of value types.
Definition: TypeRange.h:37
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
Definition: Types.h:74
bool isIndex() const
Definition: Types.cpp:54
This class provides an abstraction over the different types of ranges over Values.
Definition: ValueRange.h:387
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:96
Type getType() const
Return the type of this value.
Definition: Value.h:105
Operation * getDefiningOp() const
If this value is the result of an operation, return the operation that defines it.
Definition: Value.cpp:20
AffineBound represents a lower or upper bound in the for operation.
Definition: AffineOps.h:523
AffineDmaStartOp starts a non-blocking DMA operation that transfers data from a source memref to a de...
Definition: AffineOps.h:106
AffineDmaWaitOp blocks until the completion of a DMA operation associated with the tag element 'tag[i...
Definition: AffineOps.h:315
An AffineValueMap is an affine map plus its ML value operands and results for analysis purposes.
LogicalResult canonicalize()
Attempts to canonicalize the map and operands.
Definition: AffineOps.cpp:4045
ArrayRef< Value > getOperands() const
AffineExpr getResult(unsigned i)
unsigned getNumResults() const
Operation * getOwner() const
Return the owner of this operand.
Definition: UseDefLists.h:38
constexpr auto RecursivelySpeculatable
Speculatability
This enum is returned from the getSpeculatability method in the ConditionallySpeculatable op interfac...
constexpr auto NotSpeculatable
void buildAffineLoopNest(OpBuilder &builder, Location loc, ArrayRef< int64_t > lbs, ArrayRef< int64_t > ubs, ArrayRef< int64_t > steps, function_ref< void(OpBuilder &, Location, ValueRange)> bodyBuilderFn=nullptr)
Builds a perfect nest of affine.for loops, i.e., each loop except the innermost one contains only ano...
Definition: AffineOps.cpp:2757
void fullyComposeAffineMapAndOperands(AffineMap *map, SmallVectorImpl< Value > *operands)
Given an affine map map and its input operands, this method composes into map, maps of AffineApplyOps...
Definition: AffineOps.cpp:1158
void extractForInductionVars(ArrayRef< AffineForOp > forInsts, SmallVectorImpl< Value > *ivs)
Extracts the induction variables from a list of AffineForOps and places them in the output argument i...
Definition: AffineOps.cpp:2671
bool isValidDim(Value value)
Returns true if the given Value can be used as a dimension id in the region of the closest surroundin...
Definition: AffineOps.cpp:288
SmallVector< OpFoldResult > makeComposedFoldedMultiResultAffineApply(OpBuilder &b, Location loc, AffineMap map, ArrayRef< OpFoldResult > operands)
Variant of makeComposedFoldedAffineApply suitable for multi-result maps.
Definition: AffineOps.cpp:1264
bool isAffineInductionVar(Value val)
Returns true if the provided value is the induction variable of an AffineForOp or AffineParallelOp.
Definition: AffineOps.cpp:2643
AffineForOp getForInductionVarOwner(Value val)
Returns the loop parent of an induction variable.
Definition: AffineOps.cpp:2647
AffineApplyOp makeComposedAffineApply(OpBuilder &b, Location loc, AffineMap map, ArrayRef< OpFoldResult > operands)
Returns a composed AffineApplyOp by composing map and operands with other AffineApplyOps supplying th...
Definition: AffineOps.cpp:1168
void canonicalizeMapAndOperands(AffineMap *map, SmallVectorImpl< Value > *operands)
Modifies both map and operands in-place so as to:
Definition: AffineOps.cpp:1509
OpFoldResult makeComposedFoldedAffineMax(OpBuilder &b, Location loc, AffineMap map, ArrayRef< OpFoldResult > operands)
Constructs an AffineMinOp that computes a maximum across the results of applying map to operands,...
Definition: AffineOps.cpp:1329
bool isAffineForInductionVar(Value val)
Returns true if the provided value is the induction variable of an AffineForOp.
Definition: AffineOps.cpp:2635
OpFoldResult makeComposedFoldedAffineMin(OpBuilder &b, Location loc, AffineMap map, ArrayRef< OpFoldResult > operands)
Constructs an AffineMinOp that computes a minimum across the results of applying map to operands,...
Definition: AffineOps.cpp:1322
bool isTopLevelValue(Value value)
A utility function to check if a value is defined at the top level of an op with trait AffineScope or...
Definition: AffineOps.cpp:248
Region * getAffineAnalysisScope(Operation *op)
Returns the closest region enclosing op that is held by a non-affine operation; nullptr if there is n...
Definition: AffineOps.cpp:273
void canonicalizeSetAndOperands(IntegerSet *set, SmallVectorImpl< Value > *operands)
Canonicalizes an integer set the same way canonicalizeMapAndOperands does for affine maps.
Definition: AffineOps.cpp:1514
void extractInductionVars(ArrayRef< Operation * > affineOps, SmallVectorImpl< Value > &ivs)
Extracts the induction variables from a list of either AffineForOp or AffineParallelOp and places the...
Definition: AffineOps.cpp:2678
bool isValidSymbol(Value value)
Returns true if the given value can be used as a symbol in the region of the closest surrounding op t...
Definition: AffineOps.cpp:403
OpFoldResult makeComposedFoldedAffineApply(OpBuilder &b, Location loc, AffineMap map, ArrayRef< OpFoldResult > operands)
Constructs an AffineApplyOp that applies map to operands after composing the map with the maps of any...
Definition: AffineOps.cpp:1218
AffineParallelOp getAffineParallelInductionVarOwner(Value val)
Returns true if the provided value is among the induction variables of an AffineParallelOp.
Definition: AffineOps.cpp:2658
Region * getAffineScope(Operation *op)
Returns the closest region enclosing op that is held by an operation with trait AffineScope; nullptr ...
Definition: AffineOps.cpp:263
ParseResult parseDimAndSymbolList(OpAsmParser &parser, SmallVectorImpl< Value > &operands, unsigned &numDims)
Parses dimension and symbol list.
Definition: AffineOps.cpp:496
bool isAffineParallelInductionVar(Value val)
Returns true if val is the induction variable of an AffineParallelOp.
Definition: AffineOps.cpp:2639
AffineMinOp makeComposedAffineMin(OpBuilder &b, Location loc, AffineMap map, ArrayRef< OpFoldResult > operands)
Returns an AffineMinOp obtained by composing map and operands with AffineApplyOps supplying those ope...
Definition: AffineOps.cpp:1284
BaseMemRefType getMemRefType(Value value, const BufferizationOptions &options, MemRefLayoutAttrInterface layout={}, Attribute memorySpace=nullptr)
Return a MemRefType to which the type of the given value can be bufferized.
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Definition: Matchers.h:344
LogicalResult foldMemRefCast(Operation *op, Value inner=nullptr)
This is a common utility used for patterns of the form "someop(memref.cast) -> someop".
Definition: MemRefOps.cpp:45
QueryRef parse(llvm::StringRef line, const QuerySession &qs)
Definition: Query.cpp:20
Include the generated interface declarations.
AffineMap simplifyAffineMap(AffineMap map)
Simplifies an affine map by simplifying its underlying AffineExpr results.
Definition: AffineMap.cpp:773
bool matchPattern(Value value, const Pattern &pattern)
Entry point for matching a pattern over a Value.
Definition: Matchers.h:490
OpFoldResult getAsIndexOpFoldResult(MLIRContext *ctx, int64_t val)
Convert int64_t to integer attributes of index type and return them as OpFoldResult.
const FrozenRewritePatternSet GreedyRewriteConfig bool * changed
AffineMap removeDuplicateExprs(AffineMap map)
Returns a map with the same dimension and symbol count as map, but whose results are the unique affin...
Definition: AffineMap.cpp:783
std::optional< int64_t > getConstantIntValue(OpFoldResult ofr)
If ofr is a constant integer or an IntegerAttr, return the integer.
std::optional< int64_t > getBoundForAffineExpr(AffineExpr expr, unsigned numDims, unsigned numSymbols, ArrayRef< std::optional< int64_t >> constLowerBounds, ArrayRef< std::optional< int64_t >> constUpperBounds, bool isUpper)
Get a lower or upper (depending on isUpper) bound for expr while using the constant lower and upper b...
Type getType(OpFoldResult ofr)
Returns the int type of the integer in ofr.
Definition: Utils.cpp:305
SmallVector< int64_t > delinearize(int64_t linearIndex, ArrayRef< int64_t > strides)
Given the strides together with a linear index in the dimension space, return the vector-space offset...
InFlightDiagnostic emitError(Location loc)
Utility method to emit an error message using this location.
bool isPure(Operation *op)
Returns true if the given operation is pure, i.e., is speculatable that does not touch memory.
int64_t computeProduct(ArrayRef< int64_t > basis)
Self-explicit.
AffineExprKind
Definition: AffineExpr.h:40
@ CeilDiv
RHS of ceildiv is always a constant or a symbolic expression.
@ Mod
RHS of mod is always a constant or a symbolic expression with a positive value.
@ DimId
Dimensional identifier.
@ FloorDiv
RHS of floordiv is always a constant or a symbolic expression.
@ SymbolId
Symbolic identifier.
AffineExpr getAffineBinaryOpExpr(AffineExprKind kind, AffineExpr lhs, AffineExpr rhs)
Definition: AffineExpr.cpp:70
std::function< SmallVector< Value >(OpBuilder &b, Location loc, ArrayRef< BlockArgument > newBbArgs)> NewYieldValuesFn
A function that returns the additional yielded values during replaceWithAdditionalYields.
detail::constant_int_predicate_matcher m_Zero()
Matches a constant scalar / vector splat / tensor splat integer zero.
Definition: Matchers.h:442
const FrozenRewritePatternSet & patterns
void dispatchIndexOpFoldResults(ArrayRef< OpFoldResult > ofrs, SmallVectorImpl< Value > &dynamicVec, SmallVectorImpl< int64_t > &staticVec)
Helper function to dispatch multiple OpFoldResults according to the behavior of dispatchIndexOpFoldRe...
AffineExpr getAffineConstantExpr(int64_t constant, MLIRContext *context)
Definition: AffineExpr.cpp:645
auto get(MLIRContext *context, Ts &&...params)
Helper method that injects context only if needed, this helps unify some of the attribute constructio...
OpFoldResult getAsOpFoldResult(Value val)
Given a value, try to extract a constant Attribute.
SmallVector< OpFoldResult > getMixedValues(ArrayRef< int64_t > staticValues, ValueRange dynamicValues, MLIRContext *context)
Return a vector of OpFoldResults with the same size a staticValues, but all elements for which Shaped...
detail::constant_op_matcher m_Constant()
Matches a constant foldable operation.
Definition: Matchers.h:369
AffineExpr getAffineDimExpr(unsigned position, MLIRContext *context)
These free functions allow clients of the API to not use classes in detail.
Definition: AffineExpr.cpp:621
LogicalResult verify(Operation *op, bool verifyRecursively=true)
Perform (potentially expensive) checks of invariants, used to detect compiler bugs,...
Definition: Verifier.cpp:423
AffineMap foldAttributesIntoMap(Builder &b, AffineMap map, ArrayRef< OpFoldResult > operands, SmallVector< Value > &remainingValues)
Fold all attributes among the given operands into the affine map.
Definition: AffineMap.cpp:745
AffineExpr getAffineSymbolExpr(unsigned position, MLIRContext *context)
Definition: AffineExpr.cpp:631
Canonicalize the affine map result expression order of an affine min/max operation.
Definition: AffineOps.cpp:3558
LogicalResult matchAndRewrite(T affineOp, PatternRewriter &rewriter) const override
Definition: AffineOps.cpp:3561
LogicalResult matchAndRewrite(T affineOp, PatternRewriter &rewriter) const override
Definition: AffineOps.cpp:3575
Remove duplicated expressions in affine min/max ops.
Definition: AffineOps.cpp:3374
LogicalResult matchAndRewrite(T affineOp, PatternRewriter &rewriter) const override
Definition: AffineOps.cpp:3377
Merge an affine min/max op to its consumers if its consumer is also an affine min/max op.
Definition: AffineOps.cpp:3417
LogicalResult matchAndRewrite(T affineOp, PatternRewriter &rewriter) const override
Definition: AffineOps.cpp:3420
This is the representation of an operand reference.
This class represents a listener that may be used to hook into various actions within an OpBuilder.
Definition: Builders.h:283
OpRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting against an...
Definition: PatternMatch.h:314
OpRewritePattern(MLIRContext *context, PatternBenefit benefit=1, ArrayRef< StringRef > generatedNames={})
Patterns must specify the root operation name they match against, and can also specify the benefit of...
Definition: PatternMatch.h:319
This represents an operation in an abstracted form, suitable for use with the builder APIs.
T & getOrAddProperties()
Get (or create) a properties of the provided type to be set on the operation on creation.
SmallVector< Value, 4 > operands
void addOperands(ValueRange newOperands)
void addAttribute(StringRef name, Attribute attr)
Add an attribute with the specified name.
void addTypes(ArrayRef< Type > newTypes)
SmallVector< std::unique_ptr< Region >, 1 > regions
Regions that the op will hold.
NamedAttrList attributes
SmallVector< Type, 4 > types
Types of the results of this operation.
Region * addRegion()
Create a region that should be attached to the operation.
Eliminates variable at the specified position using Fourier-Motzkin variable elimination.