openvm_recursion_circuit/primitives/exp_bits_len/
trace.rs

1use core::borrow::BorrowMut;
2use std::sync::Mutex;
3
4use openvm_stark_backend::poly_common::Squarable;
5use openvm_stark_sdk::config::baby_bear_poseidon2::F;
6use p3_field::{PrimeCharacteristicRing, PrimeField32};
7use p3_matrix::dense::RowMajorMatrix;
8use p3_maybe_rayon::prelude::*;
9
10use super::air::ExpBitsLenCols;
11
12pub(crate) const NUM_BITS_MAX_PLUS_ONE: usize = 32;
13pub(crate) const LOW_BITS_COUNT: usize = 27;
14
15#[repr(C)]
16#[derive(Clone, Debug)]
17pub struct ExpBitsLenRecord {
18    pub num_bits: u8,
19    pub base: F,
20    pub bit_src: F,
21    pub row_offset: u32,
22    pub shift_bits: u8,
23    pub shift_mult: u32,
24}
25
26impl ExpBitsLenRecord {
27    pub(crate) fn new(
28        base: F,
29        bit_src: F,
30        num_bits: usize,
31        row_offset: u32,
32        shift_bits: usize,
33        shift_mult: u32,
34    ) -> Self {
35        debug_assert!(num_bits < NUM_BITS_MAX_PLUS_ONE);
36        Self {
37            base,
38            bit_src,
39            num_bits: u8::try_from(num_bits)
40                .expect("num_bits fits in NUM_BITS_MAX_PLUS_ONE (< 256)"),
41            row_offset,
42            shift_bits: u8::try_from(shift_bits)
43                .expect("shift_bits fits in NUM_BITS_MAX_PLUS_ONE (< 256)"),
44            shift_mult,
45        }
46    }
47
48    pub(crate) fn num_rows(&self) -> usize {
49        NUM_BITS_MAX_PLUS_ONE
50    }
51
52    pub(crate) fn end_row(&self) -> usize {
53        self.row_offset as usize + self.num_rows()
54    }
55}
56
57#[derive(Debug, Default)]
58pub struct ExpBitsLenCpuTraceGenerator {
59    pub requests: Mutex<Vec<ExpBitsLenRecord>>,
60}
61
62impl ExpBitsLenCpuTraceGenerator {
63    pub fn add_request(&self, base: F, bit_src: F, num_bits: usize) {
64        self.add_requests([(base, bit_src, num_bits)]);
65    }
66
67    pub fn add_requests<I>(&self, batch: I)
68    where
69        I: IntoIterator<Item = (F, F, usize)>,
70    {
71        self.add_requests_with_shift(batch.into_iter().map(|(x, y, z)| (x, y, z, 0, 0)));
72    }
73
74    pub fn add_requests_with_shift<I>(&self, batch: I)
75    where
76        I: IntoIterator<Item = (F, F, usize, usize, u32)>,
77    {
78        let mut records = self.requests.lock().unwrap();
79        let mut next_row_offset = records.last().map(|record| record.end_row()).unwrap_or(0);
80        for (base, bit_src, num_bits, shift_bits, shift_mult) in batch {
81            let row_offset = u32::try_from(next_row_offset).expect("row offset should fit in u32");
82            let record =
83                ExpBitsLenRecord::new(base, bit_src, num_bits, row_offset, shift_bits, shift_mult);
84            next_row_offset += record.num_rows();
85            records.push(record);
86        }
87    }
88
89    #[tracing::instrument(name = "generate_trace", level = "trace", skip_all)]
90    pub fn generate_trace_row_major(
91        self,
92        required_height: Option<usize>,
93    ) -> Option<RowMajorMatrix<F>> {
94        let records = self.requests.into_inner().unwrap();
95        let num_valid_rows = records.last().map(|record| record.end_row()).unwrap_or(0);
96        let width = ExpBitsLenCols::<F>::width();
97
98        let padded_rows = if let Some(height) = required_height {
99            if height < num_valid_rows {
100                return None;
101            }
102            height
103        } else {
104            num_valid_rows.next_power_of_two()
105        };
106        let mut trace = vec![F::ZERO; padded_rows * width];
107
108        let (data_slice, padding_slice) = trace.split_at_mut(num_valid_rows * width);
109        let mut trace_slices: Vec<&mut [F]> = Vec::with_capacity(records.len());
110        let mut remaining = data_slice;
111
112        for record in &records {
113            let chunk_size = record.num_rows() * width;
114            let (chunk, rest) = remaining.split_at_mut(chunk_size);
115            trace_slices.push(chunk);
116            remaining = rest;
117        }
118
119        tracing::info_span!("fill_valid_rows").in_scope(|| {
120            trace_slices
121                .par_iter_mut()
122                .zip(records.par_iter())
123                .for_each(|(trace_slice, request)| {
124                    fill_valid_rows(
125                        request.base,
126                        request.bit_src.as_canonical_u32(),
127                        request.num_bits,
128                        request.shift_bits,
129                        request.shift_mult,
130                        trace_slice,
131                        width,
132                    );
133                });
134        });
135
136        tracing::info_span!("fill_padding_rows").in_scope(|| {
137            padding_slice.par_chunks_exact_mut(width).for_each(|row| {
138                let cols: &mut ExpBitsLenCols<F> = row.borrow_mut();
139                cols.result = F::ONE;
140                cols.result_multiplier = F::ONE;
141            });
142        });
143
144        Some(RowMajorMatrix::new(trace, width))
145    }
146}
147
148pub(crate) fn fill_valid_rows(
149    base: F,
150    bit_src: u32,
151    n: u8,
152    shift_bits: u8,
153    shift_mult: u32,
154    trace_slice: &mut [F],
155    width: usize,
156) {
157    fill_valid_rows_with_decomp_src(base, bit_src, n, shift_bits, shift_mult, trace_slice, width);
158}
159
160pub(crate) fn fill_valid_rows_with_decomp_src(
161    base: F,
162    decomp_src: u32,
163    n: u8,
164    shift_bits: u8,
165    shift_mult: u32,
166    trace_slice: &mut [F],
167    width: usize,
168) {
169    debug_assert!(n < NUM_BITS_MAX_PLUS_ONE as u8);
170    debug_assert_eq!(trace_slice.len(), NUM_BITS_MAX_PLUS_ONE * width);
171    debug_assert!(decomp_src < (1u32 << (NUM_BITS_MAX_PLUS_ONE - 1)));
172
173    let bases: Vec<_> = base.exp_powers_of_2().take(NUM_BITS_MAX_PLUS_ONE).collect();
174    let mut results = [F::ONE; NUM_BITS_MAX_PLUS_ONE];
175    let mut acc = F::ONE;
176    for step in (0..NUM_BITS_MAX_PLUS_ONE - 1).rev() {
177        if step < n as usize && ((decomp_src >> step) & 1) == 1 {
178            acc *= bases[step];
179        }
180        results[step] = acc;
181    }
182
183    let mut low_bits_are_zero = true;
184    let mut high_bits_all_one = false;
185    for step in 0..NUM_BITS_MAX_PLUS_ONE {
186        if step == LOW_BITS_COUNT {
187            high_bits_all_one = true;
188        }
189
190        let shifted = decomp_src >> step;
191        let num_bits = (n as usize).saturating_sub(step);
192        let low_bits_left = LOW_BITS_COUNT.saturating_sub(step);
193
194        let row_offset = step * width;
195        let row = &mut trace_slice[row_offset..row_offset + width];
196        let cols: &mut ExpBitsLenCols<F> = row.borrow_mut();
197        cols.is_valid = F::ONE;
198        cols.is_first = F::from_bool(step == 0);
199        cols.bit_idx = F::from_usize(step);
200        cols.base = bases[step];
201        cols.bit_src = F::from_u32(shifted);
202        cols.num_bits = F::from_usize(num_bits);
203        cols.apply_bit = F::from_bool(num_bits != 0);
204        cols.low_bits_left = F::from_usize(low_bits_left);
205        cols.in_low_region = F::from_bool(low_bits_left != 0);
206        cols.result = results[step];
207        cols.result_multiplier = if num_bits != 0 && (shifted & 1) == 1 {
208            bases[step]
209        } else {
210            F::ONE
211        };
212        cols.bit_src_mod_2 = F::from_bool((shifted & 1) == 1);
213        cols.low_bits_are_zero = F::from_bool(low_bits_are_zero);
214        cols.high_bits_all_one = F::from_bool(high_bits_all_one);
215        cols.bit_src_original = F::from_u32(decomp_src);
216        cols.shift_mult = if shift_bits as usize == step {
217            F::from_u32(shift_mult)
218        } else {
219            F::ZERO
220        };
221
222        if step < LOW_BITS_COUNT {
223            low_bits_are_zero &= (shifted & 1) == 0;
224        } else if step + 1 < NUM_BITS_MAX_PLUS_ONE {
225            high_bits_all_one &= (shifted & 1) == 1;
226        }
227    }
228}
openvm_recursion_circuit/primitives/exp_bits_len/trace.rs

openvm_recursion_circuit/primitives/exp_bits_len/
trace.rs
