Add -mlir-reproducer-before-all (#402)

* Add -mlir-reproducer-before-all

Author: mgehre-amd

mlir/include/mlir-c/Pass.h

@@ -78,6 +78,11 @@ mlirPassManagerRunOnOp(MlirPassManager passManager, MlirOperation op);
 MLIR_CAPI_EXPORTED void
 mlirPassManagerEnableIRPrinting(MlirPassManager passManager);
 
+/// Enable lir-reproducer-before-all.
+MLIR_CAPI_EXPORTED void
+mlirPassManagerEnableReproducerBeforeAll(MlirPassManager passManager,
+                                         MlirStringRef outputDir);
+
 /// Enable / disable verify-each.
 MLIR_CAPI_EXPORTED void
 mlirPassManagerEnableVerifier(MlirPassManager passManager, bool enable);

mlir/include/mlir/Pass/PassManager.h

@@ -423,6 +423,10 @@ class PassManager : public OpPassManager {
       llvm::StringRef printTreeDir = ".pass_manager_output",
       OpPrintingFlags opPrintingFlags = OpPrintingFlags());
 
+  /// Dump a reproducer before each pass into a file in the given output
+  /// directory.
+  void enableReproducerBeforeAll(llvm::StringRef outputDir);
+
   //===--------------------------------------------------------------------===//
   // Pass Timing
 

mlir/lib/Bindings/Python/Pass.cpp

@@ -78,6 +78,14 @@ void mlir::python::populatePassManagerSubmodule(py::module &m) {
             mlirPassManagerEnableIRPrinting(passManager.get());
           },
           "Enable mlir-print-ir-after-all.")
+      .def(
+          "enable_reproducer_before_all",
+          [](PyPassManager &passManager, const std::string &outputDir) {
+            mlirPassManagerEnableReproducerBeforeAll(
+                passManager.get(),
+                mlirStringRefCreate(outputDir.data(), outputDir.size()));
+          },
+          "Enable mlir-reproducer-before-all.")
       .def(
           "enable_verifier",
           [](PyPassManager &passManager, bool enable) {

mlir/lib/CAPI/IR/Pass.cpp

@@ -48,6 +48,11 @@ void mlirPassManagerEnableIRPrinting(MlirPassManager passManager) {
   return unwrap(passManager)->enableIRPrinting();
 }
 
+void mlirPassManagerEnableReproducerBeforeAll(MlirPassManager passManager,
+                                              MlirStringRef outputDir) {
+  return unwrap(passManager)->enableReproducerBeforeAll(unwrap(outputDir));
+}
+
 void mlirPassManagerEnableVerifier(MlirPassManager passManager, bool enable) {
   unwrap(passManager)->enableVerifier(enable);
 }

mlir/lib/Pass/IRPrinting.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "PassDetail.h"
+#include "mlir/IR/AsmState.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Support/FileUtilities.h"
@@ -345,6 +346,74 @@ struct FileTreeIRPrinterConfig : public PassManager::IRPrinterConfig {
   llvm::DenseMap<Operation *, unsigned> counters;
 };
 
+/// Print a pass pipeline like `builtin.module(func.func(cse))`
+/// from a list of scopes and the pass.
+template <typename RangeT>
+void printAsPassPipeline(RangeT scopes, Pass *pass, raw_ostream &os) {
+  // Add pass scopes like 'builtin.module(emitc.tu('
+  for (OperationName scope : scopes)
+    os << scope << "(";
+  pass->printAsTextualPipeline(os);
+  for (OperationName _ : scopes)
+    os << ")";
+}
+
+/// A pass instrumentation to dump the IR before each pass into
+/// numbered files.
+/// It includes a mlir_reproducer info to rerun the pass.
+class ReproducerBeforeAll : public PassInstrumentation {
+public:
+  ReproducerBeforeAll(mlir::StringRef outputDir) : outputDir(outputDir) {}
+  void runBeforePass(Pass *pass, Operation *op) override;
+
+  std::string outputDir;
+
+  uint32_t counter = 0;
+};
+
+void ReproducerBeforeAll::runBeforePass(Pass *pass, Operation *op) {
+  // Skip adator passes (which adopt FuncOp passes to ModuleOp pass managers).
+  if (isa<OpToOpPassAdaptor>(pass))
+    return;
+
+  llvm::SmallString<128> path(outputDir);
+  if (failed(createDirectoryOrPrintErr(path)))
+    return;
+
+  // Open output file.
+  std::string fileName =
+      llvm::formatv("{0,0+2}_{1}.mlir", counter++, pass->getArgument());
+  llvm::sys::path::append(path, fileName);
+
+  std::string error;
+  std::unique_ptr<llvm::ToolOutputFile> file = openOutputFile(path, &error);
+  if (!file) {
+    llvm::errs() << "Error opening output file " << path << ": " << error
+                 << "\n";
+    return;
+  }
+
+  SmallVector<OperationName> scopes;
+  scopes.push_back(op->getName());
+  while (Operation *parentOp = op->getParentOp()) {
+    scopes.push_back(parentOp->getName());
+    op = parentOp;
+  }
+
+  std::string pipelineStr;
+  llvm::raw_string_ostream passOS(pipelineStr);
+  printAsPassPipeline(llvm::reverse(scopes), pass, passOS);
+
+  AsmState state(op);
+  state.attachResourcePrinter("mlir_reproducer",
+                              [&](Operation *op, AsmResourceBuilder &builder) {
+                                builder.buildString("pipeline", pipelineStr);
+                                builder.buildBool("disable_threading", true);
+                                builder.buildBool("verify_each", true);
+                              });
+  op->print(file->os(), state);
+  file->keep();
+}
 } // namespace
 
 /// Add an instrumentation to print the IR before and after pass execution,
@@ -383,3 +452,11 @@ void PassManager::enableIRPrintingToFileTree(
       printModuleScope, printAfterOnlyOnChange, printAfterOnlyOnFailure,
       opPrintingFlags, printTreeDir));
 }
+
+/// Add an instrumentation to print the IR before and after pass execution.
+void PassManager::enableReproducerBeforeAll(StringRef outputDir) {
+  if (getContext()->isMultithreadingEnabled())
+    llvm::report_fatal_error("IR printing can't be setup on a pass-manager "
+                             "without disabling multi-threading first.");
+  addInstrumentation(std::make_unique<ReproducerBeforeAll>(outputDir));
+}

mlir/lib/Pass/PassManagerOptions.cpp

@@ -64,6 +64,10 @@ struct PassManagerOptions {
                      "tree rooted at this directory. Use in conjunction with "
                      "mlir-print-ir-* flags")};
 
+  llvm::cl::opt<std::string> reproducerBeforeAllDir{
+      "mlir-reproducer-before-all",
+      llvm::cl::desc("Save a reproducer before each pass to this directory")};
+
   /// Add an IR printing instrumentation if enabled by any 'print-ir' flags.
   void addPrinterInstrumentation(PassManager &pm);
 
@@ -151,6 +155,9 @@ LogicalResult mlir::applyPassManagerCLOptions(PassManager &pm) {
     pm.enableCrashReproducerGeneration(options->reproducerFile,
                                        options->localReproducer);
 
+  if (!options->reproducerBeforeAllDir.empty())
+    pm.enableReproducerBeforeAll(options->reproducerBeforeAllDir);
+
   // Enable statistics dumping.
   if (options->passStatistics)
     pm.enableStatistics(options->passStatisticsDisplayMode);

mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi

@@ -17,6 +17,7 @@ class PassManager:
     def _CAPICreate(self) -> object: ...
     def _testing_release(self) -> None: ...
     def enable_ir_printing(self) -> None: ...
+    def enable_reproducer_before_all(self, output_dir: str) -> None: ...
     def enable_verifier(self, enable: bool) -> None: ...
     @staticmethod
     def parse(pipeline: str, context: Optional[_ir.Context] = None) -> PassManager: ...

mlir/test/Pass/reproducer-before-all.mlir

@@ -0,0 +1,26 @@
+// RUN: rm -rf %t || true
+// RUN: mlir-opt %s -mlir-disable-threading -mlir-reproducer-before-all=%t \
+// RUN:   -pass-pipeline='builtin.module(canonicalize,cse,func.func(canonicalize))'
+// RUN: FileCheck %s -input-file=%t/00_canonicalize.mlir --check-prefixes CHECK0
+// RUN: FileCheck %s -input-file=%t/01_cse.mlir --check-prefixes CHECK1
+// RUN: FileCheck %s -input-file=%t/02_canonicalize.mlir --check-prefixes CHECK2
+
+builtin.module @outer {
+  func.func @symA() {
+    return
+  }
+}
+
+// CHECK0:      module @outer {
+// CHECK0:      {-#
+// CHECK0-NEXT: external_resources: {
+// CHECK0-NEXT:     mlir_reproducer: {
+// CHECK0-NEXT:       pipeline: "builtin.module(canonicalize
+// CHECK0-NEXT:       disable_threading: true,
+// CHECK0-NEXT:       verify_each: true
+// CHECK0-NEXT:     }
+// CHECK0-NEXT:   }
+// CHECK0-NEXT: #-}
+
+// CHECK1:       pipeline: "builtin.module(cse
+// CHECK2:       pipeline: "builtin.module(func.func(canonicalize