1 //===------ TestDynamicPipeline.cpp --- dynamic pipeline test pass --------===//
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 // This file implements a pass to test the dynamic pipeline feature.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #include "mlir/Pass/Pass.h"
14 #include "mlir/Pass/PassManager.h"
15 
16 using namespace mlir;
17 
18 namespace {
19 
20 class TestDynamicPipelinePass
21     : public PassWrapper<TestDynamicPipelinePass, OperationPass<>> {
22 public:
23   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestDynamicPipelinePass)
24 
25   StringRef getArgument() const final { return "test-dynamic-pipeline"; }
26   StringRef getDescription() const final {
27     return "Tests the dynamic pipeline feature by applying "
28            "a pipeline on a selected set of functions";
29   }
30   void getDependentDialects(DialectRegistry &registry) const override {
31     OpPassManager pm(ModuleOp::getOperationName(),
32                      OpPassManager::Nesting::Implicit);
33     (void)parsePassPipeline(pipeline, pm, llvm::errs());
34     pm.getDependentDialects(registry);
35   }
36 
37   TestDynamicPipelinePass() = default;
38   TestDynamicPipelinePass(const TestDynamicPipelinePass &) {}
39 
40   void runOnOperation() override {
41     Operation *currentOp = getOperation();
42 
43     llvm::errs() << "Dynamic execute '" << pipeline << "' on "
44                  << currentOp->getName() << "\n";
45     if (pipeline.empty()) {
46       llvm::errs() << "Empty pipeline\n";
47       return;
48     }
49     auto symbolOp = dyn_cast<SymbolOpInterface>(currentOp);
50     if (!symbolOp) {
51       currentOp->emitWarning()
52           << "Ignoring because not implementing SymbolOpInterface\n";
53       return;
54     }
55 
56     auto opName = symbolOp.getName();
57     if (!opNames.empty() && !llvm::is_contained(opNames, opName)) {
58       llvm::errs() << "dynamic-pipeline skip op name: " << opName << "\n";
59       return;
60     }
61     OpPassManager pm(currentOp->getName().getIdentifier(),
62                      OpPassManager::Nesting::Implicit);
63     (void)parsePassPipeline(pipeline, pm, llvm::errs());
64 
65     // Check that running on the parent operation always immediately fails.
66     if (runOnParent) {
67       if (currentOp->getParentOp())
68         if (!failed(runPipeline(pm, currentOp->getParentOp())))
69           signalPassFailure();
70       return;
71     }
72 
73     if (runOnNestedOp) {
74       llvm::errs() << "Run on nested op\n";
75       currentOp->walk([&](Operation *op) {
76         if (op == currentOp || !op->hasTrait<OpTrait::IsIsolatedFromAbove>() ||
77             op->getName() != currentOp->getName())
78           return;
79         llvm::errs() << "Run on " << *op << "\n";
80         // Run on the current operation
81         if (failed(runPipeline(pm, op)))
82           signalPassFailure();
83       });
84     } else {
85       // Run on the current operation
86       if (failed(runPipeline(pm, currentOp)))
87         signalPassFailure();
88     }
89   }
90 
91   Option<bool> runOnNestedOp{
92       *this, "run-on-nested-operations",
93       llvm::cl::desc("This will apply the pipeline on nested operations under "
94                      "the visited operation.")};
95   Option<bool> runOnParent{
96       *this, "run-on-parent",
97       llvm::cl::desc("This will apply the pipeline on the parent operation if "
98                      "it exist, this is expected to fail.")};
99   Option<std::string> pipeline{
100       *this, "dynamic-pipeline",
101       llvm::cl::desc("The pipeline description that "
102                      "will run on the filtered function.")};
103   ListOption<std::string> opNames{
104       *this, "op-name",
105       llvm::cl::desc("List of function name to apply the pipeline to")};
106 };
107 } // namespace
108 
109 namespace mlir {
110 namespace test {
111 void registerTestDynamicPipelinePass() {
112   PassRegistration<TestDynamicPipelinePass>();
113 }
114 } // namespace test
115 } // namespace mlir
116