Merge pull request #168 from mooori/opt-cond-br

perf: optimize zkASM for `CondBr` if target is a label
near · Jan 4, 2024 · ca441cc · ca441cc
2 parents 8c58058 + bb03e5d
commit ca441cc
Show file tree

Hide file tree

Showing 9 changed files with 30 additions and 78 deletions.
diff --git a/cranelift/codegen/src/isa/zkasm/inst/emit.rs b/cranelift/codegen/src/isa/zkasm/inst/emit.rs
@@ -932,22 +932,24 @@ impl MachInstEmit for Inst {
                 not_taken,
                 mut kind,
             } => {
+                // TODO(#179): The handling of `CondBr` might need to be refactored.
                 kind.rs1 = allocs.next(kind.rs1);
                 kind.rs2 = allocs.next(kind.rs2);
                 // TODO(akashin): Support other types of comparisons.
                 assert!(matches!(kind.kind, IntCC::NotEqual));
                 assert_eq!(kind.rs2, zero_reg());
                 match taken {
                     BranchTarget::Label(label) => {
-                        // WARNING: next two put_string MUST be in that order,
-                        // otherwise, if kind.rs1 == A it will lead to infinite loops.
-                        // Idk how to make kind.rs1 always not equal A.
-                        // At least, both A and B marked as clobbered in
-                        // mod.rs, so it will work.
-                        put_string(&format!("{} => B\n", reg_name(kind.rs1)), sink);
-                        put_string("0 => A\n", sink);
-                        put_string("$ => A :EQ\n", sink);
-                        put_string(&format!("A :JMPZ(label_{})\n", label.index()), sink);
+                        // Due to ISLE rules for `cond_br` the comparison result is in `kind.rs1`.
+                        // See #179 for details and how they might change in the future.
+                        //
+                        // Apart from that, this zkASM remains valid as long as the branch
+                        // condition effectively checks that `kind.rs1` contains a non-zero value.
+                        // This is guaranteed by the asserts above.
+                        put_string(
+                            &format!("{} :JMPNZ(label_{})\n", reg_name(kind.rs1), label.index()),
+                            sink,
+                        );
                     }
                     BranchTarget::ResolvedOffset(offset) => {
                         assert!(offset != 0);

diff --git a/cranelift/codegen/src/isa/zkasm/inst/mod.rs b/cranelift/codegen/src/isa/zkasm/inst/mod.rs
@@ -408,11 +408,6 @@ fn zkasm_get_operands<F: Fn(VReg) -> VReg>(inst: &Inst, collector: &mut OperandC
         &Inst::Jal { .. } => {}
         &Inst::CondBr { kind, .. } => {
             collector.reg_use(kind.rs1);
-            collector.reg_use(kind.rs2);
-            let mut clobbered = PRegSet::empty();
-            clobbered.add(a0().to_real_reg().unwrap().into());
-            clobbered.add(b0().to_real_reg().unwrap().into());
-            collector.reg_clobbers(clobbered);
         }
         &Inst::LoadExtName { rd, .. } => {
             collector.reg_def(rd);

diff --git a/cranelift/zkasm_data/benchmarks/fibonacci/generated/from_rust.zkasm b/cranelift/zkasm_data/benchmarks/fibonacci/generated/from_rust.zkasm
@@ -67,10 +67,7 @@ label_1_1:
   $ => A :AND
   A => C
   A => D
-  C => B
-  0 => A
-  $ => A :EQ
-  A :JMPZ(label_1_2)
+  C :JMPNZ(label_1_2)
   :JMP(label_1_3)
 label_1_2:
   C => D

diff --git a/cranelift/zkasm_data/benchmarks/fibonacci/generated/handwritten_wat.zkasm b/cranelift/zkasm_data/benchmarks/fibonacci/generated/handwritten_wat.zkasm
@@ -26,10 +26,7 @@ label_1_1:
   A :MSTORE(SP + 8)
   10000n => B  ;; LoadConst32
   $ => A :EQ
-  A => B
-  0 => A
-  $ => A :EQ
-  A :JMPZ(label_1_3)
+  A :JMPNZ(label_1_3)
   $ => A :MLOAD(SP)
   E => B
   :JMP(label_1_1)

diff --git a/cranelift/zkasm_data/benchmarks/fibonacci/state.csv b/cranelift/zkasm_data/benchmarks/fibonacci/state.csv
@@ -1,4 +1,4 @@
 Test,Status,Cycles
-from_rust,pass,53026
+from_rust,pass,50026
 handwritten,pass,50008
-handwritten_wat,pass,170023
+handwritten_wat,pass,140023
diff --git a/cranelift/zkasm_data/benchmarks/sha256/generated/from_rust.zkasm b/cranelift/zkasm_data/benchmarks/sha256/generated/from_rust.zkasm
@@ -504,10 +504,7 @@ function_1:
   $ => B :AND
   B :MSTORE(SP)
   0n => A  ;; LoadConst32
-  A => B
-  0 => A
-  $ => A :EQ
-  A :JMPZ(label_1_1)
+  A :JMPNZ(label_1_1)
   :JMP(label_1_2)
 label_1_1:
   :JMP(label_1_3)
@@ -539,10 +536,7 @@ label_1_3:
   $ => A :XOR
   8n => B  ;; LoadConst32
   $ => A :LT
-  A => B
-  0 => A
-  $ => A :EQ
-  A :JMPZ(label_1_5)
+  A :JMPNZ(label_1_5)
   160n => B  ;; LoadConst32
   $ => A :MLOAD(SP + 88)
   $ => C :ADD
@@ -1727,10 +1721,7 @@ function_3:
   0n => B  ;; LoadConst32
   $ => A :MLOAD(SP + 16)
   $ => A :EQ
-  A => B
-  0 => A
-  $ => A :EQ
-  A :JMPZ(label_3_1)
+  A :JMPNZ(label_3_1)
   :JMP(label_3_2)
 label_3_1:
   :JMP(label_3_7)
@@ -28109,10 +28100,7 @@ label_3_3:
   $ => B :MLOAD(SP + 9416)
   $ => A :EQ
   1 - A => A
-  A => B
-  0 => A
-  $ => A :EQ
-  A :JMPZ(label_3_4)
+  A :JMPNZ(label_3_4)
   :JMP(label_3_5)
 label_3_4:
   $ => A :MLOAD(SP + 112)
@@ -28210,10 +28198,7 @@ function_5:
   B => A
   C => B
   $ => A :LT
-  A => B
-  0 => A
-  $ => A :EQ
-  A :JMPZ(label_5_2)
+  A :JMPNZ(label_5_2)
   $ => E :MLOAD(SP + 16)
   $ => C :MLOAD(SP)
   C :MSTORE(SP + 16)
@@ -28236,10 +28221,7 @@ label_5_2:
   0n => B  ;; LoadConst32
   $ => A :MLOAD(SP + 24)
   $ => A :EQ
-  A => B
-  0 => A
-  $ => A :EQ
-  A :JMPZ(label_5_3)
+  A :JMPNZ(label_5_3)
   :JMP(label_5_4)
 label_5_3:
   $ => A :MLOAD(SP + 16)
@@ -28261,10 +28243,7 @@ label_5_5:
   E => A
   $ => B :MLOAD(SP + 40)
   $ => A :LT
-  A => B
-  0 => A
-  $ => A :EQ
-  A :JMPZ(label_5_6)
+  A :JMPNZ(label_5_6)
   :JMP(label_5_7)
 label_5_6:
   E => B
@@ -28304,10 +28283,7 @@ label_5_9:
   A => B
   C => A
   $ => A :SLT
-  A => B
-  0 => A
-  $ => A :EQ
-  A :JMPZ(label_5_10)
+  A :JMPNZ(label_5_10)
   :JMP(label_5_11)
 label_5_10:
   :JMP(label_5_16)
@@ -28340,10 +28316,7 @@ label_5_12:
   E => A
   $ => B :MLOAD(SP + 16)
   $ => A :LT
-  A => B
-  0 => A
-  $ => A :EQ
-  A :JMPZ(label_5_13)
+  A :JMPNZ(label_5_13)
   :JMP(label_5_14)
 label_5_13:
   E => B
@@ -28359,10 +28332,7 @@ label_5_17:
   0n => B  ;; LoadConst32
   E => A
   $ => A :EQ
-  A => B
-  0 => A
-  $ => A :EQ
-  A :JMPZ(label_5_18)
+  A :JMPNZ(label_5_18)
   :JMP(label_5_19)
 label_5_18:
   :JMP(label_5_24)
@@ -28390,10 +28360,7 @@ label_5_20:
   E => A
   $ => B :MLOAD(SP + 24)
   $ => A :LT
-  A => B
-  0 => A
-  $ => A :EQ
-  A :JMPZ(label_5_21)
+  A :JMPNZ(label_5_21)
   :JMP(label_5_22)
 label_5_21:
   E => B

diff --git a/cranelift/zkasm_data/generated/counter.zkasm b/cranelift/zkasm_data/generated/counter.zkasm
@@ -22,10 +22,7 @@ label_1_1:
   A => E
   10n => B  ;; LoadConst32
   $ => A :EQ
-  A => B
-  0 => A
-  $ => A :EQ
-  A :JMPZ(label_1_3)
+  A :JMPNZ(label_1_3)
   :JMP(label_1_1)
 label_1_3:
   10n => B  ;; LoadConst32

diff --git a/cranelift/zkasm_data/state.csv b/cranelift/zkasm_data/state.csv
@@ -3,7 +3,7 @@ _should_fail_unreachable,runtime error,
 add,pass,19
 add_func,pass,30
 and,pass,17
-counter,pass,153
+counter,pass,123
 div,pass,26
 eqz,pass,26
 fibonacci,pass,224

diff --git a/docs/zkasm/test_summary.csv b/docs/zkasm/test_summary.csv
@@ -1,9 +1,6 @@
 Suite path,Passing count,Total count,Total cycles
-<<<<<<< HEAD,,,
-=======,,,
->>>>>>> main,,,
-cranelift/zkasm_data,27,29,1027
-cranelift/zkasm_data/benchmarks/fibonacci,3,3,273057
+cranelift/zkasm_data,27,29,997
+cranelift/zkasm_data/benchmarks/fibonacci,3,3,240057
 cranelift/zkasm_data/benchmarks/sha256,0,1,0
 cranelift/zkasm_data/spectest/conversions,12,24,180
 cranelift/zkasm_data/spectest/i32,287,364,7988