xref: /llvm-project/mlir/test/lib/Dialect/Test/TestPatterns.cpp (revision e95e94adc6bb748de015ac3053e7f0786b65f351)
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 &registry) 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(&regionOp->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 &region = *op->getParentRegion();
749     Type i32Type = rewriter.getIntegerType(32);
750     Location loc = op->getLoc();
751     rewriter.createBlock(&region, 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 &region = *op->getParentRegion();
767     Type i32Type = rewriter.getIntegerType(32);
768     Location loc = op->getLoc();
769     rewriter.createBlock(&region, 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 &region = op->getRegion(0);
856     Block *entry = &region.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(&region, 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 &registry) 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 &registry) 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 &region = op->getRegion(0);
1508     Block *entry = &region.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(&region, 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 &registry) 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