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
87