[mlir][Pass] Include anchor op in -pass-pipeline

[rkayaith]

Summary

In D134622 the printed form of a pass manager is changed to include the
name of the op that the pass manager is anchored on. This updates the
-pass-pipeline argument format to include the anchor op as well, so
that the printed form of a pipeline can be directly passed to
-pass-pipeline. In most cases this requires updating
-pass-pipeline='pipeline' to
-pass-pipeline='builtin.module(pipeline)'.

This also fixes an outdated assert that prevented running a
PassManager anchored on 'any'.

I've split out the mechanical test updates into D136193, but will
combine them into one commit before landing.

[rkayaith]

update

[rkayaith]

update description

[rkayaith]

update

[rkayaith]

add extra test

[rkayaith]

add more info to description

[rkayaith]

The approach used here currently (add a new option instead of changing pass-pipeline) has the downside that pass-pipeline used to be symmetric with the printed form, but no longer is. If we wanted to preserve that, some options I see:

• remove the current -pass-pipeline and call the new option pass-pipeline. this avoids having multiple options which do similar things
• rename -pass-pipeline to something else, e.g. -append-pipeline, and call the new option pass-pipeline. this way there's still a way to interleave the text-pipeline form with -pass-name options, for users that find that convenient.
  • upstream the main place I see this feature used is in integration tests, mostly GPU tests, e.g. https://github.com/llvm/llvm-project/blob/e3638e83db086e77243d0673a4ce8a0b4b330d42/mlir/test/Integration/GPU/CUDA/async.mlir#L1-L6

Any thoughts on how to proceed here?

[rriddle]

In D134900#3834645, @rkayaith wrote:

The approach used here currently (add a new option instead of changing pass-pipeline) has the downside that pass-pipeline used to be symmetric with the printed form, but no longer is. If we wanted to preserve that, some options I see:

• remove the current -pass-pipeline and call the new option pass-pipeline. this avoids having multiple options which do similar things
• rename -pass-pipeline to something else, e.g. -append-pipeline, and call the new option pass-pipeline. this way there's still a way to interleave the text-pipeline form with -pass-name options, for users that find that convenient.
  • upstream the main place I see this feature used is in integration tests, mostly GPU tests, e.g. https://github.com/llvm/llvm-project/blob/e3638e83db086e77243d0673a4ce8a0b4b330d42/mlir/test/Integration/GPU/CUDA/async.mlir#L1-L6

Any thoughts on how to proceed here?

I think we should proceed with changing -pass-pipeline. I'd rather have fewer things than more things. We've also already had people confused on the implicit nature of the anchor op in pass-pipeline before, and it seems better to just have that option represent the full pipeline.

[mehdi_amini]

I'm aligned with River on the goal here.

[rkayaith]

Rework patch to change the -pass-pipeline option instead. Other than
pipeline-parsing.mlir all the test updates here are mechanical
(`-pass-pipeline='pipeline' -> -pass-pipeline='builtin.module(pipeline)')

[rriddle]

Did you also scrub the docs for uses of pass-pipeline?

mlir/lib/Pass/Pass.cpp

776–781

Please avoid using if initializers with nullable things. Can you also spell out auto here?

[rkayaith]

rebase, update docs, address comment

[rkayaith]

unnecessary 'llvm::'

[rkayaith]

move mechanical test updates into a separate diff

[rkayaith]

rebase

mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h

 /// - allowUnregisteredDialects allows to parse and create operation without
 /// registering the Dialect in the MLIRContext.
 /// - preloadDialectsInContext will trigger the upfront loading of all
 ///   dialects from the global registry in the MLIRContext. This option is
 ///   deprecated and will be removed soon.
 /// - 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<llvm::MemoryBuffer> buffer,
+/// - passManager is used to initialize the pass manager that will be applied
+LogicalResult
+MlirOptMain(llvm::raw_ostream &outputStream,
+            std::unique_ptr<llvm::MemoryBuffer> buffer,
             const PassPipelineCLParser &passPipeline, DialectRegistry &registry,
             bool splitInputFile, bool verifyDiagnostics, bool verifyPasses,
-    bool allowUnregisteredDialects, bool preloadDialectsInContext = false,
-    bool emitBytecode = false, bool implicitModule = false);
+            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
 ///   apply on the loaded IR.
 LogicalResult MlirOptMain(
     llvm::raw_ostream &outputStream, std::unique_ptr<llvm::MemoryBuffer> buffer,
     PassPipelineFn passManagerSetupFn, DialectRegistry &registry,
     bool splitInputFile, bool verifyDiagnostics, bool verifyPasses,

mlir/lib/Pass/Pass.cpp

 
 PassManager::~PassManager() = default;
 
 void PassManager::enableVerifier(bool enabled) { verifyPasses = enabled; }
 
 /// 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;
   getDependentDialects(dependentDialects);
   context->appendDialectRegistry(dependentDialects);
   for (StringRef name : dependentDialects.getDialectNames())
     context->getOrLoadDialect(name);
 
   // Before running, make sure to finalize the pipeline pass list.

mlir/lib/Pass/PassRegistry.cpp

 
 FailureOr<OpPassManager> mlir::parsePassPipeline(StringRef pipeline,
                                                  raw_ostream &errorStream) {
   // Pipelines are expected to be of the form `<op-name>(<pipeline>)`.
   size_t pipelineStart = pipeline.find_first_of('(');
   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();
   }
 
   StringRef opName = pipeline.take_front(pipelineStart);
   OpPassManager pm(opName);
   if (failed(parsePassPipeline(pipeline.drop_front(1 + pipelineStart), pm)))
     return failure();
   return pm;

mlir/lib/Tools/mlir-opt/MlirOptMain.cpp

return processBuffer(os, std::move(chunkBuffer), verifyDiagnostics,
                          verifyPasses, allowUnregisteredDialects,
                          preloadDialectsInContext, emitBytecode, implicitModule,
                          passManagerSetupFn, registry, threadPool);
   };
   return splitAndProcessBuffer(std::move(buffer), chunkFn, outputStream,
                                splitInputFile, /*insertMarkerInOutput=*/true);
 }
 
-LogicalResult mlir::MlirOptMain(raw_ostream &outputStream,
-                                std::unique_ptr<MemoryBuffer> buffer,
-                                const PassPipelineCLParser &passPipeline,
-                                DialectRegistry &registry, bool splitInputFile,
-                                bool verifyDiagnostics, bool verifyPasses,
-                                bool allowUnregisteredDialects,
-                                bool preloadDialectsInContext,
-                                bool emitBytecode, bool implicitModule) {
+LogicalResult mlir::MlirOptMain(
+    raw_ostream &outputStream, std::unique_ptr<MemoryBuffer> buffer,
+    const PassPipelineCLParser &passPipeline, DialectRegistry &registry,
+    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<OpPassManager> newPm =
+          parsePassPipeline(passManager, llvm::errs());
+      if (failed(newPm))
+        return failure();
+      if (!pm.empty())
+        return errorHandler(
+            "'-pass-manager' option is overriding existing pipeline");
+      static_cast<OpPassManager &>(pm) = std::move(*newPm);
+    }
     return passPipeline.addToPipeline(pm, errorHandler);
   };
   return MlirOptMain(outputStream, std::move(buffer), passManagerSetupFn,
                      registry, splitInputFile, verifyDiagnostics, verifyPasses,
                      allowUnregisteredDialects, preloadDialectsInContext,
                      emitBytecode, implicitModule);
 }
 
static cl::opt<bool> emitBytecode(
       cl::init(false));
 
   static cl::opt<bool> noImplicitModule{
       "no-implicit-module",
       cl::desc(
           "Disable implicit addition of a top-level module op during parsing"),
       cl::init(false)};
 
+  static cl::opt<std::string> passManager(
+      "pass-manager",
+      cl::desc("Pass manager to run (format: `anchor-op(pipeline)`)"));
+
   InitLLVM y(argc, argv);
 
   // Register any command line options.
   registerAsmPrinterCLOptions();
   registerMLIRContextCLOptions();
   registerPassManagerCLOptions();
   registerDefaultTimingManagerCLOptions();
   DebugCounter::registerCLOptions();
LogicalResult mlir::MlirOptMain(int argc, char **argv, llvm::StringRef toolName,
   if (!output) {
     llvm::errs() << errorMessage << "\n";
     return failure();
   }
 
   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.
   output->keep();
   return success();
 }

mlir/test/mlir-opt/pass-manager.mlir

+//    RUN: mlir-opt %s -cse -test-dump-pipeline 2>&1 | FileCheck %s -check-prefix=CHECK0
+// CHECK0: builtin.module(cse)
+
+//    RUN: mlir-opt %s -pass-manager='any(cse)' -test-dump-pipeline 2>&1 | FileCheck %s -check-prefix=CHECK1
+// CHECK1: any(cse)
+
+//    RUN: mlir-opt %s -pass-manager='any()' -cse -test-dump-pipeline 2>&1 | FileCheck %s -check-prefix=CHECK2
+// CHECK2: any(cse)
+
+module {}