llvm-project

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