diff --git a/mlir/include/mlir-c/Pass.h b/mlir/include/mlir-c/Pass.h index 35db138305d1e2..0ab2c29bf3f777 100644 --- a/mlir/include/mlir-c/Pass.h +++ b/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); diff --git a/mlir/include/mlir/Pass/PassManager.h b/mlir/include/mlir/Pass/PassManager.h index d9bab431e2e0cc..e2d78823c835ad 100644 --- a/mlir/include/mlir/Pass/PassManager.h +++ b/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 diff --git a/mlir/lib/Bindings/Python/Pass.cpp b/mlir/lib/Bindings/Python/Pass.cpp index a68421b61641f6..e19eb450634ac1 100644 --- a/mlir/lib/Bindings/Python/Pass.cpp +++ b/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) { diff --git a/mlir/lib/CAPI/IR/Pass.cpp b/mlir/lib/CAPI/IR/Pass.cpp index d242baae99c086..0ae054c7c639f9 100644 --- a/mlir/lib/CAPI/IR/Pass.cpp +++ b/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); } diff --git a/mlir/lib/Pass/IRPrinting.cpp b/mlir/lib/Pass/IRPrinting.cpp index 9ffda6402cc07a..75f3e35b890200 100644 --- a/mlir/lib/Pass/IRPrinting.cpp +++ b/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 counters; }; +/// Print a pass pipeline like `builtin.module(func.func(cse))` +/// from a list of scopes and the pass. +template +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(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 file = openOutputFile(path, &error); + if (!file) { + llvm::errs() << "Error opening output file " << path << ": " << error + << "\n"; + return; + } + + SmallVector 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(outputDir)); +} diff --git a/mlir/lib/Pass/PassManagerOptions.cpp b/mlir/lib/Pass/PassManagerOptions.cpp index dd119a75f40696..3bdfb16336edb3 100644 --- a/mlir/lib/Pass/PassManagerOptions.cpp +++ b/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 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); diff --git a/mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi b/mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi index c072d5e0fb86f3..b91296d44a41ce 100644 --- a/mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi +++ b/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: ... diff --git a/mlir/test/Pass/reproducer-before-all.mlir b/mlir/test/Pass/reproducer-before-all.mlir new file mode 100644 index 00000000000000..bec7e257e59e67 --- /dev/null +++ b/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