openvm_rv32im_circuit/base_alu/
core.rs

1use std::{
2    array,
3    borrow::{Borrow, BorrowMut},
4    iter::zip,
5};
6
7use openvm_circuit::{
8    arch::*,
9    system::memory::{online::TracingMemory, MemoryAuxColsFactory},
10};
11use openvm_circuit_primitives::{
12    bitwise_op_lookup::{BitwiseOperationLookupBus, SharedBitwiseOperationLookupChip},
13    utils::not,
14    AlignedBytesBorrow, ColumnsAir, StructReflection, StructReflectionHelper,
15};
16use openvm_circuit_primitives_derive::AlignedBorrow;
17use openvm_instructions::{instruction::Instruction, program::DEFAULT_PC_STEP, LocalOpcode};
18use openvm_rv32im_transpiler::BaseAluOpcode;
19use openvm_stark_backend::{
20    interaction::InteractionBuilder,
21    p3_air::{AirBuilder, BaseAir},
22    p3_field::{Field, PrimeCharacteristicRing, PrimeField32},
23    BaseAirWithPublicValues,
24};
25use strum::IntoEnumIterator;
26
27#[repr(C)]
28#[derive(AlignedBorrow, StructReflection, Debug)]
29pub struct BaseAluCoreCols<T, const NUM_LIMBS: usize, const LIMB_BITS: usize> {
30    pub a: [T; NUM_LIMBS],
31    pub b: [T; NUM_LIMBS],
32    pub c: [T; NUM_LIMBS],
33
34    pub opcode_add_flag: T,
35    pub opcode_sub_flag: T,
36    pub opcode_xor_flag: T,
37    pub opcode_or_flag: T,
38    pub opcode_and_flag: T,
39}
40
41#[derive(Copy, Clone, Debug, derive_new::new, ColumnsAir)]
42#[columns_via(BaseAluCoreCols<u8, NUM_LIMBS, LIMB_BITS>)]
43pub struct BaseAluCoreAir<const NUM_LIMBS: usize, const LIMB_BITS: usize> {
44    pub bus: BitwiseOperationLookupBus,
45    pub offset: usize,
46}
47
48impl<F: Field, const NUM_LIMBS: usize, const LIMB_BITS: usize> BaseAir<F>
49    for BaseAluCoreAir<NUM_LIMBS, LIMB_BITS>
50{
51    fn width(&self) -> usize {
52        BaseAluCoreCols::<F, NUM_LIMBS, LIMB_BITS>::width()
53    }
54}
55impl<F: Field, const NUM_LIMBS: usize, const LIMB_BITS: usize> BaseAirWithPublicValues<F>
56    for BaseAluCoreAir<NUM_LIMBS, LIMB_BITS>
57{
58}
59
60impl<AB, I, const NUM_LIMBS: usize, const LIMB_BITS: usize> VmCoreAir<AB, I>
61    for BaseAluCoreAir<NUM_LIMBS, LIMB_BITS>
62where
63    AB: InteractionBuilder,
64    I: VmAdapterInterface<AB::Expr>,
65    I::Reads: From<[[AB::Expr; NUM_LIMBS]; 2]>,
66    I::Writes: From<[[AB::Expr; NUM_LIMBS]; 1]>,
67    I::ProcessedInstruction: From<MinimalInstruction<AB::Expr>>,
68{
69    fn eval(
70        &self,
71        builder: &mut AB,
72        local_core: &[AB::Var],
73        _from_pc: AB::Var,
74    ) -> AdapterAirContext<AB::Expr, I> {
75        let cols: &BaseAluCoreCols<_, NUM_LIMBS, LIMB_BITS> = local_core.borrow();
76        let flags = [
77            cols.opcode_add_flag,
78            cols.opcode_sub_flag,
79            cols.opcode_xor_flag,
80            cols.opcode_or_flag,
81            cols.opcode_and_flag,
82        ];
83
84        let is_valid = flags.iter().fold(AB::Expr::ZERO, |acc, &flag| {
85            builder.assert_bool(flag);
86            acc + flag.into()
87        });
88        builder.assert_bool(is_valid.clone());
89
90        let a = &cols.a;
91        let b = &cols.b;
92        let c = &cols.c;
93
94        // For ADD, define carry[i] = (b[i] + c[i] + carry[i - 1] - a[i]) / 2^LIMB_BITS. If
95        // each carry[i] is boolean and 0 <= a[i] < 2^LIMB_BITS, it can be proven that
96        // a[i] = (b[i] + c[i]) % 2^LIMB_BITS as necessary. The same holds for SUB when
97        // carry[i] is (a[i] + c[i] - b[i] + carry[i - 1]) / 2^LIMB_BITS.
98        let mut carry_add: [AB::Expr; NUM_LIMBS] = array::from_fn(|_| AB::Expr::ZERO);
99        let mut carry_sub: [AB::Expr; NUM_LIMBS] = array::from_fn(|_| AB::Expr::ZERO);
100        let carry_divide = AB::F::from_usize(1 << LIMB_BITS).inverse();
101
102        for i in 0..NUM_LIMBS {
103            // We explicitly separate the constraints for ADD and SUB in order to keep degree
104            // cubic. Because we constrain that the carry (which is arbitrary) is bool, if
105            // carry has degree larger than 1 the max-degree constrain could be at least 4.
106            carry_add[i] = AB::Expr::from(carry_divide)
107                * (b[i] + c[i] - a[i]
108                    + if i > 0 {
109                        carry_add[i - 1].clone()
110                    } else {
111                        AB::Expr::ZERO
112                    });
113            builder
114                .when(cols.opcode_add_flag)
115                .assert_bool(carry_add[i].clone());
116            carry_sub[i] = AB::Expr::from(carry_divide)
117                * (a[i] + c[i] - b[i]
118                    + if i > 0 {
119                        carry_sub[i - 1].clone()
120                    } else {
121                        AB::Expr::ZERO
122                    });
123            builder
124                .when(cols.opcode_sub_flag)
125                .assert_bool(carry_sub[i].clone());
126        }
127
128        // Interaction with BitwiseOperationLookup to range check a for ADD and SUB, and
129        // constrain a's correctness for XOR, OR, and AND.
130        let bitwise = cols.opcode_xor_flag + cols.opcode_or_flag + cols.opcode_and_flag;
131        for i in 0..NUM_LIMBS {
132            let x = not::<AB::Expr>(bitwise.clone()) * a[i] + bitwise.clone() * b[i];
133            let y = not::<AB::Expr>(bitwise.clone()) * a[i] + bitwise.clone() * c[i];
134            let x_xor_y = cols.opcode_xor_flag * a[i]
135                + cols.opcode_or_flag * ((AB::Expr::from_u32(2) * a[i]) - b[i] - c[i])
136                + cols.opcode_and_flag * (b[i] + c[i] - (AB::Expr::from_u32(2) * a[i]));
137            self.bus
138                .send_xor(x, y, x_xor_y)
139                .eval(builder, is_valid.clone());
140        }
141
142        let expected_opcode = VmCoreAir::<AB, I>::expr_to_global_expr(
143            self,
144            flags.iter().zip(BaseAluOpcode::iter()).fold(
145                AB::Expr::ZERO,
146                |acc, (flag, local_opcode)| {
147                    acc + (*flag).into() * AB::Expr::from_u8(local_opcode as u8)
148                },
149            ),
150        );
151
152        AdapterAirContext {
153            to_pc: None,
154            reads: [cols.b.map(Into::into), cols.c.map(Into::into)].into(),
155            writes: [cols.a.map(Into::into)].into(),
156            instruction: MinimalInstruction {
157                is_valid,
158                opcode: expected_opcode,
159            }
160            .into(),
161        }
162    }
163
164    fn start_offset(&self) -> usize {
165        self.offset
166    }
167}
168
169#[repr(C, align(4))]
170#[derive(AlignedBytesBorrow, Debug)]
171pub struct BaseAluCoreRecord<const NUM_LIMBS: usize> {
172    pub b: [u8; NUM_LIMBS],
173    pub c: [u8; NUM_LIMBS],
174    // Use u8 instead of usize for better packing
175    pub local_opcode: u8,
176}
177
178#[derive(Clone, Copy, derive_new::new)]
179pub struct BaseAluExecutor<A, const NUM_LIMBS: usize, const LIMB_BITS: usize> {
180    adapter: A,
181    pub offset: usize,
182}
183
184#[derive(derive_new::new)]
185pub struct BaseAluFiller<A, const NUM_LIMBS: usize, const LIMB_BITS: usize> {
186    adapter: A,
187    pub bitwise_lookup_chip: SharedBitwiseOperationLookupChip<LIMB_BITS>,
188    pub offset: usize,
189}
190
191impl<F, A, RA, const NUM_LIMBS: usize, const LIMB_BITS: usize> PreflightExecutor<F, RA>
192    for BaseAluExecutor<A, NUM_LIMBS, LIMB_BITS>
193where
194    F: PrimeField32,
195    A: 'static
196        + AdapterTraceExecutor<
197            F,
198            ReadData: Into<[[u8; NUM_LIMBS]; 2]>,
199            WriteData: From<[[u8; NUM_LIMBS]; 1]>,
200        >,
201    for<'buf> RA: RecordArena<
202        'buf,
203        EmptyAdapterCoreLayout<F, A>,
204        (A::RecordMut<'buf>, &'buf mut BaseAluCoreRecord<NUM_LIMBS>),
205    >,
206{
207    fn get_opcode_name(&self, opcode: usize) -> String {
208        format!("{:?}", BaseAluOpcode::from_usize(opcode - self.offset))
209    }
210
211    fn execute(
212        &self,
213        state: VmStateMut<F, TracingMemory, RA>,
214        instruction: &Instruction<F>,
215    ) -> Result<(), ExecutionError> {
216        let Instruction { opcode, .. } = instruction;
217
218        let local_opcode = BaseAluOpcode::from_usize(opcode.local_opcode_idx(self.offset));
219        let (mut adapter_record, core_record) = state.ctx.alloc(EmptyAdapterCoreLayout::new());
220
221        A::start(*state.pc, state.memory, &mut adapter_record);
222
223        [core_record.b, core_record.c] = self
224            .adapter
225            .read(state.memory, instruction, &mut adapter_record)
226            .into();
227
228        let rd = run_alu::<NUM_LIMBS, LIMB_BITS>(local_opcode, &core_record.b, &core_record.c);
229
230        core_record.local_opcode = local_opcode as u8;
231
232        self.adapter
233            .write(state.memory, instruction, [rd].into(), &mut adapter_record);
234
235        *state.pc = state.pc.wrapping_add(DEFAULT_PC_STEP);
236
237        Ok(())
238    }
239}
240
241impl<F, A, const NUM_LIMBS: usize, const LIMB_BITS: usize> TraceFiller<F>
242    for BaseAluFiller<A, NUM_LIMBS, LIMB_BITS>
243where
244    F: PrimeField32,
245    A: 'static + AdapterTraceFiller<F>,
246{
247    fn fill_trace_row(&self, mem_helper: &MemoryAuxColsFactory<F>, row_slice: &mut [F]) {
248        // SAFETY: row_slice is guaranteed by the caller to have at least A::WIDTH +
249        // BaseAluCoreCols::width() elements
250        let (adapter_row, mut core_row) = unsafe { row_slice.split_at_mut_unchecked(A::WIDTH) };
251        self.adapter.fill_trace_row(mem_helper, adapter_row);
252        // SAFETY: core_row contains a valid BaseAluCoreRecord written by the executor
253        // during trace generation
254        let record: &BaseAluCoreRecord<NUM_LIMBS> =
255            unsafe { get_record_from_slice(&mut core_row, ()) };
256        let core_row: &mut BaseAluCoreCols<F, NUM_LIMBS, LIMB_BITS> = core_row.borrow_mut();
257        // SAFETY: the following is highly unsafe. We are going to cast `core_row` to a record
258        // buffer, and then do an _overlapping_ write to the `core_row` as a row of field elements.
259        // This requires:
260        // - Cols and Record structs should be repr(C) and we write in reverse order (to ensure
261        //   non-overlapping)
262        // - Do not overwrite any reference in `record` before it has already been used or moved
263        // - alignment of `F` must be >= alignment of Record (AlignedBytesBorrow will panic
264        //   otherwise)
265
266        let local_opcode = BaseAluOpcode::from_usize(record.local_opcode as usize);
267        let a = run_alu::<NUM_LIMBS, LIMB_BITS>(local_opcode, &record.b, &record.c);
268        // PERF: needless conversion
269        core_row.opcode_and_flag = F::from_bool(local_opcode == BaseAluOpcode::AND);
270        core_row.opcode_or_flag = F::from_bool(local_opcode == BaseAluOpcode::OR);
271        core_row.opcode_xor_flag = F::from_bool(local_opcode == BaseAluOpcode::XOR);
272        core_row.opcode_sub_flag = F::from_bool(local_opcode == BaseAluOpcode::SUB);
273        core_row.opcode_add_flag = F::from_bool(local_opcode == BaseAluOpcode::ADD);
274
275        if local_opcode == BaseAluOpcode::ADD || local_opcode == BaseAluOpcode::SUB {
276            for a_val in a {
277                self.bitwise_lookup_chip
278                    .request_xor(a_val as u32, a_val as u32);
279            }
280        } else {
281            for (b_val, c_val) in zip(record.b, record.c) {
282                self.bitwise_lookup_chip
283                    .request_xor(b_val as u32, c_val as u32);
284            }
285        }
286        core_row.c = record.c.map(F::from_u8);
287        core_row.b = record.b.map(F::from_u8);
288        core_row.a = a.map(F::from_u8);
289    }
290}
291
292#[inline(always)]
293pub(super) fn run_alu<const NUM_LIMBS: usize, const LIMB_BITS: usize>(
294    opcode: BaseAluOpcode,
295    x: &[u8; NUM_LIMBS],
296    y: &[u8; NUM_LIMBS],
297) -> [u8; NUM_LIMBS] {
298    debug_assert!(LIMB_BITS <= 8, "specialize for bytes");
299    match opcode {
300        BaseAluOpcode::ADD => run_add::<NUM_LIMBS, LIMB_BITS>(x, y),
301        BaseAluOpcode::SUB => run_subtract::<NUM_LIMBS, LIMB_BITS>(x, y),
302        BaseAluOpcode::XOR => run_xor::<NUM_LIMBS>(x, y),
303        BaseAluOpcode::OR => run_or::<NUM_LIMBS>(x, y),
304        BaseAluOpcode::AND => run_and::<NUM_LIMBS>(x, y),
305    }
306}
307
308#[inline(always)]
309fn run_add<const NUM_LIMBS: usize, const LIMB_BITS: usize>(
310    x: &[u8; NUM_LIMBS],
311    y: &[u8; NUM_LIMBS],
312) -> [u8; NUM_LIMBS] {
313    let mut z = [0u8; NUM_LIMBS];
314    let mut carry = [0u8; NUM_LIMBS];
315    for i in 0..NUM_LIMBS {
316        let mut overflow =
317            (x[i] as u16) + (y[i] as u16) + if i > 0 { carry[i - 1] as u16 } else { 0 };
318        carry[i] = (overflow >> LIMB_BITS) as u8;
319        overflow &= (1u16 << LIMB_BITS) - 1;
320        z[i] = overflow as u8;
321    }
322    z
323}
324
325#[inline(always)]
326fn run_subtract<const NUM_LIMBS: usize, const LIMB_BITS: usize>(
327    x: &[u8; NUM_LIMBS],
328    y: &[u8; NUM_LIMBS],
329) -> [u8; NUM_LIMBS] {
330    let mut z = [0u8; NUM_LIMBS];
331    let mut carry = [0u8; NUM_LIMBS];
332    for i in 0..NUM_LIMBS {
333        let rhs = y[i] as u16 + if i > 0 { carry[i - 1] as u16 } else { 0 };
334        if x[i] as u16 >= rhs {
335            z[i] = x[i] - rhs as u8;
336            carry[i] = 0;
337        } else {
338            z[i] = (x[i] as u16 + (1u16 << LIMB_BITS) - rhs) as u8;
339            carry[i] = 1;
340        }
341    }
342    z
343}
344
345#[inline(always)]
346fn run_xor<const NUM_LIMBS: usize>(x: &[u8; NUM_LIMBS], y: &[u8; NUM_LIMBS]) -> [u8; NUM_LIMBS] {
347    array::from_fn(|i| x[i] ^ y[i])
348}
349
350#[inline(always)]
351fn run_or<const NUM_LIMBS: usize>(x: &[u8; NUM_LIMBS], y: &[u8; NUM_LIMBS]) -> [u8; NUM_LIMBS] {
352    array::from_fn(|i| x[i] | y[i])
353}
354
355#[inline(always)]
356fn run_and<const NUM_LIMBS: usize>(x: &[u8; NUM_LIMBS], y: &[u8; NUM_LIMBS]) -> [u8; NUM_LIMBS] {
357    array::from_fn(|i| x[i] & y[i])
358}
openvm_rv32im_circuit/base_alu/core.rs

openvm_rv32im_circuit/base_alu/
core.rs
