sp1_core_machine/control_flow/jal/
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},
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::{JalChip, JalColumns};
18
19impl<F: PrimeField32, M: TrustMode> MachineAir<F> for JalChip<M> {
20    type Record = ExecutionRecord;
21
22    type Program = Program;
23
24    fn name(&self) -> &'static str {
25        if M::IS_TRUSTED {
26            "Jal"
27        } else {
28            "JalUser"
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.jal_events.len(), input.fixed_log2_rows::<F, _>(self));
38        Some(nb_rows)
39    }
40
41    fn generate_dependencies(&self, input: &Self::Record, output: &mut Self::Record) {
42        if input.program.enable_untrusted_programs == M::IS_TRUSTED {
43            return;
44        }
45
46        let chunk_size = std::cmp::max((input.jal_events.len()) / num_cpus::get(), 1);
47        let width = <JalChip<M> as BaseAir<F>>::width(self);
48
49        let blu_batches = input
50            .jal_events
51            .chunks(chunk_size)
52            .par_bridge()
53            .map(|events| {
54                let mut blu: HashMap<ByteLookupEvent, usize> = HashMap::new();
55                events.iter().for_each(|event| {
56                    let mut row = vec![F::zero(); width];
57                    let cols: &mut JalColumns<F, M> = row.as_mut_slice().borrow_mut();
58
59                    cols.is_real = F::one();
60                    let low_limb = (event.0.pc.wrapping_add(event.0.b) & 0xFFFF) as u16;
61                    blu.add_bit_range_check(low_limb / 4, 14);
62                    cols.add_operation.populate(&mut blu, event.0.pc, event.0.b);
63                    if !event.0.op_a_0 {
64                        cols.op_a_operation.populate(&mut blu, event.0.pc, 4);
65                    }
66                    cols.state.populate(&mut blu, event.0.clk, event.0.pc);
67                    cols.adapter.populate(&mut blu, event.1);
68                });
69                blu
70            })
71            .collect::<Vec<_>>();
72
73        output.add_byte_lookup_events_from_maps(blu_batches.iter().collect_vec());
74    }
75
76    fn generate_trace_into(
77        &self,
78        input: &ExecutionRecord,
79        _output: &mut ExecutionRecord,
80        buffer: &mut [MaybeUninit<F>],
81    ) {
82        if input.program.enable_untrusted_programs == M::IS_TRUSTED {
83            return;
84        }
85
86        // Generate the rows for the trace.
87        let padded_nb_rows = <JalChip<M> as MachineAir<F>>::num_rows(self, input).unwrap();
88        let width = <JalChip<M> as BaseAir<F>>::width(self);
89
90        let chunk_size = std::cmp::max(input.jal_events.len() / num_cpus::get(), 1);
91        let num_event_rows = input.jal_events.len();
92
93        unsafe {
94            let padding_start = num_event_rows * width;
95            let padding_size = (padded_nb_rows - num_event_rows) * width;
96            if padding_size > 0 {
97                core::ptr::write_bytes(buffer[padding_start..].as_mut_ptr(), 0, padding_size);
98            }
99        }
100
101        let buffer_ptr = buffer.as_mut_ptr() as *mut F;
102        let values = unsafe { core::slice::from_raw_parts_mut(buffer_ptr, num_event_rows * width) };
103
104        values.chunks_mut(chunk_size * width).enumerate().par_bridge().for_each(|(i, rows)| {
105            let mut blu = Vec::new();
106            rows.chunks_mut(width).enumerate().for_each(|(j, row)| {
107                let idx = i * chunk_size + j;
108                if idx < input.jal_events.len() {
109                    let event = input.jal_events[idx];
110                    let cols: &mut JalColumns<F, M> = row.borrow_mut();
111                    cols.is_real = F::one();
112                    cols.add_operation.populate(&mut blu, event.0.pc, event.0.b);
113                    if !event.0.op_a_0 {
114                        cols.op_a_operation.populate(&mut blu, event.0.pc, 4);
115                    } else {
116                        cols.op_a_operation.value = Word::from(0u64);
117                    }
118                    cols.state.populate(&mut blu, event.0.clk, event.0.pc);
119                    cols.adapter.populate(&mut blu, event.1);
120                    if !M::IS_TRUSTED {
121                        let cols: &mut JalColumns<F, UserMode> = row.borrow_mut();
122                        cols.adapter_cols.is_trusted = F::from_bool(!event.1.is_untrusted);
123                    }
124                }
125            });
126        });
127    }
128
129    fn included(&self, shard: &Self::Record) -> bool {
130        if let Some(shape) = shard.shape.as_ref() {
131            shape.included::<F, _>(self)
132        } else {
133            !shard.jal_events.is_empty()
134                && (M::IS_TRUSTED != shard.program.enable_untrusted_programs)
135        }
136    }
137
138    fn column_names(&self) -> Vec<String> {
139        JalColumns::<F, M>::struct_reflection().unwrap()
140    }
141}
sp1_core_machine/control_flow/jal/trace.rs

sp1_core_machine/control_flow/jal/
trace.rs
