diff --git a/mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h b/mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h --- a/mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h +++ b/mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h @@ -53,12 +53,15 @@ /// - emitBytecode will generate bytecode output instead of text. /// - implicitModule will enable implicit addition of a top-level /// 'builtin.module' if one doesn't already exist. -LogicalResult MlirOptMain( - llvm::raw_ostream &outputStream, std::unique_ptr buffer, - const PassPipelineCLParser &passPipeline, DialectRegistry ®istry, - bool splitInputFile, bool verifyDiagnostics, bool verifyPasses, - bool allowUnregisteredDialects, bool preloadDialectsInContext = false, - bool emitBytecode = false, bool implicitModule = false); +/// - passManager is used to initialize the pass manager that will be applied +LogicalResult +MlirOptMain(llvm::raw_ostream &outputStream, + std::unique_ptr buffer, + const PassPipelineCLParser &passPipeline, DialectRegistry ®istry, + bool splitInputFile, bool verifyDiagnostics, bool verifyPasses, + bool allowUnregisteredDialects, + bool preloadDialectsInContext = false, bool emitBytecode = false, + bool implicitModule = false, StringRef passManager = ""); /// Support a callback to setup the pass manager. /// - passManagerSetupFn is the callback invoked to setup the pass manager to diff --git a/mlir/lib/Pass/Pass.cpp b/mlir/lib/Pass/Pass.cpp --- a/mlir/lib/Pass/Pass.cpp +++ b/mlir/lib/Pass/Pass.cpp @@ -773,9 +773,11 @@ /// Run the passes within this manager on the provided operation. LogicalResult PassManager::run(Operation *op) { MLIRContext *context = getContext(); - assert(op->getName() == getOpName(*context) && - "operation has a different name than the PassManager or is from a " - "different context"); + if (auto anchorOp = getOpName(*context); + anchorOp && anchorOp != op->getName()) + return emitError(op->getLoc()) + << "can't run '" << getOpAnchorName() << "' pass manager on '" + << op->getName() << "' op"; // Register all dialects for the current pipeline. DialectRegistry dependentDialects; diff --git a/mlir/lib/Pass/PassRegistry.cpp b/mlir/lib/Pass/PassRegistry.cpp --- a/mlir/lib/Pass/PassRegistry.cpp +++ b/mlir/lib/Pass/PassRegistry.cpp @@ -694,7 +694,7 @@ if (pipelineStart == 0 || pipelineStart == StringRef::npos || !pipeline.consume_back(")")) { errorStream << "expected pass pipeline to be wrapped with the anchor " - "operation type, e.g. `builtin.module(...)"; + "operation type, e.g. `builtin.module(...)`\n"; return failure(); } diff --git a/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp b/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp --- a/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp +++ b/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp @@ -187,19 +187,28 @@ splitInputFile, /*insertMarkerInOutput=*/true); } -LogicalResult mlir::MlirOptMain(raw_ostream &outputStream, - std::unique_ptr buffer, - const PassPipelineCLParser &passPipeline, - DialectRegistry ®istry, bool splitInputFile, - bool verifyDiagnostics, bool verifyPasses, - bool allowUnregisteredDialects, - bool preloadDialectsInContext, - bool emitBytecode, bool implicitModule) { +LogicalResult mlir::MlirOptMain( + raw_ostream &outputStream, std::unique_ptr buffer, + const PassPipelineCLParser &passPipeline, DialectRegistry ®istry, + bool splitInputFile, bool verifyDiagnostics, bool verifyPasses, + bool allowUnregisteredDialects, bool preloadDialectsInContext, + bool emitBytecode, bool implicitModule, StringRef passManager) { auto passManagerSetupFn = [&](PassManager &pm) { auto errorHandler = [&](const Twine &msg) { emitError(UnknownLoc::get(pm.getContext())) << msg; return failure(); }; + + if (!passManager.empty()) { + FailureOr newPm = + parsePassPipeline(passManager, llvm::errs()); + if (failed(newPm)) + return failure(); + if (!pm.empty()) + return errorHandler( + "'-pass-manager' option is overriding existing pipeline"); + static_cast(pm) = std::move(*newPm); + } return passPipeline.addToPipeline(pm, errorHandler); }; return MlirOptMain(outputStream, std::move(buffer), passManagerSetupFn, @@ -253,6 +262,10 @@ "Disable implicit addition of a top-level module op during parsing"), cl::init(false)}; + static cl::opt passManager( + "pass-manager", + cl::desc("Pass manager to run (format: `anchor-op(pipeline)`)")); + InitLLVM y(argc, argv); // Register any command line options. @@ -298,7 +311,8 @@ if (failed(MlirOptMain(output->os(), std::move(file), passPipeline, registry, splitInputFile, verifyDiagnostics, verifyPasses, allowUnregisteredDialects, preloadDialectsInContext, - emitBytecode, /*implicitModule=*/!noImplicitModule))) + emitBytecode, /*implicitModule=*/!noImplicitModule, + passManager))) return failure(); // Keep the output file if the invocation of MlirOptMain was successful. diff --git a/mlir/test/mlir-opt/pass-manager.mlir b/mlir/test/mlir-opt/pass-manager.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/mlir-opt/pass-manager.mlir @@ -0,0 +1,14 @@ +// RUN: mlir-opt %s -cse -test-dump-pipeline 2>&1 | FileCheck %s -check-prefix=CHECK0 +// CHECK0: builtin.module(cse) + +// Roundtrip: +// RUN: mlir-opt %s -pass-manager='any(cse)' -test-dump-pipeline 2>&1 | FileCheck %s -check-prefix=CHECK1 +// CHECK1: any(cse) + +// Additional pass options get appended: +// RUN: mlir-opt %s -pass-manager='any()' -cse -cse -test-dump-pipeline 2>&1 | FileCheck %s -check-prefix=CHECK2 +// CHECK2: any(cse,cse) + +// RUN: mlir-opt %s -pass-manager='wrong-op()' --verify-diagnostics +// expected-error@below {{can't run 'wrong-op' pass manager on 'builtin.module' op}} +module {}