Merge branch 'main' into viktar/update-zk-processor

near · Dec 1, 2023 · 556f9a3 · 556f9a3
2 parents 1df57d6 + bef30bd
commit 556f9a3
Show file tree

Hide file tree

Showing 35 changed files with 1,061 additions and 39 deletions.
diff --git a/cranelift/codegen/src/isa/zkasm/inst.isle b/cranelift/codegen/src/isa/zkasm/inst.isle
@@ -69,6 +69,16 @@
       (rs1 Reg)
       (rs2 Reg))
 
+    (UDivArith32
+      (rd WritableReg)
+      (rs1 Reg)
+      (rs2 Reg))
+
+    (UDivArith
+      (rd WritableReg)
+      (rs1 Reg)
+      (rs2 Reg))
+
     (RemArith
       (rd WritableReg)
       (rs1 Reg)
@@ -733,6 +743,18 @@
             (_ Unit (emit (MInst.MulArith32 dst rs1 rs2))))
         dst))
 
+(decl zk_divu (XReg XReg) XReg)
+(rule (zk_divu rs1 rs2)
+      (let ((dst WritableXReg (temp_writable_xreg))
+            (_ Unit (emit (MInst.UDivArith dst rs1 rs2))))
+        dst))
+
+(decl zk_divu_32 (XReg XReg) XReg)
+(rule (zk_divu_32 rs1 rs2)
+      (let ((dst WritableXReg (temp_writable_xreg))
+            (_ Unit (emit (MInst.UDivArith32 dst rs1 rs2))))
+        dst))
+
 (decl zk_div_32 (XReg XReg) XReg)
 (rule (zk_div_32 rs1 rs2)
       (let ((dst WritableXReg (temp_writable_xreg))

diff --git a/cranelift/codegen/src/isa/zkasm/inst/emit.rs b/cranelift/codegen/src/isa/zkasm/inst/emit.rs
@@ -679,6 +679,37 @@ impl MachInstEmit for Inst {
                 put_string("${E % B} => C\n", sink);
                 put_string("E:ARITH\n", sink);
             }
+            &Inst::UDivArith32 { rd, rs1, rs2 } => {
+                let rs1 = allocs.next(rs1);
+                let rs2 = allocs.next(rs2);
+                debug_assert_eq!(rs1, e0());
+                debug_assert_eq!(rs2, b0());
+                let rd = allocs.next_writable(rd);
+                // E / B => A
+                put_string("0 => D\n", sink);
+                put_string("${E / B} => A\n", sink);
+                put_string("${E % B} => C\n", sink);
+                put_string("E:ARITH\n", sink);
+
+                // A *= 2 ** 32
+                put_string("4294967296n => B\n", sink);
+                put_string("0 => C\n", sink);
+                // Intentionnaly ignore 64-bit overflow here
+                // (don't write _mulArith / 2 ** 64 => D)
+                // because if it is non-zero something is not OK
+                put_string("${A * B} => A :ARITH\n", sink);
+            }
+            &Inst::UDivArith { rd, rs1, rs2 } => {
+                let rs1 = allocs.next(rs1);
+                let rs2 = allocs.next(rs2);
+                debug_assert_eq!(rs1, e0());
+                debug_assert_eq!(rs2, b0());
+                let rd = allocs.next_writable(rd);
+                put_string("0 => D\n", sink);
+                put_string("${E / B} => A\n", sink);
+                put_string("${E % B} => C\n", sink);
+                put_string("E:ARITH\n", sink);
+            }
             // Rem32 is quite difficult opcode. Here we need calculate
             // ((op1 / 2**32) % (op2 / 2**32)) * 2**32
             &Inst::RemArith32 { rd, rs1, rs2 } => {

diff --git a/cranelift/codegen/src/isa/zkasm/inst/mod.rs b/cranelift/codegen/src/isa/zkasm/inst/mod.rs
@@ -346,6 +346,24 @@ fn zkasm_get_operands<F: Fn(VReg) -> VReg>(inst: &Inst, collector: &mut OperandC
             collector.reg_clobbers(clobbered);
             collector.reg_fixed_def(rd, a0());
         }
+        &Inst::UDivArith { rd, rs1, rs2, .. } => {
+            collector.reg_fixed_use(rs1, e0());
+            collector.reg_fixed_use(rs2, b0());
+            let mut clobbered = PRegSet::empty();
+            clobbered.add(c0().to_real_reg().unwrap().into());
+            clobbered.add(d0().to_real_reg().unwrap().into());
+            collector.reg_clobbers(clobbered);
+            collector.reg_fixed_def(rd, a0());
+        }
+        &Inst::UDivArith32 { rd, rs1, rs2, .. } => {
+            collector.reg_fixed_use(rs1, e0());
+            collector.reg_fixed_use(rs2, b0());
+            let mut clobbered = PRegSet::empty();
+            clobbered.add(c0().to_real_reg().unwrap().into());
+            clobbered.add(d0().to_real_reg().unwrap().into());
+            collector.reg_clobbers(clobbered);
+            collector.reg_fixed_def(rd, a0());
+        }
         &Inst::URemArith { rd, rs1, rs2, .. } => {
             collector.reg_fixed_use(rs1, e0());
             collector.reg_fixed_use(rs2, b0());
@@ -1085,6 +1103,21 @@ impl Inst {
                 let rd_s = format_reg(rd.to_reg(), allocs);
                 format!("MulArith32 rd = {}, rs1 = {}, rs2 = {}", rd_s, rs1_s, rs2_s)
             }
+            &Inst::UDivArith { rd, rs1, rs2 } => {
+                let rs1_s = format_reg(rs1, allocs);
+                let rs2_s = format_reg(rs2, allocs);
+                let rd_s = format_reg(rd.to_reg(), allocs);
+                format!("UDivArith rd = {}, rs1 = {}, rs2 = {}", rd_s, rs1_s, rs2_s)
+            }
+            &Inst::UDivArith32 { rd, rs1, rs2 } => {
+                let rs1_s = format_reg(rs1, allocs);
+                let rs2_s = format_reg(rs2, allocs);
+                let rd_s = format_reg(rd.to_reg(), allocs);
+                format!(
+                    "UDivArith32 rd = {}, rs1 = {}, rs2 = {}",
+                    rd_s, rs1_s, rs2_s
+                )
+            }
             &Inst::DivArith { rd, rs1, rs2 } => {
                 let rs1_s = format_reg(rs1, allocs);
                 let rs2_s = format_reg(rs2, allocs);

diff --git a/cranelift/codegen/src/isa/zkasm/lower.isle b/cranelift/codegen/src/isa/zkasm/lower.isle
@@ -98,22 +98,19 @@
 
 ;;;; Rules for `div` ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 
-(rule -1 (lower (has_type (fits_in_32 ty) (udiv x y)))
-  (let
-    ((y2 XReg (zext y ty $I64))
-      (_ InstOutput (gen_div_by_zero y2)))
-    (rv_divuw (zext x ty $I64) y2)))
+(rule (lower (has_type $I64 (udiv x y)))
+  (zk_divu x y))
+
+(rule (lower (has_type $I32 (udiv x y)))
+  (zk_divu_32 x y))
+
 
 (rule (lower (has_type $I64 (sdiv x y)))
   (zk_div x y))
 
 (rule (lower (has_type $I32 (sdiv x y)))
   (zk_div_32 x y))
 
-(rule (lower (has_type $I64 (udiv x y)))
-  (let
-    ((_ InstOutput (gen_div_by_zero y)))
-    (rv_divu x y)))
 
 ;;;; Rules for `rem` ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 

diff --git a/cranelift/zkasm_data/spectest/i32/generated/div_u_1.zkasm b/cranelift/zkasm_data/spectest/i32/generated/div_u_1.zkasm
@@ -0,0 +1,35 @@
+start:
+  zkPC + 2 => RR
+  :JMP(function_1)
+  :JMP(finalizeExecution)
+function_1:
+  SP + 1 => SP
+  RR :MSTORE(SP - 1)
+  SP + 4 => SP
+  C :MSTORE(SP - 1)
+  D :MSTORE(SP - 2)
+  E :MSTORE(SP - 3)
+  B :MSTORE(SP - 4)
+  4294967296n => E
+  4294967296n => B
+  0 => D
+  ${E / B} => A
+  ${E % B} => C
+  E:ARITH
+  4294967296n => B
+  0 => C
+  ${A * B} => A :ARITH
+  4294967296n => B
+  B :ASSERT
+  $ => C :MLOAD(SP - 1)
+  $ => D :MLOAD(SP - 2)
+  $ => E :MLOAD(SP - 3)
+  $ => B :MLOAD(SP - 4)
+  SP - 4 => SP
+  $ => RR :MLOAD(SP - 1)
+  SP - 1 => SP
+  :JMP(RR)
+finalizeExecution:
+  ${beforeLast()}  :JMPN(finalizeExecution)
+                   :JMP(start)
+INCLUDE "helpers/2-exp.zkasm"
diff --git a/cranelift/zkasm_data/spectest/i32/generated/div_u_10.zkasm b/cranelift/zkasm_data/spectest/i32/generated/div_u_10.zkasm
@@ -0,0 +1,35 @@
+start:
+  zkPC + 2 => RR
+  :JMP(function_1)
+  :JMP(finalizeExecution)
+function_1:
+  SP + 1 => SP
+  RR :MSTORE(SP - 1)
+  SP + 4 => SP
+  C :MSTORE(SP - 1)
+  D :MSTORE(SP - 2)
+  E :MSTORE(SP - 3)
+  B :MSTORE(SP - 4)
+  21474836480n => E
+  18446744065119617024n => B
+  0 => D
+  ${E / B} => A
+  ${E % B} => C
+  E:ARITH
+  4294967296n => B
+  0 => C
+  ${A * B} => A :ARITH
+  0n => B
+  B :ASSERT
+  $ => C :MLOAD(SP - 1)
+  $ => D :MLOAD(SP - 2)
+  $ => E :MLOAD(SP - 3)
+  $ => B :MLOAD(SP - 4)
+  SP - 4 => SP
+  $ => RR :MLOAD(SP - 1)
+  SP - 1 => SP
+  :JMP(RR)
+finalizeExecution:
+  ${beforeLast()}  :JMPN(finalizeExecution)
+                   :JMP(start)
+INCLUDE "helpers/2-exp.zkasm"
diff --git a/cranelift/zkasm_data/spectest/i32/generated/div_u_11.zkasm b/cranelift/zkasm_data/spectest/i32/generated/div_u_11.zkasm
@@ -0,0 +1,35 @@
+start:
+  zkPC + 2 => RR
+  :JMP(function_1)
+  :JMP(finalizeExecution)
+function_1:
+  SP + 1 => SP
+  RR :MSTORE(SP - 1)
+  SP + 4 => SP
+  C :MSTORE(SP - 1)
+  D :MSTORE(SP - 2)
+  E :MSTORE(SP - 3)
+  B :MSTORE(SP - 4)
+  18446744052234715136n => E
+  18446744065119617024n => B
+  0 => D
+  ${E / B} => A
+  ${E % B} => C
+  E:ARITH
+  4294967296n => B
+  0 => C
+  ${A * B} => A :ARITH
+  0n => B
+  B :ASSERT
+  $ => C :MLOAD(SP - 1)
+  $ => D :MLOAD(SP - 2)
+  $ => E :MLOAD(SP - 3)
+  $ => B :MLOAD(SP - 4)
+  SP - 4 => SP
+  $ => RR :MLOAD(SP - 1)
+  SP - 1 => SP
+  :JMP(RR)
+finalizeExecution:
+  ${beforeLast()}  :JMPN(finalizeExecution)
+                   :JMP(start)
+INCLUDE "helpers/2-exp.zkasm"
diff --git a/cranelift/zkasm_data/spectest/i32/generated/div_u_12.zkasm b/cranelift/zkasm_data/spectest/i32/generated/div_u_12.zkasm
@@ -0,0 +1,35 @@
+start:
+  zkPC + 2 => RR
+  :JMP(function_1)
+  :JMP(finalizeExecution)
+function_1:
+  SP + 1 => SP
+  RR :MSTORE(SP - 1)
+  SP + 4 => SP
+  C :MSTORE(SP - 1)
+  D :MSTORE(SP - 2)
+  E :MSTORE(SP - 3)
+  B :MSTORE(SP - 4)
+  30064771072n => E
+  12884901888n => B
+  0 => D
+  ${E / B} => A
+  ${E % B} => C
+  E:ARITH
+  4294967296n => B
+  0 => C
+  ${A * B} => A :ARITH
+  8589934592n => B
+  B :ASSERT
+  $ => C :MLOAD(SP - 1)
+  $ => D :MLOAD(SP - 2)
+  $ => E :MLOAD(SP - 3)
+  $ => B :MLOAD(SP - 4)
+  SP - 4 => SP
+  $ => RR :MLOAD(SP - 1)
+  SP - 1 => SP
+  :JMP(RR)
+finalizeExecution:
+  ${beforeLast()}  :JMPN(finalizeExecution)
+                   :JMP(start)
+INCLUDE "helpers/2-exp.zkasm"
diff --git a/cranelift/zkasm_data/spectest/i32/generated/div_u_13.zkasm b/cranelift/zkasm_data/spectest/i32/generated/div_u_13.zkasm
@@ -0,0 +1,35 @@
+start:
+  zkPC + 2 => RR
+  :JMP(function_1)
+  :JMP(finalizeExecution)
+function_1:
+  SP + 1 => SP
+  RR :MSTORE(SP - 1)
+  SP + 4 => SP
+  C :MSTORE(SP - 1)
+  D :MSTORE(SP - 2)
+  E :MSTORE(SP - 3)
+  B :MSTORE(SP - 4)
+  47244640256n => E
+  21474836480n => B
+  0 => D
+  ${E / B} => A
+  ${E % B} => C
+  E:ARITH
+  4294967296n => B
+  0 => C
+  ${A * B} => A :ARITH
+  8589934592n => B
+  B :ASSERT
+  $ => C :MLOAD(SP - 1)
+  $ => D :MLOAD(SP - 2)
+  $ => E :MLOAD(SP - 3)
+  $ => B :MLOAD(SP - 4)
+  SP - 4 => SP
+  $ => RR :MLOAD(SP - 1)
+  SP - 1 => SP
+  :JMP(RR)
+finalizeExecution:
+  ${beforeLast()}  :JMPN(finalizeExecution)
+                   :JMP(start)
+INCLUDE "helpers/2-exp.zkasm"
diff --git a/cranelift/zkasm_data/spectest/i32/generated/div_u_14.zkasm b/cranelift/zkasm_data/spectest/i32/generated/div_u_14.zkasm
@@ -0,0 +1,35 @@
+start:
+  zkPC + 2 => RR
+  :JMP(function_1)
+  :JMP(finalizeExecution)
+function_1:
+  SP + 1 => SP
+  RR :MSTORE(SP - 1)
+  SP + 4 => SP
+  C :MSTORE(SP - 1)
+  D :MSTORE(SP - 2)
+  E :MSTORE(SP - 3)
+  B :MSTORE(SP - 4)
+  73014444032n => E
+  30064771072n => B
+  0 => D
+  ${E / B} => A
+  ${E % B} => C
+  E:ARITH
+  4294967296n => B
+  0 => C
+  ${A * B} => A :ARITH
+  8589934592n => B
+  B :ASSERT
+  $ => C :MLOAD(SP - 1)
+  $ => D :MLOAD(SP - 2)
+  $ => E :MLOAD(SP - 3)
+  $ => B :MLOAD(SP - 4)
+  SP - 4 => SP
+  $ => RR :MLOAD(SP - 1)
+  SP - 1 => SP
+  :JMP(RR)
+finalizeExecution:
+  ${beforeLast()}  :JMPN(finalizeExecution)
+                   :JMP(start)
+INCLUDE "helpers/2-exp.zkasm"