Parallelize passes that prepare IR to be run in interpreter

risc0 · Sep 5, 2024 · 5e9f2d1 · 5e9f2d1
1 parent bb1eea1
commit 5e9f2d1
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 7 deletions.
diff --git a/zirgen/Dialect/ZStruct/Transforms/Passes.h b/zirgen/Dialect/ZStruct/Transforms/Passes.h
@@ -22,7 +22,7 @@ namespace zirgen::ZStruct {
 
 // Pass constructors
 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> createOptimizeLayoutPass();
-std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> createUnrollPass();
+std::unique_ptr<mlir::Pass> createUnrollPass();
 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> createExpandLayoutPass();
 std::unique_ptr<mlir::Pass> createInlineLayoutPass();
 

diff --git a/zirgen/Dialect/ZStruct/Transforms/Passes.td b/zirgen/Dialect/ZStruct/Transforms/Passes.td
@@ -23,7 +23,7 @@ def OptimizeLayout : Pass<"optimize-layout", "mlir::ModuleOp"> {
   let constructor = "zirgen::ZStruct::createOptimizeLayoutPass()";
 }
 
-def Unroll : Pass<"unroll", "mlir::ModuleOp"> {
+def Unroll : Pass<"unroll"> {
   let summary = "Unroll zhlt map and reduce";
   let description = [{
     Removes zhlt.map and zhlt.reduce instructions by unrolling the loops.

diff --git a/zirgen/Dialect/ZStruct/Transforms/Unroll.cpp b/zirgen/Dialect/ZStruct/Transforms/Unroll.cpp
@@ -36,7 +36,7 @@ struct UnrollPass : public UnrollBase<UnrollPass> {
 
 } // End namespace
 
-std::unique_ptr<OperationPass<ModuleOp>> createUnrollPass() {
+std::unique_ptr<Pass> createUnrollPass() {
   return std::make_unique<UnrollPass>();
 }
 

diff --git a/zirgen/dsl/driver.cpp b/zirgen/dsl/driver.cpp
@@ -287,15 +287,17 @@ int runTests(mlir::ModuleOp& module) {
   mlir::MLIRContext& context = *module.getContext();
   // Set all the symbols to private
   mlir::PassManager pm(&context);
+  applyDefaultTimingPassManagerCLOptions(pm);
   if (failed(applyPassManagerCLOptions(pm))) {
     llvm::errs() << "Pass manager does not agree with command line options.\n";
     return 1;
   }
   pm.enableVerifier(true);
   pm.addPass(mlir::createInlinerPass());
-  pm.addPass(zirgen::ZStruct::createUnrollPass());
-  pm.addPass(mlir::createCanonicalizerPass());
-  pm.addPass(mlir::createCSEPass());
+  mlir::OpPassManager& opm = pm.nest<zirgen::Zhlt::StepFuncOp>();
+  opm.addPass(zirgen::ZStruct::createUnrollPass());
+  opm.addPass(mlir::createCanonicalizerPass());
+  opm.addPass(mlir::createCSEPass());
   if (failed(pm.run(module))) {
     llvm::errs() << "an internal compiler error occurred while inlining the tests:\n";
     module.print(llvm::errs());
@@ -455,7 +457,6 @@ int main(int argc, char* argv[]) {
 
   llvm::SourceMgr sourceManager;
   sourceManager.setIncludeDirs(includeDirs);
-  context.disableMultithreading();
 
   mlir::SourceMgrDiagnosticHandler sourceMgrHandler(sourceManager, &context);
   openMainFile(sourceManager, inputFilename);
@@ -498,6 +499,7 @@ int main(int argc, char* argv[]) {
   }
 
   mlir::PassManager pm(&context);
+  applyDefaultTimingPassManagerCLOptions(pm);
   if (failed(applyPassManagerCLOptions(pm))) {
     llvm::errs() << "Pass manager does not agree with command line options.\n";
     return 1;