perf: Do less conversion between u32 and F (#1182)

openvm-org · Jan 8, 2025 · 5d7c7fc · 5d7c7fc
1 parent 6fdbf70
commit 5d7c7fc
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 17 deletions.
diff --git a/extensions/rv32im/circuit/src/adapters/loadstore.rs b/extensions/rv32im/circuit/src/adapters/loadstore.rs
@@ -131,7 +131,7 @@ pub struct Rv32LoadStoreReadRecord<F: Field> {
 
     pub imm: F,
     pub imm_sign: bool,
-    pub mem_ptr_limbs: [F; 2],
+    pub mem_ptr_limbs: [u32; 2],
     pub mem_as: F,
 }
 
@@ -408,7 +408,7 @@ impl<F: PrimeField32> VmAdapterChip<F> for Rv32LoadStoreAdapterChip<F> {
                 read: read_record.0,
                 imm: c,
                 imm_sign: imm_sign == 1,
-                mem_ptr_limbs: mem_ptr_limbs.map(F::from_canonical_u32),
+                mem_ptr_limbs,
                 mem_as: e,
             },
         ))
@@ -431,13 +431,8 @@ impl<F: PrimeField32> VmAdapterChip<F> for Rv32LoadStoreAdapterChip<F> {
         let (write_id, _) = match local_opcode {
             STOREW | STOREH | STOREB => {
                 let ptr = read_record.mem_ptr_limbs[0]
-                    + read_record.mem_ptr_limbs[1]
-                        * F::from_canonical_u32(1 << (RV32_CELL_BITS * 2));
-                memory.write(
-                    e,
-                    F::from_canonical_u32(ptr.as_canonical_u32() & 0xfffffffc),
-                    output.writes[0],
-                )
+                    + read_record.mem_ptr_limbs[1] * (1 << (RV32_CELL_BITS * 2));
+                memory.write(e, F::from_canonical_u32(ptr & 0xfffffffc), output.writes[0])
             }
             LOADW | LOADB | LOADH | LOADBU | LOADHU => memory.write(d, a, output.writes[0]),
         };
@@ -474,7 +469,7 @@ impl<F: PrimeField32> VmAdapterChip<F> for Rv32LoadStoreAdapterChip<F> {
             aux_cols_factory.make_read_aux_cols(memory.record_by_id(read_record.read));
         adapter_cols.imm = read_record.imm;
         adapter_cols.imm_sign = F::from_bool(read_record.imm_sign);
-        adapter_cols.mem_ptr_limbs = read_record.mem_ptr_limbs;
+        adapter_cols.mem_ptr_limbs = read_record.mem_ptr_limbs.map(F::from_canonical_u32);
         let write = memory.record_by_id(write_record.write_id);
         adapter_cols.write_base_aux = aux_cols_factory
             .make_write_aux_cols::<RV32_REGISTER_NUM_LIMBS>(write)

diff --git a/extensions/rv32im/circuit/src/jalr/core.rs b/extensions/rv32im/circuit/src/jalr/core.rs
@@ -51,7 +51,7 @@ pub struct Rv32JalrCoreRecord<F> {
     pub rs1_data: [F; RV32_REGISTER_NUM_LIMBS],
     pub rd_data: [F; RV32_REGISTER_NUM_LIMBS - 1],
     pub to_pc_least_sig_bit: F,
-    pub to_pc_limbs: [F; 2],
+    pub to_pc_limbs: [u32; 2],
     pub imm_sign: F,
 }
 
@@ -234,11 +234,9 @@ where
         let mask = (1 << 15) - 1;
         let to_pc_least_sig_bit = rs1_val.wrapping_add(imm_extended) & 1;
 
-        let to_pc_limbs = array::from_fn(|i| F::from_canonical_u32((to_pc >> (1 + i * 15)) & mask));
-        self.range_checker_chip
-            .add_count(to_pc_limbs[0].as_canonical_u32(), 15);
-        self.range_checker_chip
-            .add_count(to_pc_limbs[1].as_canonical_u32(), 14);
+        let to_pc_limbs = array::from_fn(|i| ((to_pc >> (1 + i * 15)) & mask));
+        self.range_checker_chip.add_count(to_pc_limbs[0], 15);
+        self.range_checker_chip.add_count(to_pc_limbs[1], 14);
 
         let rd_data = rd_data.map(F::from_canonical_u32);
 
@@ -270,7 +268,7 @@ where
         core_cols.rd_data = record.rd_data;
         core_cols.rs1_data = record.rs1_data;
         core_cols.to_pc_least_sig_bit = record.to_pc_least_sig_bit;
-        core_cols.to_pc_limbs = record.to_pc_limbs;
+        core_cols.to_pc_limbs = record.to_pc_limbs.map(F::from_canonical_u32);
         core_cols.imm_sign = record.imm_sign;
         core_cols.is_valid = F::ONE;
     }