zkasm: rem

near · Nov 6, 2023 · ee4c079 · ee4c079
1 parent 87a848c
commit ee4c079
Show file tree

Hide file tree

Showing 9 changed files with 114 additions and 30 deletions.
diff --git a/cranelift/codegen/src/isa/zkasm/inst.isle b/cranelift/codegen/src/isa/zkasm/inst.isle
@@ -46,6 +46,11 @@
       (rs1 Reg)
       (rs2 Reg))
 
+    (RemArith
+      (rd WritableReg)
+      (rs1 Reg)
+      (rs2 Reg))
+
     ;; An load
     (Load
       (rd WritableReg)
@@ -707,6 +712,14 @@
             (_ Unit (emit (MInst.DivArith dst rs1 rs2))))
         dst))
 
+;; Helper for emitting the `rem` instruction.
+;; rd ← rs1 % rs2
+(decl zk_rem (XReg XReg) XReg)
+(rule (zk_rem rs1 rs2)
+      (let ((dst WritableXReg (temp_writable_xreg))
+            (_ Unit (emit (MInst.RemArith dst rs1 rs2))))
+        dst))
+
 ;; Helper for emitting the `divu` ("Divide Unsigned") instruction.
 ;; rd ← rs1 ÷ rs2
 (decl rv_divu (XReg XReg) XReg)

diff --git a/cranelift/codegen/src/isa/zkasm/inst/emit.rs b/cranelift/codegen/src/isa/zkasm/inst/emit.rs
@@ -489,9 +489,24 @@ impl MachInstEmit for Inst {
                 let rd = allocs.next_writable(rd);
                 // Same as in MulArith about D register
                 put_string("0 => D\n", sink);
-                put_string("0 => C\n", sink);
                 put_string("$${var _divArith = E / B}\n", sink);
+                put_string("$${var _remArith = E % B}\n", sink);
+                put_string("${_divArith} => A\n", sink);
+                put_string("${_remArith} => C\n", sink);
+                put_string("E:ARITH\n", sink);
+            }
+            &Inst::RemArith { 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);
+                // Same as in MulArith about D register
+                put_string("0 => D\n", sink);
+                put_string("$${var _divArith = E / B}\n", sink);
+                put_string("$${var _remArith = E % B}\n", sink);
                 put_string("${_divArith} => A\n", sink);
+                put_string("${_remArith} => C\n", sink);
                 put_string("E:ARITH\n", sink);
             }
             &Inst::AluRRR {

diff --git a/cranelift/codegen/src/isa/zkasm/inst/mod.rs b/cranelift/codegen/src/isa/zkasm/inst/mod.rs
@@ -288,6 +288,15 @@ fn zkasm_get_operands<F: Fn(VReg) -> VReg>(inst: &Inst, collector: &mut OperandC
             collector.reg_clobbers(clobbered);
             collector.reg_fixed_def(rd, a0());
         }
+        &Inst::RemArith { rd, rs1, rs2, .. } => {
+            collector.reg_fixed_use(rs1, e0());
+            collector.reg_fixed_use(rs2, b0());
+            let mut clobbered = PRegSet::empty();
+            clobbered.add(a0().to_real_reg().unwrap().into());
+            clobbered.add(d0().to_real_reg().unwrap().into());
+            collector.reg_clobbers(clobbered);
+            collector.reg_fixed_def(rd, c0());
+        }
         &Inst::Load { rd, from, .. } => {
             if let Some(r) = from.get_allocatable_register() {
                 collector.reg_use(r);
@@ -995,6 +1004,12 @@ impl Inst {
                 let rd_s = format_reg(rd.to_reg(), allocs);
                 format!("DivArith rd = {}, rs1 = {}, rs2 = {}", rd_s, rs1_s, rs2_s)
             }
+            &Inst::RemArith { 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!("RemArith rd = {}, rs1 = {}, rs2 = {}", rd_s, rs1_s, rs2_s)
+            }
             Inst::AddImm32 { rd, src1, src2 } => {
                 let rd = format_reg(rd.to_reg(), allocs);
                 format!("{src1} + {src2} => {rd};")

diff --git a/cranelift/codegen/src/isa/zkasm/lower.isle b/cranelift/codegen/src/isa/zkasm/lower.isle
@@ -126,35 +126,15 @@
 
 ;;;; Rules for `rem` ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 
-(rule -1 (lower (has_type (fits_in_16 ty) (urem x y)))
-  (let
-    ((y2 XReg (zext y ty $I64))
-      (_ InstOutput (gen_div_by_zero y2)))
-    (rv_remuw (zext x ty $I64) y2)))
-
-(rule -1 (lower (has_type (fits_in_16 ty) (srem x y)))
-  (let
-    ((y2 XReg (sext y ty $I64))
-      (_ InstOutput (gen_div_by_zero y2)))
-    (rv_remw (sext x ty $I64) y2)))
-
-(rule (lower (has_type $I32 (srem x y)))
-  (let
-    ((y2 XReg (sext y $I32 $I64))
-      (_ InstOutput (gen_div_by_zero y2)))
-   (rv_remw x y2)))
+(rule (lower (srem x y))
+  (zk_rem x y))
 
 (rule (lower (has_type $I32 (urem x y)))
   (let
     ((y2 XReg (zext y $I32 $I64))
         (_ InstOutput (gen_div_by_zero y2)))
     (rv_remuw x y2)))
 
-(rule (lower (has_type $I64 (srem x y)))
-  (let
-    ((_ InstOutput (gen_div_by_zero y)))
-    (rv_rem x y)))
-
 (rule (lower (has_type $I64 (urem x y)))
   (let
     ((_ InstOutput (gen_div_by_zero y)))

diff --git a/cranelift/filetests/src/test_zkasm.rs b/cranelift/filetests/src/test_zkasm.rs
@@ -211,5 +211,6 @@ mod tests {
         _should_fail_unreachable,
         i32_const,
         i64_const,
+        rem,
     }
 }
diff --git a/cranelift/zkasm_data/div.wat b/cranelift/zkasm_data/div.wat
@@ -1,9 +1,9 @@
 (module
  (import "env" "assert_eq" (func $assert_eq (param i32) (param i32)))
  (func $main
-	i32.const 1999999999
-	i32.const 64516129
+	i32.const 5
+	i32.const 2
 	i32.div_s
-	i32.const 31
+	i32.const 2
 	call $assert_eq)
  (start $main))
diff --git a/cranelift/zkasm_data/generated/div.zkasm b/cranelift/zkasm_data/generated/div.zkasm
@@ -10,14 +10,15 @@ function_1:
   D :MSTORE(SP - 2)
   E :MSTORE(SP - 3)
   B :MSTORE(SP - 4)
-  1999999999 => E
-  64516129 => B
+  5 => E
+  2 => B
   0 => D
-  0 => C
   $${var _divArith = E / B}
+  $${var _remArith = E % B}
   ${_divArith} => A
+  ${_remArith} => C
   E:ARITH
-  31 => B
+  2 => B
   B :ASSERT
   $ => C :MLOAD(SP - 1)
   $ => D :MLOAD(SP - 2)

diff --git a/cranelift/zkasm_data/generated/rem.zkasm b/cranelift/zkasm_data/generated/rem.zkasm
@@ -0,0 +1,45 @@
+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)
+  5 => E
+  2 => B
+  0 => D
+  $${var _divArith = E / B}
+  $${var _remArith = E % B}
+  ${_divArith} => A
+  ${_remArith} => C
+  E:ARITH
+  C => A
+  1 => B
+  B :ASSERT
+  8 => E
+  3 => B
+  0 => D
+  $${var _divArith = E / B}
+  $${var _remArith = E % B}
+  ${_divArith} => A
+  ${_remArith} => C
+  E:ARITH
+  C => A
+  2 => 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)
diff --git a/cranelift/zkasm_data/rem.wat b/cranelift/zkasm_data/rem.wat
@@ -0,0 +1,14 @@
+(module
+ (import "env" "assert_eq" (func $assert_eq (param i32) (param i32)))
+ (func $main
+	i32.const 5
+	i32.const 2
+	i32.rem_s
+	i32.const 1
+	call $assert_eq
+    i32.const 8
+    i32.const 3
+    i32.rem_s
+    i32.const 2
+    call $assert_eq)
+ (start $main))