1 //===- PassManagerTest.cpp - PassManager 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/Pass/PassManager.h"
10 #include "mlir/Dialect/Func/IR/FuncOps.h"
11 #include "mlir/IR/Builders.h"
12 #include "mlir/IR/BuiltinOps.h"
13 #include "mlir/Pass/Pass.h"
14 #include "gtest/gtest.h"
15 
16 #include <memory>
17 
18 using namespace mlir;
19 using namespace mlir::detail;
20 
21 namespace {
22 /// Analysis that operates on any operation.
23 struct GenericAnalysis {
24   GenericAnalysis(Operation *op) : isFunc(isa<FuncOp>(op)) {}
25   const bool isFunc;
26 };
27 
28 /// Analysis that operates on a specific operation.
29 struct OpSpecificAnalysis {
30   OpSpecificAnalysis(FuncOp op) : isSecret(op.getName() == "secret") {}
31   const bool isSecret;
32 };
33 
34 /// Simple pass to annotate a FuncOp with the results of analysis.
35 struct AnnotateFunctionPass
36     : public PassWrapper<AnnotateFunctionPass, OperationPass<FuncOp>> {
37   void runOnOperation() override {
38     FuncOp op = getOperation();
39     Builder builder(op->getParentOfType<ModuleOp>());
40 
41     auto &ga = getAnalysis<GenericAnalysis>();
42     auto &sa = getAnalysis<OpSpecificAnalysis>();
43 
44     op->setAttr("isFunc", builder.getBoolAttr(ga.isFunc));
45     op->setAttr("isSecret", builder.getBoolAttr(sa.isSecret));
46   }
47 };
48 
49 TEST(PassManagerTest, OpSpecificAnalysis) {
50   MLIRContext context;
51   context.loadDialect<func::FuncDialect>();
52   Builder builder(&context);
53 
54   // Create a module with 2 functions.
55   OwningOpRef<ModuleOp> module(ModuleOp::create(UnknownLoc::get(&context)));
56   for (StringRef name : {"secret", "not_secret"}) {
57     FuncOp func =
58         FuncOp::create(builder.getUnknownLoc(), name,
59                        builder.getFunctionType(llvm::None, llvm::None));
60     func.setPrivate();
61     module->push_back(func);
62   }
63 
64   // Instantiate and run our pass.
65   PassManager pm(&context);
66   pm.addNestedPass<FuncOp>(std::make_unique<AnnotateFunctionPass>());
67   LogicalResult result = pm.run(module.get());
68   EXPECT_TRUE(succeeded(result));
69 
70   // Verify that each function got annotated with expected attributes.
71   for (FuncOp func : module->getOps<FuncOp>()) {
72     ASSERT_TRUE(func->getAttr("isFunc").isa<BoolAttr>());
73     EXPECT_TRUE(func->getAttr("isFunc").cast<BoolAttr>().getValue());
74 
75     bool isSecret = func.getName() == "secret";
76     ASSERT_TRUE(func->getAttr("isSecret").isa<BoolAttr>());
77     EXPECT_EQ(func->getAttr("isSecret").cast<BoolAttr>().getValue(), isSecret);
78   }
79 }
80 
81 namespace {
82 struct InvalidPass : Pass {
83   InvalidPass() : Pass(TypeID::get<InvalidPass>(), StringRef("invalid_op")) {}
84   StringRef getName() const override { return "Invalid Pass"; }
85   void runOnOperation() override {}
86   bool canScheduleOn(RegisteredOperationName opName) const override {
87     return true;
88   }
89 
90   /// A clone method to create a copy of this pass.
91   std::unique_ptr<Pass> clonePass() const override {
92     return std::make_unique<InvalidPass>(
93         *static_cast<const InvalidPass *>(this));
94   }
95 };
96 } // namespace
97 
98 TEST(PassManagerTest, InvalidPass) {
99   MLIRContext context;
100   context.allowUnregisteredDialects();
101 
102   // Create a module
103   OwningOpRef<ModuleOp> module(ModuleOp::create(UnknownLoc::get(&context)));
104 
105   // Add a single "invalid_op" operation
106   OpBuilder builder(&module->getBodyRegion());
107   OperationState state(UnknownLoc::get(&context), "invalid_op");
108   builder.insert(Operation::create(state));
109 
110   // Register a diagnostic handler to capture the diagnostic so that we can
111   // check it later.
112   std::unique_ptr<Diagnostic> diagnostic;
113   context.getDiagEngine().registerHandler([&](Diagnostic &diag) {
114     diagnostic = std::make_unique<Diagnostic>(std::move(diag));
115   });
116 
117   // Instantiate and run our pass.
118   PassManager pm(&context);
119   pm.nest("invalid_op").addPass(std::make_unique<InvalidPass>());
120   LogicalResult result = pm.run(module.get());
121   EXPECT_TRUE(failed(result));
122   ASSERT_TRUE(diagnostic.get() != nullptr);
123   EXPECT_EQ(
124       diagnostic->str(),
125       "'invalid_op' op trying to schedule a pass on an unregistered operation");
126 
127   // Check that clearing the pass manager effectively removed the pass.
128   pm.clear();
129   result = pm.run(module.get());
130   EXPECT_TRUE(succeeded(result));
131 
132   // Check that adding the pass at the top-level triggers a fatal error.
133   ASSERT_DEATH(pm.addPass(std::make_unique<InvalidPass>()), "");
134 }
135 
136 } // namespace
137