[mlir] Add callback to provide a pass pipeline to MlirOptMain

[panickal-xmos]

Summary

The callback can be used to provide a default pass pipeline.

[panickal-xmos]

Add tag to commit message

[mehdi_amini]

mlir/include/mlir/Support/MlirOptMain.h

40

58

59

I suspect we can keep a forwarding API from the old one to the new one, it can just be implemented as:

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) {
  auto passManagerSetupFn =  [&](PassManager &pm) {
     auto errorHandler = [&](const Twine &msg) {
       emitError(UnknownLoc::get(pm.getContext())) << msg;
       return failure();
    };
    return passPipeline.addToPipeline(pm, errorHandler);
  };
   return MlirOptMain(outputStream, buffer, passManagerSetupFn,
                          registry, splitInputfile, verifyDiagnostics, verifyPasses,
                          allowUnregisteredDialects, preloadDialectsInContext);
}

mlir/lib/Support/MlirOptMain.cpp

290

I wouldn't pass the errorHandler in the callback, that's hard to justify document: it seems like leaking the passPipeline.addToPipeline API to me.

Instead you should be able to write it above here and capture it:

 [&](PassManager &pm) {
  auto errorHandler = [&](const Twine &msg) {
    emitError(UnknownLoc::get(pm.getContext())) << msg;
    return failure();
  };
  return passPipeline.addToPipeline(pm, errorHandler);
},
``

But I would add this in the header (see the other comment)

[panickal-xmos]

mlir/include/mlir/Support/MlirOptMain.h

59

Forwarding it is a good idea. Should we do it in the header?

[mehdi_amini]

mlir/include/mlir/Support/MlirOptMain.h

59

If it does not need to include more header, then sure, but I'm not sure it'd work right now.

[panickal-xmos]

Forward the current API to the new one

[rriddle]

LGTM (but wait for Mehdi to check again)

mlir/include/mlir/Support/MlirOptMain.h

32

Please document this.

63

[mehdi_amini]

LG,
I was about to push this (building and running the tests), but I saw River's comment so I'll do as soon as they are addressed

[panickal-xmos]

Update comments

mlir/include/mlir/Support/MlirOptMain.h

 namespace llvm {
 class raw_ostream;
 class MemoryBuffer;
 } // end namespace llvm
 
 namespace mlir {
 class DialectRegistry;
 class PassPipelineCLParser;
+class PassManager;
+
+/// This defines the function type used to setup the pass manager. This can be
+/// used to pass in a callback to setup a default pass pipeline to be applied on
+/// the loaded IR.
+using PassPipelineFn = llvm::function_ref<LogicalResult(PassManager &pm)>;
 
 /// Perform the core processing behind `mlir-opt`:
 /// - outputStream is the stream where the resulting IR is printed.
 /// - buffer is the in-memory file to parser and process.
 /// - passPipeline is the specification of the pipeline that will be applied.
 /// - registry should contain all the dialects that can be parsed in the source.
 /// - splitInputFile will look for a "-----" marker in the input file, and load
 /// each chunk in an individual ModuleOp processed separately.
 /// - verifyDiagnostics enables a verification mode where comments starting with
 /// "expected-(error|note|remark|warning)" are parsed in the input and matched
 /// against emitted diagnostics.
 /// - verifyPasses enables the IR verifier in-between each pass in the pipeline.
 /// - 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.
 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);
 
+/// 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,
+                          bool allowUnregisteredDialects,
+                          bool preloadDialectsInContext = false);
+
 /// Implementation for tools like `mlir-opt`.
 /// - toolName is used for the header displayed by `--help`.
 /// - registry should contain all the dialects that can be parsed in the source.
 /// - 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.
 LogicalResult MlirOptMain(int argc, char **argv, llvm::StringRef toolName,
                           DialectRegistry &registry,

mlir/lib/Support/MlirOptMain.cpp

 /// within the specified context.
 ///
 /// This typically parses the main source file, runs zero or more optimization
 /// passes, then prints the output.
 ///
 static LogicalResult performActions(raw_ostream &os, bool verifyDiagnostics,
                                     bool verifyPasses, SourceMgr &sourceMgr,
                                     MLIRContext *context,
-                                    const PassPipelineCLParser &passPipeline) {
+                                    PassPipelineFn passManagerSetupFn) {
   DefaultTimingManager tm;
   applyDefaultTimingManagerCLOptions(tm);
   TimingScope timing = tm.getRootScope();
 
   // Disable multi-threading when parsing the input file. This removes the
   // unnecessary/costly context synchronization when parsing.
   bool wasThreadingEnabled = context->isMultithreadingEnabled();
   context->disableMultithreading();
 
   // Parse the input file and reset the context threading state.
   TimingScope parserTiming = timing.nest("Parser");
   OwningModuleRef module(parseSourceFile(sourceMgr, context));
   context->enableMultithreading(wasThreadingEnabled);
   if (!module)
     return failure();
   parserTiming.stop();
 
   // Apply any pass manager command line options.
   PassManager pm(context, OpPassManager::Nesting::Implicit);
   pm.enableVerifier(verifyPasses);
   applyPassManagerCLOptions(pm);
   pm.enableTiming(timing);
 
-  auto errorHandler = [&](const Twine &msg) {
-    emitError(UnknownLoc::get(context)) << msg;
-    return failure();
-  };
-
-  // Build the provided pipeline.
-  if (failed(passPipeline.addToPipeline(pm, errorHandler)))
+  // Callback to build the pipeline.
+  if (failed(passManagerSetupFn(pm)))
     return failure();
 
   // Run the pipeline.
   if (failed(pm.run(*module)))
     return failure();
 
   // Print the output.
   TimingScope outputTiming = timing.nest("Output");
   module->print(os);
   os << '\n';
   return success();
 }
 
 /// Parses the memory buffer.  If successfully, run a series of passes against
 /// it and print the result.
 static LogicalResult
 processBuffer(raw_ostream &os, std::unique_ptr<MemoryBuffer> ownedBuffer,
               bool verifyDiagnostics, bool verifyPasses,
               bool allowUnregisteredDialects, bool preloadDialectsInContext,
-              const PassPipelineCLParser &passPipeline,
-              DialectRegistry &registry, llvm::ThreadPool &threadPool) {
+              PassPipelineFn passManagerSetupFn, DialectRegistry &registry,
+              llvm::ThreadPool &threadPool) {
   // Tell sourceMgr about this buffer, which is what the parser will pick up.
   SourceMgr sourceMgr;
   sourceMgr.AddNewSourceBuffer(std::move(ownedBuffer), SMLoc());
 
   // Create a context just for the current buffer. Disable threading on creation
   // since we'll inject the thread-pool separately.
   MLIRContext context(registry, MLIRContext::Threading::DISABLED);
   context.setThreadPool(threadPool);
 
   // Parse the input file.
   if (preloadDialectsInContext)
     context.loadAllAvailableDialects();
   context.allowUnregisteredDialects(allowUnregisteredDialects);
   if (verifyDiagnostics)
     context.printOpOnDiagnostic(false);
   context.getDebugActionManager().registerActionHandler<DebugCounter>();
 
   // If we are in verify diagnostics mode then we have a lot of work to do,
   // otherwise just perform the actions without worrying about it.
   if (!verifyDiagnostics) {
     SourceMgrDiagnosticHandler sourceMgrHandler(sourceMgr, &context);
     return performActions(os, verifyDiagnostics, verifyPasses, sourceMgr,
-                          &context, passPipeline);
+                          &context, passManagerSetupFn);
   }
 
   SourceMgrDiagnosticVerifierHandler sourceMgrHandler(sourceMgr, &context);
 
   // Do any processing requested by command line flags.  We don't care whether
   // these actions succeed or fail, we only care what diagnostics they produce
   // and whether they match our expectations.
   (void)performActions(os, verifyDiagnostics, verifyPasses, sourceMgr, &context,
-                       passPipeline);
+                       passManagerSetupFn);
 
   // Verify the diagnostic handler to make sure that each of the diagnostics
   // matched.
   return sourceMgrHandler.verify();
 }
 
 LogicalResult mlir::MlirOptMain(raw_ostream &outputStream,
                                 std::unique_ptr<MemoryBuffer> buffer,
-                                const PassPipelineCLParser &passPipeline,
+                                PassPipelineFn passManagerSetupFn,
                                 DialectRegistry &registry, bool splitInputFile,
                                 bool verifyDiagnostics, bool verifyPasses,
                                 bool allowUnregisteredDialects,
                                 bool preloadDialectsInContext) {
   // The split-input-file mode is a very specific mode that slices the file
   // up into small pieces and checks each independently.
   // We use an explicit threadpool to avoid creating and joining/destroying
   // threads for each of the split.
   llvm::ThreadPool threadPool;
   if (splitInputFile)
     return splitAndProcessBuffer(
         std::move(buffer),
         [&](std::unique_ptr<MemoryBuffer> chunkBuffer, raw_ostream &os) {
           return processBuffer(os, std::move(chunkBuffer), verifyDiagnostics,
                                verifyPasses, allowUnregisteredDialects,
-                               preloadDialectsInContext, passPipeline, registry,
-                               threadPool);
+                               preloadDialectsInContext, passManagerSetupFn,
+                               registry, threadPool);
         },
         outputStream);
 
   return processBuffer(outputStream, std::move(buffer), verifyDiagnostics,
                        verifyPasses, allowUnregisteredDialects,
-                       preloadDialectsInContext, passPipeline, registry,
+                       preloadDialectsInContext, passManagerSetupFn, registry,
                        threadPool);
 }
 
+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) {
+  auto passManagerSetupFn = [&](PassManager &pm) {
+    auto errorHandler = [&](const Twine &msg) {
+      emitError(UnknownLoc::get(pm.getContext())) << msg;
+      return failure();
+    };
+    return passPipeline.addToPipeline(pm, errorHandler);
+  };
+  return MlirOptMain(outputStream, std::move(buffer), passManagerSetupFn,
+                     registry, splitInputFile, verifyDiagnostics, verifyPasses,
+                     allowUnregisteredDialects, preloadDialectsInContext);
+}
+
 LogicalResult mlir::MlirOptMain(int argc, char **argv, llvm::StringRef toolName,
                                 DialectRegistry &registry,
                                 bool preloadDialectsInContext) {
   static cl::opt<std::string> inputFilename(
       cl::Positional, cl::desc("<input file>"), cl::init("-"));
 
   static cl::opt<std::string> outputFilename("o", cl::desc("Output filename"),
                                              cl::value_desc("filename"),
LogicalResult mlir::MlirOptMain(int argc, char **argv, llvm::StringRef toolName,
 
   if (failed(MlirOptMain(output->os(), std::move(file), passPipeline, registry,
                          splitInputFile, verifyDiagnostics, verifyPasses,
                          allowUnregisteredDialects, preloadDialectsInContext)))
     return failure();
 
   // Keep the output file if the invocation of MlirOptMain was successful.
   output->keep();
   return success();
 }