sp1_core_machine/control_flow/jalr/
trace.rs

1use std::{borrow::BorrowMut, mem::MaybeUninit};
2
3use hashbrown::HashMap;
4use itertools::Itertools;
5use rayon::iter::{ParallelBridge, ParallelIterator};
6use slop_air::BaseAir;
7use slop_algebra::PrimeField32;
8use sp1_core_executor::{
9    events::{ByteLookupEvent, ByteRecord, JumpEvent},
10    ExecutionRecord, Program,
11};
12use sp1_hypercube::{air::MachineAir, Word};
13use struct_reflection::StructReflectionHelper;
14
15use crate::{utils::next_multiple_of_32, TrustMode, UserMode};
16
17use super::{JalrChip, JalrColumns};
18
19impl<F: PrimeField32, M: TrustMode> MachineAir<F> for JalrChip<M> {
20    type Record = ExecutionRecord;
21
22    type Program = Program;
23
24    fn name(&self) -> &'static str {
25        if M::IS_TRUSTED {
26            "Jalr"
27        } else {
28            "JalrUser"
29        }
30    }
31
32    fn num_rows(&self, input: &Self::Record) -> Option<usize> {
33        if input.program.enable_untrusted_programs == M::IS_TRUSTED {
34            return Some(0);
35        }
36        let nb_rows =
37            next_multiple_of_32(input.jalr_events.len(), input.fixed_log2_rows::<F, _>(self));
38        Some(nb_rows)
39    }
40
41    fn generate_trace_into(
42        &self,
43        input: &ExecutionRecord,
44        output: &mut ExecutionRecord,
45        buffer: &mut [MaybeUninit<F>],
46    ) {
47        if input.program.enable_untrusted_programs == M::IS_TRUSTED {
48            return;
49        }
50
51        let chunk_size = std::cmp::max((input.jalr_events.len()) / num_cpus::get(), 1);
52        let padded_nb_rows = <JalrChip<M> as MachineAir<F>>::num_rows(self, input).unwrap();
53        let width = <JalrChip<M> as BaseAir<F>>::width(self);
54        let num_event_rows = input.jalr_events.len();
55
56        unsafe {
57            let padding_start = num_event_rows * width;
58            let padding_size = (padded_nb_rows - num_event_rows) * width;
59            if padding_size > 0 {
60                core::ptr::write_bytes(buffer[padding_start..].as_mut_ptr(), 0, padding_size);
61            }
62        }
63
64        let buffer_ptr = buffer.as_mut_ptr() as *mut F;
65        let values = unsafe { core::slice::from_raw_parts_mut(buffer_ptr, padded_nb_rows * width) };
66
67        let blu_events = values
68            .chunks_mut(chunk_size * width)
69            .enumerate()
70            .par_bridge()
71            .map(|(i, rows)| {
72                let mut blu: HashMap<ByteLookupEvent, usize> = HashMap::new();
73                rows.chunks_mut(width).enumerate().for_each(|(j, row)| {
74                    let idx = i * chunk_size + j;
75                    let cols: &mut JalrColumns<F, M> = row.borrow_mut();
76
77                    if idx < input.jalr_events.len() {
78                        let event = &input.jalr_events[idx];
79                        self.event_to_row(&event.0, event.1.op_c, cols, &mut blu);
80                        cols.state.populate(&mut blu, event.0.clk, event.0.pc);
81                        cols.adapter.populate(&mut blu, event.1);
82                        if !M::IS_TRUSTED {
83                            let cols: &mut JalrColumns<F, UserMode> = row.borrow_mut();
84                            cols.adapter_cols.is_trusted = F::from_bool(!event.1.is_untrusted);
85                        }
86                    }
87                });
88                blu
89            })
90            .collect::<Vec<_>>();
91
92        output.add_byte_lookup_events_from_maps(blu_events.iter().collect_vec());
93    }
94
95    fn included(&self, shard: &Self::Record) -> bool {
96        if let Some(shape) = shard.shape.as_ref() {
97            shape.included::<F, _>(self)
98        } else {
99            !shard.jalr_events.is_empty()
100                && (M::IS_TRUSTED != shard.program.enable_untrusted_programs)
101        }
102    }
103
104    fn column_names(&self) -> Vec<String> {
105        JalrColumns::<F, M>::struct_reflection().unwrap()
106    }
107}
108
109impl<M: TrustMode> JalrChip<M> {
110    /// Create a row from an event.
111    fn event_to_row<F: PrimeField32>(
112        &self,
113        event: &JumpEvent,
114        imm: u64,
115        cols: &mut JalrColumns<F, M>,
116        blu: &mut HashMap<ByteLookupEvent, usize>,
117    ) {
118        cols.is_real = F::one();
119        let low_limb = (event.b.wrapping_add(imm) & 0xFFFF) as u16;
120        blu.add_bit_range_check(low_limb / 4, 14);
121        cols.lsb = F::from_canonical_u16(low_limb & 1);
122        cols.add_operation.populate(blu, event.b, imm);
123        if !event.op_a_0 {
124            cols.op_a_operation.populate(blu, event.pc, 4);
125        } else {
126            cols.op_a_operation.value = Word::from(0u64);
127        }
128    }
129}
sp1_core_machine/control_flow/jalr/trace.rs

sp1_core_machine/control_flow/jalr/
trace.rs
