openvm_rv32im_circuit/mul/
core.rs

1use std::{
2    array,
3    borrow::{Borrow, BorrowMut},
4};
5
6use openvm_circuit::{
7    arch::*,
8    system::memory::{online::TracingMemory, MemoryAuxColsFactory},
9};
10use openvm_circuit_primitives::{
11    range_tuple::{RangeTupleCheckerBus, SharedRangeTupleCheckerChip},
12    AlignedBytesBorrow, ColumnsAir, StructReflection, StructReflectionHelper,
13};
14use openvm_circuit_primitives_derive::AlignedBorrow;
15use openvm_instructions::{instruction::Instruction, program::DEFAULT_PC_STEP, LocalOpcode};
16use openvm_rv32im_transpiler::MulOpcode;
17use openvm_stark_backend::{
18    interaction::InteractionBuilder,
19    p3_air::BaseAir,
20    p3_field::{Field, PrimeCharacteristicRing, PrimeField32},
21    BaseAirWithPublicValues,
22};
23
24#[repr(C)]
25#[derive(AlignedBorrow, StructReflection)]
26pub struct MultiplicationCoreCols<T, const NUM_LIMBS: usize, const LIMB_BITS: usize> {
27    pub a: [T; NUM_LIMBS],
28    pub b: [T; NUM_LIMBS],
29    pub c: [T; NUM_LIMBS],
30    pub is_valid: T,
31}
32
33#[derive(Copy, Clone, Debug, derive_new::new, ColumnsAir)]
34#[columns_via(MultiplicationCoreCols<u8, NUM_LIMBS, LIMB_BITS>)]
35pub struct MultiplicationCoreAir<const NUM_LIMBS: usize, const LIMB_BITS: usize> {
36    pub bus: RangeTupleCheckerBus<2>,
37    pub offset: usize,
38}
39
40impl<F: Field, const NUM_LIMBS: usize, const LIMB_BITS: usize> BaseAir<F>
41    for MultiplicationCoreAir<NUM_LIMBS, LIMB_BITS>
42{
43    fn width(&self) -> usize {
44        MultiplicationCoreCols::<F, NUM_LIMBS, LIMB_BITS>::width()
45    }
46}
47impl<F: Field, const NUM_LIMBS: usize, const LIMB_BITS: usize> BaseAirWithPublicValues<F>
48    for MultiplicationCoreAir<NUM_LIMBS, LIMB_BITS>
49{
50}
51
52impl<AB, I, const NUM_LIMBS: usize, const LIMB_BITS: usize> VmCoreAir<AB, I>
53    for MultiplicationCoreAir<NUM_LIMBS, LIMB_BITS>
54where
55    AB: InteractionBuilder,
56    I: VmAdapterInterface<AB::Expr>,
57    I::Reads: From<[[AB::Expr; NUM_LIMBS]; 2]>,
58    I::Writes: From<[[AB::Expr; NUM_LIMBS]; 1]>,
59    I::ProcessedInstruction: From<MinimalInstruction<AB::Expr>>,
60{
61    fn eval(
62        &self,
63        builder: &mut AB,
64        local_core: &[AB::Var],
65        _from_pc: AB::Var,
66    ) -> AdapterAirContext<AB::Expr, I> {
67        let cols: &MultiplicationCoreCols<_, NUM_LIMBS, LIMB_BITS> = local_core.borrow();
68        builder.assert_bool(cols.is_valid);
69
70        let a = &cols.a;
71        let b = &cols.b;
72        let c = &cols.c;
73
74        // Define carry[i] = (sum_{k=0}^{i} b[k] * c[i - k] + carry[i - 1] - a[i]) / 2^LIMB_BITS.
75        // If 0 <= a[i], carry[i] < 2^LIMB_BITS, it can be proven that a[i] = sum_{k=0}^{i} (b[k] *
76        // c[i - k]) % 2^LIMB_BITS as necessary.
77        let mut carry: [AB::Expr; NUM_LIMBS] = array::from_fn(|_| AB::Expr::ZERO);
78        let carry_divide = AB::F::from_u32(1 << LIMB_BITS).inverse();
79
80        for i in 0..NUM_LIMBS {
81            let expected_limb = if i == 0 {
82                AB::Expr::ZERO
83            } else {
84                carry[i - 1].clone()
85            } + (0..=i).fold(AB::Expr::ZERO, |acc, k| acc + (b[k] * c[i - k]));
86            carry[i] = AB::Expr::from(carry_divide) * (expected_limb - a[i]);
87        }
88
89        for (a, carry) in a.iter().zip(carry.iter()) {
90            self.bus
91                .send(vec![(*a).into(), carry.clone()])
92                .eval(builder, cols.is_valid);
93        }
94
95        let expected_opcode = VmCoreAir::<AB, I>::opcode_to_global_expr(self, MulOpcode::MUL);
96
97        AdapterAirContext {
98            to_pc: None,
99            reads: [cols.b.map(Into::into), cols.c.map(Into::into)].into(),
100            writes: [cols.a.map(Into::into)].into(),
101            instruction: MinimalInstruction {
102                is_valid: cols.is_valid.into(),
103                opcode: expected_opcode,
104            }
105            .into(),
106        }
107    }
108
109    fn start_offset(&self) -> usize {
110        self.offset
111    }
112}
113
114#[repr(C)]
115#[derive(AlignedBytesBorrow, Debug)]
116pub struct MultiplicationCoreRecord<const NUM_LIMBS: usize, const LIMB_BITS: usize> {
117    pub b: [u8; NUM_LIMBS],
118    pub c: [u8; NUM_LIMBS],
119}
120
121#[derive(Clone, Copy, derive_new::new)]
122pub struct MultiplicationExecutor<A, const NUM_LIMBS: usize, const LIMB_BITS: usize> {
123    adapter: A,
124    pub offset: usize,
125}
126
127#[derive(Clone, Debug)]
128pub struct MultiplicationFiller<A, const NUM_LIMBS: usize, const LIMB_BITS: usize> {
129    adapter: A,
130    pub offset: usize,
131    pub range_tuple_chip: SharedRangeTupleCheckerChip<2>,
132}
133
134impl<A, const NUM_LIMBS: usize, const LIMB_BITS: usize>
135    MultiplicationFiller<A, NUM_LIMBS, LIMB_BITS>
136{
137    pub fn new(
138        adapter: A,
139        range_tuple_chip: SharedRangeTupleCheckerChip<2>,
140        offset: usize,
141    ) -> Self {
142        // The RangeTupleChecker is used to range check (a[i], carry[i]) pairs where 0 <= i
143        // < NUM_LIMBS. a[i] must have LIMB_BITS bits and carry[i] is the sum of i + 1 bytes
144        // (with LIMB_BITS bits).
145        debug_assert!(
146            range_tuple_chip.sizes()[0] == 1 << LIMB_BITS,
147            "First element of RangeTupleChecker must have size {}",
148            1 << LIMB_BITS
149        );
150        debug_assert!(
151            range_tuple_chip.sizes()[1] >= (1 << LIMB_BITS) * NUM_LIMBS as u32,
152            "Second element of RangeTupleChecker must have size of at least {}",
153            (1 << LIMB_BITS) * NUM_LIMBS as u32
154        );
155
156        Self {
157            adapter,
158            offset,
159            range_tuple_chip,
160        }
161    }
162}
163
164impl<F, A, RA, const NUM_LIMBS: usize, const LIMB_BITS: usize> PreflightExecutor<F, RA>
165    for MultiplicationExecutor<A, NUM_LIMBS, LIMB_BITS>
166where
167    F: PrimeField32,
168    A: 'static
169        + AdapterTraceExecutor<
170            F,
171            ReadData: Into<[[u8; NUM_LIMBS]; 2]>,
172            WriteData: From<[[u8; NUM_LIMBS]; 1]>,
173        >,
174    for<'buf> RA: RecordArena<
175        'buf,
176        EmptyAdapterCoreLayout<F, A>,
177        (
178            A::RecordMut<'buf>,
179            &'buf mut MultiplicationCoreRecord<NUM_LIMBS, LIMB_BITS>,
180        ),
181    >,
182{
183    fn get_opcode_name(&self, opcode: usize) -> String {
184        format!("{:?}", MulOpcode::from_usize(opcode - self.offset))
185    }
186
187    fn execute(
188        &self,
189        state: VmStateMut<F, TracingMemory, RA>,
190        instruction: &Instruction<F>,
191    ) -> Result<(), ExecutionError> {
192        let Instruction { opcode, .. } = instruction;
193
194        debug_assert_eq!(
195            MulOpcode::from_usize(opcode.local_opcode_idx(self.offset)),
196            MulOpcode::MUL
197        );
198        let (mut adapter_record, core_record) = state.ctx.alloc(EmptyAdapterCoreLayout::new());
199
200        A::start(*state.pc, state.memory, &mut adapter_record);
201
202        let [rs1, rs2] = self
203            .adapter
204            .read(state.memory, instruction, &mut adapter_record)
205            .into();
206
207        let (a, _) = run_mul::<NUM_LIMBS, LIMB_BITS>(&rs1, &rs2);
208
209        core_record.b = rs1;
210        core_record.c = rs2;
211
212        self.adapter
213            .write(state.memory, instruction, [a].into(), &mut adapter_record);
214
215        *state.pc = state.pc.wrapping_add(DEFAULT_PC_STEP);
216        Ok(())
217    }
218}
219
220impl<F, A, const NUM_LIMBS: usize, const LIMB_BITS: usize> TraceFiller<F>
221    for MultiplicationFiller<A, NUM_LIMBS, LIMB_BITS>
222where
223    F: PrimeField32,
224    A: 'static + AdapterTraceFiller<F>,
225{
226    fn fill_trace_row(&self, mem_helper: &MemoryAuxColsFactory<F>, row_slice: &mut [F]) {
227        // SAFETY: row_slice is guaranteed by the caller to have at least A::WIDTH +
228        // MultiplicationCoreCols::width() elements
229        let (adapter_row, mut core_row) = unsafe { row_slice.split_at_mut_unchecked(A::WIDTH) };
230        self.adapter.fill_trace_row(mem_helper, adapter_row);
231        // SAFETY: core_row contains a valid MultiplicationCoreRecord written by the executor
232        // during trace generation
233        let record: &MultiplicationCoreRecord<NUM_LIMBS, LIMB_BITS> =
234            unsafe { get_record_from_slice(&mut core_row, ()) };
235
236        let core_row: &mut MultiplicationCoreCols<F, NUM_LIMBS, LIMB_BITS> = core_row.borrow_mut();
237
238        let (a, carry) = run_mul::<NUM_LIMBS, LIMB_BITS>(&record.b, &record.c);
239
240        for (a, carry) in a.iter().zip(carry.iter()) {
241            self.range_tuple_chip.add_count(&[*a as u32, *carry]);
242        }
243
244        // write in reverse order
245        core_row.is_valid = F::ONE;
246        core_row.c = record.c.map(F::from_u8);
247        core_row.b = record.b.map(F::from_u8);
248        core_row.a = a.map(F::from_u8);
249    }
250}
251
252// returns mul, carry
253#[inline(always)]
254pub(super) fn run_mul<const NUM_LIMBS: usize, const LIMB_BITS: usize>(
255    x: &[u8; NUM_LIMBS],
256    y: &[u8; NUM_LIMBS],
257) -> ([u8; NUM_LIMBS], [u32; NUM_LIMBS]) {
258    let mut result = [0u8; NUM_LIMBS];
259    let mut carry = [0u32; NUM_LIMBS];
260    for i in 0..NUM_LIMBS {
261        let mut res = 0u32;
262        if i > 0 {
263            res = carry[i - 1];
264        }
265        for j in 0..=i {
266            res += (x[j] as u32) * (y[i - j] as u32);
267        }
268        carry[i] = res >> LIMB_BITS;
269        res %= 1u32 << LIMB_BITS;
270        result[i] = res as u8;
271    }
272    (result, carry)
273}
openvm_rv32im_circuit/mul/core.rs

openvm_rv32im_circuit/mul/
core.rs
