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