Merge pull request bytecodealliance#124 from near/viktar/i32rotl

zkasm: i32.rotl

cranelift/codegen/src/isa/zkasm/inst.isle

@@ -28,6 +28,16 @@
       (rs1 Reg)
       (rs2 Reg))
 
+    (Rotl64
+      (rd WritableReg)
+      (rs1 Reg)
+      (rs2 Reg))
+
+    (Rotl32
+      (rd WritableReg)
+      (rs1 Reg)
+      (rs2 Reg))
+
     (Shl64
       (rd WritableReg)
       (rs1 Reg)
@@ -728,6 +738,19 @@
         dst))
 
 
+(decl zk_rotl (XReg XReg) XReg)
+(rule (zk_rotl rs1 rs2)
+      (let ((dst WritableXReg (temp_writable_xreg))
+            (_ Unit (emit (MInst.Rotl64 dst rs1 rs2))))
+        dst))
+
+(decl zk_rotl_32 (XReg XReg) XReg)
+(rule (zk_rotl_32 rs1 rs2)
+      (let ((dst WritableXReg (temp_writable_xreg))
+            (_ Unit (emit (MInst.Rotl32 dst rs1 rs2))))
+        dst))
+
+
 ;; Helper for emitting the `mul` instruction.
 ;; rd ← rs1 × rs2
 (decl zk_mul (XReg XReg) XReg)

cranelift/codegen/src/isa/zkasm/inst/emit.rs

@@ -508,6 +508,69 @@ impl MachInstEmit for Inst {
                     sink,
                 );
             }
+            &Inst::Rotl64 { rd, rs1, rs2 } => {
+                unimplemented!("Rotl64");
+            }
+            &Inst::Rotl32 { rd, rs1, rs2 } => {
+                // rotl(num, amount) = num << amount | num >> (bitness - amount)
+                let rs1 = allocs.next(rs1);
+                let rs2 = allocs.next(rs2);
+                debug_assert_eq!(rs1, a0());
+                debug_assert_eq!(rs2, e0());
+
+                // a (in A) -- number, e (in E) -- rot amount.
+                // It is important to do e := e % 32. Shifts we use here do it inside.
+
+                put_string("A :MSTORE(SP)\n", sink); // a on SP
+                put_string("E :MSTORE(SP + 1)\n", sink); // e on SP + 1
+
+                // Shl and shr also use stack, so we need increase SP
+                put_string("SP + 2 => SP\n", sink);
+
+                Inst::Shl32 {
+                    rd: Writable::from_reg(e0()),
+                    rs1: a0(),
+                    rs2: e0(),
+                }
+                .emit(&[], sink, emit_info, state); // now a << e in E.
+
+                put_string(";; shl\n", sink);
+
+                put_string("SP - 2 => SP\n", sink);
+
+                put_string("$ => C :MLOAD(SP)\n", sink); // a in C
+                put_string("E :MSTORE(SP)\n", sink); // (a << e) on SP
+                put_string("$ => E :MLOAD(SP + 1)\n", sink); // e in E
+
+                put_string("137438953472n => A\n", sink);
+                put_string("E => B\n", sink);
+                put_string("$ => E :SUB\n", sink); // 32 - e in E
+
+                put_string("C => A\n", sink); // a in A
+
+                put_string("SP + 1 => SP\n", sink);
+
+                Inst::Shru32 {
+                    rd: Writable::from_reg(a0()),
+                    rs1: a0(),
+                    rs2: e0(),
+                }
+                .emit(&[], sink, emit_info, state); // now a >> (32 - e) in A
+
+                put_string("SP - 1 => SP\n", sink);
+
+                put_string(";; shr\n", sink);
+
+                put_string("$ => B :MLOAD(SP)\n", sink); // (a << e) in B
+
+                Inst::AluRRR {
+                    alu_op: AluOPRRR::Or,
+                    rd: Writable::from_reg(a0()),
+                    rs1: a0(),
+                    rs2: b0(),
+                }
+                .emit(&[], sink, emit_info, state); // now a >> (32 - e) | (a << e) in A
+            }
             &Inst::Shru64 { rd, rs1, rs2 } => {
                 let rs1 = allocs.next(rs1);
                 let rs2 = allocs.next(rs2);

cranelift/codegen/src/isa/zkasm/inst/mod.rs

@@ -267,6 +267,29 @@ fn zkasm_get_operands<F: Fn(VReg) -> VReg>(inst: &Inst, collector: &mut OperandC
             collector.reg_fixed_use(rs2, b0());
             collector.reg_def(rd);
         }
+        &Inst::Rotl64 { rd, rs1, rs2, .. } => {
+            unimplemented!("Rotl64");
+            // collector.reg_fixed_use(rs1, a0());
+            // collector.reg_fixed_use(rs2, e0());
+            // let mut clobbered = PRegSet::empty();
+            // clobbered.add(d0().to_real_reg().unwrap().into());
+            // clobbered.add(c0().to_real_reg().unwrap().into());
+            // clobbered.add(b0().to_real_reg().unwrap().into());
+            // clobbered.add(e0().to_real_reg().unwrap().into());
+            // collector.reg_clobbers(clobbered);
+            // collector.reg_fixed_def(rd, a0());
+        }
+        &Inst::Rotl32 { rd, rs1, rs2, .. } => {
+            collector.reg_fixed_use(rs1, a0());
+            collector.reg_fixed_use(rs2, e0());
+            let mut clobbered = PRegSet::empty();
+            clobbered.add(d0().to_real_reg().unwrap().into());
+            clobbered.add(c0().to_real_reg().unwrap().into());
+            clobbered.add(b0().to_real_reg().unwrap().into());
+            clobbered.add(e0().to_real_reg().unwrap().into());
+            collector.reg_clobbers(clobbered);
+            collector.reg_fixed_def(rd, a0());
+        }
         &Inst::Shru64 { rd, rs1, rs2, .. } => {
             collector.reg_fixed_use(rs1, a0());
             collector.reg_fixed_use(rs2, e0());
@@ -1067,6 +1090,18 @@ impl Inst {
                     }
                 }
             }
+            &Inst::Rotl32 { 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!("Rotl32 rd = {}, rs1 = {}, rs2 = {}", rd_s, rs1_s, rs2_s)
+            }
+            &Inst::Rotl64 { 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!("Rotl64 rd = {}, rs1 = {}, rs2 = {}", rd_s, rs1_s, rs2_s)
+            }
             &Inst::Shl32 { rd, rs1, rs2 } => {
                 let rs1_s = format_reg(rs1, allocs);
                 let rs2_s = format_reg(rs2, allocs);

cranelift/codegen/src/isa/zkasm/lower.isle

@@ -209,8 +209,11 @@
   (rv_sra x (value_regs_get y 0)))
 
 ;;;; Rules for `rotl` ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
-(rule (lower (has_type (fits_in_64 ty) (rotl x y)))
-  (lower_rotl ty (zext x ty $I64) (value_regs_get y 0)))
+(rule (lower (has_type $I64 (rotl x y)))
+  (zk_rotl x y))
+
+(rule (lower (has_type $I32 (rotl x y)))
+  (zk_rotl_32 x y))
 
 ;;;; Rules for `rotr` ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 (rule (lower (has_type (fits_in_64 ty) (rotr x y)))

cranelift/zkasm_data/spectest/i32/generated/rotl_1.zkasm

@@ -0,0 +1,103 @@
+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 => A
+  4294967296n => E
+  A :MSTORE(SP)
+  E :MSTORE(SP + 1)
+  SP + 2 => SP
+  A :MSTORE(SP)
+  0 => D
+  4294967296n => B
+  ${E / B} => A
+  ${E % B} => C
+  E:ARITH
+  A => E
+  32 => B
+  0 => D
+  ${E / B} => A
+  ${E % B} => C
+  E:ARITH
+  $ => A :MLOAD(SP)
+  C => E
+;;NEED_INCLUDE: 2-exp
+  zkPC + 2 => RR
+    :JMP(@two_power + E)
+  0 => D
+  0 => C
+  $${var _mulShlArith = A * B}
+  ${_mulShlArith / 18446744073709551616} => D
+  ${_mulShlArith % 18446744073709551616} => E :ARITH
+  ;; shl
+  SP - 2 => SP
+  $ => C :MLOAD(SP)
+  E :MSTORE(SP)
+  $ => E :MLOAD(SP + 1)
+  137438953472n => A
+  E => B
+  $ => E :SUB
+  C => A
+  SP + 1 => SP
+  E :MSTORE(SP)
+  0 => D
+  A => E
+  4294967296n => B
+  ${E / B} => A
+  ${E % B} => C
+  E:ARITH
+  $ => E :MLOAD(SP)
+  A :MSTORE(SP)
+  0 => D
+  4294967296n => B
+  ${E / B} => A
+  ${E % B} => C
+  E:ARITH
+  A => E
+  32 => B
+  0 => D
+  ${E / B} => A
+  ${E % B} => C
+  E:ARITH
+  $ => A :MLOAD(SP)
+  C => E
+;;NEED_INCLUDE: 2-exp
+  zkPC + 2 => RR
+    :JMP(@two_power + E)
+  A => E
+  0 => D
+  ${E / B} => A
+  ${E % B} => C
+  E :ARITH
+  0 => C
+  4294967296n => B
+  ${A * B} => E
+  0 => D
+  E :ARITH
+  E => A
+  SP - 1 => SP
+  ;; shr
+  $ => B :MLOAD(SP)
+  $ => A :OR
+  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"

cranelift/zkasm_data/spectest/i32/generated/rotl_10.zkasm

@@ -0,0 +1,103 @@
+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)
+  8546350908852666368n => A
+  18446743992105172992n => E
+  A :MSTORE(SP)
+  E :MSTORE(SP + 1)
+  SP + 2 => SP
+  A :MSTORE(SP)
+  0 => D
+  4294967296n => B
+  ${E / B} => A
+  ${E % B} => C
+  E:ARITH
+  A => E
+  32 => B
+  0 => D
+  ${E / B} => A
+  ${E % B} => C
+  E:ARITH
+  $ => A :MLOAD(SP)
+  C => E
+;;NEED_INCLUDE: 2-exp
+  zkPC + 2 => RR
+    :JMP(@two_power + E)
+  0 => D
+  0 => C
+  $${var _mulShlArith = A * B}
+  ${_mulShlArith / 18446744073709551616} => D
+  ${_mulShlArith % 18446744073709551616} => E :ARITH
+  ;; shl
+  SP - 2 => SP
+  $ => C :MLOAD(SP)
+  E :MSTORE(SP)
+  $ => E :MLOAD(SP + 1)
+  137438953472n => A
+  E => B
+  $ => E :SUB
+  C => A
+  SP + 1 => SP
+  E :MSTORE(SP)
+  0 => D
+  A => E
+  4294967296n => B
+  ${E / B} => A
+  ${E % B} => C
+  E:ARITH
+  $ => E :MLOAD(SP)
+  A :MSTORE(SP)
+  0 => D
+  4294967296n => B
+  ${E / B} => A
+  ${E % B} => C
+  E:ARITH
+  A => E
+  32 => B
+  0 => D
+  ${E / B} => A
+  ${E % B} => C
+  E:ARITH
+  $ => A :MLOAD(SP)
+  C => E
+;;NEED_INCLUDE: 2-exp
+  zkPC + 2 => RR
+    :JMP(@two_power + E)
+  A => E
+  0 => D
+  ${E / B} => A
+  ${E % B} => C
+  E :ARITH
+  0 => C
+  4294967296n => B
+  ${A * B} => E
+  0 => D
+  E :ARITH
+  E => A
+  SP - 1 => SP
+  ;; shr
+  $ => B :MLOAD(SP)
+  $ => A :OR
+  6312901892695392256n => 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"