openvm_rv32im_circuit/loadstore/
core.rs

1use std::{
2    array,
3    borrow::{Borrow, BorrowMut},
4    fmt::Debug,
5};
6
7use openvm_circuit::{
8    arch::*,
9    system::memory::{online::TracingMemory, MemoryAuxColsFactory},
10};
11use openvm_circuit_primitives::{
12    AlignedBorrow, AlignedBytesBorrow, ColumnsAir, StructReflection, StructReflectionHelper,
13};
14use openvm_instructions::{
15    instruction::Instruction, program::DEFAULT_PC_STEP, riscv::RV32_REGISTER_NUM_LIMBS, LocalOpcode,
16};
17use openvm_rv32im_transpiler::Rv32LoadStoreOpcode::{self, *};
18use openvm_stark_backend::{
19    interaction::InteractionBuilder,
20    p3_air::{AirBuilder, BaseAir},
21    p3_field::{Field, PrimeCharacteristicRing, PrimeField32},
22    BaseAirWithPublicValues,
23};
24
25use crate::adapters::{LoadStoreInstruction, Rv32LoadStoreAdapterFiller};
26
27#[derive(Debug, Clone, Copy)]
28enum InstructionOpcode {
29    LoadW0,
30    LoadHu0,
31    LoadHu2,
32    LoadBu0,
33    LoadBu1,
34    LoadBu2,
35    LoadBu3,
36    StoreW0,
37    StoreH0,
38    StoreH2,
39    StoreB0,
40    StoreB1,
41    StoreB2,
42    StoreB3,
43}
44
45use InstructionOpcode::*;
46
47/// LoadStore Core Chip handles byte/halfword into word conversions and unsigned extends
48/// This chip uses read_data and prev_data to constrain the write_data
49/// It also handles the shifting in case of not 4 byte aligned instructions
50/// This chips treats each (opcode, shift) pair as a separate instruction
51#[repr(C)]
52#[derive(Debug, Clone, AlignedBorrow, StructReflection)]
53pub struct LoadStoreCoreCols<T, const NUM_CELLS: usize> {
54    pub flags: [T; 4],
55    /// we need to keep the degree of is_valid and is_load to 1
56    pub is_valid: T,
57    pub is_load: T,
58
59    pub read_data: [T; NUM_CELLS],
60    pub prev_data: [T; NUM_CELLS],
61    /// write_data will be constrained against read_data and prev_data
62    /// depending on the opcode and the shift amount
63    pub write_data: [T; NUM_CELLS],
64}
65
66#[derive(Debug, Clone, derive_new::new, ColumnsAir)]
67#[columns_via(LoadStoreCoreCols<u8, NUM_CELLS>)]
68pub struct LoadStoreCoreAir<const NUM_CELLS: usize> {
69    pub offset: usize,
70}
71
72impl<F: Field, const NUM_CELLS: usize> BaseAir<F> for LoadStoreCoreAir<NUM_CELLS> {
73    fn width(&self) -> usize {
74        LoadStoreCoreCols::<F, NUM_CELLS>::width()
75    }
76}
77
78impl<F: Field, const NUM_CELLS: usize> BaseAirWithPublicValues<F> for LoadStoreCoreAir<NUM_CELLS> {}
79
80impl<AB, I, const NUM_CELLS: usize> VmCoreAir<AB, I> for LoadStoreCoreAir<NUM_CELLS>
81where
82    AB: InteractionBuilder,
83    I: VmAdapterInterface<AB::Expr>,
84    I::Reads: From<([AB::Var; NUM_CELLS], [AB::Expr; NUM_CELLS])>,
85    I::Writes: From<[[AB::Expr; NUM_CELLS]; 1]>,
86    I::ProcessedInstruction: From<LoadStoreInstruction<AB::Expr>>,
87{
88    fn eval(
89        &self,
90        builder: &mut AB,
91        local_core: &[AB::Var],
92        _from_pc: AB::Var,
93    ) -> AdapterAirContext<AB::Expr, I> {
94        let cols: &LoadStoreCoreCols<AB::Var, NUM_CELLS> = (*local_core).borrow();
95        let LoadStoreCoreCols::<AB::Var, NUM_CELLS> {
96            read_data,
97            prev_data,
98            write_data,
99            flags,
100            is_valid,
101            is_load,
102        } = *cols;
103
104        let get_expr_12 = |x: &AB::Expr| (x.clone() - AB::Expr::ONE) * (x.clone() - AB::Expr::TWO);
105
106        builder.assert_bool(is_valid);
107        let sum = flags.iter().fold(AB::Expr::ZERO, |acc, &flag| {
108            builder.assert_zero(flag * get_expr_12(&flag.into()));
109            acc + flag
110        });
111        builder.assert_zero(sum.clone() * get_expr_12(&sum));
112        // when sum is 0, is_valid must be 0
113        builder.when(get_expr_12(&sum)).assert_zero(is_valid);
114
115        // We will use the InstructionOpcode enum to encode the opcodes
116        // the appended digit to each opcode is the shift amount
117        let inv_2 = AB::F::from_u32(2).inverse();
118        let mut opcode_flags = vec![];
119        for flag in flags {
120            opcode_flags.push(flag * (flag - AB::F::ONE) * inv_2);
121        }
122        for flag in flags {
123            opcode_flags.push(flag * (sum.clone() - AB::F::TWO) * AB::F::NEG_ONE);
124        }
125        (0..4).for_each(|i| {
126            ((i + 1)..4).for_each(|j| opcode_flags.push(flags[i] * flags[j]));
127        });
128
129        let opcode_when = |idxs: &[InstructionOpcode]| -> AB::Expr {
130            idxs.iter().fold(AB::Expr::ZERO, |acc, &idx| {
131                acc + opcode_flags[idx as usize].clone()
132            })
133        };
134
135        // Constrain that is_load matches the opcode
136        builder.assert_eq(
137            is_load,
138            opcode_when(&[LoadW0, LoadHu0, LoadHu2, LoadBu0, LoadBu1, LoadBu2, LoadBu3]),
139        );
140        builder.when(is_load).assert_one(is_valid);
141
142        // There are three parts to write_data:
143        // - 1st limb is always read_data
144        // - 2nd to (NUM_CELLS/2)th limbs are:
145        //   - read_data if loadw/loadhu/storew/storeh
146        //   - prev_data if storeb
147        //   - zero if loadbu
148        // - (NUM_CELLS/2 + 1)th to last limbs are:
149        //   - read_data if loadw/storew
150        //   - prev_data if storeb/storeh
151        //   - zero if loadbu/loadhu
152        // Shifting needs to be carefully handled in case by case basis
153        // refer to [run_write_data] for the expected behavior in each case
154        for (i, cell) in write_data.iter().enumerate() {
155            // handling loads, expected_load_val = 0 if a store operation is happening
156            let expected_load_val = if i == 0 {
157                opcode_when(&[LoadW0, LoadHu0, LoadBu0]) * read_data[0]
158                    + opcode_when(&[LoadBu1]) * read_data[1]
159                    + opcode_when(&[LoadHu2, LoadBu2]) * read_data[2]
160                    + opcode_when(&[LoadBu3]) * read_data[3]
161            } else if i < NUM_CELLS / 2 {
162                opcode_when(&[LoadW0, LoadHu0]) * read_data[i]
163                    + opcode_when(&[LoadHu2]) * read_data[i + 2]
164            } else {
165                opcode_when(&[LoadW0]) * read_data[i]
166            };
167
168            // handling stores, expected_store_val = 0 if a load operation is happening
169            let expected_store_val = if i == 0 {
170                opcode_when(&[StoreW0, StoreH0, StoreB0]) * read_data[i]
171                    + opcode_when(&[StoreH2, StoreB1, StoreB2, StoreB3]) * prev_data[i]
172            } else if i == 1 {
173                opcode_when(&[StoreB1]) * read_data[i - 1]
174                    + opcode_when(&[StoreW0, StoreH0]) * read_data[i]
175                    + opcode_when(&[StoreH2, StoreB0, StoreB2, StoreB3]) * prev_data[i]
176            } else if i == 2 {
177                opcode_when(&[StoreH2, StoreB2]) * read_data[i - 2]
178                    + opcode_when(&[StoreW0]) * read_data[i]
179                    + opcode_when(&[StoreH0, StoreB0, StoreB1, StoreB3]) * prev_data[i]
180            } else if i == 3 {
181                opcode_when(&[StoreB3]) * read_data[i - 3]
182                    + opcode_when(&[StoreH2]) * read_data[i - 2]
183                    + opcode_when(&[StoreW0]) * read_data[i]
184                    + opcode_when(&[StoreH0, StoreB0, StoreB1, StoreB2]) * prev_data[i]
185            } else {
186                opcode_when(&[StoreW0]) * read_data[i]
187                    + opcode_when(&[StoreB0, StoreB1, StoreB2, StoreB3]) * prev_data[i]
188                    + opcode_when(&[StoreH0])
189                        * if i < NUM_CELLS / 2 {
190                            read_data[i]
191                        } else {
192                            prev_data[i]
193                        }
194                    + opcode_when(&[StoreH2])
195                        * if i - 2 < NUM_CELLS / 2 {
196                            read_data[i - 2]
197                        } else {
198                            prev_data[i]
199                        }
200            };
201            let expected_val = expected_load_val + expected_store_val;
202            builder.assert_eq(*cell, expected_val);
203        }
204
205        let expected_opcode = opcode_when(&[LoadW0]) * AB::Expr::from_u8(LOADW as u8)
206            + opcode_when(&[LoadHu0, LoadHu2]) * AB::Expr::from_u8(LOADHU as u8)
207            + opcode_when(&[LoadBu0, LoadBu1, LoadBu2, LoadBu3]) * AB::Expr::from_u8(LOADBU as u8)
208            + opcode_when(&[StoreW0]) * AB::Expr::from_u8(STOREW as u8)
209            + opcode_when(&[StoreH0, StoreH2]) * AB::Expr::from_u8(STOREH as u8)
210            + opcode_when(&[StoreB0, StoreB1, StoreB2, StoreB3]) * AB::Expr::from_u8(STOREB as u8);
211        let expected_opcode = VmCoreAir::<AB, I>::expr_to_global_expr(self, expected_opcode);
212
213        let load_shift_amount = opcode_when(&[LoadBu1]) * AB::Expr::ONE
214            + opcode_when(&[LoadHu2, LoadBu2]) * AB::Expr::TWO
215            + opcode_when(&[LoadBu3]) * AB::Expr::from_u32(3);
216
217        let store_shift_amount = opcode_when(&[StoreB1]) * AB::Expr::ONE
218            + opcode_when(&[StoreH2, StoreB2]) * AB::Expr::TWO
219            + opcode_when(&[StoreB3]) * AB::Expr::from_u32(3);
220
221        AdapterAirContext {
222            to_pc: None,
223            reads: (prev_data, read_data.map(|x| x.into())).into(),
224            writes: [write_data.map(|x| x.into())].into(),
225            instruction: LoadStoreInstruction {
226                is_valid: is_valid.into(),
227                opcode: expected_opcode,
228                is_load: is_load.into(),
229                load_shift_amount,
230                store_shift_amount,
231            }
232            .into(),
233        }
234    }
235
236    fn start_offset(&self) -> usize {
237        self.offset
238    }
239}
240
241#[repr(C)]
242#[derive(AlignedBytesBorrow, Debug)]
243pub struct LoadStoreCoreRecord<const NUM_CELLS: usize> {
244    pub local_opcode: u8,
245    pub shift_amount: u8,
246    pub read_data: [u8; NUM_CELLS],
247    // Note: `prev_data` can be a field, so we need to use u32
248    pub prev_data: [u32; NUM_CELLS],
249}
250
251#[derive(Clone, Copy, derive_new::new)]
252pub struct LoadStoreExecutor<A, const NUM_CELLS: usize> {
253    adapter: A,
254    pub offset: usize,
255}
256
257#[derive(Clone, derive_new::new)]
258pub struct LoadStoreFiller<
259    A = Rv32LoadStoreAdapterFiller,
260    const NUM_CELLS: usize = RV32_REGISTER_NUM_LIMBS,
261> {
262    adapter: A,
263    pub offset: usize,
264}
265
266impl<F, A, RA, const NUM_CELLS: usize> PreflightExecutor<F, RA> for LoadStoreExecutor<A, NUM_CELLS>
267where
268    F: PrimeField32,
269    A: 'static
270        + AdapterTraceExecutor<
271            F,
272            ReadData = (([u32; NUM_CELLS], [u8; NUM_CELLS]), u8),
273            WriteData = [u32; NUM_CELLS],
274        >,
275    for<'buf> RA: RecordArena<
276        'buf,
277        EmptyAdapterCoreLayout<F, A>,
278        (A::RecordMut<'buf>, &'buf mut LoadStoreCoreRecord<NUM_CELLS>),
279    >,
280{
281    fn get_opcode_name(&self, opcode: usize) -> String {
282        format!(
283            "{:?}",
284            Rv32LoadStoreOpcode::from_usize(opcode - self.offset)
285        )
286    }
287
288    fn execute(
289        &self,
290        state: VmStateMut<F, TracingMemory, RA>,
291        instruction: &Instruction<F>,
292    ) -> Result<(), ExecutionError> {
293        let Instruction { opcode, .. } = instruction;
294
295        let (mut adapter_record, core_record) = state.ctx.alloc(EmptyAdapterCoreLayout::new());
296
297        A::start(*state.pc, state.memory, &mut adapter_record);
298
299        (
300            (core_record.prev_data, core_record.read_data),
301            core_record.shift_amount,
302        ) = self
303            .adapter
304            .read(state.memory, instruction, &mut adapter_record);
305
306        let local_opcode = Rv32LoadStoreOpcode::from_usize(opcode.local_opcode_idx(self.offset));
307        core_record.local_opcode = local_opcode as u8;
308
309        let write_data = run_write_data(
310            local_opcode,
311            core_record.read_data,
312            core_record.prev_data,
313            core_record.shift_amount as usize,
314        );
315        self.adapter
316            .write(state.memory, instruction, write_data, &mut adapter_record);
317
318        *state.pc = state.pc.wrapping_add(DEFAULT_PC_STEP);
319
320        Ok(())
321    }
322}
323
324impl<F, A, const NUM_CELLS: usize> TraceFiller<F> for LoadStoreFiller<A, NUM_CELLS>
325where
326    F: PrimeField32,
327    A: 'static + AdapterTraceFiller<F>,
328{
329    fn fill_trace_row(&self, mem_helper: &MemoryAuxColsFactory<F>, row_slice: &mut [F]) {
330        // SAFETY: row_slice is guaranteed by the caller to have at least A::WIDTH +
331        // LoadStoreCoreCols::width() elements
332        let (adapter_row, mut core_row) = unsafe { row_slice.split_at_mut_unchecked(A::WIDTH) };
333        self.adapter.fill_trace_row(mem_helper, adapter_row);
334        // SAFETY: core_row contains a valid LoadStoreCoreRecord written by the executor
335        // during trace generation
336        let record: &LoadStoreCoreRecord<NUM_CELLS> =
337            unsafe { get_record_from_slice(&mut core_row, ()) };
338        let core_row: &mut LoadStoreCoreCols<F, NUM_CELLS> = core_row.borrow_mut();
339
340        let opcode = Rv32LoadStoreOpcode::from_usize(record.local_opcode as usize);
341        let shift = record.shift_amount;
342
343        let write_data = run_write_data(opcode, record.read_data, record.prev_data, shift as usize);
344        // Writing in reverse order
345        core_row.write_data = write_data.map(F::from_u32);
346        core_row.prev_data = record.prev_data.map(F::from_u32);
347        core_row.read_data = record.read_data.map(F::from_u8);
348        core_row.is_load = F::from_bool([LOADW, LOADHU, LOADBU].contains(&opcode));
349        core_row.is_valid = F::ONE;
350        let flags = &mut core_row.flags;
351        *flags = [F::ZERO; 4];
352        match (opcode, shift) {
353            (LOADW, 0) => flags[0] = F::TWO,
354            (LOADHU, 0) => flags[1] = F::TWO,
355            (LOADHU, 2) => flags[2] = F::TWO,
356            (LOADBU, 0) => flags[3] = F::TWO,
357
358            (LOADBU, 1) => flags[0] = F::ONE,
359            (LOADBU, 2) => flags[1] = F::ONE,
360            (LOADBU, 3) => flags[2] = F::ONE,
361            (STOREW, 0) => flags[3] = F::ONE,
362
363            (STOREH, 0) => (flags[0], flags[1]) = (F::ONE, F::ONE),
364            (STOREH, 2) => (flags[0], flags[2]) = (F::ONE, F::ONE),
365            (STOREB, 0) => (flags[0], flags[3]) = (F::ONE, F::ONE),
366            (STOREB, 1) => (flags[1], flags[2]) = (F::ONE, F::ONE),
367            (STOREB, 2) => (flags[1], flags[3]) = (F::ONE, F::ONE),
368            (STOREB, 3) => (flags[2], flags[3]) = (F::ONE, F::ONE),
369            _ => unreachable!(),
370        };
371    }
372}
373
374// Returns the write data
375#[inline(always)]
376pub(super) fn run_write_data<const NUM_CELLS: usize>(
377    opcode: Rv32LoadStoreOpcode,
378    read_data: [u8; NUM_CELLS],
379    prev_data: [u32; NUM_CELLS],
380    shift: usize,
381) -> [u32; NUM_CELLS] {
382    match (opcode, shift) {
383        (LOADW, 0) => {
384            read_data.map(|x| x as u32)
385        },
386        (LOADBU, 0) | (LOADBU, 1) | (LOADBU, 2) | (LOADBU, 3) => {
387           let mut wrie_data = [0; NUM_CELLS];
388           wrie_data[0] = read_data[shift] as u32;
389           wrie_data
390        }
391        (LOADHU, 0) | (LOADHU, 2) => {
392            let mut write_data = [0; NUM_CELLS];
393            for (i, cell) in write_data.iter_mut().take(NUM_CELLS / 2).enumerate() {
394                *cell = read_data[i + shift] as u32;
395            }
396            write_data
397        }
398        (STOREW, 0) => {
399            read_data.map(|x| x as u32)
400        },
401        (STOREB, 0) | (STOREB, 1) | (STOREB, 2) | (STOREB, 3) => {
402            let mut write_data = prev_data;
403            write_data[shift] = read_data[0] as u32;
404            write_data
405        }
406        (STOREH, 0) | (STOREH, 2) => {
407            array::from_fn(|i| {
408                if i >= shift && i < (NUM_CELLS / 2 + shift){
409                    read_data[i - shift] as u32
410                } else {
411                    prev_data[i]
412                }
413            })
414        }
415        // Currently the adapter AIR requires `ptr_val` to be aligned to the data size in bytes.
416        // The circuit requires that `shift = ptr_val % 4` so that `ptr_val - shift` is a multiple of 4.
417        // This requirement is non-trivial to remove, because we use it to ensure that `ptr_val - shift + 4 <= 2^pointer_max_bits`.
418        _ => unreachable!(
419            "unaligned memory access not supported by this execution environment: {opcode:?}, shift: {shift}"
420        ),
421    }
422}
openvm_rv32im_circuit/loadstore/core.rs

openvm_rv32im_circuit/loadstore/
core.rs
