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