[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

39

54–56

57

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

274

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

57

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

[mehdi_amini]

mlir/include/mlir/Support/MlirOptMain.h

57

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.

61

[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;
+
+using PassPipelineFn = llvm::function_ref<LogicalResult(
+    PassManager &pm,
+    llvm::function_ref<LogicalResult(const Twine &)> errorHandler)>;
 
 /// 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.
+/// - pipelineFn is the callback to the pass 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,
+                          PassPipelineFn pipelineFn, 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,
                           bool preloadDialectsInContext = false);
 
 /// Helper wrapper to return the result of MlirOptMain directly from main.

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 pipelineFn) {
   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();
static LogicalResult performActions(raw_ostream &os, bool verifyDiagnostics,
   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(pipelineFn(pm, errorHandler)))
     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 pipelineFn, 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, pipelineFn);
   }
 
   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);
+                       pipelineFn);
 
   // 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 pipelineFn,
                                 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,
+                               preloadDialectsInContext, pipelineFn, registry,
                                threadPool);
         },
         outputStream);
 
   return processBuffer(outputStream, std::move(buffer), verifyDiagnostics,
                        verifyPasses, allowUnregisteredDialects,
-                       preloadDialectsInContext, passPipeline, registry,
+                       preloadDialectsInContext, pipelineFn, registry,
                        threadPool);
 }
 
 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("-"));
LogicalResult mlir::MlirOptMain(int argc, char **argv, llvm::StringRef toolName,
   }
 
   auto output = openOutputFile(outputFilename, &errorMessage);
   if (!output) {
     llvm::errs() << errorMessage << "\n";
     return failure();
   }
 
-  if (failed(MlirOptMain(output->os(), std::move(file), passPipeline, registry,
-                         splitInputFile, verifyDiagnostics, verifyPasses,
-                         allowUnregisteredDialects, preloadDialectsInContext)))
+  if (failed(MlirOptMain(
+          output->os(), std::move(file),
+          [&](PassManager &pm,
+              function_ref<LogicalResult(const Twine &)> errorHandler) {
+            return passPipeline.addToPipeline(pm, errorHandler);
+          },
+          registry, splitInputFile, verifyDiagnostics, verifyPasses,
     return failure();
 
   // Keep the output file if the invocation of MlirOptMain was successful.
   output->keep();
   return success();
 }