mlir/unittests/Transforms/Canonicalizer.cpp

   1 //===- DialectConversion.cpp - Dialect conversion unit tests --------------===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8
   9 #include "mlir/IR/PatternMatch.h"
  10 #include "mlir/Parser/Parser.h"
  11 #include "mlir/Pass/PassManager.h"
  12 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
  13 #include "mlir/Transforms/Passes.h"
  14 #include "gtest/gtest.h"
  15
  16 using namespace mlir;
  17
  18 namespace {
  19
  20 struct DisabledPattern : public RewritePattern {
  21   DisabledPattern(MLIRContext *context)
  22       : RewritePattern("test.foo", /*benefit=*/0, context,
  23                        /*generatedNamed=*/{}) {
  24     setDebugName("DisabledPattern");
  25   }
  26
  27   LogicalResult matchAndRewrite(Operation *op,
  28                                 PatternRewriter &rewriter) const override {
  29     if (op->getNumResults() != 1)
  30       return failure();
  31     rewriter.eraseOp(op);
  32     return success();
  33   }
  34 };
  35
  36 struct EnabledPattern : public RewritePattern {
  37   EnabledPattern(MLIRContext *context)
  38       : RewritePattern("test.foo", /*benefit=*/0, context,
  39                        /*generatedNamed=*/{}) {
  40     setDebugName("EnabledPattern");
  41   }
  42
  43   LogicalResult matchAndRewrite(Operation *op,
  44                                 PatternRewriter &rewriter) const override {
  45     if (op->getNumResults() == 1)
  46       return failure();
  47     rewriter.eraseOp(op);
  48     return success();
  49   }
  50 };
  51
  52 struct TestDialect : public Dialect {
  53   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestDialect)
  54
  55   static StringRef getDialectNamespace() { return "test"; }
  56
  57   TestDialect(MLIRContext *context)
  58       : Dialect(getDialectNamespace(), context, TypeID::get<TestDialect>()) {
  59     allowUnknownOperations();
  60   }
  61
  62   void getCanonicalizationPatterns(RewritePatternSet &results) const override {
  63     results.add<DisabledPattern, EnabledPattern>(results.getContext());
  64   }
  65 };
  66
  67 TEST(CanonicalizerTest, TestDisablePatterns) {
  68   MLIRContext context;
  69   context.getOrLoadDialect<TestDialect>();
  70   PassManager mgr(&context);
  71   mgr.addPass(
  72       createCanonicalizerPass(GreedyRewriteConfig(), {"DisabledPattern"}));
  73
  74   const char *const code = R"mlir(
  75     %0:2 = "test.foo"() {sym_name = "A"} : () -> (i32, i32)
  76     %1 = "test.foo"() {sym_name = "B"} : () -> (f32)
  77   )mlir";
  78
  79   OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(code, &context);
  80   ASSERT_TRUE(succeeded(mgr.run(*module)));
  81
  82   EXPECT_TRUE(module->lookupSymbol("B"));
  83   EXPECT_FALSE(module->lookupSymbol("A"));
  84 }
  85
  86 } // end anonymous namespace