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