openvm_rv32im_circuit/auipc/
core.rs

1use std::{
2    array::{self, from_fn},
3    borrow::{Borrow, BorrowMut},
4};
5
6use openvm_circuit::{
7    arch::*,
8    system::memory::{online::TracingMemory, MemoryAuxColsFactory},
9};
10use openvm_circuit_primitives::{
11    bitwise_op_lookup::{BitwiseOperationLookupBus, SharedBitwiseOperationLookupChip},
12    AlignedBytesBorrow, ColumnsAir, StructReflection, StructReflectionHelper,
13};
14use openvm_circuit_primitives_derive::AlignedBorrow;
15use openvm_instructions::{
16    instruction::Instruction,
17    program::{DEFAULT_PC_STEP, PC_BITS},
18    LocalOpcode,
19};
20use openvm_rv32im_transpiler::Rv32AuipcOpcode::{self, *};
21use openvm_stark_backend::{
22    interaction::InteractionBuilder,
23    p3_air::{AirBuilder, BaseAir},
24    p3_field::{Field, PrimeCharacteristicRing, PrimeField32},
25    BaseAirWithPublicValues,
26};
27
28use crate::adapters::{
29    Rv32RdWriteAdapterExecutor, Rv32RdWriteAdapterFiller, RV32_CELL_BITS, RV32_REGISTER_NUM_LIMBS,
30};
31
32#[repr(C)]
33#[derive(Debug, Clone, AlignedBorrow, StructReflection)]
34pub struct Rv32AuipcCoreCols<T> {
35    pub is_valid: T,
36    // The limbs of the immediate except the least significant limb since it is always 0
37    pub imm_limbs: [T; RV32_REGISTER_NUM_LIMBS - 1],
38    // The limbs of the PC except the most significant and the least significant limbs
39    pub pc_limbs: [T; RV32_REGISTER_NUM_LIMBS - 2],
40    pub rd_data: [T; RV32_REGISTER_NUM_LIMBS],
41}
42
43#[derive(Debug, Clone, Copy, derive_new::new, ColumnsAir)]
44#[columns_via(Rv32AuipcCoreCols<u8>)]
45pub struct Rv32AuipcCoreAir {
46    pub bus: BitwiseOperationLookupBus,
47}
48
49impl<F: Field> BaseAir<F> for Rv32AuipcCoreAir {
50    fn width(&self) -> usize {
51        Rv32AuipcCoreCols::<F>::width()
52    }
53}
54
55impl<F: Field> BaseAirWithPublicValues<F> for Rv32AuipcCoreAir {}
56
57impl<AB, I> VmCoreAir<AB, I> for Rv32AuipcCoreAir
58where
59    AB: InteractionBuilder,
60    I: VmAdapterInterface<AB::Expr>,
61    I::Reads: From<[[AB::Expr; 0]; 0]>,
62    I::Writes: From<[[AB::Expr; RV32_REGISTER_NUM_LIMBS]; 1]>,
63    I::ProcessedInstruction: From<ImmInstruction<AB::Expr>>,
64{
65    fn eval(
66        &self,
67        builder: &mut AB,
68        local_core: &[AB::Var],
69        from_pc: AB::Var,
70    ) -> AdapterAirContext<AB::Expr, I> {
71        let cols: &Rv32AuipcCoreCols<AB::Var> = (*local_core).borrow();
72
73        let Rv32AuipcCoreCols {
74            is_valid,
75            imm_limbs,
76            pc_limbs,
77            rd_data,
78        } = *cols;
79        builder.assert_bool(is_valid);
80
81        // We want to constrain rd = pc + imm (i32 add) where:
82        // - rd_data represents limbs of rd
83        // - pc_limbs are limbs of pc except the most and least significant limbs
84        // - imm_limbs are limbs of imm except the least significant limb
85
86        // We know that rd_data[0] is equal to the least significant limb of PC
87        // Thus, the intermediate value will be equal to PC without its most significant limb:
88        let intermed_val = rd_data[0]
89            + pc_limbs
90                .iter()
91                .enumerate()
92                .fold(AB::Expr::ZERO, |acc, (i, &val)| {
93                    acc + val * AB::Expr::from_u32(1 << ((i + 1) * RV32_CELL_BITS))
94                });
95
96        // Compute the most significant limb of PC
97        let pc_msl = (from_pc - intermed_val)
98            * AB::F::from_usize(1 << (RV32_CELL_BITS * (RV32_REGISTER_NUM_LIMBS - 1))).inverse();
99
100        // The vector pc_limbs contains the actual limbs of PC in little endian order
101        let pc_limbs = [rd_data[0]]
102            .iter()
103            .chain(pc_limbs.iter())
104            .map(|x| (*x).into())
105            .chain([pc_msl])
106            .collect::<Vec<AB::Expr>>();
107
108        let mut carry: [AB::Expr; RV32_REGISTER_NUM_LIMBS] = array::from_fn(|_| AB::Expr::ZERO);
109        let carry_divide = AB::F::from_usize(1 << RV32_CELL_BITS).inverse();
110
111        // Don't need to constrain the least significant limb of the addition
112        // since we already know that rd_data[0] = pc_limbs[0] and the least significant limb of imm
113        // is 0 Note: imm_limbs doesn't include the least significant limb so imm_limbs[i -
114        // 1] means the i-th limb of imm
115        for i in 1..RV32_REGISTER_NUM_LIMBS {
116            carry[i] = AB::Expr::from(carry_divide)
117                * (pc_limbs[i].clone() + imm_limbs[i - 1] - rd_data[i] + carry[i - 1].clone());
118            builder.when(is_valid).assert_bool(carry[i].clone());
119        }
120
121        // Range checking of rd_data entries to RV32_CELL_BITS bits
122        for i in 0..(RV32_REGISTER_NUM_LIMBS / 2) {
123            self.bus
124                .send_range(rd_data[i * 2], rd_data[i * 2 + 1])
125                .eval(builder, is_valid);
126        }
127
128        // The immediate and PC limbs need range checking to ensure they're within [0,
129        // 2^RV32_CELL_BITS) Since we range check two items at a time, doing this way helps
130        // efficiently divide the limbs into groups of 2 Note: range checking the limbs of
131        // immediate and PC separately would result in additional range checks       since
132        // they both have odd number of limbs that need to be range checked
133        let mut need_range_check: Vec<AB::Expr> = Vec::new();
134        for limb in imm_limbs {
135            need_range_check.push(limb.into());
136        }
137
138        assert_eq!(pc_limbs.len(), RV32_REGISTER_NUM_LIMBS);
139        // use enumerate to match pc_limbs[0] => i = 0, pc_limbs[1] => i = 1, ...
140        // pc_limbs[0] is already range checked through rd_data[0], so we skip it
141        for (i, limb) in pc_limbs.iter().enumerate().skip(1) {
142            // the most significant limb is pc_limbs[3] => i = 3
143            if i == pc_limbs.len() - 1 {
144                // Range check the most significant limb of pc to be in [0,
145                // 2^{PC_BITS-(RV32_REGISTER_NUM_LIMBS-1)*RV32_CELL_BITS})
146                need_range_check.push(
147                    (*limb).clone()
148                        * AB::Expr::from_usize(1 << (pc_limbs.len() * RV32_CELL_BITS - PC_BITS)),
149                );
150            } else {
151                need_range_check.push((*limb).clone());
152            }
153        }
154
155        // need_range_check contains (RV32_REGISTER_NUM_LIMBS - 1) elements from imm_limbs
156        // and (RV32_REGISTER_NUM_LIMBS - 1) elements from pc_limbs
157        // Hence, is of even length 2*RV32_REGISTER_NUM_LIMBS - 2
158        assert_eq!(need_range_check.len() % 2, 0);
159        for pair in need_range_check.chunks_exact(2) {
160            self.bus
161                .send_range(pair[0].clone(), pair[1].clone())
162                .eval(builder, is_valid);
163        }
164
165        let imm = imm_limbs
166            .iter()
167            .enumerate()
168            .fold(AB::Expr::ZERO, |acc, (i, &val)| {
169                acc + val * AB::Expr::from_u32(1 << (i * RV32_CELL_BITS))
170            });
171        let expected_opcode = VmCoreAir::<AB, I>::opcode_to_global_expr(self, AUIPC);
172        AdapterAirContext {
173            to_pc: None,
174            reads: [].into(),
175            writes: [rd_data.map(|x| x.into())].into(),
176            instruction: ImmInstruction {
177                is_valid: is_valid.into(),
178                opcode: expected_opcode,
179                immediate: imm,
180            }
181            .into(),
182        }
183    }
184
185    fn start_offset(&self) -> usize {
186        Rv32AuipcOpcode::CLASS_OFFSET
187    }
188}
189
190#[repr(C)]
191#[derive(AlignedBytesBorrow, Debug, Clone)]
192pub struct Rv32AuipcCoreRecord {
193    pub from_pc: u32,
194    pub imm: u32,
195}
196
197#[derive(Clone, Copy, derive_new::new)]
198pub struct Rv32AuipcExecutor<A = Rv32RdWriteAdapterExecutor> {
199    adapter: A,
200}
201
202#[derive(Clone, derive_new::new)]
203pub struct Rv32AuipcFiller<A = Rv32RdWriteAdapterFiller> {
204    adapter: A,
205    pub bitwise_lookup_chip: SharedBitwiseOperationLookupChip<RV32_CELL_BITS>,
206}
207
208impl<F, A, RA> PreflightExecutor<F, RA> for Rv32AuipcExecutor<A>
209where
210    F: PrimeField32,
211    A: 'static + AdapterTraceExecutor<F, ReadData = (), WriteData = [u8; RV32_REGISTER_NUM_LIMBS]>,
212    for<'buf> RA: RecordArena<
213        'buf,
214        EmptyAdapterCoreLayout<F, A>,
215        (A::RecordMut<'buf>, &'buf mut Rv32AuipcCoreRecord),
216    >,
217{
218    fn get_opcode_name(&self, _: usize) -> String {
219        format!("{AUIPC:?}")
220    }
221
222    fn execute(
223        &self,
224        state: VmStateMut<F, TracingMemory, RA>,
225        instruction: &Instruction<F>,
226    ) -> Result<(), ExecutionError> {
227        let (mut adapter_record, core_record) = state.ctx.alloc(EmptyAdapterCoreLayout::new());
228
229        A::start(*state.pc, state.memory, &mut adapter_record);
230
231        core_record.from_pc = *state.pc;
232        core_record.imm = instruction.c.as_canonical_u32();
233
234        let rd = run_auipc(*state.pc, core_record.imm);
235
236        self.adapter
237            .write(state.memory, instruction, rd, &mut adapter_record);
238
239        *state.pc = state.pc.wrapping_add(DEFAULT_PC_STEP);
240
241        Ok(())
242    }
243}
244
245impl<F, A> TraceFiller<F> for Rv32AuipcFiller<A>
246where
247    F: PrimeField32,
248    A: 'static + AdapterTraceFiller<F>,
249{
250    fn fill_trace_row(&self, mem_helper: &MemoryAuxColsFactory<F>, row_slice: &mut [F]) {
251        // SAFETY: row_slice is guaranteed by the caller to have at least A::WIDTH +
252        // Rv32AuipcCoreCols::width() elements
253        let (adapter_row, mut core_row) = unsafe { row_slice.split_at_mut_unchecked(A::WIDTH) };
254        self.adapter.fill_trace_row(mem_helper, adapter_row);
255        // SAFETY: core_row contains a valid Rv32AuipcCoreRecord written by the executor
256        // during trace generation
257        let record: &Rv32AuipcCoreRecord = unsafe { get_record_from_slice(&mut core_row, ()) };
258
259        let core_row: &mut Rv32AuipcCoreCols<F> = core_row.borrow_mut();
260
261        let imm_limbs = record.imm.to_le_bytes();
262        let pc_limbs = record.from_pc.to_le_bytes();
263        let rd_data = run_auipc(record.from_pc, record.imm);
264        debug_assert_eq!(imm_limbs[3], 0);
265
266        // range checks:
267        // hardcoding for performance: first 3 limbs of imm_limbs, last 3 limbs of pc_limbs where
268        // most significant limb of pc_limbs is shifted up
269        self.bitwise_lookup_chip
270            .request_range(imm_limbs[0] as u32, imm_limbs[1] as u32);
271        self.bitwise_lookup_chip
272            .request_range(imm_limbs[2] as u32, pc_limbs[1] as u32);
273        let msl_shift = RV32_REGISTER_NUM_LIMBS * RV32_CELL_BITS - PC_BITS;
274        self.bitwise_lookup_chip
275            .request_range(pc_limbs[2] as u32, (pc_limbs[3] as u32) << msl_shift);
276        for pair in rd_data.chunks_exact(2) {
277            self.bitwise_lookup_chip
278                .request_range(pair[0] as u32, pair[1] as u32);
279        }
280        // Writing in reverse order
281        core_row.rd_data = rd_data.map(F::from_u8);
282        // only the middle 2 limbs:
283        core_row.pc_limbs = from_fn(|i| F::from_u8(pc_limbs[i + 1]));
284        core_row.imm_limbs = from_fn(|i| F::from_u8(imm_limbs[i]));
285
286        core_row.is_valid = F::ONE;
287    }
288}
289
290// returns rd_data
291#[inline(always)]
292pub(super) fn run_auipc(pc: u32, imm: u32) -> [u8; RV32_REGISTER_NUM_LIMBS] {
293    let rd = pc.wrapping_add(imm << RV32_CELL_BITS);
294    rd.to_le_bytes()
295}
openvm_rv32im_circuit/auipc/core.rs

openvm_rv32im_circuit/auipc/
core.rs
