openvm_rv32im_circuit/jalr/
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    bitwise_op_lookup::{BitwiseOperationLookupBus, SharedBitwiseOperationLookupChip},
12    var_range::{SharedVariableRangeCheckerChip, VariableRangeCheckerBus},
13    AlignedBytesBorrow, ColumnsAir, StructReflection, StructReflectionHelper,
14};
15use openvm_circuit_primitives_derive::AlignedBorrow;
16use openvm_instructions::{
17    instruction::Instruction,
18    program::{DEFAULT_PC_STEP, PC_BITS},
19    LocalOpcode,
20};
21use openvm_rv32im_transpiler::Rv32JalrOpcode::{self, *};
22use openvm_stark_backend::{
23    interaction::InteractionBuilder,
24    p3_air::{AirBuilder, BaseAir},
25    p3_field::{Field, PrimeCharacteristicRing, PrimeField32},
26    BaseAirWithPublicValues,
27};
28
29use crate::adapters::{
30    Rv32JalrAdapterExecutor, Rv32JalrAdapterFiller, RV32_CELL_BITS, RV32_REGISTER_NUM_LIMBS,
31};
32
33#[repr(C)]
34#[derive(Debug, Clone, AlignedBorrow, StructReflection)]
35pub struct Rv32JalrCoreCols<T> {
36    pub imm: T,
37    pub rs1_data: [T; RV32_REGISTER_NUM_LIMBS],
38    // To save a column, we only store the 3 most significant limbs of `rd_data`
39    // the least significant limb can be derived using from_pc and the other limbs
40    pub rd_data: [T; RV32_REGISTER_NUM_LIMBS - 1],
41    pub is_valid: T,
42
43    pub to_pc_least_sig_bit: T,
44    /// These are the limbs of `to_pc * 2`.
45    pub to_pc_limbs: [T; 2],
46    pub imm_sign: T,
47}
48
49#[derive(Debug, Clone, derive_new::new, ColumnsAir)]
50#[columns_via(Rv32JalrCoreCols<u8>)]
51pub struct Rv32JalrCoreAir {
52    pub bitwise_lookup_bus: BitwiseOperationLookupBus,
53    pub range_bus: VariableRangeCheckerBus,
54}
55
56impl<F: Field> BaseAir<F> for Rv32JalrCoreAir {
57    fn width(&self) -> usize {
58        Rv32JalrCoreCols::<F>::width()
59    }
60}
61
62impl<F: Field> BaseAirWithPublicValues<F> for Rv32JalrCoreAir {}
63
64impl<AB, I> VmCoreAir<AB, I> for Rv32JalrCoreAir
65where
66    AB: InteractionBuilder,
67    I: VmAdapterInterface<AB::Expr>,
68    I::Reads: From<[[AB::Expr; RV32_REGISTER_NUM_LIMBS]; 1]>,
69    I::Writes: From<[[AB::Expr; RV32_REGISTER_NUM_LIMBS]; 1]>,
70    I::ProcessedInstruction: From<SignedImmInstruction<AB::Expr>>,
71{
72    fn eval(
73        &self,
74        builder: &mut AB,
75        local_core: &[AB::Var],
76        from_pc: AB::Var,
77    ) -> AdapterAirContext<AB::Expr, I> {
78        let cols: &Rv32JalrCoreCols<AB::Var> = (*local_core).borrow();
79        let Rv32JalrCoreCols::<AB::Var> {
80            imm,
81            rs1_data: rs1,
82            rd_data: rd,
83            is_valid,
84            imm_sign,
85            to_pc_least_sig_bit,
86            to_pc_limbs,
87        } = *cols;
88
89        builder.assert_bool(is_valid);
90
91        // composed is the composition of 3 most significant limbs of rd
92        let composed = rd
93            .iter()
94            .enumerate()
95            .fold(AB::Expr::ZERO, |acc, (i, &val)| {
96                acc + val * AB::Expr::from_u32(1 << ((i + 1) * RV32_CELL_BITS))
97            });
98
99        let least_sig_limb = from_pc + AB::F::from_u32(DEFAULT_PC_STEP) - composed;
100
101        // rd_data is the final decomposition of `from_pc + DEFAULT_PC_STEP` we need.
102        // The range check on `least_sig_limb` also ensures that `rd_data` correctly represents
103        // `from_pc + DEFAULT_PC_STEP`. Specifically, if `rd_data` does not match the
104        // expected limb, then `least_sig_limb` becomes the real `least_sig_limb` plus the
105        // difference between `composed` and the three most significant limbs of `from_pc +
106        // DEFAULT_PC_STEP`. In that case, `least_sig_limb` >= 2^RV32_CELL_BITS.
107        let rd_data = array::from_fn(|i| {
108            if i == 0 {
109                least_sig_limb.clone()
110            } else {
111                rd[i - 1].into().clone()
112            }
113        });
114
115        // Constrain rd_data
116        // Assumes only from_pc in [0,2^PC_BITS) is allowed by program bus
117        self.bitwise_lookup_bus
118            .send_range(rd_data[0].clone(), rd_data[1].clone())
119            .eval(builder, is_valid);
120        self.range_bus
121            .range_check(rd_data[2].clone(), RV32_CELL_BITS)
122            .eval(builder, is_valid);
123        self.range_bus
124            .range_check(rd_data[3].clone(), PC_BITS - RV32_CELL_BITS * 3)
125            .eval(builder, is_valid);
126
127        builder.assert_bool(imm_sign);
128
129        // Constrain to_pc_least_sig_bit + 2 * to_pc_limbs = rs1 + imm as a i32 addition with 2
130        // limbs RISC-V spec explicitly sets the least significant bit of `to_pc` to 0
131        let rs1_limbs_01 = rs1[0] + rs1[1] * AB::F::from_u32(1 << RV32_CELL_BITS);
132        let rs1_limbs_23 = rs1[2] + rs1[3] * AB::F::from_u32(1 << RV32_CELL_BITS);
133        let inv = AB::F::from_u32(1 << 16).inverse();
134
135        builder.assert_bool(to_pc_least_sig_bit);
136        let carry = (rs1_limbs_01 + imm - to_pc_limbs[0] * AB::F::TWO - to_pc_least_sig_bit) * inv;
137        builder.when(is_valid).assert_bool(carry.clone());
138
139        let imm_extend_limb = imm_sign * AB::F::from_u32((1 << 16) - 1);
140        let carry = (rs1_limbs_23 + imm_extend_limb + carry - to_pc_limbs[1]) * inv;
141        builder.when(is_valid).assert_bool(carry);
142
143        // preventing to_pc overflow
144        self.range_bus
145            .range_check(to_pc_limbs[1], PC_BITS - 16)
146            .eval(builder, is_valid);
147        self.range_bus
148            .range_check(to_pc_limbs[0], 15)
149            .eval(builder, is_valid);
150        let to_pc = to_pc_limbs[0] * AB::F::TWO + to_pc_limbs[1] * AB::F::from_u32(1 << 16);
151
152        let expected_opcode = VmCoreAir::<AB, I>::opcode_to_global_expr(self, JALR);
153
154        AdapterAirContext {
155            to_pc: Some(to_pc),
156            reads: [rs1.map(|x| x.into())].into(),
157            writes: [rd_data].into(),
158            instruction: SignedImmInstruction {
159                is_valid: is_valid.into(),
160                opcode: expected_opcode,
161                immediate: imm.into(),
162                imm_sign: imm_sign.into(),
163            }
164            .into(),
165        }
166    }
167
168    fn start_offset(&self) -> usize {
169        Rv32JalrOpcode::CLASS_OFFSET
170    }
171}
172
173#[repr(C)]
174#[derive(AlignedBytesBorrow, Debug)]
175pub struct Rv32JalrCoreRecord {
176    pub imm: u16,
177    pub from_pc: u32,
178    pub rs1_val: u32,
179    pub imm_sign: bool,
180}
181
182#[derive(Clone, Copy, derive_new::new)]
183pub struct Rv32JalrExecutor<A = Rv32JalrAdapterExecutor> {
184    adapter: A,
185}
186
187#[derive(Clone)]
188pub struct Rv32JalrFiller<A = Rv32JalrAdapterFiller> {
189    adapter: A,
190    pub bitwise_lookup_chip: SharedBitwiseOperationLookupChip<RV32_CELL_BITS>,
191    pub range_checker_chip: SharedVariableRangeCheckerChip,
192}
193
194impl<A> Rv32JalrFiller<A> {
195    pub fn new(
196        adapter: A,
197        bitwise_lookup_chip: SharedBitwiseOperationLookupChip<RV32_CELL_BITS>,
198        range_checker_chip: SharedVariableRangeCheckerChip,
199    ) -> Self {
200        assert!(range_checker_chip.range_max_bits() >= 16);
201        Self {
202            adapter,
203            bitwise_lookup_chip,
204            range_checker_chip,
205        }
206    }
207}
208
209impl<F, A, RA> PreflightExecutor<F, RA> for Rv32JalrExecutor<A>
210where
211    F: PrimeField32,
212    A: 'static
213        + AdapterTraceExecutor<
214            F,
215            ReadData = [u8; RV32_REGISTER_NUM_LIMBS],
216            WriteData = [u8; RV32_REGISTER_NUM_LIMBS],
217        >,
218    for<'buf> RA: RecordArena<
219        'buf,
220        EmptyAdapterCoreLayout<F, A>,
221        (A::RecordMut<'buf>, &'buf mut Rv32JalrCoreRecord),
222    >,
223{
224    fn get_opcode_name(&self, opcode: usize) -> String {
225        format!(
226            "{:?}",
227            Rv32JalrOpcode::from_usize(opcode - Rv32JalrOpcode::CLASS_OFFSET)
228        )
229    }
230
231    fn execute(
232        &self,
233        state: VmStateMut<F, TracingMemory, RA>,
234        instruction: &Instruction<F>,
235    ) -> Result<(), ExecutionError> {
236        let Instruction { opcode, c, g, .. } = *instruction;
237
238        debug_assert_eq!(
239            opcode.local_opcode_idx(Rv32JalrOpcode::CLASS_OFFSET),
240            JALR as usize
241        );
242
243        let (mut adapter_record, core_record) = state.ctx.alloc(EmptyAdapterCoreLayout::new());
244
245        A::start(*state.pc, state.memory, &mut adapter_record);
246
247        core_record.rs1_val = u32::from_le_bytes(self.adapter.read(
248            state.memory,
249            instruction,
250            &mut adapter_record,
251        ));
252
253        core_record.imm = c.as_canonical_u32() as u16;
254        core_record.imm_sign = g.is_one();
255        core_record.from_pc = *state.pc;
256
257        let (to_pc, rd_data) = run_jalr(
258            core_record.from_pc,
259            core_record.rs1_val,
260            core_record.imm,
261            core_record.imm_sign,
262        );
263
264        self.adapter
265            .write(state.memory, instruction, rd_data, &mut adapter_record);
266
267        // RISC-V spec explicitly sets the least significant bit of `to_pc` to 0
268        *state.pc = to_pc & !1;
269
270        Ok(())
271    }
272}
273impl<F, A> TraceFiller<F> for Rv32JalrFiller<A>
274where
275    F: PrimeField32,
276    A: 'static + AdapterTraceFiller<F>,
277{
278    fn fill_trace_row(&self, mem_helper: &MemoryAuxColsFactory<F>, row_slice: &mut [F]) {
279        // SAFETY: row_slice is guaranteed by the caller to have at least A::WIDTH +
280        // Rv32JalrCoreCols::width() elements
281        let (adapter_row, mut core_row) = unsafe { row_slice.split_at_mut_unchecked(A::WIDTH) };
282        self.adapter.fill_trace_row(mem_helper, adapter_row);
283        // SAFETY: core_row contains a valid Rv32JalrCoreRecord written by the executor
284        // during trace generation
285        let record: &Rv32JalrCoreRecord = unsafe { get_record_from_slice(&mut core_row, ()) };
286
287        let core_row: &mut Rv32JalrCoreCols<F> = core_row.borrow_mut();
288
289        let (to_pc, rd_data) =
290            run_jalr(record.from_pc, record.rs1_val, record.imm, record.imm_sign);
291        let to_pc_limbs = [(to_pc & ((1 << 16) - 1)) >> 1, to_pc >> 16];
292        self.range_checker_chip.add_count(to_pc_limbs[0], 15);
293        self.range_checker_chip
294            .add_count(to_pc_limbs[1], PC_BITS - 16);
295        self.bitwise_lookup_chip
296            .request_range(rd_data[0] as u32, rd_data[1] as u32);
297
298        self.range_checker_chip
299            .add_count(rd_data[2] as u32, RV32_CELL_BITS);
300        self.range_checker_chip
301            .add_count(rd_data[3] as u32, PC_BITS - RV32_CELL_BITS * 3);
302
303        // Write in reverse order
304        core_row.imm_sign = F::from_bool(record.imm_sign);
305        core_row.to_pc_limbs = to_pc_limbs.map(F::from_u32);
306        core_row.to_pc_least_sig_bit = F::from_bool(to_pc & 1 == 1);
307        // fill_trace_row is called only on valid rows
308        core_row.is_valid = F::ONE;
309        core_row.rs1_data = record.rs1_val.to_le_bytes().map(F::from_u8);
310        core_row
311            .rd_data
312            .iter_mut()
313            .rev()
314            .zip(rd_data.iter().skip(1).rev())
315            .for_each(|(dst, src)| {
316                *dst = F::from_u8(*src);
317            });
318        core_row.imm = F::from_u16(record.imm);
319    }
320}
321
322// returns (to_pc, rd_data)
323#[inline(always)]
324pub(super) fn run_jalr(pc: u32, rs1: u32, imm: u16, imm_sign: bool) -> (u32, [u8; 4]) {
325    let to_pc = rs1.wrapping_add(imm as u32 + (imm_sign as u32 * 0xffff0000));
326    assert!(to_pc < (1 << PC_BITS));
327    (to_pc, pc.wrapping_add(DEFAULT_PC_STEP).to_le_bytes())
328}
openvm_rv32im_circuit/jalr/core.rs

openvm_rv32im_circuit/jalr/
core.rs
