From 2505128202ed12c42ea495d2c1bcfc6264541cdc Mon Sep 17 00:00:00 2001
From: nils m <nils@risczero.com>
Date: Tue, 4 Feb 2025 14:16:55 -0800
Subject: [PATCH] 5 runs of PO2=19 completed in 8.994s, avg=1.799s, 291481.380
 cycles/sec, 1457.162 keccak/sec

  __device__ static constexpr size_t getFpBufferSize() { return 1554; }
---
 zirgen/Dialect/ByteCode/Transforms/Schedule.cpp    | 3 ++-
 zirgen/Dialect/ByteCode/Transforms/ScheduleZll.cpp | 4 ++--
 zirgen/Dialect/Zll/Transforms/BUILD.bazel          | 1 +
 zirgen/Dialect/Zll/Transforms/Passes.h             | 1 +
 zirgen/Dialect/Zll/Transforms/Passes.td            | 9 +++++++++
 zirgen/Main/gen_zirgen.cpp                         | 3 ++-
 6 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/zirgen/Dialect/ByteCode/Transforms/Schedule.cpp b/zirgen/Dialect/ByteCode/Transforms/Schedule.cpp
index 4f747f0b..4b1835e5 100644
--- a/zirgen/Dialect/ByteCode/Transforms/Schedule.cpp
+++ b/zirgen/Dialect/ByteCode/Transforms/Schedule.cpp
@@ -327,7 +327,7 @@ const OpInfo& Scheduler::getOpInfo(Operation* topOp) {
       return lhs.second->queuePriority() < rhs.second->queuePriority();
     });
 
-    // Registers that have mee saved so far by operands making use of them.
+    // Registers that have been saved so far by operands making use of them.
     ssize_t adjustedRegs = 0;
 
     size_t maxLive = 0;
@@ -404,3 +404,4 @@ void scheduleBlock(Block* block, ScheduleInterface& scheduleInterface) {
 }
 
 } // namespace zirgen::ByteCode
+
diff --git a/zirgen/Dialect/ByteCode/Transforms/ScheduleZll.cpp b/zirgen/Dialect/ByteCode/Transforms/ScheduleZll.cpp
index ac35fe75..7c62165f 100644
--- a/zirgen/Dialect/ByteCode/Transforms/ScheduleZll.cpp
+++ b/zirgen/Dialect/ByteCode/Transforms/ScheduleZll.cpp
@@ -484,9 +484,9 @@ struct CloneActiveZllPass : public impl::CloneActiveZllBase<CloneActiveZllPass>
             activeVal.print(llvm::dbgs(), asmState);
             llvm::dbgs() << "\n";
           });
-/*          llvm::errs() << "Spilling ";
+          llvm::errs() << "Spilling ";
           activeVal.print(llvm::errs(), asmState);
-          llvm::errs() << "\n";*/
+          llvm::errs() << "\n";
           active.update(it, ActiveState::SpilledDoNotClone);
           ++opSpillVals;
         } else if (state == ActiveState::Used) {
diff --git a/zirgen/Dialect/Zll/Transforms/BUILD.bazel b/zirgen/Dialect/Zll/Transforms/BUILD.bazel
index 7bd090c8..5f50389f 100644
--- a/zirgen/Dialect/Zll/Transforms/BUILD.bazel
+++ b/zirgen/Dialect/Zll/Transforms/BUILD.bazel
@@ -42,6 +42,7 @@ cc_library(
         "MakePolynomial.cpp",
         "MultiplyToIf.cpp",
         "PassDetail.h",
+        "ReorderConstraints.cpp",
         "SortForReproducibility.cpp",
         "SplitStage.cpp",
     ],
diff --git a/zirgen/Dialect/Zll/Transforms/Passes.h b/zirgen/Dialect/Zll/Transforms/Passes.h
index dd6daf61..74406ce0 100644
--- a/zirgen/Dialect/Zll/Transforms/Passes.h
+++ b/zirgen/Dialect/Zll/Transforms/Passes.h
@@ -34,6 +34,7 @@ std::unique_ptr<mlir::Pass> createBalancedSplitPass(size_t maxOps = 1000);
 std::unique_ptr<mlir::Pass> createSortForReproducibilityPass();
 std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>> createExtractPolyMixPass();
 std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>> createAnnotatePolyMixPass();
+std::unique_ptr<mlir::Pass> createReorderConstraintsPass();
 
 // Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
diff --git a/zirgen/Dialect/Zll/Transforms/Passes.td b/zirgen/Dialect/Zll/Transforms/Passes.td
index 847a0364..b38e4f39 100644
--- a/zirgen/Dialect/Zll/Transforms/Passes.td
+++ b/zirgen/Dialect/Zll/Transforms/Passes.td
@@ -106,4 +106,13 @@ def AnnotatePolyMix : Pass<"annotate-poly-mix", "mlir::func::FuncOp"> {
   let constructor = "zirgen::Zll::createAnnotatePolyMixPass()";
 }
 
+def ReorderConstraints : Pass<"reorder-constraints"> {
+  let summary = "Optimize by reordering zll.eqz constraints";
+  let description = [{
+     ReorderConstraints reorders constraints, and puts groups of constraints
+     inside a `if(1)` such that they can be better optimized by CSE.
+  }];
+  let constructor = "zirgen::Zll::createReorderConstraintsPass()";
+}
+
 #endif // ZLL_TRANSFORM_PASSES
diff --git a/zirgen/Main/gen_zirgen.cpp b/zirgen/Main/gen_zirgen.cpp
index 34f27aa1..9a7d8598 100644
--- a/zirgen/Main/gen_zirgen.cpp
+++ b/zirgen/Main/gen_zirgen.cpp
@@ -1,4 +1,4 @@
-b// Copyright 2025 RISC Zero, Inc.
+// Copyright 2025 RISC Zero, Inc.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -406,6 +406,7 @@ int main(int argc, char* argv[]) {
     checkPasses.addPass(mlir::createCanonicalizerPass());
     checkPasses.addPass(zirgen::dsl::createTopologicalShufflePass());
   }
+  checkPasses.addPass(zirgen::Zll::createReorderConstraintsPass());
 
   if (failed(pm.run(typedModule.value()))) {
     llvm::errs() << "an internal compiler error occurred while lowering this module:\n";