117 lines · cpp
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-exception6//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 17using namespace mlir;18 19namespace {20 21class TestDynamicPipelinePass22 : public PassWrapper<TestDynamicPipelinePass, OperationPass<>> {23public: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, llvm::errs());35 pm.getDependentDialects(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, llvm::errs());65 66 // Check that running on the parent operation always immediately fails.67 if (runOnParent) {68 if (currentOp->getParentOp())69 if (!failed(runPipeline(pm, currentOp->getParentOp())))70 signalPassFailure();71 return;72 }73 74 if (runOnNestedOp) {75 llvm::errs() << "Run on nested op\n";76 currentOp->walk([&](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 operation82 if (failed(runPipeline(pm, op)))83 signalPassFailure();84 });85 } else {86 // Run on the current operation87 if (failed(runPipeline(pm, 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} // namespace109 110namespace mlir {111namespace test {112void registerTestDynamicPipelinePass() {113 PassRegistration<TestDynamicPipelinePass>();114}115} // namespace test116} // namespace mlir117