feat: add LOADW4 STOREW4 SHINTW4

extensions/native/circuit/src/adapters/loadstore_native_adapter.rs

@@ -253,20 +253,22 @@ impl<F: PrimeField32, const NUM_CELLS: usize> VmAdapterChip<F>
 
         let (data_read_as, data_write_as) = {
             match local_opcode {
-                LOADW => (e, d),
-                STOREW | SHINTW => (d, e),
+                LOADW | LOADW4 => (e, d),
+                STOREW | STOREW4 | SHINTW | SHINTW4 => (d, e),
             }
         };
         let (data_read_ptr, data_write_ptr) = {
             match local_opcode {
-                LOADW => (read_cell.1 + b, a),
-                STOREW | SHINTW => (a, read_cell.1 + b),
+                LOADW | LOADW4 => (read_cell.1 + b, a),
+                STOREW | STOREW4 | SHINTW | SHINTW4 => (a, read_cell.1 + b),
             }
         };
 
         let data_read = match local_opcode {
-            SHINTW => None,
-            LOADW | STOREW => Some(memory.read::<NUM_CELLS>(data_read_as, data_read_ptr)),
+            SHINTW | SHINTW4 => None,
+            LOADW | LOADW4 | STOREW | STOREW4 => {
+                Some(memory.read::<NUM_CELLS>(data_read_as, data_read_ptr))
+            }
         };
         let record = NativeLoadStoreReadRecord {
             pointer_read: read_cell.0,

extensions/native/circuit/src/extension.rs

@@ -72,6 +72,7 @@ pub struct Native;
 #[derive(ChipUsageGetter, Chip, InstructionExecutor, From, AnyEnum)]
 pub enum NativeExecutor<F: PrimeField32> {
     LoadStore(NativeLoadStoreChip<F, 1>),
+    BlockLoadStore(NativeLoadStoreChip<F, 4>),
     BranchEqual(NativeBranchEqChip<F>),
     Jal(NativeJalChip<F>),
     FieldArithmetic(FieldArithmeticChip<F>),
@@ -115,7 +116,38 @@ impl<F: PrimeField32> VmExtension<F> for Native {
 
         inventory.add_executor(
             load_store_chip,
-            NativeLoadStoreOpcode::iter().map(VmOpcode::with_default_offset),
+            [
+                NativeLoadStoreOpcode::LOADW,
+                NativeLoadStoreOpcode::STOREW,
+                NativeLoadStoreOpcode::SHINTW,
+            ]
+            .iter()
+            .map(|&opcode| VmOpcode::with_default_offset(opcode)),
+        )?;
+
+        let mut block_load_store_chip = NativeLoadStoreChip::<F, 4>::new(
+            NativeLoadStoreAdapterChip::new(
+                execution_bus,
+                program_bus,
+                memory_bridge,
+                NativeLoadStoreOpcode::default_offset(),
+            ),
+            NativeLoadStoreCoreChip::new(NativeLoadStoreOpcode::default_offset()),
+            offline_memory.clone(),
+        );
+        block_load_store_chip
+            .core
+            .set_streams(builder.streams().clone());
+
+        inventory.add_executor(
+            block_load_store_chip,
+            [
+                NativeLoadStoreOpcode::LOADW4,
+                NativeLoadStoreOpcode::STOREW4,
+                NativeLoadStoreOpcode::SHINTW4,
+            ]
+            .iter()
+            .map(|&opcode| VmOpcode::with_default_offset(opcode)),
         )?;
 
         let branch_equal_chip = NativeBranchEqChip::new(

extensions/native/circuit/src/loadstore/core.rs

@@ -19,7 +19,6 @@ use openvm_stark_backend::{
 };
 use serde::{Deserialize, Serialize};
 use serde_big_array::BigArray;
-use strum::IntoEnumIterator;
 
 use super::super::adapters::loadstore_native_adapter::NativeLoadStoreInstruction;
 
@@ -84,12 +83,27 @@ where
         });
         builder.assert_bool(is_valid.clone());
 
-        let expected_opcode = flags.iter().zip(NativeLoadStoreOpcode::iter()).fold(
-            AB::Expr::ZERO,
-            |acc, (flag, opcode)| {
+        let expected_opcode = flags
+            .iter()
+            .zip(if NUM_CELLS == 1 {
+                [
+                    NativeLoadStoreOpcode::LOADW,
+                    NativeLoadStoreOpcode::STOREW,
+                    NativeLoadStoreOpcode::SHINTW,
+                ]
+            } else if NUM_CELLS == 4 {
+                [
+                    NativeLoadStoreOpcode::LOADW4,
+                    NativeLoadStoreOpcode::STOREW4,
+                    NativeLoadStoreOpcode::SHINTW4,
+                ]
+            } else {
+                panic!("Unsupported number of cells: {}", NUM_CELLS);
+            })
+            .fold(AB::Expr::ZERO, |acc, (flag, opcode)| {
                 acc + (*flag).into() * AB::Expr::from_canonical_usize(opcode.as_usize())
-            },
-        ) + AB::Expr::from_canonical_usize(self.offset);
+            })
+            + AB::Expr::from_canonical_usize(self.offset);
 
         AdapterAirContext {
             to_pc: None,
@@ -145,7 +159,9 @@ where
             NativeLoadStoreOpcode::from_usize(opcode.local_opcode_idx(self.air.offset));
         let (pointer_read, data_read) = reads.into();
 
-        let data_write = if local_opcode == NativeLoadStoreOpcode::SHINTW {
+        let data_write = if (NUM_CELLS == 1 && local_opcode == NativeLoadStoreOpcode::SHINTW)
+            || (NUM_CELLS == 4 && local_opcode == NativeLoadStoreOpcode::SHINTW4)
+        {
             let mut streams = self.streams.get().unwrap().lock().unwrap();
             if streams.hint_stream.len() < NUM_CELLS {
                 return Err(ExecutionError::HintOutOfBounds { pc: from_pc });
@@ -174,9 +190,17 @@ where
 
     fn generate_trace_row(&self, row_slice: &mut [F], record: Self::Record) {
         let cols: &mut NativeLoadStoreCoreCols<_, NUM_CELLS> = row_slice.borrow_mut();
-        cols.is_loadw = F::from_bool(record.opcode == NativeLoadStoreOpcode::LOADW);
-        cols.is_storew = F::from_bool(record.opcode == NativeLoadStoreOpcode::STOREW);
-        cols.is_shintw = F::from_bool(record.opcode == NativeLoadStoreOpcode::SHINTW);
+        if NUM_CELLS == 1 {
+            cols.is_loadw = F::from_bool(record.opcode == NativeLoadStoreOpcode::LOADW);
+            cols.is_storew = F::from_bool(record.opcode == NativeLoadStoreOpcode::STOREW);
+            cols.is_shintw = F::from_bool(record.opcode == NativeLoadStoreOpcode::SHINTW);
+        } else if NUM_CELLS == 4 {
+            cols.is_loadw = F::from_bool(record.opcode == NativeLoadStoreOpcode::LOADW4);
+            cols.is_storew = F::from_bool(record.opcode == NativeLoadStoreOpcode::STOREW4);
+            cols.is_shintw = F::from_bool(record.opcode == NativeLoadStoreOpcode::SHINTW4);
+        } else {
+            panic!("Unsupported number of cells: {}", NUM_CELLS);
+        }
 
         cols.pointer_read = record.pointer_read;
         cols.data_read = record.data_read;

extensions/native/compiler/src/asm/compiler.rs

@@ -1104,33 +1104,17 @@ impl<F: PrimeField32 + TwoAdicField, EF: ExtensionField<F> + TwoAdicField> AsmCo
     }
 
     fn load_ext(&mut self, val: Ext<F, EF>, addr: i32, offset: F, debug_info: Option<DebugInfo>) {
-        for i in 0..EF::D {
-            self.push(
-                AsmInstruction::LoadFI(
-                    val.fp() + i as i32,
-                    addr,
-                    F::from_canonical_usize(i),
-                    F::ONE,
-                    offset,
-                ),
-                debug_info.clone(),
-            )
-        }
+        self.push(
+            AsmInstruction::LoadEI(val.fp(), addr, F::ZERO, F::ONE, offset),
+            debug_info.clone(),
+        );
     }
 
     fn store_ext(&mut self, val: Ext<F, EF>, addr: i32, offset: F, debug_info: Option<DebugInfo>) {
-        for i in 0..EF::D {
-            self.push(
-                AsmInstruction::StoreFI(
-                    val.fp() + i as i32,
-                    addr,
-                    F::from_canonical_usize(i),
-                    F::ONE,
-                    offset,
-                ),
-                debug_info.clone(),
-            )
-        }
+        self.push(
+            AsmInstruction::StoreEI(val.fp(), addr, F::ZERO, F::ONE, offset),
+            debug_info.clone(),
+        );
     }
 
     fn add_ext_exti(

extensions/native/compiler/src/asm/instruction.rs

@@ -12,11 +12,21 @@ pub enum AsmInstruction<F, EF> {
     /// Load a value from the address stored at src(fp) into dst(fp) with given index and offset.
     LoadFI(i32, i32, F, F, F),
 
+    /// Load extension word (dst, src, var_index, size, offset).
+    ///
+    /// Load an extension from the address stored at src(fp) into dst(fp) with given index and offset.
+    LoadEI(i32, i32, F, F, F),
+
     /// Store word (val, addr, var_index, size, offset)
     ///
     /// Store a value from val(fp) into the address stored at addr(fp) with given index and offset.
     StoreFI(i32, i32, F, F, F),
 
+    /// Store extension word (val, addr, var_index, size, offset)
+    ///
+    /// Store an extension from val(fp) into the address stored at addr(fp) with given index and offset.
+    StoreEI(i32, i32, F, F, F),
+
     /// Set dst = imm.
     ImmF(i32, F),
 
@@ -132,6 +142,9 @@ pub enum AsmInstruction<F, EF> {
     /// Stores the next hint stream word into value stored at addr + value.
     StoreHintWordI(i32, F),
 
+    /// Stores the next hint stream ext into value stored at addr + value.
+    StoreHintExtI(i32, F),
+
     /// Publish(val, index).
     Publish(i32, i32),
 
@@ -154,13 +167,27 @@ impl<F: PrimeField32, EF: ExtensionField<F>> AsmInstruction<F, EF> {
                     dst, src, var_index, size, offset
                 )
             }
+            AsmInstruction::LoadEI(dst, src, var_index, size, offset) => {
+                write!(
+                    f,
+                    "lei   ({})fp, ({})fp, {}, {}, {}",
+                    dst, src, var_index, size, offset
+                )
+            }
             AsmInstruction::StoreFI(dst, src, var_index, size, offset) => {
                 write!(
                     f,
                     "swi   ({})fp, ({})fp, {}, {}, {}",
                     dst, src, var_index, size, offset
                 )
             }
+            AsmInstruction::StoreEI(dst, src, var_index, size, offset) => {
+                write!(
+                    f,
+                    "sei   ({})fp, ({})fp, {}, {}, {}",
+                    dst, src, var_index, size, offset
+                )
+            }
             AsmInstruction::ImmF(dst, src) => {
                 write!(f, "imm   ({})fp, ({})", dst, src)
             }
@@ -315,6 +342,9 @@ impl<F: PrimeField32, EF: ExtensionField<F>> AsmInstruction<F, EF> {
             AsmInstruction::StoreHintWordI(dst, offset) => {
                 write!(f, "shintw ({})fp {}", dst, offset)
             }
+            AsmInstruction::StoreHintExtI(dst, offset) => {
+                write!(f, "shinte ({})fp {}", dst, offset)
+            }
             AsmInstruction::Publish(val, index) => {
                 write!(f, "commit ({})fp ({})fp", val, index)
             }

extensions/native/compiler/src/conversion/mod.rs

@@ -342,6 +342,17 @@ fn convert_instruction<F: PrimeField32, EF: ExtensionField<F>>(
                 AS::Native,
             ),
         ],
+        AsmInstruction::LoadEI(dst, src, index, size, offset) => vec![
+            // mem[dst] <- mem[mem[src] + index * size + offset]
+            inst(
+                options.opcode_with_offset(NativeLoadStoreOpcode::LOADW4),
+                i32_f(dst),
+                index * size + offset,
+                i32_f(src),
+                AS::Native,
+                AS::Native,
+            ),
+        ],
         AsmInstruction::StoreFI(val, addr, index, size, offset) => vec![
             // mem[mem[addr] + index * size + offset] <- mem[val]
             inst(
@@ -353,6 +364,17 @@ fn convert_instruction<F: PrimeField32, EF: ExtensionField<F>>(
                 AS::Native,
             ),
         ],
+        AsmInstruction::StoreEI(val, addr, index, size, offset) => vec![
+            // mem[mem[addr] + index * size + offset] <- mem[val]
+            inst(
+                options.opcode_with_offset(NativeLoadStoreOpcode::STOREW4),
+                i32_f(val),
+                index * size + offset,
+                i32_f(addr),
+                AS::Native,
+                AS::Native,
+            ),
+        ],
         AsmInstruction::Jump(dst, label) => {
             vec![
                 // pc <- labels[label], mem[dst] <- pc
@@ -496,6 +518,14 @@ fn convert_instruction<F: PrimeField32, EF: ExtensionField<F>>(
             AS::Native,
             AS::Native,
         )],
+        AsmInstruction::StoreHintExtI(val, offset) => vec![inst(
+            options.opcode_with_offset(NativeLoadStoreOpcode::SHINTW4),
+            F::ZERO,
+            offset,
+            i32_f(val),
+            AS::Native,
+            AS::Native,
+        )],
         AsmInstruction::PrintV(..) | AsmInstruction::PrintF(..) | AsmInstruction::PrintE(..) => {
             if options.compile_prints {
                 convert_print_instruction(instruction, options)

extensions/native/compiler/src/lib.rs

@@ -47,6 +47,9 @@ pub enum NativeLoadStoreOpcode {
     STOREW,
     /// Instruction to write the next hint word into memory.
     SHINTW,
+    LOADW4,
+    STOREW4,
+    SHINTW4,
 }
 
 #[derive(Copy, Clone, Debug, UsizeOpcode)]