1 //===- TestPatterns.cpp - Test dialect pattern driver ---------------------===// 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 "TestDialect.h" 10 #include "TestOps.h" 11 #include "TestTypes.h" 12 #include "mlir/Dialect/Arith/IR/Arith.h" 13 #include "mlir/Dialect/Func/IR/FuncOps.h" 14 #include "mlir/Dialect/Func/Transforms/FuncConversions.h" 15 #include "mlir/Dialect/Tensor/IR/Tensor.h" 16 #include "mlir/IR/Matchers.h" 17 #include "mlir/Pass/Pass.h" 18 #include "mlir/Transforms/DialectConversion.h" 19 #include "mlir/Transforms/FoldUtils.h" 20 #include "mlir/Transforms/GreedyPatternRewriteDriver.h" 21 #include "llvm/ADT/ScopeExit.h" 22 23 using namespace mlir; 24 using namespace test; 25 26 // Native function for testing NativeCodeCall 27 static Value chooseOperand(Value input1, Value input2, BoolAttr choice) { 28 return choice.getValue() ? input1 : input2; 29 } 30 31 static void createOpI(PatternRewriter &rewriter, Location loc, Value input) { 32 rewriter.create<OpI>(loc, input); 33 } 34 35 static void handleNoResultOp(PatternRewriter &rewriter, 36 OpSymbolBindingNoResult op) { 37 // Turn the no result op to a one-result op. 38 rewriter.create<OpSymbolBindingB>(op.getLoc(), op.getOperand().getType(), 39 op.getOperand()); 40 } 41 42 static bool getFirstI32Result(Operation *op, Value &value) { 43 if (!Type(op->getResult(0).getType()).isSignlessInteger(32)) 44 return false; 45 value = op->getResult(0); 46 return true; 47 } 48 49 static Value bindNativeCodeCallResult(Value value) { return value; } 50 51 static SmallVector<Value, 2> bindMultipleNativeCodeCallResult(Value input1, 52 Value input2) { 53 return SmallVector<Value, 2>({input2, input1}); 54 } 55 56 // Test that natives calls are only called once during rewrites. 57 // OpM_Test will return Pi, increased by 1 for each subsequent calls. 58 // This let us check the number of times OpM_Test was called by inspecting 59 // the returned value in the MLIR output. 60 static int64_t opMIncreasingValue = 314159265; 61 static Attribute opMTest(PatternRewriter &rewriter, Value val) { 62 int64_t i = opMIncreasingValue++; 63 return rewriter.getIntegerAttr(rewriter.getIntegerType(32), i); 64 } 65 66 namespace { 67 #include "TestPatterns.inc" 68 } // namespace 69 70 //===----------------------------------------------------------------------===// 71 // Test Reduce Pattern Interface 72 //===----------------------------------------------------------------------===// 73 74 void test::populateTestReductionPatterns(RewritePatternSet &patterns) { 75 populateWithGenerated(patterns); 76 } 77 78 //===----------------------------------------------------------------------===// 79 // Canonicalizer Driver. 80 //===----------------------------------------------------------------------===// 81 82 namespace { 83 struct FoldingPattern : public RewritePattern { 84 public: 85 FoldingPattern(MLIRContext *context) 86 : RewritePattern(TestOpInPlaceFoldAnchor::getOperationName(), 87 /*benefit=*/1, context) {} 88 89 LogicalResult matchAndRewrite(Operation *op, 90 PatternRewriter &rewriter) const override { 91 // Exercise createOrFold API for a single-result operation that is folded 92 // upon construction. The operation being created has an in-place folder, 93 // and it should be still present in the output. Furthermore, the folder 94 // should not crash when attempting to recover the (unchanged) operation 95 // result. 96 Value result = rewriter.createOrFold<TestOpInPlaceFold>( 97 op->getLoc(), rewriter.getIntegerType(32), op->getOperand(0)); 98 assert(result); 99 rewriter.replaceOp(op, result); 100 return success(); 101 } 102 }; 103 104 /// This pattern creates a foldable operation at the entry point of the block. 105 /// This tests the situation where the operation folder will need to replace an 106 /// operation with a previously created constant that does not initially 107 /// dominate the operation to replace. 108 struct FolderInsertBeforePreviouslyFoldedConstantPattern 109 : public OpRewritePattern<TestCastOp> { 110 public: 111 using OpRewritePattern<TestCastOp>::OpRewritePattern; 112 113 LogicalResult matchAndRewrite(TestCastOp op, 114 PatternRewriter &rewriter) const override { 115 if (!op->hasAttr("test_fold_before_previously_folded_op")) 116 return failure(); 117 rewriter.setInsertionPointToStart(op->getBlock()); 118 119 auto constOp = rewriter.create<arith::ConstantOp>( 120 op.getLoc(), rewriter.getBoolAttr(true)); 121 rewriter.replaceOpWithNewOp<TestCastOp>(op, rewriter.getI32Type(), 122 Value(constOp)); 123 return success(); 124 } 125 }; 126 127 /// This pattern matches test.op_commutative2 with the first operand being 128 /// another test.op_commutative2 with a constant on the right side and fold it 129 /// away by propagating it as its result. This is intend to check that patterns 130 /// are applied after the commutative property moves constant to the right. 131 struct FolderCommutativeOp2WithConstant 132 : public OpRewritePattern<TestCommutative2Op> { 133 public: 134 using OpRewritePattern<TestCommutative2Op>::OpRewritePattern; 135 136 LogicalResult matchAndRewrite(TestCommutative2Op op, 137 PatternRewriter &rewriter) const override { 138 auto operand = 139 dyn_cast_or_null<TestCommutative2Op>(op->getOperand(0).getDefiningOp()); 140 if (!operand) 141 return failure(); 142 Attribute constInput; 143 if (!matchPattern(operand->getOperand(1), m_Constant(&constInput))) 144 return failure(); 145 rewriter.replaceOp(op, operand->getOperand(1)); 146 return success(); 147 } 148 }; 149 150 /// This pattern matches test.any_attr_of_i32_str ops. In case of an integer 151 /// attribute with value smaller than MaxVal, it increments the value by 1. 152 template <int MaxVal> 153 struct IncrementIntAttribute : public OpRewritePattern<AnyAttrOfOp> { 154 using OpRewritePattern<AnyAttrOfOp>::OpRewritePattern; 155 156 LogicalResult matchAndRewrite(AnyAttrOfOp op, 157 PatternRewriter &rewriter) const override { 158 auto intAttr = dyn_cast<IntegerAttr>(op.getAttr()); 159 if (!intAttr) 160 return failure(); 161 int64_t val = intAttr.getInt(); 162 if (val >= MaxVal) 163 return failure(); 164 rewriter.modifyOpInPlace( 165 op, [&]() { op.setAttrAttr(rewriter.getI32IntegerAttr(val + 1)); }); 166 return success(); 167 } 168 }; 169 170 /// This patterns adds an "eligible" attribute to "foo.maybe_eligible_op". 171 struct MakeOpEligible : public RewritePattern { 172 MakeOpEligible(MLIRContext *context) 173 : RewritePattern("foo.maybe_eligible_op", /*benefit=*/1, context) {} 174 175 LogicalResult matchAndRewrite(Operation *op, 176 PatternRewriter &rewriter) const override { 177 if (op->hasAttr("eligible")) 178 return failure(); 179 rewriter.modifyOpInPlace( 180 op, [&]() { op->setAttr("eligible", rewriter.getUnitAttr()); }); 181 return success(); 182 } 183 }; 184 185 /// This pattern hoists eligible ops out of a "test.one_region_op". 186 struct HoistEligibleOps : public OpRewritePattern<test::OneRegionOp> { 187 using OpRewritePattern<test::OneRegionOp>::OpRewritePattern; 188 189 LogicalResult matchAndRewrite(test::OneRegionOp op, 190 PatternRewriter &rewriter) const override { 191 Operation *terminator = op.getRegion().front().getTerminator(); 192 Operation *toBeHoisted = terminator->getOperands()[0].getDefiningOp(); 193 if (toBeHoisted->getParentOp() != op) 194 return failure(); 195 if (!toBeHoisted->hasAttr("eligible")) 196 return failure(); 197 rewriter.moveOpBefore(toBeHoisted, op); 198 return success(); 199 } 200 }; 201 202 /// This pattern moves "test.move_before_parent_op" before the parent op. 203 struct MoveBeforeParentOp : public RewritePattern { 204 MoveBeforeParentOp(MLIRContext *context) 205 : RewritePattern("test.move_before_parent_op", /*benefit=*/1, context) {} 206 207 LogicalResult matchAndRewrite(Operation *op, 208 PatternRewriter &rewriter) const override { 209 // Do not hoist past functions. 210 if (isa<FunctionOpInterface>(op->getParentOp())) 211 return failure(); 212 rewriter.moveOpBefore(op, op->getParentOp()); 213 return success(); 214 } 215 }; 216 217 /// This pattern inlines blocks that are nested in 218 /// "test.inline_blocks_into_parent" into the parent block. 219 struct InlineBlocksIntoParent : public RewritePattern { 220 InlineBlocksIntoParent(MLIRContext *context) 221 : RewritePattern("test.inline_blocks_into_parent", /*benefit=*/1, 222 context) {} 223 224 LogicalResult matchAndRewrite(Operation *op, 225 PatternRewriter &rewriter) const override { 226 bool changed = false; 227 for (Region &r : op->getRegions()) { 228 while (!r.empty()) { 229 rewriter.inlineBlockBefore(&r.front(), op); 230 changed = true; 231 } 232 } 233 return success(changed); 234 } 235 }; 236 237 /// This pattern splits blocks at "test.split_block_here" and replaces the op 238 /// with a new op (to prevent an infinite loop of block splitting). 239 struct SplitBlockHere : public RewritePattern { 240 SplitBlockHere(MLIRContext *context) 241 : RewritePattern("test.split_block_here", /*benefit=*/1, context) {} 242 243 LogicalResult matchAndRewrite(Operation *op, 244 PatternRewriter &rewriter) const override { 245 rewriter.splitBlock(op->getBlock(), op->getIterator()); 246 Operation *newOp = rewriter.create( 247 op->getLoc(), 248 OperationName("test.new_op", op->getContext()).getIdentifier(), 249 op->getOperands(), op->getResultTypes()); 250 rewriter.replaceOp(op, newOp); 251 return success(); 252 } 253 }; 254 255 /// This pattern clones "test.clone_me" ops. 256 struct CloneOp : public RewritePattern { 257 CloneOp(MLIRContext *context) 258 : RewritePattern("test.clone_me", /*benefit=*/1, context) {} 259 260 LogicalResult matchAndRewrite(Operation *op, 261 PatternRewriter &rewriter) const override { 262 // Do not clone already cloned ops to avoid going into an infinite loop. 263 if (op->hasAttr("was_cloned")) 264 return failure(); 265 Operation *cloned = rewriter.clone(*op); 266 cloned->setAttr("was_cloned", rewriter.getUnitAttr()); 267 return success(); 268 } 269 }; 270 271 /// This pattern clones regions of "test.clone_region_before" ops before the 272 /// parent block. 273 struct CloneRegionBeforeOp : public RewritePattern { 274 CloneRegionBeforeOp(MLIRContext *context) 275 : RewritePattern("test.clone_region_before", /*benefit=*/1, context) {} 276 277 LogicalResult matchAndRewrite(Operation *op, 278 PatternRewriter &rewriter) const override { 279 // Do not clone already cloned ops to avoid going into an infinite loop. 280 if (op->hasAttr("was_cloned")) 281 return failure(); 282 for (Region &r : op->getRegions()) 283 rewriter.cloneRegionBefore(r, op->getBlock()); 284 op->setAttr("was_cloned", rewriter.getUnitAttr()); 285 return success(); 286 } 287 }; 288 289 struct TestPatternDriver 290 : public PassWrapper<TestPatternDriver, OperationPass<>> { 291 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestPatternDriver) 292 293 TestPatternDriver() = default; 294 TestPatternDriver(const TestPatternDriver &other) : PassWrapper(other) {} 295 296 StringRef getArgument() const final { return "test-patterns"; } 297 StringRef getDescription() const final { return "Run test dialect patterns"; } 298 void runOnOperation() override { 299 mlir::RewritePatternSet patterns(&getContext()); 300 populateWithGenerated(patterns); 301 302 // Verify named pattern is generated with expected name. 303 patterns.add<FoldingPattern, TestNamedPatternRule, 304 FolderInsertBeforePreviouslyFoldedConstantPattern, 305 FolderCommutativeOp2WithConstant, HoistEligibleOps, 306 MakeOpEligible>(&getContext()); 307 308 // Additional patterns for testing the GreedyPatternRewriteDriver. 309 patterns.insert<IncrementIntAttribute<3>>(&getContext()); 310 311 GreedyRewriteConfig config; 312 config.useTopDownTraversal = this->useTopDownTraversal; 313 config.maxIterations = this->maxIterations; 314 (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns), 315 config); 316 } 317 318 Option<bool> useTopDownTraversal{ 319 *this, "top-down", 320 llvm::cl::desc("Seed the worklist in general top-down order"), 321 llvm::cl::init(GreedyRewriteConfig().useTopDownTraversal)}; 322 Option<int> maxIterations{ 323 *this, "max-iterations", 324 llvm::cl::desc("Max. iterations in the GreedyRewriteConfig"), 325 llvm::cl::init(GreedyRewriteConfig().maxIterations)}; 326 }; 327 328 struct DumpNotifications : public RewriterBase::Listener { 329 void notifyBlockInserted(Block *block, Region *previous, 330 Region::iterator previousIt) override { 331 llvm::outs() << "notifyBlockInserted"; 332 if (block->getParentOp()) { 333 llvm::outs() << " into " << block->getParentOp()->getName() << ": "; 334 } else { 335 llvm::outs() << " into unknown op: "; 336 } 337 if (previous == nullptr) { 338 llvm::outs() << "was unlinked\n"; 339 } else { 340 llvm::outs() << "was linked\n"; 341 } 342 } 343 void notifyOperationInserted(Operation *op, 344 OpBuilder::InsertPoint previous) override { 345 llvm::outs() << "notifyOperationInserted: " << op->getName(); 346 if (!previous.isSet()) { 347 llvm::outs() << ", was unlinked\n"; 348 } else { 349 if (!previous.getPoint().getNodePtr()) { 350 llvm::outs() << ", was linked, exact position unknown\n"; 351 } else if (previous.getPoint() == previous.getBlock()->end()) { 352 llvm::outs() << ", was last in block\n"; 353 } else { 354 llvm::outs() << ", previous = " << previous.getPoint()->getName() 355 << "\n"; 356 } 357 } 358 } 359 void notifyBlockErased(Block *block) override { 360 llvm::outs() << "notifyBlockErased\n"; 361 } 362 void notifyOperationErased(Operation *op) override { 363 llvm::outs() << "notifyOperationErased: " << op->getName() << "\n"; 364 } 365 void notifyOperationModified(Operation *op) override { 366 llvm::outs() << "notifyOperationModified: " << op->getName() << "\n"; 367 } 368 void notifyOperationReplaced(Operation *op, ValueRange values) override { 369 llvm::outs() << "notifyOperationReplaced: " << op->getName() << "\n"; 370 } 371 }; 372 373 struct TestStrictPatternDriver 374 : public PassWrapper<TestStrictPatternDriver, OperationPass<func::FuncOp>> { 375 public: 376 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestStrictPatternDriver) 377 378 TestStrictPatternDriver() = default; 379 TestStrictPatternDriver(const TestStrictPatternDriver &other) 380 : PassWrapper(other) { 381 strictMode = other.strictMode; 382 } 383 384 StringRef getArgument() const final { return "test-strict-pattern-driver"; } 385 StringRef getDescription() const final { 386 return "Test strict mode of pattern driver"; 387 } 388 389 void runOnOperation() override { 390 MLIRContext *ctx = &getContext(); 391 mlir::RewritePatternSet patterns(ctx); 392 patterns.add< 393 // clang-format off 394 ChangeBlockOp, 395 CloneOp, 396 CloneRegionBeforeOp, 397 EraseOp, 398 ImplicitChangeOp, 399 InlineBlocksIntoParent, 400 InsertSameOp, 401 MoveBeforeParentOp, 402 ReplaceWithNewOp, 403 SplitBlockHere 404 // clang-format on 405 >(ctx); 406 SmallVector<Operation *> ops; 407 getOperation()->walk([&](Operation *op) { 408 StringRef opName = op->getName().getStringRef(); 409 if (opName == "test.insert_same_op" || opName == "test.change_block_op" || 410 opName == "test.replace_with_new_op" || opName == "test.erase_op" || 411 opName == "test.move_before_parent_op" || 412 opName == "test.inline_blocks_into_parent" || 413 opName == "test.split_block_here" || opName == "test.clone_me" || 414 opName == "test.clone_region_before") { 415 ops.push_back(op); 416 } 417 }); 418 419 DumpNotifications dumpNotifications; 420 GreedyRewriteConfig config; 421 config.listener = &dumpNotifications; 422 if (strictMode == "AnyOp") { 423 config.strictMode = GreedyRewriteStrictness::AnyOp; 424 } else if (strictMode == "ExistingAndNewOps") { 425 config.strictMode = GreedyRewriteStrictness::ExistingAndNewOps; 426 } else if (strictMode == "ExistingOps") { 427 config.strictMode = GreedyRewriteStrictness::ExistingOps; 428 } else { 429 llvm_unreachable("invalid strictness option"); 430 } 431 432 // Check if these transformations introduce visiting of operations that 433 // are not in the `ops` set (The new created ops are valid). An invalid 434 // operation will trigger the assertion while processing. 435 bool changed = false; 436 bool allErased = false; 437 (void)applyOpPatternsAndFold(ArrayRef(ops), std::move(patterns), config, 438 &changed, &allErased); 439 Builder b(ctx); 440 getOperation()->setAttr("pattern_driver_changed", b.getBoolAttr(changed)); 441 getOperation()->setAttr("pattern_driver_all_erased", 442 b.getBoolAttr(allErased)); 443 } 444 445 Option<std::string> strictMode{ 446 *this, "strictness", 447 llvm::cl::desc("Can be {AnyOp, ExistingAndNewOps, ExistingOps}"), 448 llvm::cl::init("AnyOp")}; 449 450 private: 451 // New inserted operation is valid for further transformation. 452 class InsertSameOp : public RewritePattern { 453 public: 454 InsertSameOp(MLIRContext *context) 455 : RewritePattern("test.insert_same_op", /*benefit=*/1, context) {} 456 457 LogicalResult matchAndRewrite(Operation *op, 458 PatternRewriter &rewriter) const override { 459 if (op->hasAttr("skip")) 460 return failure(); 461 462 Operation *newOp = 463 rewriter.create(op->getLoc(), op->getName().getIdentifier(), 464 op->getOperands(), op->getResultTypes()); 465 rewriter.modifyOpInPlace( 466 op, [&]() { op->setAttr("skip", rewriter.getBoolAttr(true)); }); 467 newOp->setAttr("skip", rewriter.getBoolAttr(true)); 468 469 return success(); 470 } 471 }; 472 473 // Replace an operation may introduce the re-visiting of its users. 474 class ReplaceWithNewOp : public RewritePattern { 475 public: 476 ReplaceWithNewOp(MLIRContext *context) 477 : RewritePattern("test.replace_with_new_op", /*benefit=*/1, context) {} 478 479 LogicalResult matchAndRewrite(Operation *op, 480 PatternRewriter &rewriter) const override { 481 Operation *newOp; 482 if (op->hasAttr("create_erase_op")) { 483 newOp = rewriter.create( 484 op->getLoc(), 485 OperationName("test.erase_op", op->getContext()).getIdentifier(), 486 ValueRange(), TypeRange()); 487 } else { 488 newOp = rewriter.create( 489 op->getLoc(), 490 OperationName("test.new_op", op->getContext()).getIdentifier(), 491 op->getOperands(), op->getResultTypes()); 492 } 493 // "replaceOp" could be used instead of "replaceAllOpUsesWith"+"eraseOp". 494 // A "notifyOperationReplaced" callback is triggered in either case. 495 rewriter.replaceAllOpUsesWith(op, newOp->getResults()); 496 rewriter.eraseOp(op); 497 return success(); 498 } 499 }; 500 501 // Remove an operation may introduce the re-visiting of its operands. 502 class EraseOp : public RewritePattern { 503 public: 504 EraseOp(MLIRContext *context) 505 : RewritePattern("test.erase_op", /*benefit=*/1, context) {} 506 LogicalResult matchAndRewrite(Operation *op, 507 PatternRewriter &rewriter) const override { 508 rewriter.eraseOp(op); 509 return success(); 510 } 511 }; 512 513 // The following two patterns test RewriterBase::replaceAllUsesWith. 514 // 515 // That function replaces all usages of a Block (or a Value) with another one 516 // *and tracks these changes in the rewriter.* The GreedyPatternRewriteDriver 517 // with GreedyRewriteStrictness::AnyOp uses that tracking to construct its 518 // worklist: when an op is modified, it is added to the worklist. The two 519 // patterns below make the tracking observable: ChangeBlockOp replaces all 520 // usages of a block and that pattern is applied because the corresponding ops 521 // are put on the initial worklist (see above). ImplicitChangeOp does an 522 // unrelated change but ops of the corresponding type are *not* on the initial 523 // worklist, so the effect of the second pattern is only visible if the 524 // tracking and subsequent adding to the worklist actually works. 525 526 // Replace all usages of the first successor with the second successor. 527 class ChangeBlockOp : public RewritePattern { 528 public: 529 ChangeBlockOp(MLIRContext *context) 530 : RewritePattern("test.change_block_op", /*benefit=*/1, context) {} 531 LogicalResult matchAndRewrite(Operation *op, 532 PatternRewriter &rewriter) const override { 533 if (op->getNumSuccessors() < 2) 534 return failure(); 535 Block *firstSuccessor = op->getSuccessor(0); 536 Block *secondSuccessor = op->getSuccessor(1); 537 if (firstSuccessor == secondSuccessor) 538 return failure(); 539 // This is the function being tested: 540 rewriter.replaceAllUsesWith(firstSuccessor, secondSuccessor); 541 // Using the following line instead would make the test fail: 542 // firstSuccessor->replaceAllUsesWith(secondSuccessor); 543 return success(); 544 } 545 }; 546 547 // Changes the successor to the parent block. 548 class ImplicitChangeOp : public RewritePattern { 549 public: 550 ImplicitChangeOp(MLIRContext *context) 551 : RewritePattern("test.implicit_change_op", /*benefit=*/1, context) {} 552 LogicalResult matchAndRewrite(Operation *op, 553 PatternRewriter &rewriter) const override { 554 if (op->getNumSuccessors() < 1 || op->getSuccessor(0) == op->getBlock()) 555 return failure(); 556 rewriter.modifyOpInPlace(op, 557 [&]() { op->setSuccessor(op->getBlock(), 0); }); 558 return success(); 559 } 560 }; 561 }; 562 563 } // namespace 564 565 //===----------------------------------------------------------------------===// 566 // ReturnType Driver. 567 //===----------------------------------------------------------------------===// 568 569 namespace { 570 // Generate ops for each instance where the type can be successfully inferred. 571 template <typename OpTy> 572 static void invokeCreateWithInferredReturnType(Operation *op) { 573 auto *context = op->getContext(); 574 auto fop = op->getParentOfType<func::FuncOp>(); 575 auto location = UnknownLoc::get(context); 576 OpBuilder b(op); 577 b.setInsertionPointAfter(op); 578 579 // Use permutations of 2 args as operands. 580 assert(fop.getNumArguments() >= 2); 581 for (int i = 0, e = fop.getNumArguments(); i < e; ++i) { 582 for (int j = 0; j < e; ++j) { 583 std::array<Value, 2> values = {{fop.getArgument(i), fop.getArgument(j)}}; 584 SmallVector<Type, 2> inferredReturnTypes; 585 if (succeeded(OpTy::inferReturnTypes( 586 context, std::nullopt, values, op->getDiscardableAttrDictionary(), 587 op->getPropertiesStorage(), op->getRegions(), 588 inferredReturnTypes))) { 589 OperationState state(location, OpTy::getOperationName()); 590 // TODO: Expand to regions. 591 OpTy::build(b, state, values, op->getAttrs()); 592 (void)b.create(state); 593 } 594 } 595 } 596 } 597 598 static void reifyReturnShape(Operation *op) { 599 OpBuilder b(op); 600 601 // Use permutations of 2 args as operands. 602 auto shapedOp = cast<OpWithShapedTypeInferTypeInterfaceOp>(op); 603 SmallVector<Value, 2> shapes; 604 if (failed(shapedOp.reifyReturnTypeShapes(b, op->getOperands(), shapes)) || 605 !llvm::hasSingleElement(shapes)) 606 return; 607 for (const auto &it : llvm::enumerate(shapes)) { 608 op->emitRemark() << "value " << it.index() << ": " 609 << it.value().getDefiningOp(); 610 } 611 } 612 613 struct TestReturnTypeDriver 614 : public PassWrapper<TestReturnTypeDriver, OperationPass<func::FuncOp>> { 615 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestReturnTypeDriver) 616 617 void getDependentDialects(DialectRegistry ®istry) const override { 618 registry.insert<tensor::TensorDialect>(); 619 } 620 StringRef getArgument() const final { return "test-return-type"; } 621 StringRef getDescription() const final { return "Run return type functions"; } 622 623 void runOnOperation() override { 624 if (getOperation().getName() == "testCreateFunctions") { 625 std::vector<Operation *> ops; 626 // Collect ops to avoid triggering on inserted ops. 627 for (auto &op : getOperation().getBody().front()) 628 ops.push_back(&op); 629 // Generate test patterns for each, but skip terminator. 630 for (auto *op : llvm::ArrayRef(ops).drop_back()) { 631 // Test create method of each of the Op classes below. The resultant 632 // output would be in reverse order underneath `op` from which 633 // the attributes and regions are used. 634 invokeCreateWithInferredReturnType<OpWithInferTypeInterfaceOp>(op); 635 invokeCreateWithInferredReturnType<OpWithInferTypeAdaptorInterfaceOp>( 636 op); 637 invokeCreateWithInferredReturnType< 638 OpWithShapedTypeInferTypeInterfaceOp>(op); 639 }; 640 return; 641 } 642 if (getOperation().getName() == "testReifyFunctions") { 643 std::vector<Operation *> ops; 644 // Collect ops to avoid triggering on inserted ops. 645 for (auto &op : getOperation().getBody().front()) 646 if (isa<OpWithShapedTypeInferTypeInterfaceOp>(op)) 647 ops.push_back(&op); 648 // Generate test patterns for each, but skip terminator. 649 for (auto *op : ops) 650 reifyReturnShape(op); 651 } 652 } 653 }; 654 } // namespace 655 656 namespace { 657 struct TestDerivedAttributeDriver 658 : public PassWrapper<TestDerivedAttributeDriver, 659 OperationPass<func::FuncOp>> { 660 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestDerivedAttributeDriver) 661 662 StringRef getArgument() const final { return "test-derived-attr"; } 663 StringRef getDescription() const final { 664 return "Run test derived attributes"; 665 } 666 void runOnOperation() override; 667 }; 668 } // namespace 669 670 void TestDerivedAttributeDriver::runOnOperation() { 671 getOperation().walk([](DerivedAttributeOpInterface dOp) { 672 auto dAttr = dOp.materializeDerivedAttributes(); 673 if (!dAttr) 674 return; 675 for (auto d : dAttr) 676 dOp.emitRemark() << d.getName().getValue() << " = " << d.getValue(); 677 }); 678 } 679 680 //===----------------------------------------------------------------------===// 681 // Legalization Driver. 682 //===----------------------------------------------------------------------===// 683 684 namespace { 685 //===----------------------------------------------------------------------===// 686 // Region-Block Rewrite Testing 687 688 /// This pattern is a simple pattern that inlines the first region of a given 689 /// operation into the parent region. 690 struct TestRegionRewriteBlockMovement : public ConversionPattern { 691 TestRegionRewriteBlockMovement(MLIRContext *ctx) 692 : ConversionPattern("test.region", 1, ctx) {} 693 694 LogicalResult 695 matchAndRewrite(Operation *op, ArrayRef<Value> operands, 696 ConversionPatternRewriter &rewriter) const final { 697 // Inline this region into the parent region. 698 auto &parentRegion = *op->getParentRegion(); 699 auto &opRegion = op->getRegion(0); 700 if (op->getDiscardableAttr("legalizer.should_clone")) 701 rewriter.cloneRegionBefore(opRegion, parentRegion, parentRegion.end()); 702 else 703 rewriter.inlineRegionBefore(opRegion, parentRegion, parentRegion.end()); 704 705 if (op->getDiscardableAttr("legalizer.erase_old_blocks")) { 706 while (!opRegion.empty()) 707 rewriter.eraseBlock(&opRegion.front()); 708 } 709 710 // Drop this operation. 711 rewriter.eraseOp(op); 712 return success(); 713 } 714 }; 715 /// This pattern is a simple pattern that generates a region containing an 716 /// illegal operation. 717 struct TestRegionRewriteUndo : public RewritePattern { 718 TestRegionRewriteUndo(MLIRContext *ctx) 719 : RewritePattern("test.region_builder", 1, ctx) {} 720 721 LogicalResult matchAndRewrite(Operation *op, 722 PatternRewriter &rewriter) const final { 723 // Create the region operation with an entry block containing arguments. 724 OperationState newRegion(op->getLoc(), "test.region"); 725 newRegion.addRegion(); 726 auto *regionOp = rewriter.create(newRegion); 727 auto *entryBlock = rewriter.createBlock(®ionOp->getRegion(0)); 728 entryBlock->addArgument(rewriter.getIntegerType(64), 729 rewriter.getUnknownLoc()); 730 731 // Add an explicitly illegal operation to ensure the conversion fails. 732 rewriter.create<ILLegalOpF>(op->getLoc(), rewriter.getIntegerType(32)); 733 rewriter.create<TestValidOp>(op->getLoc(), ArrayRef<Value>()); 734 735 // Drop this operation. 736 rewriter.eraseOp(op); 737 return success(); 738 } 739 }; 740 /// A simple pattern that creates a block at the end of the parent region of the 741 /// matched operation. 742 struct TestCreateBlock : public RewritePattern { 743 TestCreateBlock(MLIRContext *ctx) 744 : RewritePattern("test.create_block", /*benefit=*/1, ctx) {} 745 746 LogicalResult matchAndRewrite(Operation *op, 747 PatternRewriter &rewriter) const final { 748 Region ®ion = *op->getParentRegion(); 749 Type i32Type = rewriter.getIntegerType(32); 750 Location loc = op->getLoc(); 751 rewriter.createBlock(®ion, region.end(), {i32Type, i32Type}, {loc, loc}); 752 rewriter.create<TerminatorOp>(loc); 753 rewriter.eraseOp(op); 754 return success(); 755 } 756 }; 757 758 /// A simple pattern that creates a block containing an invalid operation in 759 /// order to trigger the block creation undo mechanism. 760 struct TestCreateIllegalBlock : public RewritePattern { 761 TestCreateIllegalBlock(MLIRContext *ctx) 762 : RewritePattern("test.create_illegal_block", /*benefit=*/1, ctx) {} 763 764 LogicalResult matchAndRewrite(Operation *op, 765 PatternRewriter &rewriter) const final { 766 Region ®ion = *op->getParentRegion(); 767 Type i32Type = rewriter.getIntegerType(32); 768 Location loc = op->getLoc(); 769 rewriter.createBlock(®ion, region.end(), {i32Type, i32Type}, {loc, loc}); 770 // Create an illegal op to ensure the conversion fails. 771 rewriter.create<ILLegalOpF>(loc, i32Type); 772 rewriter.create<TerminatorOp>(loc); 773 rewriter.eraseOp(op); 774 return success(); 775 } 776 }; 777 778 /// A simple pattern that tests the undo mechanism when replacing the uses of a 779 /// block argument. 780 struct TestUndoBlockArgReplace : public ConversionPattern { 781 TestUndoBlockArgReplace(MLIRContext *ctx) 782 : ConversionPattern("test.undo_block_arg_replace", /*benefit=*/1, ctx) {} 783 784 LogicalResult 785 matchAndRewrite(Operation *op, ArrayRef<Value> operands, 786 ConversionPatternRewriter &rewriter) const final { 787 auto illegalOp = 788 rewriter.create<ILLegalOpF>(op->getLoc(), rewriter.getF32Type()); 789 rewriter.replaceUsesOfBlockArgument(op->getRegion(0).getArgument(0), 790 illegalOp->getResult(0)); 791 rewriter.modifyOpInPlace(op, [] {}); 792 return success(); 793 } 794 }; 795 796 /// This pattern hoists ops out of a "test.hoist_me" and then fails conversion. 797 /// This is to test the rollback logic. 798 struct TestUndoMoveOpBefore : public ConversionPattern { 799 TestUndoMoveOpBefore(MLIRContext *ctx) 800 : ConversionPattern("test.hoist_me", /*benefit=*/1, ctx) {} 801 802 LogicalResult 803 matchAndRewrite(Operation *op, ArrayRef<Value> operands, 804 ConversionPatternRewriter &rewriter) const override { 805 rewriter.moveOpBefore(op, op->getParentOp()); 806 // Replace with an illegal op to ensure the conversion fails. 807 rewriter.replaceOpWithNewOp<ILLegalOpF>(op, rewriter.getF32Type()); 808 return success(); 809 } 810 }; 811 812 /// A rewrite pattern that tests the undo mechanism when erasing a block. 813 struct TestUndoBlockErase : public ConversionPattern { 814 TestUndoBlockErase(MLIRContext *ctx) 815 : ConversionPattern("test.undo_block_erase", /*benefit=*/1, ctx) {} 816 817 LogicalResult 818 matchAndRewrite(Operation *op, ArrayRef<Value> operands, 819 ConversionPatternRewriter &rewriter) const final { 820 Block *secondBlock = &*std::next(op->getRegion(0).begin()); 821 rewriter.setInsertionPointToStart(secondBlock); 822 rewriter.create<ILLegalOpF>(op->getLoc(), rewriter.getF32Type()); 823 rewriter.eraseBlock(secondBlock); 824 rewriter.modifyOpInPlace(op, [] {}); 825 return success(); 826 } 827 }; 828 829 /// A pattern that modifies a property in-place, but keeps the op illegal. 830 struct TestUndoPropertiesModification : public ConversionPattern { 831 TestUndoPropertiesModification(MLIRContext *ctx) 832 : ConversionPattern("test.with_properties", /*benefit=*/1, ctx) {} 833 LogicalResult 834 matchAndRewrite(Operation *op, ArrayRef<Value> operands, 835 ConversionPatternRewriter &rewriter) const final { 836 if (!op->hasAttr("modify_inplace")) 837 return failure(); 838 rewriter.modifyOpInPlace( 839 op, [&]() { cast<TestOpWithProperties>(op).getProperties().setA(42); }); 840 return success(); 841 } 842 }; 843 844 //===----------------------------------------------------------------------===// 845 // Type-Conversion Rewrite Testing 846 847 /// This patterns erases a region operation that has had a type conversion. 848 struct TestDropOpSignatureConversion : public ConversionPattern { 849 TestDropOpSignatureConversion(MLIRContext *ctx, 850 const TypeConverter &converter) 851 : ConversionPattern(converter, "test.drop_region_op", 1, ctx) {} 852 LogicalResult 853 matchAndRewrite(Operation *op, ArrayRef<Value> operands, 854 ConversionPatternRewriter &rewriter) const override { 855 Region ®ion = op->getRegion(0); 856 Block *entry = ®ion.front(); 857 858 // Convert the original entry arguments. 859 const TypeConverter &converter = *getTypeConverter(); 860 TypeConverter::SignatureConversion result(entry->getNumArguments()); 861 if (failed(converter.convertSignatureArgs(entry->getArgumentTypes(), 862 result)) || 863 failed(rewriter.convertRegionTypes(®ion, converter, &result))) 864 return failure(); 865 866 // Convert the region signature and just drop the operation. 867 rewriter.eraseOp(op); 868 return success(); 869 } 870 }; 871 /// This pattern simply updates the operands of the given operation. 872 struct TestPassthroughInvalidOp : public ConversionPattern { 873 TestPassthroughInvalidOp(MLIRContext *ctx) 874 : ConversionPattern("test.invalid", 1, ctx) {} 875 LogicalResult 876 matchAndRewrite(Operation *op, ArrayRef<Value> operands, 877 ConversionPatternRewriter &rewriter) const final { 878 rewriter.replaceOpWithNewOp<TestValidOp>(op, std::nullopt, operands, 879 std::nullopt); 880 return success(); 881 } 882 }; 883 /// This pattern handles the case of a split return value. 884 struct TestSplitReturnType : public ConversionPattern { 885 TestSplitReturnType(MLIRContext *ctx) 886 : ConversionPattern("test.return", 1, ctx) {} 887 LogicalResult 888 matchAndRewrite(Operation *op, ArrayRef<Value> operands, 889 ConversionPatternRewriter &rewriter) const final { 890 // Check for a return of F32. 891 if (op->getNumOperands() != 1 || !op->getOperand(0).getType().isF32()) 892 return failure(); 893 894 // Check if the first operation is a cast operation, if it is we use the 895 // results directly. 896 auto *defOp = operands[0].getDefiningOp(); 897 if (auto packerOp = 898 llvm::dyn_cast_or_null<UnrealizedConversionCastOp>(defOp)) { 899 rewriter.replaceOpWithNewOp<TestReturnOp>(op, packerOp.getOperands()); 900 return success(); 901 } 902 903 // Otherwise, fail to match. 904 return failure(); 905 } 906 }; 907 908 //===----------------------------------------------------------------------===// 909 // Multi-Level Type-Conversion Rewrite Testing 910 struct TestChangeProducerTypeI32ToF32 : public ConversionPattern { 911 TestChangeProducerTypeI32ToF32(MLIRContext *ctx) 912 : ConversionPattern("test.type_producer", 1, ctx) {} 913 LogicalResult 914 matchAndRewrite(Operation *op, ArrayRef<Value> operands, 915 ConversionPatternRewriter &rewriter) const final { 916 // If the type is I32, change the type to F32. 917 if (!Type(*op->result_type_begin()).isSignlessInteger(32)) 918 return failure(); 919 rewriter.replaceOpWithNewOp<TestTypeProducerOp>(op, rewriter.getF32Type()); 920 return success(); 921 } 922 }; 923 struct TestChangeProducerTypeF32ToF64 : public ConversionPattern { 924 TestChangeProducerTypeF32ToF64(MLIRContext *ctx) 925 : ConversionPattern("test.type_producer", 1, ctx) {} 926 LogicalResult 927 matchAndRewrite(Operation *op, ArrayRef<Value> operands, 928 ConversionPatternRewriter &rewriter) const final { 929 // If the type is F32, change the type to F64. 930 if (!Type(*op->result_type_begin()).isF32()) 931 return rewriter.notifyMatchFailure(op, "expected single f32 operand"); 932 rewriter.replaceOpWithNewOp<TestTypeProducerOp>(op, rewriter.getF64Type()); 933 return success(); 934 } 935 }; 936 struct TestChangeProducerTypeF32ToInvalid : public ConversionPattern { 937 TestChangeProducerTypeF32ToInvalid(MLIRContext *ctx) 938 : ConversionPattern("test.type_producer", 10, ctx) {} 939 LogicalResult 940 matchAndRewrite(Operation *op, ArrayRef<Value> operands, 941 ConversionPatternRewriter &rewriter) const final { 942 // Always convert to B16, even though it is not a legal type. This tests 943 // that values are unmapped correctly. 944 rewriter.replaceOpWithNewOp<TestTypeProducerOp>(op, rewriter.getBF16Type()); 945 return success(); 946 } 947 }; 948 struct TestUpdateConsumerType : public ConversionPattern { 949 TestUpdateConsumerType(MLIRContext *ctx) 950 : ConversionPattern("test.type_consumer", 1, ctx) {} 951 LogicalResult 952 matchAndRewrite(Operation *op, ArrayRef<Value> operands, 953 ConversionPatternRewriter &rewriter) const final { 954 // Verify that the incoming operand has been successfully remapped to F64. 955 if (!operands[0].getType().isF64()) 956 return failure(); 957 rewriter.replaceOpWithNewOp<TestTypeConsumerOp>(op, operands[0]); 958 return success(); 959 } 960 }; 961 962 //===----------------------------------------------------------------------===// 963 // Non-Root Replacement Rewrite Testing 964 /// This pattern generates an invalid operation, but replaces it before the 965 /// pattern is finished. This checks that we don't need to legalize the 966 /// temporary op. 967 struct TestNonRootReplacement : public RewritePattern { 968 TestNonRootReplacement(MLIRContext *ctx) 969 : RewritePattern("test.replace_non_root", 1, ctx) {} 970 971 LogicalResult matchAndRewrite(Operation *op, 972 PatternRewriter &rewriter) const final { 973 auto resultType = *op->result_type_begin(); 974 auto illegalOp = rewriter.create<ILLegalOpF>(op->getLoc(), resultType); 975 auto legalOp = rewriter.create<LegalOpB>(op->getLoc(), resultType); 976 977 rewriter.replaceOp(illegalOp, legalOp); 978 rewriter.replaceOp(op, illegalOp); 979 return success(); 980 } 981 }; 982 983 //===----------------------------------------------------------------------===// 984 // Recursive Rewrite Testing 985 /// This pattern is applied to the same operation multiple times, but has a 986 /// bounded recursion. 987 struct TestBoundedRecursiveRewrite 988 : public OpRewritePattern<TestRecursiveRewriteOp> { 989 using OpRewritePattern<TestRecursiveRewriteOp>::OpRewritePattern; 990 991 void initialize() { 992 // The conversion target handles bounding the recursion of this pattern. 993 setHasBoundedRewriteRecursion(); 994 } 995 996 LogicalResult matchAndRewrite(TestRecursiveRewriteOp op, 997 PatternRewriter &rewriter) const final { 998 // Decrement the depth of the op in-place. 999 rewriter.modifyOpInPlace(op, [&] { 1000 op->setAttr("depth", rewriter.getI64IntegerAttr(op.getDepth() - 1)); 1001 }); 1002 return success(); 1003 } 1004 }; 1005 1006 struct TestNestedOpCreationUndoRewrite 1007 : public OpRewritePattern<IllegalOpWithRegionAnchor> { 1008 using OpRewritePattern<IllegalOpWithRegionAnchor>::OpRewritePattern; 1009 1010 LogicalResult matchAndRewrite(IllegalOpWithRegionAnchor op, 1011 PatternRewriter &rewriter) const final { 1012 // rewriter.replaceOpWithNewOp<IllegalOpWithRegion>(op); 1013 rewriter.replaceOpWithNewOp<IllegalOpWithRegion>(op); 1014 return success(); 1015 }; 1016 }; 1017 1018 // This pattern matches `test.blackhole` and delete this op and its producer. 1019 struct TestReplaceEraseOp : public OpRewritePattern<BlackHoleOp> { 1020 using OpRewritePattern<BlackHoleOp>::OpRewritePattern; 1021 1022 LogicalResult matchAndRewrite(BlackHoleOp op, 1023 PatternRewriter &rewriter) const final { 1024 Operation *producer = op.getOperand().getDefiningOp(); 1025 // Always erase the user before the producer, the framework should handle 1026 // this correctly. 1027 rewriter.eraseOp(op); 1028 rewriter.eraseOp(producer); 1029 return success(); 1030 }; 1031 }; 1032 1033 // This pattern replaces explicitly illegal op with explicitly legal op, 1034 // but in addition creates unregistered operation. 1035 struct TestCreateUnregisteredOp : public OpRewritePattern<ILLegalOpG> { 1036 using OpRewritePattern<ILLegalOpG>::OpRewritePattern; 1037 1038 LogicalResult matchAndRewrite(ILLegalOpG op, 1039 PatternRewriter &rewriter) const final { 1040 IntegerAttr attr = rewriter.getI32IntegerAttr(0); 1041 Value val = rewriter.create<arith::ConstantOp>(op->getLoc(), attr); 1042 rewriter.replaceOpWithNewOp<LegalOpC>(op, val); 1043 return success(); 1044 }; 1045 }; 1046 } // namespace 1047 1048 namespace { 1049 struct TestTypeConverter : public TypeConverter { 1050 using TypeConverter::TypeConverter; 1051 TestTypeConverter() { 1052 addConversion(convertType); 1053 addArgumentMaterialization(materializeCast); 1054 addSourceMaterialization(materializeCast); 1055 } 1056 1057 static LogicalResult convertType(Type t, SmallVectorImpl<Type> &results) { 1058 // Drop I16 types. 1059 if (t.isSignlessInteger(16)) 1060 return success(); 1061 1062 // Convert I64 to F64. 1063 if (t.isSignlessInteger(64)) { 1064 results.push_back(FloatType::getF64(t.getContext())); 1065 return success(); 1066 } 1067 1068 // Convert I42 to I43. 1069 if (t.isInteger(42)) { 1070 results.push_back(IntegerType::get(t.getContext(), 43)); 1071 return success(); 1072 } 1073 1074 // Split F32 into F16,F16. 1075 if (t.isF32()) { 1076 results.assign(2, FloatType::getF16(t.getContext())); 1077 return success(); 1078 } 1079 1080 // Otherwise, convert the type directly. 1081 results.push_back(t); 1082 return success(); 1083 } 1084 1085 /// Hook for materializing a conversion. This is necessary because we generate 1086 /// 1->N type mappings. 1087 static std::optional<Value> materializeCast(OpBuilder &builder, 1088 Type resultType, 1089 ValueRange inputs, Location loc) { 1090 return builder.create<TestCastOp>(loc, resultType, inputs).getResult(); 1091 } 1092 }; 1093 1094 struct TestLegalizePatternDriver 1095 : public PassWrapper<TestLegalizePatternDriver, OperationPass<>> { 1096 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestLegalizePatternDriver) 1097 1098 StringRef getArgument() const final { return "test-legalize-patterns"; } 1099 StringRef getDescription() const final { 1100 return "Run test dialect legalization patterns"; 1101 } 1102 /// The mode of conversion to use with the driver. 1103 enum class ConversionMode { Analysis, Full, Partial }; 1104 1105 TestLegalizePatternDriver(ConversionMode mode) : mode(mode) {} 1106 1107 void getDependentDialects(DialectRegistry ®istry) const override { 1108 registry.insert<func::FuncDialect, test::TestDialect>(); 1109 } 1110 1111 void runOnOperation() override { 1112 TestTypeConverter converter; 1113 mlir::RewritePatternSet patterns(&getContext()); 1114 populateWithGenerated(patterns); 1115 patterns 1116 .add<TestRegionRewriteBlockMovement, TestRegionRewriteUndo, 1117 TestCreateBlock, TestCreateIllegalBlock, TestUndoBlockArgReplace, 1118 TestUndoBlockErase, TestPassthroughInvalidOp, TestSplitReturnType, 1119 TestChangeProducerTypeI32ToF32, TestChangeProducerTypeF32ToF64, 1120 TestChangeProducerTypeF32ToInvalid, TestUpdateConsumerType, 1121 TestNonRootReplacement, TestBoundedRecursiveRewrite, 1122 TestNestedOpCreationUndoRewrite, TestReplaceEraseOp, 1123 TestCreateUnregisteredOp, TestUndoMoveOpBefore, 1124 TestUndoPropertiesModification>(&getContext()); 1125 patterns.add<TestDropOpSignatureConversion>(&getContext(), converter); 1126 mlir::populateAnyFunctionOpInterfaceTypeConversionPattern(patterns, 1127 converter); 1128 mlir::populateCallOpTypeConversionPattern(patterns, converter); 1129 1130 // Define the conversion target used for the test. 1131 ConversionTarget target(getContext()); 1132 target.addLegalOp<ModuleOp>(); 1133 target.addLegalOp<LegalOpA, LegalOpB, LegalOpC, TestCastOp, TestValidOp, 1134 TerminatorOp, OneRegionOp>(); 1135 target 1136 .addIllegalOp<ILLegalOpF, TestRegionBuilderOp, TestOpWithRegionFold>(); 1137 target.addDynamicallyLegalOp<TestReturnOp>([](TestReturnOp op) { 1138 // Don't allow F32 operands. 1139 return llvm::none_of(op.getOperandTypes(), 1140 [](Type type) { return type.isF32(); }); 1141 }); 1142 target.addDynamicallyLegalOp<func::FuncOp>([&](func::FuncOp op) { 1143 return converter.isSignatureLegal(op.getFunctionType()) && 1144 converter.isLegal(&op.getBody()); 1145 }); 1146 target.addDynamicallyLegalOp<func::CallOp>( 1147 [&](func::CallOp op) { return converter.isLegal(op); }); 1148 1149 // TestCreateUnregisteredOp creates `arith.constant` operation, 1150 // which was not added to target intentionally to test 1151 // correct error code from conversion driver. 1152 target.addDynamicallyLegalOp<ILLegalOpG>([](ILLegalOpG) { return false; }); 1153 1154 // Expect the type_producer/type_consumer operations to only operate on f64. 1155 target.addDynamicallyLegalOp<TestTypeProducerOp>( 1156 [](TestTypeProducerOp op) { return op.getType().isF64(); }); 1157 target.addDynamicallyLegalOp<TestTypeConsumerOp>([](TestTypeConsumerOp op) { 1158 return op.getOperand().getType().isF64(); 1159 }); 1160 1161 // Check support for marking certain operations as recursively legal. 1162 target.markOpRecursivelyLegal<func::FuncOp, ModuleOp>([](Operation *op) { 1163 return static_cast<bool>( 1164 op->getAttrOfType<UnitAttr>("test.recursively_legal")); 1165 }); 1166 1167 // Mark the bound recursion operation as dynamically legal. 1168 target.addDynamicallyLegalOp<TestRecursiveRewriteOp>( 1169 [](TestRecursiveRewriteOp op) { return op.getDepth() == 0; }); 1170 1171 // Handle a partial conversion. 1172 if (mode == ConversionMode::Partial) { 1173 DenseSet<Operation *> unlegalizedOps; 1174 ConversionConfig config; 1175 DumpNotifications dumpNotifications; 1176 config.listener = &dumpNotifications; 1177 config.unlegalizedOps = &unlegalizedOps; 1178 if (failed(applyPartialConversion(getOperation(), target, 1179 std::move(patterns), config))) { 1180 getOperation()->emitRemark() << "applyPartialConversion failed"; 1181 } 1182 // Emit remarks for each legalizable operation. 1183 for (auto *op : unlegalizedOps) 1184 op->emitRemark() << "op '" << op->getName() << "' is not legalizable"; 1185 return; 1186 } 1187 1188 // Handle a full conversion. 1189 if (mode == ConversionMode::Full) { 1190 // Check support for marking unknown operations as dynamically legal. 1191 target.markUnknownOpDynamicallyLegal([](Operation *op) { 1192 return (bool)op->getAttrOfType<UnitAttr>("test.dynamically_legal"); 1193 }); 1194 1195 ConversionConfig config; 1196 DumpNotifications dumpNotifications; 1197 config.listener = &dumpNotifications; 1198 if (failed(applyFullConversion(getOperation(), target, 1199 std::move(patterns), config))) { 1200 getOperation()->emitRemark() << "applyFullConversion failed"; 1201 } 1202 return; 1203 } 1204 1205 // Otherwise, handle an analysis conversion. 1206 assert(mode == ConversionMode::Analysis); 1207 1208 // Analyze the convertible operations. 1209 DenseSet<Operation *> legalizedOps; 1210 ConversionConfig config; 1211 config.legalizableOps = &legalizedOps; 1212 if (failed(applyAnalysisConversion(getOperation(), target, 1213 std::move(patterns), config))) 1214 return signalPassFailure(); 1215 1216 // Emit remarks for each legalizable operation. 1217 for (auto *op : legalizedOps) 1218 op->emitRemark() << "op '" << op->getName() << "' is legalizable"; 1219 } 1220 1221 /// The mode of conversion to use. 1222 ConversionMode mode; 1223 }; 1224 } // namespace 1225 1226 static llvm::cl::opt<TestLegalizePatternDriver::ConversionMode> 1227 legalizerConversionMode( 1228 "test-legalize-mode", 1229 llvm::cl::desc("The legalization mode to use with the test driver"), 1230 llvm::cl::init(TestLegalizePatternDriver::ConversionMode::Partial), 1231 llvm::cl::values( 1232 clEnumValN(TestLegalizePatternDriver::ConversionMode::Analysis, 1233 "analysis", "Perform an analysis conversion"), 1234 clEnumValN(TestLegalizePatternDriver::ConversionMode::Full, "full", 1235 "Perform a full conversion"), 1236 clEnumValN(TestLegalizePatternDriver::ConversionMode::Partial, 1237 "partial", "Perform a partial conversion"))); 1238 1239 //===----------------------------------------------------------------------===// 1240 // ConversionPatternRewriter::getRemappedValue testing. This method is used 1241 // to get the remapped value of an original value that was replaced using 1242 // ConversionPatternRewriter. 1243 namespace { 1244 struct TestRemapValueTypeConverter : public TypeConverter { 1245 using TypeConverter::TypeConverter; 1246 1247 TestRemapValueTypeConverter() { 1248 addConversion( 1249 [](Float32Type type) { return Float64Type::get(type.getContext()); }); 1250 addConversion([](Type type) { return type; }); 1251 } 1252 }; 1253 1254 /// Converter that replaces a one-result one-operand OneVResOneVOperandOp1 with 1255 /// a one-operand two-result OneVResOneVOperandOp1 by replicating its original 1256 /// operand twice. 1257 /// 1258 /// Example: 1259 /// %1 = test.one_variadic_out_one_variadic_in1"(%0) 1260 /// is replaced with: 1261 /// %1 = test.one_variadic_out_one_variadic_in1"(%0, %0) 1262 struct OneVResOneVOperandOp1Converter 1263 : public OpConversionPattern<OneVResOneVOperandOp1> { 1264 using OpConversionPattern<OneVResOneVOperandOp1>::OpConversionPattern; 1265 1266 LogicalResult 1267 matchAndRewrite(OneVResOneVOperandOp1 op, OpAdaptor adaptor, 1268 ConversionPatternRewriter &rewriter) const override { 1269 auto origOps = op.getOperands(); 1270 assert(std::distance(origOps.begin(), origOps.end()) == 1 && 1271 "One operand expected"); 1272 Value origOp = *origOps.begin(); 1273 SmallVector<Value, 2> remappedOperands; 1274 // Replicate the remapped original operand twice. Note that we don't used 1275 // the remapped 'operand' since the goal is testing 'getRemappedValue'. 1276 remappedOperands.push_back(rewriter.getRemappedValue(origOp)); 1277 remappedOperands.push_back(rewriter.getRemappedValue(origOp)); 1278 1279 rewriter.replaceOpWithNewOp<OneVResOneVOperandOp1>(op, op.getResultTypes(), 1280 remappedOperands); 1281 return success(); 1282 } 1283 }; 1284 1285 /// A rewriter pattern that tests that blocks can be merged. 1286 struct TestRemapValueInRegion 1287 : public OpConversionPattern<TestRemappedValueRegionOp> { 1288 using OpConversionPattern<TestRemappedValueRegionOp>::OpConversionPattern; 1289 1290 LogicalResult 1291 matchAndRewrite(TestRemappedValueRegionOp op, OpAdaptor adaptor, 1292 ConversionPatternRewriter &rewriter) const final { 1293 Block &block = op.getBody().front(); 1294 Operation *terminator = block.getTerminator(); 1295 1296 // Merge the block into the parent region. 1297 Block *parentBlock = op->getBlock(); 1298 Block *finalBlock = rewriter.splitBlock(parentBlock, op->getIterator()); 1299 rewriter.mergeBlocks(&block, parentBlock, ValueRange()); 1300 rewriter.mergeBlocks(finalBlock, parentBlock, ValueRange()); 1301 1302 // Replace the results of this operation with the remapped terminator 1303 // values. 1304 SmallVector<Value> terminatorOperands; 1305 if (failed(rewriter.getRemappedValues(terminator->getOperands(), 1306 terminatorOperands))) 1307 return failure(); 1308 1309 rewriter.eraseOp(terminator); 1310 rewriter.replaceOp(op, terminatorOperands); 1311 return success(); 1312 } 1313 }; 1314 1315 struct TestRemappedValue 1316 : public mlir::PassWrapper<TestRemappedValue, OperationPass<>> { 1317 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestRemappedValue) 1318 1319 StringRef getArgument() const final { return "test-remapped-value"; } 1320 StringRef getDescription() const final { 1321 return "Test public remapped value mechanism in ConversionPatternRewriter"; 1322 } 1323 void runOnOperation() override { 1324 TestRemapValueTypeConverter typeConverter; 1325 1326 mlir::RewritePatternSet patterns(&getContext()); 1327 patterns.add<OneVResOneVOperandOp1Converter>(&getContext()); 1328 patterns.add<TestChangeProducerTypeF32ToF64, TestUpdateConsumerType>( 1329 &getContext()); 1330 patterns.add<TestRemapValueInRegion>(typeConverter, &getContext()); 1331 1332 mlir::ConversionTarget target(getContext()); 1333 target.addLegalOp<ModuleOp, func::FuncOp, TestReturnOp>(); 1334 1335 // Expect the type_producer/type_consumer operations to only operate on f64. 1336 target.addDynamicallyLegalOp<TestTypeProducerOp>( 1337 [](TestTypeProducerOp op) { return op.getType().isF64(); }); 1338 target.addDynamicallyLegalOp<TestTypeConsumerOp>([](TestTypeConsumerOp op) { 1339 return op.getOperand().getType().isF64(); 1340 }); 1341 1342 // We make OneVResOneVOperandOp1 legal only when it has more that one 1343 // operand. This will trigger the conversion that will replace one-operand 1344 // OneVResOneVOperandOp1 with two-operand OneVResOneVOperandOp1. 1345 target.addDynamicallyLegalOp<OneVResOneVOperandOp1>( 1346 [](Operation *op) { return op->getNumOperands() > 1; }); 1347 1348 if (failed(mlir::applyFullConversion(getOperation(), target, 1349 std::move(patterns)))) { 1350 signalPassFailure(); 1351 } 1352 } 1353 }; 1354 } // namespace 1355 1356 //===----------------------------------------------------------------------===// 1357 // Test patterns without a specific root operation kind 1358 //===----------------------------------------------------------------------===// 1359 1360 namespace { 1361 /// This pattern matches and removes any operation in the test dialect. 1362 struct RemoveTestDialectOps : public RewritePattern { 1363 RemoveTestDialectOps(MLIRContext *context) 1364 : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {} 1365 1366 LogicalResult matchAndRewrite(Operation *op, 1367 PatternRewriter &rewriter) const override { 1368 if (!isa<TestDialect>(op->getDialect())) 1369 return failure(); 1370 rewriter.eraseOp(op); 1371 return success(); 1372 } 1373 }; 1374 1375 struct TestUnknownRootOpDriver 1376 : public mlir::PassWrapper<TestUnknownRootOpDriver, OperationPass<>> { 1377 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestUnknownRootOpDriver) 1378 1379 StringRef getArgument() const final { 1380 return "test-legalize-unknown-root-patterns"; 1381 } 1382 StringRef getDescription() const final { 1383 return "Test public remapped value mechanism in ConversionPatternRewriter"; 1384 } 1385 void runOnOperation() override { 1386 mlir::RewritePatternSet patterns(&getContext()); 1387 patterns.add<RemoveTestDialectOps>(&getContext()); 1388 1389 mlir::ConversionTarget target(getContext()); 1390 target.addIllegalDialect<TestDialect>(); 1391 if (failed(applyPartialConversion(getOperation(), target, 1392 std::move(patterns)))) 1393 signalPassFailure(); 1394 } 1395 }; 1396 } // namespace 1397 1398 //===----------------------------------------------------------------------===// 1399 // Test patterns that uses operations and types defined at runtime 1400 //===----------------------------------------------------------------------===// 1401 1402 namespace { 1403 /// This pattern matches dynamic operations 'test.one_operand_two_results' and 1404 /// replace them with dynamic operations 'test.generic_dynamic_op'. 1405 struct RewriteDynamicOp : public RewritePattern { 1406 RewriteDynamicOp(MLIRContext *context) 1407 : RewritePattern("test.dynamic_one_operand_two_results", /*benefit=*/1, 1408 context) {} 1409 1410 LogicalResult matchAndRewrite(Operation *op, 1411 PatternRewriter &rewriter) const override { 1412 assert(op->getName().getStringRef() == 1413 "test.dynamic_one_operand_two_results" && 1414 "rewrite pattern should only match operations with the right name"); 1415 1416 OperationState state(op->getLoc(), "test.dynamic_generic", 1417 op->getOperands(), op->getResultTypes(), 1418 op->getAttrs()); 1419 auto *newOp = rewriter.create(state); 1420 rewriter.replaceOp(op, newOp->getResults()); 1421 return success(); 1422 } 1423 }; 1424 1425 struct TestRewriteDynamicOpDriver 1426 : public PassWrapper<TestRewriteDynamicOpDriver, OperationPass<>> { 1427 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestRewriteDynamicOpDriver) 1428 1429 void getDependentDialects(DialectRegistry ®istry) const override { 1430 registry.insert<TestDialect>(); 1431 } 1432 StringRef getArgument() const final { return "test-rewrite-dynamic-op"; } 1433 StringRef getDescription() const final { 1434 return "Test rewritting on dynamic operations"; 1435 } 1436 void runOnOperation() override { 1437 RewritePatternSet patterns(&getContext()); 1438 patterns.add<RewriteDynamicOp>(&getContext()); 1439 1440 ConversionTarget target(getContext()); 1441 target.addIllegalOp( 1442 OperationName("test.dynamic_one_operand_two_results", &getContext())); 1443 target.addLegalOp(OperationName("test.dynamic_generic", &getContext())); 1444 if (failed(applyPartialConversion(getOperation(), target, 1445 std::move(patterns)))) 1446 signalPassFailure(); 1447 } 1448 }; 1449 } // end anonymous namespace 1450 1451 //===----------------------------------------------------------------------===// 1452 // Test type conversions 1453 //===----------------------------------------------------------------------===// 1454 1455 namespace { 1456 struct TestTypeConversionProducer 1457 : public OpConversionPattern<TestTypeProducerOp> { 1458 using OpConversionPattern<TestTypeProducerOp>::OpConversionPattern; 1459 LogicalResult 1460 matchAndRewrite(TestTypeProducerOp op, OpAdaptor adaptor, 1461 ConversionPatternRewriter &rewriter) const final { 1462 Type resultType = op.getType(); 1463 Type convertedType = getTypeConverter() 1464 ? getTypeConverter()->convertType(resultType) 1465 : resultType; 1466 if (isa<FloatType>(resultType)) 1467 resultType = rewriter.getF64Type(); 1468 else if (resultType.isInteger(16)) 1469 resultType = rewriter.getIntegerType(64); 1470 else if (isa<test::TestRecursiveType>(resultType) && 1471 convertedType != resultType) 1472 resultType = convertedType; 1473 else 1474 return failure(); 1475 1476 rewriter.replaceOpWithNewOp<TestTypeProducerOp>(op, resultType); 1477 return success(); 1478 } 1479 }; 1480 1481 /// Call signature conversion and then fail the rewrite to trigger the undo 1482 /// mechanism. 1483 struct TestSignatureConversionUndo 1484 : public OpConversionPattern<TestSignatureConversionUndoOp> { 1485 using OpConversionPattern<TestSignatureConversionUndoOp>::OpConversionPattern; 1486 1487 LogicalResult 1488 matchAndRewrite(TestSignatureConversionUndoOp op, OpAdaptor adaptor, 1489 ConversionPatternRewriter &rewriter) const final { 1490 (void)rewriter.convertRegionTypes(&op->getRegion(0), *getTypeConverter()); 1491 return failure(); 1492 } 1493 }; 1494 1495 /// Call signature conversion without providing a type converter to handle 1496 /// materializations. 1497 struct TestTestSignatureConversionNoConverter 1498 : public OpConversionPattern<TestSignatureConversionNoConverterOp> { 1499 TestTestSignatureConversionNoConverter(const TypeConverter &converter, 1500 MLIRContext *context) 1501 : OpConversionPattern<TestSignatureConversionNoConverterOp>(context), 1502 converter(converter) {} 1503 1504 LogicalResult 1505 matchAndRewrite(TestSignatureConversionNoConverterOp op, OpAdaptor adaptor, 1506 ConversionPatternRewriter &rewriter) const final { 1507 Region ®ion = op->getRegion(0); 1508 Block *entry = ®ion.front(); 1509 1510 // Convert the original entry arguments. 1511 TypeConverter::SignatureConversion result(entry->getNumArguments()); 1512 if (failed( 1513 converter.convertSignatureArgs(entry->getArgumentTypes(), result))) 1514 return failure(); 1515 rewriter.modifyOpInPlace( 1516 op, [&] { rewriter.applySignatureConversion(®ion, result); }); 1517 return success(); 1518 } 1519 1520 const TypeConverter &converter; 1521 }; 1522 1523 /// Just forward the operands to the root op. This is essentially a no-op 1524 /// pattern that is used to trigger target materialization. 1525 struct TestTypeConsumerForward 1526 : public OpConversionPattern<TestTypeConsumerOp> { 1527 using OpConversionPattern<TestTypeConsumerOp>::OpConversionPattern; 1528 1529 LogicalResult 1530 matchAndRewrite(TestTypeConsumerOp op, OpAdaptor adaptor, 1531 ConversionPatternRewriter &rewriter) const final { 1532 rewriter.modifyOpInPlace(op, 1533 [&] { op->setOperands(adaptor.getOperands()); }); 1534 return success(); 1535 } 1536 }; 1537 1538 struct TestTypeConversionAnotherProducer 1539 : public OpRewritePattern<TestAnotherTypeProducerOp> { 1540 using OpRewritePattern<TestAnotherTypeProducerOp>::OpRewritePattern; 1541 1542 LogicalResult matchAndRewrite(TestAnotherTypeProducerOp op, 1543 PatternRewriter &rewriter) const final { 1544 rewriter.replaceOpWithNewOp<TestTypeProducerOp>(op, op.getType()); 1545 return success(); 1546 } 1547 }; 1548 1549 struct TestTypeConversionDriver 1550 : public PassWrapper<TestTypeConversionDriver, OperationPass<>> { 1551 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestTypeConversionDriver) 1552 1553 void getDependentDialects(DialectRegistry ®istry) const override { 1554 registry.insert<TestDialect>(); 1555 } 1556 StringRef getArgument() const final { 1557 return "test-legalize-type-conversion"; 1558 } 1559 StringRef getDescription() const final { 1560 return "Test various type conversion functionalities in DialectConversion"; 1561 } 1562 1563 void runOnOperation() override { 1564 // Initialize the type converter. 1565 SmallVector<Type, 2> conversionCallStack; 1566 TypeConverter converter; 1567 1568 /// Add the legal set of type conversions. 1569 converter.addConversion([](Type type) -> Type { 1570 // Treat F64 as legal. 1571 if (type.isF64()) 1572 return type; 1573 // Allow converting BF16/F16/F32 to F64. 1574 if (type.isBF16() || type.isF16() || type.isF32()) 1575 return FloatType::getF64(type.getContext()); 1576 // Otherwise, the type is illegal. 1577 return nullptr; 1578 }); 1579 converter.addConversion([](IntegerType type, SmallVectorImpl<Type> &) { 1580 // Drop all integer types. 1581 return success(); 1582 }); 1583 converter.addConversion( 1584 // Convert a recursive self-referring type into a non-self-referring 1585 // type named "outer_converted_type" that contains a SimpleAType. 1586 [&](test::TestRecursiveType type, 1587 SmallVectorImpl<Type> &results) -> std::optional<LogicalResult> { 1588 // If the type is already converted, return it to indicate that it is 1589 // legal. 1590 if (type.getName() == "outer_converted_type") { 1591 results.push_back(type); 1592 return success(); 1593 } 1594 1595 conversionCallStack.push_back(type); 1596 auto popConversionCallStack = llvm::make_scope_exit( 1597 [&conversionCallStack]() { conversionCallStack.pop_back(); }); 1598 1599 // If the type is on the call stack more than once (it is there at 1600 // least once because of the _current_ call, which is always the last 1601 // element on the stack), we've hit the recursive case. Just return 1602 // SimpleAType here to create a non-recursive type as a result. 1603 if (llvm::is_contained(ArrayRef(conversionCallStack).drop_back(), 1604 type)) { 1605 results.push_back(test::SimpleAType::get(type.getContext())); 1606 return success(); 1607 } 1608 1609 // Convert the body recursively. 1610 auto result = test::TestRecursiveType::get(type.getContext(), 1611 "outer_converted_type"); 1612 if (failed(result.setBody(converter.convertType(type.getBody())))) 1613 return failure(); 1614 results.push_back(result); 1615 return success(); 1616 }); 1617 1618 /// Add the legal set of type materializations. 1619 converter.addSourceMaterialization([](OpBuilder &builder, Type resultType, 1620 ValueRange inputs, 1621 Location loc) -> Value { 1622 // Allow casting from F64 back to F32. 1623 if (!resultType.isF16() && inputs.size() == 1 && 1624 inputs[0].getType().isF64()) 1625 return builder.create<TestCastOp>(loc, resultType, inputs).getResult(); 1626 // Allow producing an i32 or i64 from nothing. 1627 if ((resultType.isInteger(32) || resultType.isInteger(64)) && 1628 inputs.empty()) 1629 return builder.create<TestTypeProducerOp>(loc, resultType); 1630 // Allow producing an i64 from an integer. 1631 if (isa<IntegerType>(resultType) && inputs.size() == 1 && 1632 isa<IntegerType>(inputs[0].getType())) 1633 return builder.create<TestCastOp>(loc, resultType, inputs).getResult(); 1634 // Otherwise, fail. 1635 return nullptr; 1636 }); 1637 1638 // Initialize the conversion target. 1639 mlir::ConversionTarget target(getContext()); 1640 target.addDynamicallyLegalOp<TestTypeProducerOp>([](TestTypeProducerOp op) { 1641 auto recursiveType = dyn_cast<test::TestRecursiveType>(op.getType()); 1642 return op.getType().isF64() || op.getType().isInteger(64) || 1643 (recursiveType && 1644 recursiveType.getName() == "outer_converted_type"); 1645 }); 1646 target.addDynamicallyLegalOp<func::FuncOp>([&](func::FuncOp op) { 1647 return converter.isSignatureLegal(op.getFunctionType()) && 1648 converter.isLegal(&op.getBody()); 1649 }); 1650 target.addDynamicallyLegalOp<TestCastOp>([&](TestCastOp op) { 1651 // Allow casts from F64 to F32. 1652 return (*op.operand_type_begin()).isF64() && op.getType().isF32(); 1653 }); 1654 target.addDynamicallyLegalOp<TestSignatureConversionNoConverterOp>( 1655 [&](TestSignatureConversionNoConverterOp op) { 1656 return converter.isLegal(op.getRegion().front().getArgumentTypes()); 1657 }); 1658 1659 // Initialize the set of rewrite patterns. 1660 RewritePatternSet patterns(&getContext()); 1661 patterns.add<TestTypeConsumerForward, TestTypeConversionProducer, 1662 TestSignatureConversionUndo, 1663 TestTestSignatureConversionNoConverter>(converter, 1664 &getContext()); 1665 patterns.add<TestTypeConversionAnotherProducer>(&getContext()); 1666 mlir::populateAnyFunctionOpInterfaceTypeConversionPattern(patterns, 1667 converter); 1668 1669 if (failed(applyPartialConversion(getOperation(), target, 1670 std::move(patterns)))) 1671 signalPassFailure(); 1672 } 1673 }; 1674 } // namespace 1675 1676 //===----------------------------------------------------------------------===// 1677 // Test Target Materialization With No Uses 1678 //===----------------------------------------------------------------------===// 1679 1680 namespace { 1681 struct ForwardOperandPattern : public OpConversionPattern<TestTypeChangerOp> { 1682 using OpConversionPattern<TestTypeChangerOp>::OpConversionPattern; 1683 1684 LogicalResult 1685 matchAndRewrite(TestTypeChangerOp op, OpAdaptor adaptor, 1686 ConversionPatternRewriter &rewriter) const final { 1687 rewriter.replaceOp(op, adaptor.getOperands()); 1688 return success(); 1689 } 1690 }; 1691 1692 struct TestTargetMaterializationWithNoUses 1693 : public PassWrapper<TestTargetMaterializationWithNoUses, OperationPass<>> { 1694 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID( 1695 TestTargetMaterializationWithNoUses) 1696 1697 StringRef getArgument() const final { 1698 return "test-target-materialization-with-no-uses"; 1699 } 1700 StringRef getDescription() const final { 1701 return "Test a special case of target materialization in DialectConversion"; 1702 } 1703 1704 void runOnOperation() override { 1705 TypeConverter converter; 1706 converter.addConversion([](Type t) { return t; }); 1707 converter.addConversion([](IntegerType intTy) -> Type { 1708 if (intTy.getWidth() == 16) 1709 return IntegerType::get(intTy.getContext(), 64); 1710 return intTy; 1711 }); 1712 converter.addTargetMaterialization( 1713 [](OpBuilder &builder, Type type, ValueRange inputs, Location loc) { 1714 return builder.create<TestCastOp>(loc, type, inputs).getResult(); 1715 }); 1716 1717 ConversionTarget target(getContext()); 1718 target.addIllegalOp<TestTypeChangerOp>(); 1719 1720 RewritePatternSet patterns(&getContext()); 1721 patterns.add<ForwardOperandPattern>(converter, &getContext()); 1722 1723 if (failed(applyPartialConversion(getOperation(), target, 1724 std::move(patterns)))) 1725 signalPassFailure(); 1726 } 1727 }; 1728 } // namespace 1729 1730 //===----------------------------------------------------------------------===// 1731 // Test Block Merging 1732 //===----------------------------------------------------------------------===// 1733 1734 namespace { 1735 /// A rewriter pattern that tests that blocks can be merged. 1736 struct TestMergeBlock : public OpConversionPattern<TestMergeBlocksOp> { 1737 using OpConversionPattern<TestMergeBlocksOp>::OpConversionPattern; 1738 1739 LogicalResult 1740 matchAndRewrite(TestMergeBlocksOp op, OpAdaptor adaptor, 1741 ConversionPatternRewriter &rewriter) const final { 1742 Block &firstBlock = op.getBody().front(); 1743 Operation *branchOp = firstBlock.getTerminator(); 1744 Block *secondBlock = &*(std::next(op.getBody().begin())); 1745 auto succOperands = branchOp->getOperands(); 1746 SmallVector<Value, 2> replacements(succOperands); 1747 rewriter.eraseOp(branchOp); 1748 rewriter.mergeBlocks(secondBlock, &firstBlock, replacements); 1749 rewriter.modifyOpInPlace(op, [] {}); 1750 return success(); 1751 } 1752 }; 1753 1754 /// A rewrite pattern to tests the undo mechanism of blocks being merged. 1755 struct TestUndoBlocksMerge : public ConversionPattern { 1756 TestUndoBlocksMerge(MLIRContext *ctx) 1757 : ConversionPattern("test.undo_blocks_merge", /*benefit=*/1, ctx) {} 1758 LogicalResult 1759 matchAndRewrite(Operation *op, ArrayRef<Value> operands, 1760 ConversionPatternRewriter &rewriter) const final { 1761 Block &firstBlock = op->getRegion(0).front(); 1762 Operation *branchOp = firstBlock.getTerminator(); 1763 Block *secondBlock = &*(std::next(op->getRegion(0).begin())); 1764 rewriter.setInsertionPointToStart(secondBlock); 1765 rewriter.create<ILLegalOpF>(op->getLoc(), rewriter.getF32Type()); 1766 auto succOperands = branchOp->getOperands(); 1767 SmallVector<Value, 2> replacements(succOperands); 1768 rewriter.eraseOp(branchOp); 1769 rewriter.mergeBlocks(secondBlock, &firstBlock, replacements); 1770 rewriter.modifyOpInPlace(op, [] {}); 1771 return success(); 1772 } 1773 }; 1774 1775 /// A rewrite mechanism to inline the body of the op into its parent, when both 1776 /// ops can have a single block. 1777 struct TestMergeSingleBlockOps 1778 : public OpConversionPattern<SingleBlockImplicitTerminatorOp> { 1779 using OpConversionPattern< 1780 SingleBlockImplicitTerminatorOp>::OpConversionPattern; 1781 1782 LogicalResult 1783 matchAndRewrite(SingleBlockImplicitTerminatorOp op, OpAdaptor adaptor, 1784 ConversionPatternRewriter &rewriter) const final { 1785 SingleBlockImplicitTerminatorOp parentOp = 1786 op->getParentOfType<SingleBlockImplicitTerminatorOp>(); 1787 if (!parentOp) 1788 return failure(); 1789 Block &innerBlock = op.getRegion().front(); 1790 TerminatorOp innerTerminator = 1791 cast<TerminatorOp>(innerBlock.getTerminator()); 1792 rewriter.inlineBlockBefore(&innerBlock, op); 1793 rewriter.eraseOp(innerTerminator); 1794 rewriter.eraseOp(op); 1795 return success(); 1796 } 1797 }; 1798 1799 struct TestMergeBlocksPatternDriver 1800 : public PassWrapper<TestMergeBlocksPatternDriver, OperationPass<>> { 1801 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestMergeBlocksPatternDriver) 1802 1803 StringRef getArgument() const final { return "test-merge-blocks"; } 1804 StringRef getDescription() const final { 1805 return "Test Merging operation in ConversionPatternRewriter"; 1806 } 1807 void runOnOperation() override { 1808 MLIRContext *context = &getContext(); 1809 mlir::RewritePatternSet patterns(context); 1810 patterns.add<TestMergeBlock, TestUndoBlocksMerge, TestMergeSingleBlockOps>( 1811 context); 1812 ConversionTarget target(*context); 1813 target.addLegalOp<func::FuncOp, ModuleOp, TerminatorOp, TestBranchOp, 1814 TestTypeConsumerOp, TestTypeProducerOp, TestReturnOp>(); 1815 target.addIllegalOp<ILLegalOpF>(); 1816 1817 /// Expect the op to have a single block after legalization. 1818 target.addDynamicallyLegalOp<TestMergeBlocksOp>( 1819 [&](TestMergeBlocksOp op) -> bool { 1820 return llvm::hasSingleElement(op.getBody()); 1821 }); 1822 1823 /// Only allow `test.br` within test.merge_blocks op. 1824 target.addDynamicallyLegalOp<TestBranchOp>([&](TestBranchOp op) -> bool { 1825 return op->getParentOfType<TestMergeBlocksOp>(); 1826 }); 1827 1828 /// Expect that all nested test.SingleBlockImplicitTerminator ops are 1829 /// inlined. 1830 target.addDynamicallyLegalOp<SingleBlockImplicitTerminatorOp>( 1831 [&](SingleBlockImplicitTerminatorOp op) -> bool { 1832 return !op->getParentOfType<SingleBlockImplicitTerminatorOp>(); 1833 }); 1834 1835 DenseSet<Operation *> unlegalizedOps; 1836 ConversionConfig config; 1837 config.unlegalizedOps = &unlegalizedOps; 1838 (void)applyPartialConversion(getOperation(), target, std::move(patterns), 1839 config); 1840 for (auto *op : unlegalizedOps) 1841 op->emitRemark() << "op '" << op->getName() << "' is not legalizable"; 1842 } 1843 }; 1844 } // namespace 1845 1846 //===----------------------------------------------------------------------===// 1847 // Test Selective Replacement 1848 //===----------------------------------------------------------------------===// 1849 1850 namespace { 1851 /// A rewrite mechanism to inline the body of the op into its parent, when both 1852 /// ops can have a single block. 1853 struct TestSelectiveOpReplacementPattern : public OpRewritePattern<TestCastOp> { 1854 using OpRewritePattern<TestCastOp>::OpRewritePattern; 1855 1856 LogicalResult matchAndRewrite(TestCastOp op, 1857 PatternRewriter &rewriter) const final { 1858 if (op.getNumOperands() != 2) 1859 return failure(); 1860 OperandRange operands = op.getOperands(); 1861 1862 // Replace non-terminator uses with the first operand. 1863 rewriter.replaceUsesWithIf(op, operands[0], [](OpOperand &operand) { 1864 return operand.getOwner()->hasTrait<OpTrait::IsTerminator>(); 1865 }); 1866 // Replace everything else with the second operand if the operation isn't 1867 // dead. 1868 rewriter.replaceOp(op, op.getOperand(1)); 1869 return success(); 1870 } 1871 }; 1872 1873 struct TestSelectiveReplacementPatternDriver 1874 : public PassWrapper<TestSelectiveReplacementPatternDriver, 1875 OperationPass<>> { 1876 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID( 1877 TestSelectiveReplacementPatternDriver) 1878 1879 StringRef getArgument() const final { 1880 return "test-pattern-selective-replacement"; 1881 } 1882 StringRef getDescription() const final { 1883 return "Test selective replacement in the PatternRewriter"; 1884 } 1885 void runOnOperation() override { 1886 MLIRContext *context = &getContext(); 1887 mlir::RewritePatternSet patterns(context); 1888 patterns.add<TestSelectiveOpReplacementPattern>(context); 1889 (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns)); 1890 } 1891 }; 1892 } // namespace 1893 1894 //===----------------------------------------------------------------------===// 1895 // PassRegistration 1896 //===----------------------------------------------------------------------===// 1897 1898 namespace mlir { 1899 namespace test { 1900 void registerPatternsTestPass() { 1901 PassRegistration<TestReturnTypeDriver>(); 1902 1903 PassRegistration<TestDerivedAttributeDriver>(); 1904 1905 PassRegistration<TestPatternDriver>(); 1906 PassRegistration<TestStrictPatternDriver>(); 1907 1908 PassRegistration<TestLegalizePatternDriver>([] { 1909 return std::make_unique<TestLegalizePatternDriver>(legalizerConversionMode); 1910 }); 1911 1912 PassRegistration<TestRemappedValue>(); 1913 1914 PassRegistration<TestUnknownRootOpDriver>(); 1915 1916 PassRegistration<TestTypeConversionDriver>(); 1917 PassRegistration<TestTargetMaterializationWithNoUses>(); 1918 1919 PassRegistration<TestRewriteDynamicOpDriver>(); 1920 1921 PassRegistration<TestMergeBlocksPatternDriver>(); 1922 PassRegistration<TestSelectiveReplacementPatternDriver>(); 1923 } 1924 } // namespace test 1925 } // namespace mlir 1926