brintos

brintos / llvm-project-archived public Read only

0
0
Text · 97.1 KiB · 7eabaae Raw
2492 lines · cpp
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-exception6//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/CommonFolders.h"14#include "mlir/Dialect/ControlFlow/Transforms/StructuralTypeConversions.h"15#include "mlir/Dialect/Func/IR/FuncOps.h"16#include "mlir/Dialect/Func/Transforms/FuncConversions.h"17#include "mlir/Dialect/SCF/Transforms/Patterns.h"18#include "mlir/Dialect/Tensor/IR/Tensor.h"19#include "mlir/IR/BuiltinAttributes.h"20#include "mlir/IR/Matchers.h"21#include "mlir/IR/PatternMatch.h"22#include "mlir/IR/Visitors.h"23#include "mlir/Pass/Pass.h"24#include "mlir/Transforms/DialectConversion.h"25#include "mlir/Transforms/FoldUtils.h"26#include "mlir/Transforms/GreedyPatternRewriteDriver.h"27#include "mlir/Transforms/WalkPatternRewriteDriver.h"28#include "llvm/ADT/ScopeExit.h"29#include <cstdint>30 31using namespace mlir;32using namespace test;33 34// Native function for testing NativeCodeCall35static Value chooseOperand(Value input1, Value input2, BoolAttr choice) {36  return choice.getValue() ? input1 : input2;37}38 39static void createOpI(PatternRewriter &rewriter, Location loc, Value input) {40  OpI::create(rewriter, loc, input);41}42 43static void handleNoResultOp(PatternRewriter &rewriter,44                             OpSymbolBindingNoResult op) {45  // Turn the no result op to a one-result op.46  OpSymbolBindingB::create(rewriter, op.getLoc(), op.getOperand().getType(),47                           op.getOperand());48}49 50static bool getFirstI32Result(Operation *op, Value &value) {51  if (!Type(op->getResult(0).getType()).isSignlessInteger(32))52    return false;53  value = op->getResult(0);54  return true;55}56 57static Value bindNativeCodeCallResult(Value value) { return value; }58 59static SmallVector<Value, 2> bindMultipleNativeCodeCallResult(Value input1,60                                                              Value input2) {61  return SmallVector<Value, 2>({input2, input1});62}63 64// Test that natives calls are only called once during rewrites.65// OpM_Test will return Pi, increased by 1 for each subsequent calls.66// This let us check the number of times OpM_Test was called by inspecting67// the returned value in the MLIR output.68static int64_t opMIncreasingValue = 314159265;69static Attribute opMTest(PatternRewriter &rewriter, Value val) {70  int64_t i = opMIncreasingValue++;71  return rewriter.getIntegerAttr(rewriter.getIntegerType(32), i);72}73 74static bool assertBinOpEqualArgsAndReturnTrue(Value v) {75  Operation *operation = v.getDefiningOp();76  if (operation->getOperand(0) != operation->getOperand(1)) {77    // Name binding equality check must happen before user-defined constraints,78    // thus this must not be triggered.79    llvm::report_fatal_error("Arguments are not equal");80  }81  return true;82}83 84namespace {85#include "TestPatterns.inc"86} // namespace87 88//===----------------------------------------------------------------------===//89// Test Reduce Pattern Interface90//===----------------------------------------------------------------------===//91 92void test::populateTestReductionPatterns(RewritePatternSet &patterns) {93  populateWithGenerated(patterns);94}95 96//===----------------------------------------------------------------------===//97// Canonicalizer Driver.98//===----------------------------------------------------------------------===//99 100namespace {101struct FoldingPattern : public RewritePattern {102public:103  FoldingPattern(MLIRContext *context)104      : RewritePattern(TestOpInPlaceFoldAnchor::getOperationName(),105                       /*benefit=*/1, context) {}106 107  LogicalResult matchAndRewrite(Operation *op,108                                PatternRewriter &rewriter) const override {109    // Exercise createOrFold API for a single-result operation that is folded110    // upon construction. The operation being created has an in-place folder,111    // and it should be still present in the output. Furthermore, the folder112    // should not crash when attempting to recover the (unchanged) operation113    // result.114    Value result = rewriter.createOrFold<TestOpInPlaceFold>(115        op->getLoc(), rewriter.getIntegerType(32), op->getOperand(0));116    assert(result);117    rewriter.replaceOp(op, result);118    return success();119  }120};121 122/// This pattern creates a foldable operation at the entry point of the block.123/// This tests the situation where the operation folder will need to replace an124/// operation with a previously created constant that does not initially125/// dominate the operation to replace.126struct FolderInsertBeforePreviouslyFoldedConstantPattern127    : public OpRewritePattern<TestCastOp> {128public:129  static_assert(std::is_same_v<Base, OpRewritePattern<TestCastOp>>);130  using Base::Base;131 132  LogicalResult matchAndRewrite(TestCastOp op,133                                PatternRewriter &rewriter) const override {134    if (!op->hasAttr("test_fold_before_previously_folded_op"))135      return failure();136    rewriter.setInsertionPointToStart(op->getBlock());137 138    auto constOp = arith::ConstantOp::create(rewriter, op.getLoc(),139                                             rewriter.getBoolAttr(true));140    rewriter.replaceOpWithNewOp<TestCastOp>(op, rewriter.getI32Type(),141                                            Value(constOp));142    return success();143  }144};145 146/// This pattern matches test.op_commutative2 with the first operand being147/// another test.op_commutative2 with a constant on the right side and fold it148/// away by propagating it as its result. This is intend to check that patterns149/// are applied after the commutative property moves constant to the right.150struct FolderCommutativeOp2WithConstant151    : public OpRewritePattern<TestCommutative2Op> {152public:153  using OpRewritePattern<TestCommutative2Op>::OpRewritePattern;154 155  LogicalResult matchAndRewrite(TestCommutative2Op op,156                                PatternRewriter &rewriter) const override {157    auto operand = op->getOperand(0).getDefiningOp<TestCommutative2Op>();158    if (!operand)159      return failure();160    Attribute constInput;161    if (!matchPattern(operand->getOperand(1), m_Constant(&constInput)))162      return failure();163    rewriter.replaceOp(op, operand->getOperand(1));164    return success();165  }166};167 168/// This pattern matches test.any_attr_of_i32_str ops. In case of an integer169/// attribute with value smaller than MaxVal, it increments the value by 1.170template <int MaxVal>171struct IncrementIntAttribute : public OpRewritePattern<AnyAttrOfOp> {172  using OpRewritePattern<AnyAttrOfOp>::OpRewritePattern;173 174  LogicalResult matchAndRewrite(AnyAttrOfOp op,175                                PatternRewriter &rewriter) const override {176    auto intAttr = dyn_cast<IntegerAttr>(op.getAttr());177    if (!intAttr)178      return failure();179    int64_t val = intAttr.getInt();180    if (val >= MaxVal)181      return failure();182    rewriter.modifyOpInPlace(183        op, [&]() { op.setAttrAttr(rewriter.getI32IntegerAttr(val + 1)); });184    return success();185  }186};187 188/// This patterns adds an "eligible" attribute to "foo.maybe_eligible_op".189struct MakeOpEligible : public RewritePattern {190  MakeOpEligible(MLIRContext *context)191      : RewritePattern("foo.maybe_eligible_op", /*benefit=*/1, context) {}192 193  LogicalResult matchAndRewrite(Operation *op,194                                PatternRewriter &rewriter) const override {195    if (op->hasAttr("eligible"))196      return failure();197    rewriter.modifyOpInPlace(198        op, [&]() { op->setAttr("eligible", rewriter.getUnitAttr()); });199    return success();200  }201};202 203/// This pattern hoists eligible ops out of a "test.one_region_op".204struct HoistEligibleOps : public OpRewritePattern<test::OneRegionOp> {205  using OpRewritePattern<test::OneRegionOp>::OpRewritePattern;206 207  LogicalResult matchAndRewrite(test::OneRegionOp op,208                                PatternRewriter &rewriter) const override {209    Operation *terminator = op.getRegion().front().getTerminator();210    Operation *toBeHoisted = terminator->getOperands()[0].getDefiningOp();211    if (toBeHoisted->getParentOp() != op)212      return failure();213    if (!toBeHoisted->hasAttr("eligible"))214      return failure();215    rewriter.moveOpBefore(toBeHoisted, op);216    return success();217  }218};219 220struct FoldSignOpF32ToSI32 : public OpRewritePattern<test::SignOp> {221  using OpRewritePattern<test::SignOp>::OpRewritePattern;222 223  LogicalResult matchAndRewrite(test::SignOp op,224                                PatternRewriter &rewriter) const override {225    if (op->getNumOperands() != 1 || op->getNumResults() != 1)226      return failure();227 228    TypedAttr operandAttr;229    matchPattern(op->getOperand(0), m_Constant(&operandAttr));230    if (!operandAttr)231      return failure();232 233    TypedAttr res = cast_or_null<TypedAttr>(234        constFoldUnaryOp<FloatAttr, FloatAttr::ValueType, void, IntegerAttr>(235            operandAttr, op.getType(), [](APFloat operand) -> APSInt {236              static const APFloat zero(0.0f);237              int operandSign = 0;238              if (operand != zero)239                operandSign = (operand < zero) ? -1 : +1;240              return APSInt(APInt(32, operandSign), false);241            }));242    if (!res)243      return failure();244 245    rewriter.replaceOpWithNewOp<arith::ConstantOp>(op, res);246    return success();247  }248};249 250struct FoldLessThanOpF32ToI1 : public OpRewritePattern<test::LessThanOp> {251  using OpRewritePattern<test::LessThanOp>::OpRewritePattern;252 253  LogicalResult matchAndRewrite(test::LessThanOp op,254                                PatternRewriter &rewriter) const override {255    if (op->getNumOperands() != 2 || op->getNumResults() != 1)256      return failure();257 258    TypedAttr lhsAttr;259    TypedAttr rhsAttr;260    matchPattern(op->getOperand(0), m_Constant(&lhsAttr));261    matchPattern(op->getOperand(1), m_Constant(&rhsAttr));262 263    if (!lhsAttr || !rhsAttr)264      return failure();265 266    Attribute operandAttrs[2] = {lhsAttr, rhsAttr};267    TypedAttr res = cast_or_null<TypedAttr>(268        constFoldBinaryOp<FloatAttr, FloatAttr::ValueType, void, IntegerAttr>(269            operandAttrs, op.getType(), [](APFloat lhs, APFloat rhs) -> APInt {270              return APInt(1, lhs < rhs);271            }));272    if (!res)273      return failure();274 275    rewriter.replaceOpWithNewOp<arith::ConstantOp>(op, res);276    return success();277  }278};279 280/// This pattern moves "test.move_before_parent_op" before the parent op.281struct MoveBeforeParentOp : public RewritePattern {282  MoveBeforeParentOp(MLIRContext *context)283      : RewritePattern("test.move_before_parent_op", /*benefit=*/1, context) {}284 285  LogicalResult matchAndRewrite(Operation *op,286                                PatternRewriter &rewriter) const override {287    // Do not hoist past functions.288    if (isa<FunctionOpInterface>(op->getParentOp()))289      return failure();290    rewriter.moveOpBefore(op, op->getParentOp());291    return success();292  }293};294 295/// This pattern moves "test.move_after_parent_op" after the parent op.296struct MoveAfterParentOp : public RewritePattern {297  MoveAfterParentOp(MLIRContext *context)298      : RewritePattern("test.move_after_parent_op", /*benefit=*/1, context) {}299 300  LogicalResult matchAndRewrite(Operation *op,301                                PatternRewriter &rewriter) const override {302    // Do not hoist past functions.303    if (isa<FunctionOpInterface>(op->getParentOp()))304      return failure();305 306    int64_t moveForwardBy = 0;307    if (auto advanceBy = op->getAttrOfType<IntegerAttr>("advance"))308      moveForwardBy = advanceBy.getInt();309 310    Operation *moveAfter = op->getParentOp();311    for (int64_t i = 0; i < moveForwardBy; ++i)312      moveAfter = moveAfter->getNextNode();313 314    rewriter.moveOpAfter(op, moveAfter);315    return success();316  }317};318 319/// This pattern inlines blocks that are nested in320/// "test.inline_blocks_into_parent" into the parent block.321struct InlineBlocksIntoParent : public RewritePattern {322  InlineBlocksIntoParent(MLIRContext *context)323      : RewritePattern("test.inline_blocks_into_parent", /*benefit=*/1,324                       context) {}325 326  LogicalResult matchAndRewrite(Operation *op,327                                PatternRewriter &rewriter) const override {328    bool changed = false;329    for (Region &r : op->getRegions()) {330      while (!r.empty()) {331        rewriter.inlineBlockBefore(&r.front(), op);332        changed = true;333      }334    }335    return success(changed);336  }337};338 339/// This pattern splits blocks at "test.split_block_here" and replaces the op340/// with a new op (to prevent an infinite loop of block splitting).341struct SplitBlockHere : public RewritePattern {342  SplitBlockHere(MLIRContext *context)343      : RewritePattern("test.split_block_here", /*benefit=*/1, context) {}344 345  LogicalResult matchAndRewrite(Operation *op,346                                PatternRewriter &rewriter) const override {347    rewriter.splitBlock(op->getBlock(), op->getIterator());348    Operation *newOp = rewriter.create(349        op->getLoc(),350        OperationName("test.new_op", op->getContext()).getIdentifier(),351        op->getOperands(), op->getResultTypes());352    rewriter.replaceOp(op, newOp);353    return success();354  }355};356 357/// This pattern clones "test.clone_me" ops.358struct CloneOp : public RewritePattern {359  CloneOp(MLIRContext *context)360      : RewritePattern("test.clone_me", /*benefit=*/1, context) {}361 362  LogicalResult matchAndRewrite(Operation *op,363                                PatternRewriter &rewriter) const override {364    // Do not clone already cloned ops to avoid going into an infinite loop.365    if (op->hasAttr("was_cloned"))366      return failure();367    Operation *cloned = rewriter.clone(*op);368    cloned->setAttr("was_cloned", rewriter.getUnitAttr());369    return success();370  }371};372 373/// This pattern clones regions of "test.clone_region_before" ops before the374/// parent block.375struct CloneRegionBeforeOp : public RewritePattern {376  CloneRegionBeforeOp(MLIRContext *context)377      : RewritePattern("test.clone_region_before", /*benefit=*/1, context) {}378 379  LogicalResult matchAndRewrite(Operation *op,380                                PatternRewriter &rewriter) const override {381    // Do not clone already cloned ops to avoid going into an infinite loop.382    if (op->hasAttr("was_cloned"))383      return failure();384    for (Region &r : op->getRegions())385      rewriter.cloneRegionBefore(r, op->getBlock());386    op->setAttr("was_cloned", rewriter.getUnitAttr());387    return success();388  }389};390 391/// Replace an operation may introduce the re-visiting of its users.392class ReplaceWithNewOp : public RewritePattern {393public:394  ReplaceWithNewOp(MLIRContext *context)395      : RewritePattern("test.replace_with_new_op", /*benefit=*/1, context) {}396 397  LogicalResult matchAndRewrite(Operation *op,398                                PatternRewriter &rewriter) const override {399    Operation *newOp;400    if (op->hasAttr("create_erase_op")) {401      newOp = rewriter.create(402          op->getLoc(),403          OperationName("test.erase_op", op->getContext()).getIdentifier(),404          ValueRange(), TypeRange());405    } else {406      newOp = rewriter.create(407          op->getLoc(),408          OperationName("test.new_op", op->getContext()).getIdentifier(),409          op->getOperands(), op->getResultTypes());410    }411    // "replaceOp" could be used instead of "replaceAllOpUsesWith"+"eraseOp".412    // A "notifyOperationReplaced" callback is triggered in either case.413    rewriter.replaceAllOpUsesWith(op, newOp->getResults());414    rewriter.eraseOp(op);415    return success();416  }417};418 419/// Erases the first child block of the matched "test.erase_first_block"420/// operation.421class EraseFirstBlock : public RewritePattern {422public:423  EraseFirstBlock(MLIRContext *context)424      : RewritePattern("test.erase_first_block", /*benefit=*/1, context) {}425 426  LogicalResult matchAndRewrite(Operation *op,427                                PatternRewriter &rewriter) const override {428    for (Region &r : op->getRegions()) {429      for (Block &b : r.getBlocks()) {430        rewriter.eraseBlock(&b);431        return success();432      }433    }434 435    return failure();436  }437};438 439struct TestGreedyPatternDriver440    : public PassWrapper<TestGreedyPatternDriver, OperationPass<>> {441  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestGreedyPatternDriver)442 443  TestGreedyPatternDriver() = default;444  TestGreedyPatternDriver(const TestGreedyPatternDriver &other)445      : PassWrapper(other) {}446 447  StringRef getArgument() const final { return "test-greedy-patterns"; }448  StringRef getDescription() const final { return "Run test dialect patterns"; }449  void runOnOperation() override {450    mlir::RewritePatternSet patterns(&getContext());451    populateWithGenerated(patterns);452 453    // Verify named pattern is generated with expected name.454    patterns.add<FoldingPattern, TestNamedPatternRule,455                 FolderInsertBeforePreviouslyFoldedConstantPattern,456                 FolderCommutativeOp2WithConstant, HoistEligibleOps,457                 MakeOpEligible>(&getContext());458 459    // Additional patterns for testing the GreedyPatternRewriteDriver.460    patterns.insert<IncrementIntAttribute<3>>(&getContext());461 462    GreedyRewriteConfig config;463    config.setUseTopDownTraversal(useTopDownTraversal)464        .setMaxIterations(this->maxIterations)465        .enableFolding(this->fold)466        .enableConstantCSE(this->cseConstants);467    (void)applyPatternsGreedily(getOperation(), std::move(patterns), config);468  }469 470  Option<bool> useTopDownTraversal{471      *this, "top-down",472      llvm::cl::desc("Seed the worklist in general top-down order"),473      llvm::cl::init(GreedyRewriteConfig().getUseTopDownTraversal())};474  Option<int> maxIterations{475      *this, "max-iterations",476      llvm::cl::desc("Max. iterations in the GreedyRewriteConfig"),477      llvm::cl::init(GreedyRewriteConfig().getMaxIterations())};478  Option<bool> fold{*this, "fold", llvm::cl::desc("Whether to fold"),479                    llvm::cl::init(GreedyRewriteConfig().isFoldingEnabled())};480  Option<bool> cseConstants{481      *this, "cse-constants", llvm::cl::desc("Whether to CSE constants"),482      llvm::cl::init(GreedyRewriteConfig().isConstantCSEEnabled())};483};484 485struct DumpNotifications : public RewriterBase::Listener {486  void notifyBlockInserted(Block *block, Region *previous,487                           Region::iterator previousIt) override {488    llvm::outs() << "notifyBlockInserted";489    if (block->getParentOp()) {490      llvm::outs() << " into " << block->getParentOp()->getName() << ": ";491    } else {492      llvm::outs() << " into unknown op: ";493    }494    if (previous == nullptr) {495      llvm::outs() << "was unlinked\n";496    } else {497      llvm::outs() << "was linked\n";498    }499  }500  void notifyOperationInserted(Operation *op,501                               OpBuilder::InsertPoint previous) override {502    llvm::outs() << "notifyOperationInserted: " << op->getName();503    if (!previous.isSet()) {504      llvm::outs() << ", was unlinked\n";505    } else {506      if (!previous.getPoint().getNodePtr()) {507        llvm::outs() << ", was linked, exact position unknown\n";508      } else if (previous.getPoint() == previous.getBlock()->end()) {509        llvm::outs() << ", was last in block\n";510      } else {511        llvm::outs() << ", previous = " << previous.getPoint()->getName()512                     << "\n";513      }514    }515  }516  void notifyBlockErased(Block *block) override {517    llvm::outs() << "notifyBlockErased\n";518  }519  void notifyOperationErased(Operation *op) override {520    llvm::outs() << "notifyOperationErased: " << op->getName() << "\n";521  }522  void notifyOperationModified(Operation *op) override {523    llvm::outs() << "notifyOperationModified: " << op->getName() << "\n";524  }525  void notifyOperationReplaced(Operation *op, ValueRange values) override {526    llvm::outs() << "notifyOperationReplaced: " << op->getName() << "\n";527  }528};529 530struct TestStrictPatternDriver531    : public PassWrapper<TestStrictPatternDriver, OperationPass<func::FuncOp>> {532public:533  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestStrictPatternDriver)534 535  TestStrictPatternDriver() = default;536  TestStrictPatternDriver(const TestStrictPatternDriver &other)537      : PassWrapper(other) {538    strictMode = other.strictMode;539  }540 541  StringRef getArgument() const final { return "test-strict-pattern-driver"; }542  StringRef getDescription() const final {543    return "Test strict mode of pattern driver";544  }545 546  void runOnOperation() override {547    MLIRContext *ctx = &getContext();548    mlir::RewritePatternSet patterns(ctx);549    patterns.add<550        // clang-format off551        ChangeBlockOp,552        CloneOp,553        CloneRegionBeforeOp,554        EraseOp,555        ImplicitChangeOp,556        InlineBlocksIntoParent,557        InsertSameOp,558        MoveBeforeParentOp,559        ReplaceWithNewOp,560        SplitBlockHere561        // clang-format on562        >(ctx);563    SmallVector<Operation *> ops;564    getOperation()->walk([&](Operation *op) {565      StringRef opName = op->getName().getStringRef();566      if (opName == "test.insert_same_op" || opName == "test.change_block_op" ||567          opName == "test.replace_with_new_op" || opName == "test.erase_op" ||568          opName == "test.move_before_parent_op" ||569          opName == "test.inline_blocks_into_parent" ||570          opName == "test.split_block_here" || opName == "test.clone_me" ||571          opName == "test.clone_region_before") {572        ops.push_back(op);573      }574    });575 576    DumpNotifications dumpNotifications;577    GreedyRewriteConfig config;578    config.setListener(&dumpNotifications);579    if (strictMode == "AnyOp") {580      config.setStrictness(GreedyRewriteStrictness::AnyOp);581    } else if (strictMode == "ExistingAndNewOps") {582      config.setStrictness(GreedyRewriteStrictness::ExistingAndNewOps);583    } else if (strictMode == "ExistingOps") {584      config.setStrictness(GreedyRewriteStrictness::ExistingOps);585    } else {586      llvm_unreachable("invalid strictness option");587    }588 589    // Check if these transformations introduce visiting of operations that590    // are not in the `ops` set (The new created ops are valid). An invalid591    // operation will trigger the assertion while processing.592    bool changed = false;593    bool allErased = false;594    (void)applyOpPatternsGreedily(ArrayRef(ops), std::move(patterns), config,595                                  &changed, &allErased);596    Builder b(ctx);597    getOperation()->setAttr("pattern_driver_changed", b.getBoolAttr(changed));598    getOperation()->setAttr("pattern_driver_all_erased",599                            b.getBoolAttr(allErased));600  }601 602  Option<std::string> strictMode{603      *this, "strictness",604      llvm::cl::desc("Can be {AnyOp, ExistingAndNewOps, ExistingOps}"),605      llvm::cl::init("AnyOp")};606 607private:608  // New inserted operation is valid for further transformation.609  class InsertSameOp : public RewritePattern {610  public:611    InsertSameOp(MLIRContext *context)612        : RewritePattern("test.insert_same_op", /*benefit=*/1, context) {}613 614    LogicalResult matchAndRewrite(Operation *op,615                                  PatternRewriter &rewriter) const override {616      if (op->hasAttr("skip"))617        return failure();618 619      Operation *newOp =620          rewriter.create(op->getLoc(), op->getName().getIdentifier(),621                          op->getOperands(), op->getResultTypes());622      rewriter.modifyOpInPlace(623          op, [&]() { op->setAttr("skip", rewriter.getBoolAttr(true)); });624      newOp->setAttr("skip", rewriter.getBoolAttr(true));625 626      return success();627    }628  };629 630  // Remove an operation may introduce the re-visiting of its operands.631  class EraseOp : public RewritePattern {632  public:633    EraseOp(MLIRContext *context)634        : RewritePattern("test.erase_op", /*benefit=*/1, context) {}635    LogicalResult matchAndRewrite(Operation *op,636                                  PatternRewriter &rewriter) const override {637      rewriter.eraseOp(op);638      return success();639    }640  };641 642  // The following two patterns test RewriterBase::replaceAllUsesWith.643  //644  // That function replaces all usages of a Block (or a Value) with another one645  // *and tracks these changes in the rewriter.* The GreedyPatternRewriteDriver646  // with GreedyRewriteStrictness::AnyOp uses that tracking to construct its647  // worklist: when an op is modified, it is added to the worklist. The two648  // patterns below make the tracking observable: ChangeBlockOp replaces all649  // usages of a block and that pattern is applied because the corresponding ops650  // are put on the initial worklist (see above). ImplicitChangeOp does an651  // unrelated change but ops of the corresponding type are *not* on the initial652  // worklist, so the effect of the second pattern is only visible if the653  // tracking and subsequent adding to the worklist actually works.654 655  // Replace all usages of the first successor with the second successor.656  class ChangeBlockOp : public RewritePattern {657  public:658    ChangeBlockOp(MLIRContext *context)659        : RewritePattern("test.change_block_op", /*benefit=*/1, context) {}660    LogicalResult matchAndRewrite(Operation *op,661                                  PatternRewriter &rewriter) const override {662      if (op->getNumSuccessors() < 2)663        return failure();664      Block *firstSuccessor = op->getSuccessor(0);665      Block *secondSuccessor = op->getSuccessor(1);666      if (firstSuccessor == secondSuccessor)667        return failure();668      // This is the function being tested:669      rewriter.replaceAllUsesWith(firstSuccessor, secondSuccessor);670      // Using the following line instead would make the test fail:671      // firstSuccessor->replaceAllUsesWith(secondSuccessor);672      return success();673    }674  };675 676  // Changes the successor to the parent block.677  class ImplicitChangeOp : public RewritePattern {678  public:679    ImplicitChangeOp(MLIRContext *context)680        : RewritePattern("test.implicit_change_op", /*benefit=*/1, context) {}681    LogicalResult matchAndRewrite(Operation *op,682                                  PatternRewriter &rewriter) const override {683      if (op->getNumSuccessors() < 1 || op->getSuccessor(0) == op->getBlock())684        return failure();685      rewriter.modifyOpInPlace(op,686                               [&]() { op->setSuccessor(op->getBlock(), 0); });687      return success();688    }689  };690};691 692struct TestWalkPatternDriver final693    : PassWrapper<TestWalkPatternDriver, OperationPass<>> {694  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestWalkPatternDriver)695 696  TestWalkPatternDriver() = default;697  TestWalkPatternDriver(const TestWalkPatternDriver &other)698      : PassWrapper(other) {}699 700  StringRef getArgument() const override {701    return "test-walk-pattern-rewrite-driver";702  }703  StringRef getDescription() const override {704    return "Run test walk pattern rewrite driver";705  }706  void runOnOperation() override {707    mlir::RewritePatternSet patterns(&getContext());708 709    // Patterns for testing the WalkPatternRewriteDriver.710    patterns.add<IncrementIntAttribute<3>, MoveBeforeParentOp,711                 MoveAfterParentOp, CloneOp, ReplaceWithNewOp, EraseFirstBlock>(712        &getContext());713 714    DumpNotifications dumpListener;715    walkAndApplyPatterns(getOperation(), std::move(patterns),716                         dumpNotifications ? &dumpListener : nullptr);717  }718 719  Option<bool> dumpNotifications{720      *this, "dump-notifications",721      llvm::cl::desc("Print rewrite listener notifications"),722      llvm::cl::init(false)};723};724 725} // namespace726 727//===----------------------------------------------------------------------===//728// ReturnType Driver.729//===----------------------------------------------------------------------===//730 731namespace {732// Generate ops for each instance where the type can be successfully inferred.733template <typename OpTy>734static void invokeCreateWithInferredReturnType(Operation *op) {735  auto *context = op->getContext();736  auto fop = op->getParentOfType<func::FuncOp>();737  auto location = UnknownLoc::get(context);738  OpBuilder b(op);739  b.setInsertionPointAfter(op);740 741  // Use permutations of 2 args as operands.742  assert(fop.getNumArguments() >= 2);743  for (int i = 0, e = fop.getNumArguments(); i < e; ++i) {744    for (int j = 0; j < e; ++j) {745      std::array<Value, 2> values = {{fop.getArgument(i), fop.getArgument(j)}};746      SmallVector<Type, 2> inferredReturnTypes;747      if (succeeded(OpTy::inferReturnTypes(748              context, std::nullopt, values, op->getDiscardableAttrDictionary(),749              op->getPropertiesStorage(), op->getRegions(),750              inferredReturnTypes))) {751        OperationState state(location, OpTy::getOperationName());752        // TODO: Expand to regions.753        OpTy::build(b, state, values, op->getAttrs());754        (void)b.create(state);755      }756    }757  }758}759 760static void reifyReturnShape(Operation *op) {761  OpBuilder b(op);762 763  // Use permutations of 2 args as operands.764  auto shapedOp = cast<OpWithShapedTypeInferTypeInterfaceOp>(op);765  SmallVector<Value, 2> shapes;766  if (failed(shapedOp.reifyReturnTypeShapes(b, op->getOperands(), shapes)) ||767      !llvm::hasSingleElement(shapes))768    return;769  for (const auto &it : llvm::enumerate(shapes)) {770    op->emitRemark() << "value " << it.index() << ": "771                     << it.value().getDefiningOp();772  }773}774 775struct TestReturnTypeDriver776    : public PassWrapper<TestReturnTypeDriver, OperationPass<func::FuncOp>> {777  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestReturnTypeDriver)778 779  void getDependentDialects(DialectRegistry &registry) const override {780    registry.insert<tensor::TensorDialect>();781  }782  StringRef getArgument() const final { return "test-return-type"; }783  StringRef getDescription() const final { return "Run return type functions"; }784 785  void runOnOperation() override {786    if (getOperation().getName() == "testCreateFunctions") {787      std::vector<Operation *> ops;788      // Collect ops to avoid triggering on inserted ops.789      for (auto &op : getOperation().getBody().front())790        ops.push_back(&op);791      // Generate test patterns for each, but skip terminator.792      for (auto *op : llvm::ArrayRef(ops).drop_back()) {793        // Test create method of each of the Op classes below. The resultant794        // output would be in reverse order underneath `op` from which795        // the attributes and regions are used.796        invokeCreateWithInferredReturnType<OpWithInferTypeInterfaceOp>(op);797        invokeCreateWithInferredReturnType<OpWithInferTypeAdaptorInterfaceOp>(798            op);799        invokeCreateWithInferredReturnType<800            OpWithShapedTypeInferTypeInterfaceOp>(op);801      };802      return;803    }804    if (getOperation().getName() == "testReifyFunctions") {805      std::vector<Operation *> ops;806      // Collect ops to avoid triggering on inserted ops.807      for (auto &op : getOperation().getBody().front())808        if (isa<OpWithShapedTypeInferTypeInterfaceOp>(op))809          ops.push_back(&op);810      // Generate test patterns for each, but skip terminator.811      for (auto *op : ops)812        reifyReturnShape(op);813    }814  }815};816} // namespace817 818namespace {819struct TestDerivedAttributeDriver820    : public PassWrapper<TestDerivedAttributeDriver,821                         OperationPass<func::FuncOp>> {822  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestDerivedAttributeDriver)823 824  StringRef getArgument() const final { return "test-derived-attr"; }825  StringRef getDescription() const final {826    return "Run test derived attributes";827  }828  void runOnOperation() override;829};830} // namespace831 832void TestDerivedAttributeDriver::runOnOperation() {833  getOperation().walk([](DerivedAttributeOpInterface dOp) {834    auto dAttr = dOp.materializeDerivedAttributes();835    if (!dAttr)836      return;837    for (auto d : dAttr)838      dOp.emitRemark() << d.getName().getValue() << " = " << d.getValue();839  });840}841 842//===----------------------------------------------------------------------===//843// Legalization Driver.844//===----------------------------------------------------------------------===//845 846namespace {847//===----------------------------------------------------------------------===//848// Region-Block Rewrite Testing849//===----------------------------------------------------------------------===//850 851/// This pattern applies a signature conversion to a block inside a detached852/// region.853struct TestDetachedSignatureConversion : public ConversionPattern {854  TestDetachedSignatureConversion(MLIRContext *ctx)855      : ConversionPattern("test.detached_signature_conversion", /*benefit=*/1,856                          ctx) {}857 858  LogicalResult859  matchAndRewrite(Operation *op, ArrayRef<Value> operands,860                  ConversionPatternRewriter &rewriter) const final {861    if (op->getNumRegions() != 1)862      return failure();863    OperationState state(op->getLoc(), "test.legal_op", operands,864                         op->getResultTypes(), {}, BlockRange());865    Region *newRegion = state.addRegion();866    rewriter.inlineRegionBefore(op->getRegion(0), *newRegion,867                                newRegion->begin());868    TypeConverter::SignatureConversion result(newRegion->getNumArguments());869    for (unsigned i = 0, e = newRegion->getNumArguments(); i < e; ++i)870      result.addInputs(i, rewriter.getF64Type());871    rewriter.applySignatureConversion(&newRegion->front(), result);872    Operation *newOp = rewriter.create(state);873    rewriter.replaceOp(op, newOp->getResults());874    return success();875  }876};877 878/// This pattern is a simple pattern that inlines the first region of a given879/// operation into the parent region.880struct TestRegionRewriteBlockMovement : public ConversionPattern {881  TestRegionRewriteBlockMovement(MLIRContext *ctx)882      : ConversionPattern("test.region", 1, ctx) {}883 884  LogicalResult885  matchAndRewrite(Operation *op, ArrayRef<Value> operands,886                  ConversionPatternRewriter &rewriter) const final {887    // Inline this region into the parent region.888    auto &parentRegion = *op->getParentRegion();889    auto &opRegion = op->getRegion(0);890    if (op->getDiscardableAttr("legalizer.should_clone"))891      rewriter.cloneRegionBefore(opRegion, parentRegion, parentRegion.end());892    else893      rewriter.inlineRegionBefore(opRegion, parentRegion, parentRegion.end());894 895    if (op->getDiscardableAttr("legalizer.erase_old_blocks")) {896      while (!opRegion.empty())897        rewriter.eraseBlock(&opRegion.front());898    }899 900    // Drop this operation.901    rewriter.eraseOp(op);902    return success();903  }904};905/// This pattern is a simple pattern that generates a region containing an906/// illegal operation.907struct TestRegionRewriteUndo : public RewritePattern {908  TestRegionRewriteUndo(MLIRContext *ctx)909      : RewritePattern("test.region_builder", 1, ctx) {}910 911  LogicalResult matchAndRewrite(Operation *op,912                                PatternRewriter &rewriter) const final {913    // Create the region operation with an entry block containing arguments.914    OperationState newRegion(op->getLoc(), "test.region");915    newRegion.addRegion();916    auto *regionOp = rewriter.create(newRegion);917    auto *entryBlock = rewriter.createBlock(&regionOp->getRegion(0));918    entryBlock->addArgument(rewriter.getIntegerType(64),919                            rewriter.getUnknownLoc());920 921    // Add an explicitly illegal operation to ensure the conversion fails.922    ILLegalOpF::create(rewriter, op->getLoc(), rewriter.getIntegerType(32));923    TestValidOp::create(rewriter, op->getLoc(), ArrayRef<Value>());924 925    // Drop this operation.926    rewriter.eraseOp(op);927    return success();928  }929};930/// A simple pattern that creates a block at the end of the parent region of the931/// matched operation.932struct TestCreateBlock : public RewritePattern {933  TestCreateBlock(MLIRContext *ctx)934      : RewritePattern("test.create_block", /*benefit=*/1, ctx) {}935 936  LogicalResult matchAndRewrite(Operation *op,937                                PatternRewriter &rewriter) const final {938    Region &region = *op->getParentRegion();939    Type i32Type = rewriter.getIntegerType(32);940    Location loc = op->getLoc();941    rewriter.createBlock(&region, region.end(), {i32Type, i32Type}, {loc, loc});942    TerminatorOp::create(rewriter, loc);943    rewriter.eraseOp(op);944    return success();945  }946};947 948/// A simple pattern that creates a block containing an invalid operation in949/// order to trigger the block creation undo mechanism.950struct TestCreateIllegalBlock : public RewritePattern {951  TestCreateIllegalBlock(MLIRContext *ctx)952      : RewritePattern("test.create_illegal_block", /*benefit=*/1, ctx) {}953 954  LogicalResult matchAndRewrite(Operation *op,955                                PatternRewriter &rewriter) const final {956    Region &region = *op->getParentRegion();957    Type i32Type = rewriter.getIntegerType(32);958    Location loc = op->getLoc();959    rewriter.createBlock(&region, region.end(), {i32Type, i32Type}, {loc, loc});960    // Create an illegal op to ensure the conversion fails.961    ILLegalOpF::create(rewriter, loc, i32Type);962    TerminatorOp::create(rewriter, loc);963    rewriter.eraseOp(op);964    return success();965  }966};967 968/// A simple pattern that tests the "replaceAllUsesWith" API.969struct TestValueReplace : public ConversionPattern {970  TestValueReplace(MLIRContext *ctx, const TypeConverter &converter)971      : ConversionPattern(converter, "test.value_replace", /*benefit=*/1, ctx) {972  }973 974  LogicalResult975  matchAndRewrite(Operation *op, ArrayRef<ValueRange> operands,976                  ConversionPatternRewriter &rewriter) const final {977    // Replace the first operand with 2x the second operand.978    Value from = op->getOperand(0);979    Value repl = op->getOperand(1);980    if (op->hasAttr("conditional")) {981      rewriter.replaceUsesWithIf(from, {repl, repl}, [=](OpOperand &use) {982        return use.getOwner()->hasAttr("replace_uses");983      });984    } else {985      rewriter.replaceAllUsesWith(from, {repl, repl});986    }987    rewriter.modifyOpInPlace(op, [&] {988      // If the "trigger_rollback" attribute is set, keep the op illegal, so989      // that a rollback is triggered.990      if (!op->hasAttr("trigger_rollback"))991        op->setAttr("is_legal", rewriter.getUnitAttr());992    });993    return success();994  }995};996 997/// This pattern hoists ops out of a "test.hoist_me" and then fails conversion.998/// This is to test the rollback logic.999struct TestUndoMoveOpBefore : public ConversionPattern {1000  TestUndoMoveOpBefore(MLIRContext *ctx)1001      : ConversionPattern("test.hoist_me", /*benefit=*/1, ctx) {}1002 1003  LogicalResult1004  matchAndRewrite(Operation *op, ArrayRef<Value> operands,1005                  ConversionPatternRewriter &rewriter) const override {1006    rewriter.moveOpBefore(op, op->getParentOp());1007    // Replace with an illegal op to ensure the conversion fails.1008    rewriter.replaceOpWithNewOp<ILLegalOpF>(op, rewriter.getF32Type());1009    return success();1010  }1011};1012 1013/// A rewrite pattern that tests the undo mechanism when erasing a block.1014struct TestUndoBlockErase : public ConversionPattern {1015  TestUndoBlockErase(MLIRContext *ctx)1016      : ConversionPattern("test.undo_block_erase", /*benefit=*/1, ctx) {}1017 1018  LogicalResult1019  matchAndRewrite(Operation *op, ArrayRef<Value> operands,1020                  ConversionPatternRewriter &rewriter) const final {1021    Block *secondBlock = &*std::next(op->getRegion(0).begin());1022    rewriter.setInsertionPointToStart(secondBlock);1023    ILLegalOpF::create(rewriter, op->getLoc(), rewriter.getF32Type());1024    rewriter.eraseBlock(secondBlock);1025    rewriter.modifyOpInPlace(op, [] {});1026    return success();1027  }1028};1029 1030/// A pattern that modifies a property in-place, but keeps the op illegal.1031struct TestUndoPropertiesModification : public ConversionPattern {1032  TestUndoPropertiesModification(MLIRContext *ctx)1033      : ConversionPattern("test.with_properties", /*benefit=*/1, ctx) {}1034  LogicalResult1035  matchAndRewrite(Operation *op, ArrayRef<Value> operands,1036                  ConversionPatternRewriter &rewriter) const final {1037    if (!op->hasAttr("modify_inplace"))1038      return failure();1039    rewriter.modifyOpInPlace(1040        op, [&]() { cast<TestOpWithProperties>(op).getProperties().setA(42); });1041    return success();1042  }1043};1044 1045//===----------------------------------------------------------------------===//1046// Type-Conversion Rewrite Testing1047//===----------------------------------------------------------------------===//1048 1049/// This patterns erases a region operation that has had a type conversion.1050struct TestDropOpSignatureConversion : public ConversionPattern {1051  TestDropOpSignatureConversion(MLIRContext *ctx,1052                                const TypeConverter &converter)1053      : ConversionPattern(converter, "test.drop_region_op", 1, ctx) {}1054  LogicalResult1055  matchAndRewrite(Operation *op, ArrayRef<Value> operands,1056                  ConversionPatternRewriter &rewriter) const override {1057    Region &region = op->getRegion(0);1058    Block *entry = &region.front();1059 1060    // Convert the original entry arguments.1061    const TypeConverter &converter = *getTypeConverter();1062    TypeConverter::SignatureConversion result(entry->getNumArguments());1063    if (failed(converter.convertSignatureArgs(entry->getArgumentTypes(),1064                                              result)) ||1065        failed(rewriter.convertRegionTypes(&region, converter, &result)))1066      return failure();1067 1068    // Convert the region signature and just drop the operation.1069    rewriter.eraseOp(op);1070    return success();1071  }1072};1073/// This pattern simply updates the operands of the given operation.1074struct TestPassthroughInvalidOp : public ConversionPattern {1075  TestPassthroughInvalidOp(MLIRContext *ctx, const TypeConverter &converter)1076      : ConversionPattern(converter, "test.invalid", 1, ctx) {}1077  LogicalResult1078  matchAndRewrite(Operation *op, ArrayRef<ValueRange> operands,1079                  ConversionPatternRewriter &rewriter) const final {1080    SmallVector<Value> flattened;1081    for (auto it : llvm::enumerate(operands)) {1082      ValueRange range = it.value();1083      if (range.size() == 1) {1084        flattened.push_back(range.front());1085        continue;1086      }1087 1088      // This is a 1:N replacement. Insert a test.cast op. (That's what the1089      // argument materialization used to do.)1090      flattened.push_back(1091          TestCastOp::create(rewriter, op->getLoc(),1092                             op->getOperand(it.index()).getType(), range)1093              .getResult());1094    }1095    rewriter.replaceOpWithNewOp<TestValidOp>(op, TypeRange(), flattened,1096                                             ArrayRef<NamedAttribute>());1097    return success();1098  }1099};1100/// Replace with valid op, but simply drop the operands. This is used in a1101/// regression where we used to generate circular unrealized_conversion_cast1102/// ops.1103struct TestDropAndReplaceInvalidOp : public ConversionPattern {1104  TestDropAndReplaceInvalidOp(MLIRContext *ctx, const TypeConverter &converter)1105      : ConversionPattern(converter,1106                          "test.drop_operands_and_replace_with_valid", 1, ctx) {1107  }1108  LogicalResult1109  matchAndRewrite(Operation *op, ArrayRef<Value> operands,1110                  ConversionPatternRewriter &rewriter) const final {1111    rewriter.replaceOpWithNewOp<TestValidOp>(op, TypeRange(), ValueRange(),1112                                             ArrayRef<NamedAttribute>());1113    return success();1114  }1115};1116/// This pattern handles the case of a split return value.1117struct TestSplitReturnType : public ConversionPattern {1118  TestSplitReturnType(MLIRContext *ctx)1119      : ConversionPattern("test.return", 1, ctx) {}1120  LogicalResult1121  matchAndRewrite(Operation *op, ArrayRef<ValueRange> operands,1122                  ConversionPatternRewriter &rewriter) const final {1123    // Check for a return of F32.1124    if (op->getNumOperands() != 1 || !op->getOperand(0).getType().isF32())1125      return failure();1126    rewriter.replaceOpWithNewOp<TestReturnOp>(op, operands[0]);1127    return success();1128  }1129};1130 1131//===----------------------------------------------------------------------===//1132// Multi-Level Type-Conversion Rewrite Testing1133struct TestChangeProducerTypeI32ToF32 : public ConversionPattern {1134  TestChangeProducerTypeI32ToF32(MLIRContext *ctx)1135      : ConversionPattern("test.type_producer", 1, ctx) {}1136  LogicalResult1137  matchAndRewrite(Operation *op, ArrayRef<Value> operands,1138                  ConversionPatternRewriter &rewriter) const final {1139    // If the type is I32, change the type to F32.1140    if (!Type(*op->result_type_begin()).isSignlessInteger(32))1141      return failure();1142    rewriter.replaceOpWithNewOp<TestTypeProducerOp>(op, rewriter.getF32Type());1143    return success();1144  }1145};1146struct TestChangeProducerTypeF32ToF64 : public ConversionPattern {1147  TestChangeProducerTypeF32ToF64(MLIRContext *ctx)1148      : ConversionPattern("test.type_producer", 1, ctx) {}1149  LogicalResult1150  matchAndRewrite(Operation *op, ArrayRef<Value> operands,1151                  ConversionPatternRewriter &rewriter) const final {1152    // If the type is F32, change the type to F64.1153    if (!Type(*op->result_type_begin()).isF32())1154      return rewriter.notifyMatchFailure(op, "expected single f32 operand");1155    rewriter.replaceOpWithNewOp<TestTypeProducerOp>(op, rewriter.getF64Type());1156    return success();1157  }1158};1159struct TestChangeProducerTypeF32ToInvalid : public ConversionPattern {1160  TestChangeProducerTypeF32ToInvalid(MLIRContext *ctx)1161      : ConversionPattern("test.type_producer", 10, ctx) {}1162  LogicalResult1163  matchAndRewrite(Operation *op, ArrayRef<Value> operands,1164                  ConversionPatternRewriter &rewriter) const final {1165    // Always convert to B16, even though it is not a legal type. This tests1166    // that values are unmapped correctly.1167    rewriter.replaceOpWithNewOp<TestTypeProducerOp>(op, rewriter.getBF16Type());1168    return success();1169  }1170};1171struct TestUpdateConsumerType : public ConversionPattern {1172  TestUpdateConsumerType(MLIRContext *ctx)1173      : ConversionPattern("test.type_consumer", 1, ctx) {}1174  LogicalResult1175  matchAndRewrite(Operation *op, ArrayRef<Value> operands,1176                  ConversionPatternRewriter &rewriter) const final {1177    // Verify that the incoming operand has been successfully remapped to F64.1178    if (!operands[0].getType().isF64())1179      return failure();1180    rewriter.replaceOpWithNewOp<TestTypeConsumerOp>(op, operands[0]);1181    return success();1182  }1183};1184 1185//===----------------------------------------------------------------------===//1186// Non-Root Replacement Rewrite Testing1187/// This pattern generates an invalid operation, but replaces it before the1188/// pattern is finished. This checks that we don't need to legalize the1189/// temporary op.1190struct TestNonRootReplacement : public RewritePattern {1191  TestNonRootReplacement(MLIRContext *ctx)1192      : RewritePattern("test.replace_non_root", 1, ctx) {}1193 1194  LogicalResult matchAndRewrite(Operation *op,1195                                PatternRewriter &rewriter) const final {1196    auto resultType = *op->result_type_begin();1197    auto illegalOp = ILLegalOpF::create(rewriter, op->getLoc(), resultType);1198    auto legalOp = LegalOpB::create(rewriter, op->getLoc(), resultType);1199 1200    rewriter.replaceOp(op, illegalOp);1201    rewriter.replaceOp(illegalOp, legalOp);1202    return success();1203  }1204};1205 1206//===----------------------------------------------------------------------===//1207// Recursive Rewrite Testing1208/// This pattern is applied to the same operation multiple times, but has a1209/// bounded recursion.1210struct TestBoundedRecursiveRewrite1211    : public OpRewritePattern<TestRecursiveRewriteOp> {1212  using OpRewritePattern<TestRecursiveRewriteOp>::OpRewritePattern;1213 1214  void initialize() {1215    // The conversion target handles bounding the recursion of this pattern.1216    setHasBoundedRewriteRecursion();1217  }1218 1219  LogicalResult matchAndRewrite(TestRecursiveRewriteOp op,1220                                PatternRewriter &rewriter) const final {1221    // Decrement the depth of the op in-place.1222    rewriter.modifyOpInPlace(op, [&] {1223      op->setAttr("depth", rewriter.getI64IntegerAttr(op.getDepth() - 1));1224    });1225    return success();1226  }1227};1228 1229struct TestNestedOpCreationUndoRewrite1230    : public OpRewritePattern<IllegalOpWithRegionAnchor> {1231  using OpRewritePattern<IllegalOpWithRegionAnchor>::OpRewritePattern;1232 1233  LogicalResult matchAndRewrite(IllegalOpWithRegionAnchor op,1234                                PatternRewriter &rewriter) const final {1235    // rewriter.replaceOpWithNewOp<IllegalOpWithRegion>(op);1236    rewriter.replaceOpWithNewOp<IllegalOpWithRegion>(op);1237    return success();1238  };1239};1240 1241// This pattern matches `test.blackhole` and delete this op and its producer.1242struct TestReplaceEraseOp : public OpRewritePattern<BlackHoleOp> {1243  using OpRewritePattern<BlackHoleOp>::OpRewritePattern;1244 1245  LogicalResult matchAndRewrite(BlackHoleOp op,1246                                PatternRewriter &rewriter) const final {1247    Operation *producer = op.getOperand().getDefiningOp();1248    // Always erase the user before the producer, the framework should handle1249    // this correctly.1250    rewriter.eraseOp(op);1251    rewriter.eraseOp(producer);1252    return success();1253  };1254};1255 1256// This pattern replaces explicitly illegal op with explicitly legal op,1257// but in addition creates unregistered operation.1258struct TestCreateUnregisteredOp : public OpRewritePattern<ILLegalOpG> {1259  using OpRewritePattern<ILLegalOpG>::OpRewritePattern;1260 1261  LogicalResult matchAndRewrite(ILLegalOpG op,1262                                PatternRewriter &rewriter) const final {1263    IntegerAttr attr = rewriter.getI32IntegerAttr(0);1264    Value val = arith::ConstantOp::create(rewriter, op->getLoc(), attr);1265    rewriter.replaceOpWithNewOp<LegalOpC>(op, val);1266    return success();1267  };1268};1269 1270class TestEraseOp : public ConversionPattern {1271public:1272  TestEraseOp(MLIRContext *ctx) : ConversionPattern("test.erase_op", 1, ctx) {}1273  LogicalResult1274  matchAndRewrite(Operation *op, ArrayRef<Value> operands,1275                  ConversionPatternRewriter &rewriter) const final {1276    // Erase op without replacements.1277    rewriter.eraseOp(op);1278    return success();1279  }1280};1281 1282/// Pattern that replaces test.replace_with_valid_producer with1283/// test.valid_producer and the specified type.1284class TestReplaceWithValidProducer : public ConversionPattern {1285public:1286  TestReplaceWithValidProducer(MLIRContext *ctx)1287      : ConversionPattern("test.replace_with_valid_producer", 1, ctx) {}1288  LogicalResult1289  matchAndRewrite(Operation *op, ArrayRef<Value> operands,1290                  ConversionPatternRewriter &rewriter) const final {1291    auto attr = op->getAttrOfType<TypeAttr>("type");1292    if (!attr)1293      return failure();1294    rewriter.replaceOpWithNewOp<TestValidProducerOp>(op, attr.getValue());1295    return success();1296  }1297};1298 1299/// Pattern that replaces test.replace_with_valid_consumer with1300/// test.valid_consumer. Can be used with and without a type converter.1301class TestReplaceWithValidConsumer : public ConversionPattern {1302public:1303  TestReplaceWithValidConsumer(MLIRContext *ctx, const TypeConverter &converter)1304      : ConversionPattern(converter, "test.replace_with_valid_consumer", 1,1305                          ctx) {}1306  TestReplaceWithValidConsumer(MLIRContext *ctx)1307      : ConversionPattern("test.replace_with_valid_consumer", 1, ctx) {}1308 1309  LogicalResult1310  matchAndRewrite(Operation *op, ArrayRef<Value> operands,1311                  ConversionPatternRewriter &rewriter) const final {1312    // with_converter present: pattern must have been initialized with a type1313    // converter.1314    // with_converter absent: pattern must have been initialized without a type1315    // converter.1316    if (op->hasAttr("with_converter") != static_cast<bool>(getTypeConverter()))1317      return failure();1318    rewriter.replaceOpWithNewOp<TestValidConsumerOp>(op, operands[0]);1319    return success();1320  }1321};1322 1323/// This pattern matches a test.convert_block_args op. It either:1324/// a) Duplicates all block arguments,1325/// b) or: drops all block arguments and replaces each with 2x the first1326///    operand.1327class TestConvertBlockArgs : public OpConversionPattern<ConvertBlockArgsOp> {1328  static_assert(std::is_same_v<Base, OpConversionPattern<ConvertBlockArgsOp>>);1329  using Base::Base;1330 1331  LogicalResult1332  matchAndRewrite(ConvertBlockArgsOp op, OpAdaptor adaptor,1333                  ConversionPatternRewriter &rewriter) const override {1334    if (op.getIsLegal())1335      return failure();1336    Block *body = &op.getBody().front();1337    TypeConverter::SignatureConversion result(body->getNumArguments());1338    for (auto it : llvm::enumerate(body->getArgumentTypes())) {1339      if (op.getReplaceWithOperand()) {1340        result.remapInput(it.index(), {adaptor.getVal(), adaptor.getVal()});1341      } else if (op.getDuplicate()) {1342        result.addInputs(it.index(), {it.value(), it.value()});1343      } else {1344        // No action specified. Pattern does not apply.1345        return failure();1346      }1347    }1348    rewriter.startOpModification(op);1349    rewriter.applySignatureConversion(body, result, getTypeConverter());1350    op.setIsLegal(true);1351    rewriter.finalizeOpModification(op);1352    return success();1353  }1354};1355 1356/// This pattern replaces test.repetitive_1_to_n_consumer ops with a test.valid1357/// op. The pattern supports 1:N replacements and forwards the replacement1358/// values of the single operand as test.valid operands.1359class TestRepetitive1ToNConsumer : public ConversionPattern {1360public:1361  TestRepetitive1ToNConsumer(MLIRContext *ctx)1362      : ConversionPattern("test.repetitive_1_to_n_consumer", 1, ctx) {}1363  LogicalResult1364  matchAndRewrite(Operation *op, ArrayRef<ValueRange> operands,1365                  ConversionPatternRewriter &rewriter) const final {1366    // A single operand is expected.1367    if (op->getNumOperands() != 1)1368      return failure();1369    rewriter.replaceOpWithNewOp<TestValidOp>(op, operands.front());1370    return success();1371  }1372};1373 1374/// A pattern that tests two back-to-back 1 -> 2 op replacements.1375class TestMultiple1ToNReplacement : public ConversionPattern {1376public:1377  TestMultiple1ToNReplacement(MLIRContext *ctx, const TypeConverter &converter)1378      : ConversionPattern(converter, "test.multiple_1_to_n_replacement", 1,1379                          ctx) {}1380  LogicalResult1381  matchAndRewrite(Operation *op, ArrayRef<ValueRange> operands,1382                  ConversionPatternRewriter &rewriter) const final {1383    // Helper function that replaces the given op with a new op of the given1384    // name and doubles each result (1 -> 2 replacement of each result).1385    auto replaceWithDoubleResults = [&](Operation *op, StringRef name) {1386      rewriter.setInsertionPointAfter(op);1387      SmallVector<Type> types;1388      for (Type t : op->getResultTypes()) {1389        types.push_back(t);1390        types.push_back(t);1391      }1392      OperationState state(op->getLoc(), name,1393                           /*operands=*/{}, types, op->getAttrs());1394      auto *newOp = rewriter.create(state);1395      SmallVector<ValueRange> repls;1396      for (size_t i = 0, e = op->getNumResults(); i < e; ++i)1397        repls.push_back(newOp->getResults().slice(2 * i, 2));1398      rewriter.replaceOpWithMultiple(op, repls);1399      return newOp;1400    };1401 1402    // Replace test.multiple_1_to_n_replacement with test.step_1.1403    Operation *repl1 = replaceWithDoubleResults(op, "test.step_1");1404    // Now replace test.step_1 with test.legal_op.1405    replaceWithDoubleResults(repl1, "test.legal_op");1406    return success();1407  }1408};1409 1410/// Pattern that erases 'test.type_consumers' iff the input operand is the1411/// result of a 1:1 type conversion.1412/// Used to test correct skipping of 1:1 patterns in the 1:N case.1413class TestTypeConsumerOpPattern1414    : public OpConversionPattern<TestTypeConsumerOp> {1415public:1416  TestTypeConsumerOpPattern(MLIRContext *ctx, const TypeConverter &converter)1417      : OpConversionPattern<TestTypeConsumerOp>(converter, ctx) {}1418 1419  LogicalResult1420  matchAndRewrite(TestTypeConsumerOp op, OpAdaptor operands,1421                  ConversionPatternRewriter &rewriter) const final {1422    rewriter.eraseOp(op);1423    return success();1424  }1425};1426 1427class TestPostOrderLegalization : public ConversionPattern {1428public:1429  TestPostOrderLegalization(MLIRContext *ctx, const TypeConverter &converter)1430      : ConversionPattern(converter, "test.post_order_legalization", 1, ctx) {}1431  LogicalResult1432  matchAndRewrite(Operation *op, ArrayRef<ValueRange> operands,1433                  ConversionPatternRewriter &rewriter) const final {1434    for (Region &r : op->getRegions())1435      if (failed(rewriter.legalize(&r)))1436        return failure();1437    rewriter.modifyOpInPlace(1438        op, [&]() { op->setAttr("is_legal", rewriter.getUnitAttr()); });1439    return success();1440  }1441};1442 1443/// Test unambiguous overload resolution of replaceOpWithMultiple. This1444/// function is just to trigger compiler errors. It is never executed.1445[[maybe_unused]] void testReplaceOpWithMultipleOverloads(1446    ConversionPatternRewriter &rewriter, Operation *op, ArrayRef<ValueRange> r1,1447    SmallVector<ValueRange> r2, ArrayRef<SmallVector<Value>> r3,1448    SmallVector<SmallVector<Value>> r4, ArrayRef<ArrayRef<Value>> r5,1449    SmallVector<ArrayRef<Value>> r6, SmallVector<SmallVector<Value>> &&r7,1450    Value v, ValueRange vr, ArrayRef<Value> ar) {1451  rewriter.replaceOpWithMultiple(op, r1);1452  rewriter.replaceOpWithMultiple(op, r2);1453  rewriter.replaceOpWithMultiple(op, r3);1454  rewriter.replaceOpWithMultiple(op, r4);1455  rewriter.replaceOpWithMultiple(op, r5);1456  rewriter.replaceOpWithMultiple(op, r6);1457  rewriter.replaceOpWithMultiple(op, std::move(r7));1458  rewriter.replaceOpWithMultiple(op, {vr});1459  rewriter.replaceOpWithMultiple(op, {ar});1460  rewriter.replaceOpWithMultiple(op, {{v}});1461  rewriter.replaceOpWithMultiple(op, {{v, v}});1462  rewriter.replaceOpWithMultiple(op, {{v, v}, vr});1463  rewriter.replaceOpWithMultiple(op, {{v, v}, ar});1464  rewriter.replaceOpWithMultiple(op, {ar, {v, v}, vr});1465}1466} // namespace1467 1468namespace {1469struct TestTypeConverter : public TypeConverter {1470  static_assert(std::is_same_v<Base, TypeConverter>);1471  using Base::Base;1472 1473  TestTypeConverter() {1474    addConversion(convertType);1475    addSourceMaterialization(materializeCast);1476    addTargetMaterialization(materializeCast);1477  }1478 1479  static LogicalResult convertType(Type t, SmallVectorImpl<Type> &results) {1480    // Drop I16 types.1481    if (t.isSignlessInteger(16))1482      return success();1483 1484    // Convert I64 to F64.1485    if (t.isSignlessInteger(64)) {1486      results.push_back(Float64Type::get(t.getContext()));1487      return success();1488    }1489 1490    // Convert I42 to I43.1491    if (t.isInteger(42)) {1492      results.push_back(IntegerType::get(t.getContext(), 43));1493      return success();1494    }1495 1496    // Split F32 into F16,F16.1497    if (t.isF32()) {1498      results.assign(2, Float16Type::get(t.getContext()));1499      return success();1500    }1501 1502    // Drop I24 types.1503    if (t.isInteger(24)) {1504      return success();1505    }1506 1507    // Otherwise, convert the type directly.1508    results.push_back(t);1509    return success();1510  }1511 1512  /// Hook for materializing a conversion. This is necessary because we generate1513  /// 1->N type mappings.1514  static Value materializeCast(OpBuilder &builder, Type resultType,1515                               ValueRange inputs, Location loc) {1516    return TestCastOp::create(builder, loc, resultType, inputs).getResult();1517  }1518};1519 1520struct TestLegalizePatternDriver1521    : public PassWrapper<TestLegalizePatternDriver, OperationPass<>> {1522  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestLegalizePatternDriver)1523 1524  TestLegalizePatternDriver() = default;1525  TestLegalizePatternDriver(const TestLegalizePatternDriver &other)1526      : PassWrapper(other) {}1527 1528  StringRef getArgument() const final { return "test-legalize-patterns"; }1529  StringRef getDescription() const final {1530    return "Run test dialect legalization patterns";1531  }1532  /// The mode of conversion to use with the driver.1533  enum class ConversionMode { Analysis, Full, Partial };1534 1535  void getDependentDialects(DialectRegistry &registry) const override {1536    registry.insert<func::FuncDialect, test::TestDialect>();1537  }1538 1539  void runOnOperation() override {1540    TestTypeConverter converter;1541    mlir::RewritePatternSet patterns(&getContext());1542    populateWithGenerated(patterns);1543    patterns.add<1544        TestRegionRewriteBlockMovement, TestDetachedSignatureConversion,1545        TestRegionRewriteUndo, TestCreateBlock, TestCreateIllegalBlock,1546        TestUndoBlockErase, TestSplitReturnType, TestChangeProducerTypeI32ToF32,1547        TestChangeProducerTypeF32ToF64, TestChangeProducerTypeF32ToInvalid,1548        TestUpdateConsumerType, TestNonRootReplacement,1549        TestBoundedRecursiveRewrite, TestNestedOpCreationUndoRewrite,1550        TestReplaceEraseOp, TestCreateUnregisteredOp, TestUndoMoveOpBefore,1551        TestUndoPropertiesModification, TestEraseOp,1552        TestReplaceWithValidProducer, TestReplaceWithValidConsumer,1553        TestRepetitive1ToNConsumer>(&getContext());1554    patterns.add<TestDropOpSignatureConversion, TestDropAndReplaceInvalidOp,1555                 TestPassthroughInvalidOp, TestMultiple1ToNReplacement,1556                 TestValueReplace, TestReplaceWithValidConsumer,1557                 TestTypeConsumerOpPattern, TestPostOrderLegalization>(1558        &getContext(), converter);1559    patterns.add<TestConvertBlockArgs>(converter, &getContext());1560    mlir::populateAnyFunctionOpInterfaceTypeConversionPattern(patterns,1561                                                              converter);1562    mlir::populateCallOpTypeConversionPattern(patterns, converter);1563 1564    // Define the conversion target used for the test.1565    ConversionTarget target(getContext());1566    target.addLegalOp<ModuleOp>();1567    target.addLegalOp<LegalOpA, LegalOpB, LegalOpC, TestCastOp, TestValidOp,1568                      TerminatorOp, TestOpConstant, OneRegionOp,1569                      TestValidProducerOp, TestValidConsumerOp>();1570    target.addLegalOp(OperationName("test.legal_op", &getContext()));1571    target1572        .addIllegalOp<ILLegalOpF, TestRegionBuilderOp, TestOpWithRegionFold>();1573    target.addDynamicallyLegalOp<TestReturnOp>([](TestReturnOp op) {1574      // Don't allow F32 operands.1575      return llvm::none_of(op.getOperandTypes(),1576                           [](Type type) { return type.isF32(); });1577    });1578    target.addDynamicallyLegalOp<func::FuncOp>([&](func::FuncOp op) {1579      return converter.isSignatureLegal(op.getFunctionType());1580    });1581    target.addDynamicallyLegalOp<func::CallOp>(1582        [&](func::CallOp op) { return converter.isLegal(op); });1583    target.addDynamicallyLegalOp(1584        OperationName("test.value_replace", &getContext()),1585        [](Operation *op) { return op->hasAttr("is_legal"); });1586    target.addDynamicallyLegalOp(1587        OperationName("test.post_order_legalization", &getContext()),1588        [](Operation *op) { return op->hasAttr("is_legal"); });1589 1590    // TestCreateUnregisteredOp creates `arith.constant` operation,1591    // which was not added to target intentionally to test1592    // correct error code from conversion driver.1593    target.addDynamicallyLegalOp<ILLegalOpG>([](ILLegalOpG) { return false; });1594 1595    // Expect the type_producer/type_consumer operations to only operate on f64.1596    target.addDynamicallyLegalOp<TestTypeProducerOp>(1597        [](TestTypeProducerOp op) { return op.getType().isF64(); });1598    target.addDynamicallyLegalOp<TestTypeConsumerOp>([](TestTypeConsumerOp op) {1599      return op.getOperand().getType().isF64();1600    });1601 1602    // Check support for marking certain operations as recursively legal.1603    target.markOpRecursivelyLegal<func::FuncOp, ModuleOp>([](Operation *op) {1604      return static_cast<bool>(1605          op->getAttrOfType<UnitAttr>("test.recursively_legal"));1606    });1607 1608    // Mark the bound recursion operation as dynamically legal.1609    target.addDynamicallyLegalOp<TestRecursiveRewriteOp>(1610        [](TestRecursiveRewriteOp op) { return op.getDepth() == 0; });1611 1612    // Create a dynamically legal rule that can only be legalized by folding it.1613    target.addDynamicallyLegalOp<TestOpInPlaceSelfFold>(1614        [](TestOpInPlaceSelfFold op) { return op.getFolded(); });1615 1616    target.addDynamicallyLegalOp<ConvertBlockArgsOp>(1617        [](ConvertBlockArgsOp op) { return op.getIsLegal(); });1618 1619    // Set up configuration.1620    ConversionConfig config;1621    config.allowPatternRollback = allowPatternRollback;1622    config.foldingMode = foldingMode;1623    config.buildMaterializations = buildMaterializations;1624    config.attachDebugMaterializationKind = attachDebugMaterializationKind;1625    DumpNotifications dumpNotifications;1626    config.listener = &dumpNotifications;1627 1628    // Handle a partial conversion.1629    if (mode == ConversionMode::Partial) {1630      DenseSet<Operation *> unlegalizedOps;1631      config.unlegalizedOps = &unlegalizedOps;1632      if (failed(applyPartialConversion(getOperation(), target,1633                                        std::move(patterns), config))) {1634        getOperation()->emitRemark() << "applyPartialConversion failed";1635      }1636      // Emit remarks for each legalizable operation.1637      for (auto *op : unlegalizedOps)1638        op->emitRemark() << "op '" << op->getName() << "' is not legalizable";1639      return;1640    }1641 1642    // Handle a full conversion.1643    if (mode == ConversionMode::Full) {1644      // Check support for marking unknown operations as dynamically legal.1645      target.markUnknownOpDynamicallyLegal([](Operation *op) {1646        return (bool)op->getAttrOfType<UnitAttr>("test.dynamically_legal");1647      });1648 1649      if (failed(applyFullConversion(getOperation(), target,1650                                     std::move(patterns), config))) {1651        getOperation()->emitRemark() << "applyFullConversion failed";1652      }1653      return;1654    }1655 1656    // Otherwise, handle an analysis conversion.1657    assert(mode == ConversionMode::Analysis);1658 1659    // Analyze the convertible operations.1660    DenseSet<Operation *> legalizedOps;1661    config.legalizableOps = &legalizedOps;1662    if (failed(applyAnalysisConversion(getOperation(), target,1663                                       std::move(patterns), config)))1664      return signalPassFailure();1665 1666    // Emit remarks for each legalizable operation.1667    for (auto *op : legalizedOps)1668      op->emitRemark() << "op '" << op->getName() << "' is legalizable";1669  }1670 1671  Option<ConversionMode> mode{1672      *this, "test-legalize-mode",1673      llvm::cl::desc("The legalization mode to use with the test driver"),1674      llvm::cl::init(ConversionMode::Partial),1675      llvm::cl::values(1676          clEnumValN(ConversionMode::Analysis, "analysis",1677                     "Perform an analysis conversion"),1678          clEnumValN(ConversionMode::Full, "full", "Perform a full conversion"),1679          clEnumValN(ConversionMode::Partial, "partial",1680                     "Perform a partial conversion"))};1681 1682  Option<DialectConversionFoldingMode> foldingMode{1683      *this, "test-legalize-folding-mode",1684      llvm::cl::desc("The folding mode to use with the test driver"),1685      llvm::cl::init(DialectConversionFoldingMode::BeforePatterns),1686      llvm::cl::values(clEnumValN(DialectConversionFoldingMode::Never, "never",1687                                  "Never attempt to fold"),1688                       clEnumValN(DialectConversionFoldingMode::BeforePatterns,1689                                  "before-patterns",1690                                  "Only attempt to fold not legal operations "1691                                  "before applying patterns"),1692                       clEnumValN(DialectConversionFoldingMode::AfterPatterns,1693                                  "after-patterns",1694                                  "Only attempt to fold not legal operations "1695                                  "after applying patterns"))};1696  Option<bool> allowPatternRollback{*this, "allow-pattern-rollback",1697                                    llvm::cl::desc("Allow pattern rollback"),1698                                    llvm::cl::init(true)};1699  Option<bool> attachDebugMaterializationKind{1700      *this, "attach-debug-materialization-kind",1701      llvm::cl::desc(1702          "Attach materialization kind to unrealized_conversion_cast ops"),1703      llvm::cl::init(false)};1704  Option<bool> buildMaterializations{1705      *this, "build-materializations",1706      llvm::cl::desc(1707          "If set to 'false', leave unrealized_conversion_cast ops in place"),1708      llvm::cl::init(true)};1709};1710} // namespace1711 1712//===----------------------------------------------------------------------===//1713// ConversionPatternRewriter::getRemappedValue testing. This method is used1714// to get the remapped value of an original value that was replaced using1715// ConversionPatternRewriter.1716namespace {1717struct TestRemapValueTypeConverter : public TypeConverter {1718  using TypeConverter::TypeConverter;1719 1720  TestRemapValueTypeConverter() {1721    addConversion(1722        [](Float32Type type) { return Float64Type::get(type.getContext()); });1723    addConversion([](Type type) { return type; });1724  }1725};1726 1727/// Converter that replaces a one-result one-operand OneVResOneVOperandOp1 with1728/// a one-operand two-result OneVResOneVOperandOp1 by replicating its original1729/// operand twice.1730///1731/// Example:1732///   %1 = test.one_variadic_out_one_variadic_in1"(%0)1733/// is replaced with:1734///   %1 = test.one_variadic_out_one_variadic_in1"(%0, %0)1735struct OneVResOneVOperandOp1Converter1736    : public OpConversionPattern<OneVResOneVOperandOp1> {1737  using OpConversionPattern<OneVResOneVOperandOp1>::OpConversionPattern;1738 1739  LogicalResult1740  matchAndRewrite(OneVResOneVOperandOp1 op, OpAdaptor adaptor,1741                  ConversionPatternRewriter &rewriter) const override {1742    auto origOps = op.getOperands();1743    assert(std::distance(origOps.begin(), origOps.end()) == 1 &&1744           "One operand expected");1745    Value origOp = *origOps.begin();1746    SmallVector<Value, 2> remappedOperands;1747    // Replicate the remapped original operand twice. Note that we don't used1748    // the remapped 'operand' since the goal is testing 'getRemappedValue'.1749    remappedOperands.push_back(rewriter.getRemappedValue(origOp));1750    remappedOperands.push_back(rewriter.getRemappedValue(origOp));1751 1752    rewriter.replaceOpWithNewOp<OneVResOneVOperandOp1>(op, op.getResultTypes(),1753                                                       remappedOperands);1754    return success();1755  }1756};1757 1758/// A rewriter pattern that tests that blocks can be merged.1759struct TestRemapValueInRegion1760    : public OpConversionPattern<TestRemappedValueRegionOp> {1761  using OpConversionPattern<TestRemappedValueRegionOp>::OpConversionPattern;1762 1763  LogicalResult1764  matchAndRewrite(TestRemappedValueRegionOp op, OpAdaptor adaptor,1765                  ConversionPatternRewriter &rewriter) const final {1766    Block &block = op.getBody().front();1767    Operation *terminator = block.getTerminator();1768 1769    // Merge the block into the parent region.1770    Block *parentBlock = op->getBlock();1771    Block *finalBlock = rewriter.splitBlock(parentBlock, op->getIterator());1772    rewriter.mergeBlocks(&block, parentBlock, ValueRange());1773    rewriter.mergeBlocks(finalBlock, parentBlock, ValueRange());1774 1775    // Replace the results of this operation with the remapped terminator1776    // values.1777    SmallVector<Value> terminatorOperands;1778    if (failed(rewriter.getRemappedValues(terminator->getOperands(),1779                                          terminatorOperands)))1780      return failure();1781 1782    rewriter.eraseOp(terminator);1783    rewriter.replaceOp(op, terminatorOperands);1784    return success();1785  }1786};1787 1788struct TestRemappedValue1789    : public mlir::PassWrapper<TestRemappedValue, OperationPass<>> {1790  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestRemappedValue)1791 1792  StringRef getArgument() const final { return "test-remapped-value"; }1793  StringRef getDescription() const final {1794    return "Test public remapped value mechanism in ConversionPatternRewriter";1795  }1796  void runOnOperation() override {1797    TestRemapValueTypeConverter typeConverter;1798 1799    mlir::RewritePatternSet patterns(&getContext());1800    patterns.add<OneVResOneVOperandOp1Converter>(&getContext());1801    patterns.add<TestChangeProducerTypeF32ToF64, TestUpdateConsumerType>(1802        &getContext());1803    patterns.add<TestRemapValueInRegion>(typeConverter, &getContext());1804 1805    mlir::ConversionTarget target(getContext());1806    target.addLegalOp<ModuleOp, func::FuncOp, TestReturnOp>();1807 1808    // Expect the type_producer/type_consumer operations to only operate on f64.1809    target.addDynamicallyLegalOp<TestTypeProducerOp>(1810        [](TestTypeProducerOp op) { return op.getType().isF64(); });1811    target.addDynamicallyLegalOp<TestTypeConsumerOp>([](TestTypeConsumerOp op) {1812      return op.getOperand().getType().isF64();1813    });1814 1815    // We make OneVResOneVOperandOp1 legal only when it has more that one1816    // operand. This will trigger the conversion that will replace one-operand1817    // OneVResOneVOperandOp1 with two-operand OneVResOneVOperandOp1.1818    target.addDynamicallyLegalOp<OneVResOneVOperandOp1>(1819        [](Operation *op) { return op->getNumOperands() > 1; });1820 1821    if (failed(mlir::applyFullConversion(getOperation(), target,1822                                         std::move(patterns)))) {1823      signalPassFailure();1824    }1825  }1826};1827} // namespace1828 1829//===----------------------------------------------------------------------===//1830// Test patterns without a specific root operation kind1831//===----------------------------------------------------------------------===//1832 1833namespace {1834/// This pattern matches and removes any operation in the test dialect.1835struct RemoveTestDialectOps : public RewritePattern {1836  RemoveTestDialectOps(MLIRContext *context)1837      : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}1838 1839  LogicalResult matchAndRewrite(Operation *op,1840                                PatternRewriter &rewriter) const override {1841    if (!isa<TestDialect>(op->getDialect()))1842      return failure();1843    rewriter.eraseOp(op);1844    return success();1845  }1846};1847 1848struct TestUnknownRootOpDriver1849    : public mlir::PassWrapper<TestUnknownRootOpDriver, OperationPass<>> {1850  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestUnknownRootOpDriver)1851 1852  StringRef getArgument() const final {1853    return "test-legalize-unknown-root-patterns";1854  }1855  StringRef getDescription() const final {1856    return "Test public remapped value mechanism in ConversionPatternRewriter";1857  }1858  void runOnOperation() override {1859    mlir::RewritePatternSet patterns(&getContext());1860    patterns.add<RemoveTestDialectOps>(&getContext());1861 1862    mlir::ConversionTarget target(getContext());1863    target.addIllegalDialect<TestDialect>();1864    if (failed(applyPartialConversion(getOperation(), target,1865                                      std::move(patterns))))1866      signalPassFailure();1867  }1868};1869} // namespace1870 1871//===----------------------------------------------------------------------===//1872// Test patterns that uses operations and types defined at runtime1873//===----------------------------------------------------------------------===//1874 1875namespace {1876/// This pattern matches dynamic operations 'test.one_operand_two_results' and1877/// replace them with dynamic operations 'test.generic_dynamic_op'.1878struct RewriteDynamicOp : public RewritePattern {1879  RewriteDynamicOp(MLIRContext *context)1880      : RewritePattern("test.dynamic_one_operand_two_results", /*benefit=*/1,1881                       context) {}1882 1883  LogicalResult matchAndRewrite(Operation *op,1884                                PatternRewriter &rewriter) const override {1885    assert(op->getName().getStringRef() ==1886               "test.dynamic_one_operand_two_results" &&1887           "rewrite pattern should only match operations with the right name");1888 1889    OperationState state(op->getLoc(), "test.dynamic_generic",1890                         op->getOperands(), op->getResultTypes(),1891                         op->getAttrs());1892    auto *newOp = rewriter.create(state);1893    rewriter.replaceOp(op, newOp->getResults());1894    return success();1895  }1896};1897 1898struct TestRewriteDynamicOpDriver1899    : public PassWrapper<TestRewriteDynamicOpDriver, OperationPass<>> {1900  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestRewriteDynamicOpDriver)1901 1902  void getDependentDialects(DialectRegistry &registry) const override {1903    registry.insert<TestDialect>();1904  }1905  StringRef getArgument() const final { return "test-rewrite-dynamic-op"; }1906  StringRef getDescription() const final {1907    return "Test rewritting on dynamic operations";1908  }1909  void runOnOperation() override {1910    RewritePatternSet patterns(&getContext());1911    patterns.add<RewriteDynamicOp>(&getContext());1912 1913    ConversionTarget target(getContext());1914    target.addIllegalOp(1915        OperationName("test.dynamic_one_operand_two_results", &getContext()));1916    target.addLegalOp(OperationName("test.dynamic_generic", &getContext()));1917    if (failed(applyPartialConversion(getOperation(), target,1918                                      std::move(patterns))))1919      signalPassFailure();1920  }1921};1922} // end anonymous namespace1923 1924//===----------------------------------------------------------------------===//1925// Test type conversions1926//===----------------------------------------------------------------------===//1927 1928namespace {1929struct TestTypeConversionProducer1930    : public OpConversionPattern<TestTypeProducerOp> {1931  using OpConversionPattern<TestTypeProducerOp>::OpConversionPattern;1932  LogicalResult1933  matchAndRewrite(TestTypeProducerOp op, OpAdaptor adaptor,1934                  ConversionPatternRewriter &rewriter) const final {1935    Type resultType = op.getType();1936    Type convertedType = getTypeConverter()1937                             ? getTypeConverter()->convertType(resultType)1938                             : resultType;1939    if (isa<FloatType>(resultType))1940      resultType = rewriter.getF64Type();1941    else if (resultType.isInteger(16))1942      resultType = rewriter.getIntegerType(64);1943    else if (isa<test::TestRecursiveType>(resultType) &&1944             convertedType != resultType)1945      resultType = convertedType;1946    else1947      return failure();1948 1949    rewriter.replaceOpWithNewOp<TestTypeProducerOp>(op, resultType);1950    return success();1951  }1952};1953 1954/// Call signature conversion and then fail the rewrite to trigger the undo1955/// mechanism.1956struct TestSignatureConversionUndo1957    : public OpConversionPattern<TestSignatureConversionUndoOp> {1958  using OpConversionPattern<TestSignatureConversionUndoOp>::OpConversionPattern;1959 1960  LogicalResult1961  matchAndRewrite(TestSignatureConversionUndoOp op, OpAdaptor adaptor,1962                  ConversionPatternRewriter &rewriter) const final {1963    (void)rewriter.convertRegionTypes(&op->getRegion(0), *getTypeConverter());1964    return failure();1965  }1966};1967 1968/// Call signature conversion without providing a type converter to handle1969/// materializations.1970struct TestTestSignatureConversionNoConverter1971    : public OpConversionPattern<TestSignatureConversionNoConverterOp> {1972  TestTestSignatureConversionNoConverter(const TypeConverter &converter,1973                                         MLIRContext *context)1974      : OpConversionPattern<TestSignatureConversionNoConverterOp>(context),1975        converter(converter) {}1976 1977  LogicalResult1978  matchAndRewrite(TestSignatureConversionNoConverterOp op, OpAdaptor adaptor,1979                  ConversionPatternRewriter &rewriter) const final {1980    Region &region = op->getRegion(0);1981    Block *entry = &region.front();1982 1983    // Convert the original entry arguments.1984    TypeConverter::SignatureConversion result(entry->getNumArguments());1985    if (failed(1986            converter.convertSignatureArgs(entry->getArgumentTypes(), result)))1987      return failure();1988    rewriter.modifyOpInPlace(op, [&] {1989      rewriter.applySignatureConversion(&region.front(), result);1990    });1991    return success();1992  }1993 1994  const TypeConverter &converter;1995};1996 1997/// Just forward the operands to the root op. This is essentially a no-op1998/// pattern that is used to trigger target materialization.1999struct TestTypeConsumerForward2000    : public OpConversionPattern<TestTypeConsumerOp> {2001  using OpConversionPattern<TestTypeConsumerOp>::OpConversionPattern;2002 2003  LogicalResult2004  matchAndRewrite(TestTypeConsumerOp op, OpAdaptor adaptor,2005                  ConversionPatternRewriter &rewriter) const final {2006    rewriter.modifyOpInPlace(op,2007                             [&] { op->setOperands(adaptor.getOperands()); });2008    return success();2009  }2010};2011 2012struct TestTypeConversionAnotherProducer2013    : public OpRewritePattern<TestAnotherTypeProducerOp> {2014  using OpRewritePattern<TestAnotherTypeProducerOp>::OpRewritePattern;2015 2016  LogicalResult matchAndRewrite(TestAnotherTypeProducerOp op,2017                                PatternRewriter &rewriter) const final {2018    rewriter.replaceOpWithNewOp<TestTypeProducerOp>(op, op.getType());2019    return success();2020  }2021};2022 2023struct TestReplaceWithLegalOp : public ConversionPattern {2024  TestReplaceWithLegalOp(const TypeConverter &converter, MLIRContext *ctx)2025      : ConversionPattern(converter, "test.replace_with_legal_op",2026                          /*benefit=*/1, ctx) {}2027  LogicalResult2028  matchAndRewrite(Operation *op, ArrayRef<ValueRange> operands,2029                  ConversionPatternRewriter &rewriter) const final {2030    rewriter.replaceOpWithNewOp<LegalOpD>(op, operands[0].front());2031    return success();2032  }2033};2034 2035struct TestTypeConversionDriver2036    : public PassWrapper<TestTypeConversionDriver, OperationPass<>> {2037  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestTypeConversionDriver)2038 2039  TestTypeConversionDriver() = default;2040  TestTypeConversionDriver(const TestTypeConversionDriver &other)2041      : PassWrapper(other) {}2042 2043  void getDependentDialects(DialectRegistry &registry) const override {2044    registry.insert<TestDialect>();2045  }2046  StringRef getArgument() const final {2047    return "test-legalize-type-conversion";2048  }2049  StringRef getDescription() const final {2050    return "Test various type conversion functionalities in DialectConversion";2051  }2052 2053  void runOnOperation() override {2054    // Initialize the type converter.2055    SmallVector<Type, 2> conversionCallStack;2056    TypeConverter converter;2057 2058    /// Add the legal set of type conversions.2059    converter.addConversion([](Type type) -> Type {2060      // Treat F64 as legal.2061      if (type.isF64())2062        return type;2063      // Allow converting BF16/F16/F32 to F64.2064      if (type.isBF16() || type.isF16() || type.isF32())2065        return Float64Type::get(type.getContext());2066      // Otherwise, the type is illegal.2067      return nullptr;2068    });2069    converter.addConversion([](IndexType type) { return type; });2070    converter.addConversion([](IntegerType type, SmallVectorImpl<Type> &types) {2071      if (type.isInteger(1)) {2072        // i1 is legal.2073        types.push_back(type);2074      }2075      if (type.isInteger(38)) {2076        // i38 is legal.2077        types.push_back(type);2078      }2079      // Drop all other integer types.2080      return success();2081    });2082    converter.addConversion(2083        // Convert a recursive self-referring type into a non-self-referring2084        // type named "outer_converted_type" that contains a SimpleAType.2085        [&](test::TestRecursiveType type,2086            SmallVectorImpl<Type> &results) -> std::optional<LogicalResult> {2087          // If the type is already converted, return it to indicate that it is2088          // legal.2089          if (type.getName() == "outer_converted_type") {2090            results.push_back(type);2091            return success();2092          }2093 2094          conversionCallStack.push_back(type);2095          auto popConversionCallStack = llvm::make_scope_exit(2096              [&conversionCallStack]() { conversionCallStack.pop_back(); });2097 2098          // If the type is on the call stack more than once (it is there at2099          // least once because of the _current_ call, which is always the last2100          // element on the stack), we've hit the recursive case. Just return2101          // SimpleAType here to create a non-recursive type as a result.2102          if (llvm::is_contained(ArrayRef(conversionCallStack).drop_back(),2103                                 type)) {2104            results.push_back(test::SimpleAType::get(type.getContext()));2105            return success();2106          }2107 2108          // Convert the body recursively.2109          auto result = test::TestRecursiveType::get(type.getContext(),2110                                                     "outer_converted_type");2111          if (failed(result.setBody(converter.convertType(type.getBody()))))2112            return failure();2113          results.push_back(result);2114          return success();2115        });2116    converter.addConversion([](Value v) -> std::optional<Type> {2117      // Context-aware type conversion rule that converts i37 to2118      // i(37 + increment). The increment is taken from the enclosing2119      // function.2120      auto intType = dyn_cast<IntegerType>(v.getType());2121      if (!intType || intType.getWidth() != 37)2122        return std::nullopt;2123      Region *r = v.getParentRegion();2124      if (!r) {2125        // No enclosing region found. This can happen when running with2126        // allow-pattern-rollback = true. Context-aware type conversions are2127        // not fully supported when running in rollback mode.2128        return Type();2129      }2130      Operation *op = r->getParentOp();2131      if (!op)2132        return Type();2133      if (!isa<FunctionOpInterface>(op))2134        op = op->getParentOfType<FunctionOpInterface>();2135      if (!op)2136        return Type();2137      auto incrementAttr = op->getAttrOfType<IntegerAttr>("increment");2138      if (!incrementAttr)2139        return Type();2140      return IntegerType::get(v.getContext(),2141                              intType.getWidth() + incrementAttr.getInt());2142    });2143 2144    /// Add the legal set of type materializations.2145    converter.addSourceMaterialization([](OpBuilder &builder, Type resultType,2146                                          ValueRange inputs,2147                                          Location loc) -> Value {2148      // Allow casting from F64 back to F32.2149      if (!resultType.isF16() && inputs.size() == 1 &&2150          inputs[0].getType().isF64())2151        return TestCastOp::create(builder, loc, resultType, inputs).getResult();2152      // Allow producing an i32 or i64 from nothing.2153      if ((resultType.isInteger(32) || resultType.isInteger(64)) &&2154          inputs.empty())2155        return TestTypeProducerOp::create(builder, loc, resultType);2156      // Allow producing an i64 from an integer.2157      if (isa<IntegerType>(resultType) && inputs.size() == 1 &&2158          isa<IntegerType>(inputs[0].getType()))2159        return TestCastOp::create(builder, loc, resultType, inputs).getResult();2160      // Otherwise, fail.2161      return nullptr;2162    });2163    // Materialize i37 to any desired type with unrealized_conversion_cast.2164    converter.addTargetMaterialization([](OpBuilder &builder, Type type,2165                                          ValueRange inputs,2166                                          Location loc) -> Value {2167      if (inputs.size() != 1 || !inputs[0].getType().isInteger(37))2168        return Value();2169      return UnrealizedConversionCastOp::create(builder, loc, type, inputs)2170          .getResult(0);2171    });2172 2173    // Initialize the conversion target.2174    mlir::ConversionTarget target(getContext());2175    target.addLegalOp(OperationName("test.context_op", &getContext()));2176    target.addLegalOp<LegalOpD>();2177    target.addDynamicallyLegalOp<TestTypeProducerOp>([](TestTypeProducerOp op) {2178      auto recursiveType = dyn_cast<test::TestRecursiveType>(op.getType());2179      return op.getType().isF64() || op.getType().isInteger(64) ||2180             (recursiveType &&2181              recursiveType.getName() == "outer_converted_type");2182    });2183    target.addDynamicallyLegalOp<func::FuncOp>([&](func::FuncOp op) {2184      return converter.isSignatureLegal(op.getFunctionType());2185    });2186    target.addDynamicallyLegalOp<TestCastOp>([&](TestCastOp op) {2187      // Allow casts from F64 to F32.2188      return (*op.operand_type_begin()).isF64() && op.getType().isF32();2189    });2190    target.addDynamicallyLegalOp<TestSignatureConversionNoConverterOp>(2191        [&](TestSignatureConversionNoConverterOp op) {2192          return converter.isLegal(op.getRegion().front().getArgumentTypes());2193        });2194 2195    // Initialize the set of rewrite patterns.2196    RewritePatternSet patterns(&getContext());2197    patterns2198        .add<TestTypeConsumerForward, TestTypeConversionProducer,2199             TestSignatureConversionUndo,2200             TestTestSignatureConversionNoConverter, TestReplaceWithLegalOp>(2201            converter, &getContext());2202    patterns.add<TestTypeConversionAnotherProducer>(&getContext());2203    mlir::populateAnyFunctionOpInterfaceTypeConversionPattern(patterns,2204                                                              converter);2205    mlir::scf::populateSCFStructuralTypeConversionsAndLegality(2206        converter, patterns, target);2207    mlir::cf::populateCFStructuralTypeConversionsAndLegality(converter,2208                                                             patterns, target);2209 2210    ConversionConfig config;2211    config.allowPatternRollback = allowPatternRollback;2212    if (failed(applyPartialConversion(getOperation(), target,2213                                      std::move(patterns), config)))2214      signalPassFailure();2215  }2216 2217  Option<bool> allowPatternRollback{*this, "allow-pattern-rollback",2218                                    llvm::cl::desc("Allow pattern rollback"),2219                                    llvm::cl::init(true)};2220};2221} // namespace2222 2223//===----------------------------------------------------------------------===//2224// Test Target Materialization With No Uses2225//===----------------------------------------------------------------------===//2226 2227namespace {2228struct ForwardOperandPattern : public OpConversionPattern<TestTypeChangerOp> {2229  using OpConversionPattern<TestTypeChangerOp>::OpConversionPattern;2230 2231  LogicalResult2232  matchAndRewrite(TestTypeChangerOp op, OpAdaptor adaptor,2233                  ConversionPatternRewriter &rewriter) const final {2234    rewriter.replaceOp(op, adaptor.getOperands());2235    return success();2236  }2237};2238 2239struct TestTargetMaterializationWithNoUses2240    : public PassWrapper<TestTargetMaterializationWithNoUses, OperationPass<>> {2241  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(2242      TestTargetMaterializationWithNoUses)2243 2244  StringRef getArgument() const final {2245    return "test-target-materialization-with-no-uses";2246  }2247  StringRef getDescription() const final {2248    return "Test a special case of target materialization in DialectConversion";2249  }2250 2251  void runOnOperation() override {2252    TypeConverter converter;2253    converter.addConversion([](Type t) { return t; });2254    converter.addConversion([](IntegerType intTy) -> Type {2255      if (intTy.getWidth() == 16)2256        return IntegerType::get(intTy.getContext(), 64);2257      return intTy;2258    });2259    converter.addTargetMaterialization(2260        [](OpBuilder &builder, Type type, ValueRange inputs, Location loc) {2261          return TestCastOp::create(builder, loc, type, inputs).getResult();2262        });2263 2264    ConversionTarget target(getContext());2265    target.addIllegalOp<TestTypeChangerOp>();2266 2267    RewritePatternSet patterns(&getContext());2268    patterns.add<ForwardOperandPattern>(converter, &getContext());2269 2270    if (failed(applyPartialConversion(getOperation(), target,2271                                      std::move(patterns))))2272      signalPassFailure();2273  }2274};2275} // namespace2276 2277//===----------------------------------------------------------------------===//2278// Test Block Merging2279//===----------------------------------------------------------------------===//2280 2281namespace {2282/// A rewriter pattern that tests that blocks can be merged.2283struct TestMergeBlock : public OpConversionPattern<TestMergeBlocksOp> {2284  using OpConversionPattern<TestMergeBlocksOp>::OpConversionPattern;2285 2286  LogicalResult2287  matchAndRewrite(TestMergeBlocksOp op, OpAdaptor adaptor,2288                  ConversionPatternRewriter &rewriter) const final {2289    Block &firstBlock = op.getBody().front();2290    Operation *branchOp = firstBlock.getTerminator();2291    Block *secondBlock = &*(std::next(op.getBody().begin()));2292    auto succOperands = branchOp->getOperands();2293    SmallVector<Value, 2> replacements(succOperands);2294    rewriter.eraseOp(branchOp);2295    rewriter.mergeBlocks(secondBlock, &firstBlock, replacements);2296    rewriter.modifyOpInPlace(op, [] {});2297    return success();2298  }2299};2300 2301/// A rewrite pattern to tests the undo mechanism of blocks being merged.2302struct TestUndoBlocksMerge : public ConversionPattern {2303  TestUndoBlocksMerge(MLIRContext *ctx)2304      : ConversionPattern("test.undo_blocks_merge", /*benefit=*/1, ctx) {}2305  LogicalResult2306  matchAndRewrite(Operation *op, ArrayRef<Value> operands,2307                  ConversionPatternRewriter &rewriter) const final {2308    Block &firstBlock = op->getRegion(0).front();2309    Operation *branchOp = firstBlock.getTerminator();2310    Block *secondBlock = &*(std::next(op->getRegion(0).begin()));2311    rewriter.setInsertionPointToStart(secondBlock);2312    ILLegalOpF::create(rewriter, op->getLoc(), rewriter.getF32Type());2313    auto succOperands = branchOp->getOperands();2314    SmallVector<Value, 2> replacements(succOperands);2315    rewriter.eraseOp(branchOp);2316    rewriter.mergeBlocks(secondBlock, &firstBlock, replacements);2317    rewriter.modifyOpInPlace(op, [] {});2318    return success();2319  }2320};2321 2322/// A rewrite mechanism to inline the body of the op into its parent, when both2323/// ops can have a single block.2324struct TestMergeSingleBlockOps2325    : public OpConversionPattern<SingleBlockImplicitTerminatorOp> {2326  using OpConversionPattern<2327      SingleBlockImplicitTerminatorOp>::OpConversionPattern;2328 2329  LogicalResult2330  matchAndRewrite(SingleBlockImplicitTerminatorOp op, OpAdaptor adaptor,2331                  ConversionPatternRewriter &rewriter) const final {2332    SingleBlockImplicitTerminatorOp parentOp =2333        op->getParentOfType<SingleBlockImplicitTerminatorOp>();2334    if (!parentOp)2335      return failure();2336    Block &innerBlock = op.getRegion().front();2337    TerminatorOp innerTerminator =2338        cast<TerminatorOp>(innerBlock.getTerminator());2339    rewriter.inlineBlockBefore(&innerBlock, op);2340    rewriter.eraseOp(innerTerminator);2341    rewriter.eraseOp(op);2342    return success();2343  }2344};2345 2346struct TestMergeBlocksPatternDriver2347    : public PassWrapper<TestMergeBlocksPatternDriver, OperationPass<>> {2348  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestMergeBlocksPatternDriver)2349 2350  StringRef getArgument() const final { return "test-merge-blocks"; }2351  StringRef getDescription() const final {2352    return "Test Merging operation in ConversionPatternRewriter";2353  }2354  void runOnOperation() override {2355    MLIRContext *context = &getContext();2356    mlir::RewritePatternSet patterns(context);2357    patterns.add<TestMergeBlock, TestUndoBlocksMerge, TestMergeSingleBlockOps>(2358        context);2359    ConversionTarget target(*context);2360    target.addLegalOp<func::FuncOp, ModuleOp, TerminatorOp, TestBranchOp,2361                      TestTypeConsumerOp, TestTypeProducerOp, TestReturnOp>();2362    target.addIllegalOp<ILLegalOpF>();2363 2364    /// Expect the op to have a single block after legalization.2365    target.addDynamicallyLegalOp<TestMergeBlocksOp>(2366        [&](TestMergeBlocksOp op) -> bool {2367          return op.getBody().hasOneBlock();2368        });2369 2370    /// Only allow `test.br` within test.merge_blocks op.2371    target.addDynamicallyLegalOp<TestBranchOp>([&](TestBranchOp op) -> bool {2372      return op->getParentOfType<TestMergeBlocksOp>();2373    });2374 2375    /// Expect that all nested test.SingleBlockImplicitTerminator ops are2376    /// inlined.2377    target.addDynamicallyLegalOp<SingleBlockImplicitTerminatorOp>(2378        [&](SingleBlockImplicitTerminatorOp op) -> bool {2379          return !op->getParentOfType<SingleBlockImplicitTerminatorOp>();2380        });2381 2382    DenseSet<Operation *> unlegalizedOps;2383    ConversionConfig config;2384    config.unlegalizedOps = &unlegalizedOps;2385    (void)applyPartialConversion(getOperation(), target, std::move(patterns),2386                                 config);2387    for (auto *op : unlegalizedOps)2388      op->emitRemark() << "op '" << op->getName() << "' is not legalizable";2389  }2390};2391} // namespace2392 2393//===----------------------------------------------------------------------===//2394// Test Selective Replacement2395//===----------------------------------------------------------------------===//2396 2397namespace {2398/// A rewrite mechanism to inline the body of the op into its parent, when both2399/// ops can have a single block.2400struct TestSelectiveOpReplacementPattern : public OpRewritePattern<TestCastOp> {2401  using OpRewritePattern<TestCastOp>::OpRewritePattern;2402 2403  LogicalResult matchAndRewrite(TestCastOp op,2404                                PatternRewriter &rewriter) const final {2405    if (op.getNumOperands() != 2)2406      return failure();2407    OperandRange operands = op.getOperands();2408 2409    // Replace non-terminator uses with the first operand.2410    rewriter.replaceUsesWithIf(op, operands[0], [](OpOperand &operand) {2411      return operand.getOwner()->hasTrait<OpTrait::IsTerminator>();2412    });2413    // Replace everything else with the second operand if the operation isn't2414    // dead.2415    rewriter.replaceOp(op, op.getOperand(1));2416    return success();2417  }2418};2419 2420struct TestSelectiveReplacementPatternDriver2421    : public PassWrapper<TestSelectiveReplacementPatternDriver,2422                         OperationPass<>> {2423  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(2424      TestSelectiveReplacementPatternDriver)2425 2426  StringRef getArgument() const final {2427    return "test-pattern-selective-replacement";2428  }2429  StringRef getDescription() const final {2430    return "Test selective replacement in the PatternRewriter";2431  }2432  void runOnOperation() override {2433    MLIRContext *context = &getContext();2434    mlir::RewritePatternSet patterns(context);2435    patterns.add<TestSelectiveOpReplacementPattern>(context);2436    (void)applyPatternsGreedily(getOperation(), std::move(patterns));2437  }2438};2439 2440struct TestFoldTypeConvertingOp2441    : public PassWrapper<TestFoldTypeConvertingOp, OperationPass<>> {2442  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestFoldTypeConvertingOp)2443 2444  StringRef getArgument() const final { return "test-fold-type-converting-op"; }2445  StringRef getDescription() const final {2446    return "Test helper functions for folding ops whose input and output types "2447           "differ, e.g. float comparisons of the form `(f32, f32) -> i1`.";2448  }2449  void runOnOperation() override {2450    MLIRContext *context = &getContext();2451    mlir::RewritePatternSet patterns(context);2452    patterns.add<FoldSignOpF32ToSI32, FoldLessThanOpF32ToI1>(context);2453    if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))2454      signalPassFailure();2455  }2456};2457} // namespace2458 2459//===----------------------------------------------------------------------===//2460// PassRegistration2461//===----------------------------------------------------------------------===//2462 2463namespace mlir {2464namespace test {2465void registerPatternsTestPass() {2466  PassRegistration<TestReturnTypeDriver>();2467 2468  PassRegistration<TestDerivedAttributeDriver>();2469 2470  PassRegistration<TestGreedyPatternDriver>();2471  PassRegistration<TestStrictPatternDriver>();2472  PassRegistration<TestWalkPatternDriver>();2473 2474  PassRegistration<TestLegalizePatternDriver>();2475 2476  PassRegistration<TestRemappedValue>();2477 2478  PassRegistration<TestUnknownRootOpDriver>();2479 2480  PassRegistration<TestTypeConversionDriver>();2481  PassRegistration<TestTargetMaterializationWithNoUses>();2482 2483  PassRegistration<TestRewriteDynamicOpDriver>();2484 2485  PassRegistration<TestMergeBlocksPatternDriver>();2486  PassRegistration<TestSelectiveReplacementPatternDriver>();2487 2488  PassRegistration<TestFoldTypeConvertingOp>();2489}2490} // namespace test2491} // namespace mlir2492