Fix SHA256 benchmark

- Implements memory loads and stores - Add a test with many args
near · Jan 26, 2024 · 0b18571 · 0b18571
1 parent ca0d6dc
commit 0b18571
Show file tree

Hide file tree

Showing 23 changed files with 1,805 additions and 1,791 deletions.
diff --git a/cranelift/codegen/src/isa/zkasm/inst/args.rs b/cranelift/codegen/src/isa/zkasm/inst/args.rs
@@ -802,6 +802,19 @@ impl LoadOP {
             Self::Fld => 0b011,
         }
     }
+    pub(crate) fn width(self) -> u32 {
+        match self {
+            Self::I8 => 1,
+            Self::I16 => 2,
+            Self::I32 => 4,
+            Self::U32 => 4,
+            Self::U8 => 1,
+            Self::U16 => 2,
+            Self::U64 => 8,
+            Self::Flw => unimplemented!(),
+            Self::Fld => unimplemented!(),
+        }
+    }
 }
 
 impl StoreOP {
@@ -843,6 +856,16 @@ impl StoreOP {
             Self::Fsd => 0b011,
         }
     }
+    pub(crate) fn width(self) -> u32 {
+        match self {
+            Self::I8 => 1,
+            Self::I16 => 2,
+            Self::I32 => 4,
+            Self::I64 => 8,
+            Self::Fsw => unimplemented!(),
+            Self::Fsd => unimplemented!(),
+        }
+    }
 }
 
 impl IntSelectOP {

diff --git a/cranelift/codegen/src/isa/zkasm/inst/emit.rs b/cranelift/codegen/src/isa/zkasm/inst/emit.rs
@@ -590,15 +590,52 @@ impl MachInstEmit for Inst {
                 let rd = allocs.next_writable(rd);
                 match from {
                     AMode::RegOffset(r, ..) => {
+                        let rem = offset % 8;
+                        let width = op.width() as i64;
+                        if rem + width > 8 {
+                            put_string(&format!(";; FIXME Read spans two slots\n"), sink);
+                        }
+
                         debug_assert_eq!(r, e0());
+                        put_string(
+                            &format!(
+                                "${{ ({}) / 8 }} => {}\n",
+                                access_reg_with_offset(r, offset),
+                                reg_name(r)
+                            ),
+                            sink,
+                        );
                         put_string(
                             &format!(
                                 "$ => {} :MLOAD(MEM:{})\n",
                                 reg_name(rd.to_reg()),
-                                access_reg_with_offset(r, offset)
+                                reg_name(r)
                             ),
                             sink,
                         );
+                        if rem > 0 {
+                            put_string(
+                                &format!(
+                                    "${{ {} >> {} }} => {}\n",
+                                    reg_name(rd.to_reg()),
+                                    8 * rem,
+                                    reg_name(rd.to_reg())
+                                ),
+                                sink,
+                            );
+                        }
+
+                        if width < 8 {
+                            put_string(
+                                &format!(
+                                    "${{ {} & ((1 << {}) - 1) }} => {}\n",
+                                    reg_name(rd.to_reg()),
+                                    8 * width,
+                                    reg_name(rd.to_reg()),
+                                ),
+                                sink,
+                            );
+                        }
                     }
                     AMode::SPOffset(..) | AMode::NominalSPOffset(..) | AMode::FPOffset(..) => {
                         assert_eq!(offset % 8, 0);
@@ -631,12 +668,77 @@ impl MachInstEmit for Inst {
                         debug_assert_eq!(r, e0());
                         put_string(
                             &format!(
-                                "{} :MSTORE(MEM:{})\n",
+                                "${{ ({}) / 8 }} => {}\n",
+                                access_reg_with_offset(r, offset),
+                                reg_name(r)
+                            ),
+                            sink,
+                        );
+                        let rem = offset % 8;
+                        let width = op.width() as i64;
+
+                        if op == StoreOP::I64 && rem == 0 {
+                            put_string(
+                                &format!("{} :MSTORE(MEM:{})\n", reg_name(src), reg_name(r)),
+                                sink,
+                            );
+                            return;
+                        }
+
+                        // MEM = 1111111122222222
+                        // VAL = XXXXXXXX
+                        // Store:
+                        //      1111111122222222
+                        //        XXXXXXXX
+                        //      ->offset
+                        if width < 8 {
+                            put_string(
+                                &format!(
+                                    "${{ {} & ((1 << {}) - 1) }} => {}\n",
+                                    reg_name(src),
+                                    8 * width,
+                                    reg_name(src),
+                                ),
+                                sink,
+                            );
+                        }
+                        put_string(
+                            &format!("$ => {} :MLOAD(MEM:{})\n", reg_name(d0()), reg_name(r)),
+                            sink,
+                        );
+                        put_string(
+                            &format!(
+                                "${{ (D & ~((1 << {} - 1) << {})) | ({} << {}) }} :MSTORE(MEM:{})\n",
+                                8 * width,
+                                8 * rem,
                                 reg_name(src),
-                                access_reg_with_offset(r, offset)
+                                8 * rem,
+                                reg_name(r),
                             ),
                             sink,
                         );
+                        // Such writes needs to be split across two slots.
+                        if rem + width > 8 {
+                            let rem = rem + width - 8;
+                            put_string(
+                                &format!(
+                                    "$ => {} :MLOAD(MEM:{})\n",
+                                    reg_name(d0()),
+                                    access_reg_with_offset(r, 1)
+                                ),
+                                sink,
+                            );
+                            put_string(
+                                &format!(
+                                    "${{ (D & ~(1 << {} - 1)) | ({} & (1 << {} - 1)) }} :MSTORE(MEM:{})\n",
+                                    8 * rem,
+                                    reg_name(src),
+                                    8 * rem,
+                                    access_reg_with_offset(r, 1),
+                                ),
+                                sink,
+                            );
+                        }
                     }
                     AMode::SPOffset(..) | AMode::NominalSPOffset(..) | AMode::FPOffset(..) => {
                         assert_eq!(offset % 8, 0);

diff --git a/cranelift/codegen/src/isa/zkasm/inst/mod.rs b/cranelift/codegen/src/isa/zkasm/inst/mod.rs
@@ -342,7 +342,10 @@ fn zkasm_get_operands<F: Fn(VReg) -> VReg>(inst: &Inst, collector: &mut OperandC
             if let Some(r) = to.get_allocatable_register() {
                 collector.reg_fixed_use(r, e0());
             }
-            collector.reg_use(src);
+            let mut clobbered = PRegSet::empty();
+            clobbered.add(d0().to_real_reg().unwrap().into());
+            collector.reg_clobbers(clobbered);
+            collector.reg_late_use(src);
         }
         &Inst::Args { ref args } => {
             for arg in args {

diff --git a/cranelift/filetests/src/test_zkasm.rs b/cranelift/filetests/src/test_zkasm.rs
@@ -60,7 +60,7 @@ mod tests {
 
         // Generate const data segments definitions.
         for (offset, data) in data_segments {
-            program.push(format!("  {offset} => E"));
+            program.push(format!("  {} => E", offset / 8));
             // Each slot stores 8 consecutive u8 numbers, with earlier addresses stored in lower
             // bits.
             for (i, chunk) in data.chunks(8).enumerate() {
@@ -449,5 +449,7 @@ mod tests {
         i64_mul_overflows,
         i64_rem,
         memory_i32,
+        store_i32,
+        many_args,
     }
 }