miden_processor/trace/
utils.rs

1#[cfg(test)]
2use alloc::vec::Vec;
3
4use miden_air::{MIDEN_AIR_COUNT, PcsParams, memory, trace::MIN_TRACE_LEN};
5
6use super::chiplets::Chiplets;
7use crate::{Felt, ONE};
8#[cfg(test)]
9use crate::{operation::Operation, utils::ToElements};
10
11// ROW-MAJOR TRACE WRITER
12// ================================================================================================
13
14/// Row-major flat buffer writer (`write_row` is a single `copy_from_slice`).
15///
16/// `payload` is the number of leading columns written per row; `stride` is the physical row
17/// width of the backing buffer. When `stride > payload`, the trailing `stride - payload`
18/// columns of each row are left untouched (callers rely on them staying zero-initialized).
19#[derive(Debug)]
20pub struct RowMajorTraceWriter<'a, E> {
21    data: &'a mut [E],
22    payload: usize,
23    stride: usize,
24}
25
26impl<'a, E: Copy> RowMajorTraceWriter<'a, E> {
27    /// Creates a writer whose physical row width equals the per-row payload.
28    #[cfg(test)]
29    pub(super) fn new(data: &'a mut [E], width: usize) -> Self {
30        Self::with_stride(data, width, width)
31    }
32
33    /// Creates a writer that writes `payload` columns per row into a buffer with physical row
34    /// width `stride` (`stride >= payload`).
35    pub fn with_stride(data: &'a mut [E], payload: usize, stride: usize) -> Self {
36        debug_assert!(stride >= payload, "stride must be >= payload");
37        debug_assert_eq!(data.len() % stride, 0, "buffer length must be a multiple of stride");
38        Self { data, payload, stride }
39    }
40
41    /// Writes one row's payload; `values.len()` must equal `payload`.
42    #[inline(always)]
43    pub fn write_row(&mut self, row: usize, values: &[E]) {
44        debug_assert_eq!(values.len(), self.payload);
45        let start = row * self.stride;
46        self.data[start..start + self.payload].copy_from_slice(values);
47    }
48}
49
50// TRACE FRAGMENT
51// ================================================================================================
52
53/// A writable, row-major view over one chiplet's region of the chiplets trace.
54///
55/// A chiplet occupies a contiguous band of rows and a contiguous band of columns
56/// `[col_start, col_start + num_cols)`. [`Self::copy_rows_from`] also writes the per-row
57/// `prefix_one_cols` selectors and, when there is room, the trailing `chip_clk` column.
58pub struct ChipletTraceFragment<'a> {
59    /// Contiguous `num_rows * stride` row-major slice (this chiplet's rows).
60    band: &'a mut [Felt],
61    stride: usize,
62    col_start: usize,
63    num_rows: usize,
64    num_cols: usize,
65    /// Global row offset of `band[0]` in the chiplets trace; used to compute `chip_clk`.
66    row_offset: usize,
67    /// Columns to set to ONE on every row in this band.
68    prefix_one_cols: &'static [usize],
69    /// Whether to write `chip_clk` to the trailing column.
70    write_chip_clk: bool,
71}
72
73impl<'a> ChipletTraceFragment<'a> {
74    /// Bare fragment with no prefix selectors or `chip_clk`. For chiplet-level unit tests.
75    pub fn row_major(
76        band: &'a mut [Felt],
77        stride: usize,
78        col_start: usize,
79        num_cols: usize,
80    ) -> Self {
81        Self::new(band, stride, col_start, num_cols, 0, &[], false)
82    }
83
84    /// Adds the chiplets-trace overheads: per-row ONEs at `prefix_one_cols` and `chip_clk` at
85    /// column `stride - 1` when the fragment leaves a trailing column for it.
86    pub fn with_overheads(
87        band: &'a mut [Felt],
88        stride: usize,
89        col_start: usize,
90        num_cols: usize,
91        row_offset: usize,
92        prefix_one_cols: &'static [usize],
93    ) -> Self {
94        Self::new(
95            band,
96            stride,
97            col_start,
98            num_cols,
99            row_offset,
100            prefix_one_cols,
101            stride > col_start + num_cols,
102        )
103    }
104
105    fn new(
106        band: &'a mut [Felt],
107        stride: usize,
108        col_start: usize,
109        num_cols: usize,
110        row_offset: usize,
111        prefix_one_cols: &'static [usize],
112        write_chip_clk: bool,
113    ) -> Self {
114        debug_assert_eq!(band.len() % stride, 0, "band length must be a multiple of stride");
115        debug_assert!(col_start + num_cols <= stride, "column band overruns the row stride");
116        debug_assert!(
117            prefix_one_cols.iter().all(|&col| col < col_start),
118            "prefix_one_cols must lie before col_start",
119        );
120        let num_rows = band.len() / stride;
121        Self {
122            band,
123            stride,
124            col_start,
125            num_rows,
126            num_cols,
127            row_offset,
128            prefix_one_cols,
129            write_chip_clk,
130        }
131    }
132
133    // PUBLIC ACCESSORS
134    // --------------------------------------------------------------------------------------------
135
136    /// Returns the number of columns in this execution trace fragment.
137    pub fn width(&self) -> usize {
138        self.num_cols
139    }
140
141    /// Returns the number of rows in this execution trace fragment.
142    pub fn len(&self) -> usize {
143        self.num_rows
144    }
145
146    // DATA MUTATORS
147    // --------------------------------------------------------------------------------------------
148
149    /// Copies a chiplet's row-major buffer (`num_cols` cells per row) into this fragment's
150    /// band, fusing the per-row prefix-selector ONEs and trailing `chip_clk` when configured.
151    pub fn copy_rows_from(&mut self, src: &[Felt]) {
152        debug_assert_eq!(src.len(), self.num_rows * self.num_cols, "source buffer size mismatch");
153        self.copy_rows_into(0, src);
154    }
155
156    /// Copies `src.len() / num_cols` rows starting at `row_offset` into this fragment's band,
157    /// fusing the per-row prefix-selector ONEs and the `chip_clk` column when configured.
158    pub fn copy_rows_into(&mut self, row_offset: usize, src: &[Felt]) {
159        debug_assert_eq!(src.len() % self.num_cols, 0, "source buffer size not row-aligned");
160        let chunk_rows = src.len() / self.num_cols;
161        debug_assert!(
162            row_offset + chunk_rows <= self.num_rows,
163            "chunk overruns fragment row range",
164        );
165        let clk_col = self.stride - 1;
166        for r in 0..chunk_rows {
167            let dst_row = row_offset + r;
168            let row_start = dst_row * self.stride;
169            let row = &mut self.band[row_start..row_start + self.stride];
170            for &col in self.prefix_one_cols {
171                row[col] = ONE;
172            }
173            let src_row = &src[r * self.num_cols..(r + 1) * self.num_cols];
174            row[self.col_start..self.col_start + self.num_cols].copy_from_slice(src_row);
175            if self.write_chip_clk {
176                row[clk_col] = Felt::from_u32((self.row_offset + dst_row + 1) as u32);
177            }
178        }
179    }
180}
181
182// TRACE LENGTH SUMMARY
183// ================================================================================================
184
185/// Contains the unpadded lengths of the trace parts.
186///
187/// - `core_trace_len` contains the length of the core trace (system + decoder + stack).
188/// - `range_trace_len` contains the length of the range checker trace.
189/// - `chiplets_trace_len` contains the chiplets-trace component lengths.
190/// - `poseidon2_permutation_trace_len` contains the Poseidon2 permutation AIR length.
191#[derive(Debug, Default, Eq, PartialEq, Clone, Copy)]
192pub struct TraceLenSummary {
193    core_trace_len: usize,
194    range_trace_len: usize,
195    chiplets_trace_len: ChipletsLengths,
196    poseidon2_permutation_trace_len: usize,
197    /// Set by the trace builder when known, in [`miden_air::AIRS`] order. `None` falls back to
198    /// deriving [`Self::padded_trace_len`] from the unpadded component lengths via
199    /// `next_power_of_two`.
200    padded_heights: Option<[usize; MIDEN_AIR_COUNT]>,
201}
202
203impl TraceLenSummary {
204    pub fn new(
205        core_trace_len: usize,
206        range_trace_len: usize,
207        chiplets_trace_len: ChipletsLengths,
208    ) -> Self {
209        TraceLenSummary {
210            core_trace_len,
211            range_trace_len,
212            chiplets_trace_len,
213            poseidon2_permutation_trace_len: 0,
214            padded_heights: None,
215        }
216    }
217
218    /// Builds a summary after the trace builder has computed the padded per-AIR heights, in
219    /// [`miden_air::AIRS`] order.
220    pub fn new_with_padded(
221        core_trace_len: usize,
222        range_trace_len: usize,
223        chiplets_trace_len: ChipletsLengths,
224        poseidon2_permutation_trace_len: usize,
225        padded_heights: [usize; MIDEN_AIR_COUNT],
226    ) -> Self {
227        TraceLenSummary {
228            core_trace_len,
229            range_trace_len,
230            chiplets_trace_len,
231            poseidon2_permutation_trace_len,
232            padded_heights: Some(padded_heights),
233        }
234    }
235
236    /// Returns length of the core trace (system + decoder + stack).
237    pub fn core_trace_len(&self) -> usize {
238        self.core_trace_len
239    }
240
241    /// Returns length of the range checker trace.
242    pub fn range_trace_len(&self) -> usize {
243        self.range_trace_len
244    }
245
246    /// Returns the chiplets-trace component lengths.
247    pub fn chiplets_trace_len(&self) -> ChipletsLengths {
248        self.chiplets_trace_len
249    }
250
251    /// Returns the Poseidon2 permutation AIR trace length.
252    pub fn poseidon2_permutation_trace_len(&self) -> usize {
253        self.poseidon2_permutation_trace_len
254    }
255
256    /// Returns the maximum of all component lengths.
257    pub fn trace_len(&self) -> usize {
258        self.range_trace_len
259            .max(self.core_trace_len)
260            .max(self.chiplets_trace_len.trace_len())
261            .max(self.poseidon2_permutation_trace_len)
262    }
263
264    /// Returns `trace_len` rounded up to the next power of two, clamped to `MIN_TRACE_LEN`.
265    pub fn padded_trace_len(&self) -> usize {
266        self.padded_heights
267            .map(|heights| heights.into_iter().max().expect("heights is non-empty"))
268            .unwrap_or_else(|| self.trace_len().next_power_of_two().max(MIN_TRACE_LEN))
269    }
270
271    /// Returns the padded per-AIR heights, in [`miden_air::AIRS`] order, if known.
272    pub fn padded_heights(&self) -> Option<&[usize; MIDEN_AIR_COUNT]> {
273        self.padded_heights.as_ref()
274    }
275
276    /// Returns the modelled peak prover memory, in bytes, for the padded per-AIR heights, if
277    /// known. See [`miden_air::memory::prover_peak_bytes`] for what this does and does not cover.
278    pub fn prover_memory_bytes(&self, params: &PcsParams) -> Option<u64> {
279        memory::prover_peak_bytes(self.padded_heights()?, params)
280    }
281
282    /// Returns the percent (0 - 100) of rows added by padding.
283    pub fn padding_percentage(&self) -> usize {
284        (self.padded_trace_len() - self.trace_len()) * 100 / self.padded_trace_len()
285    }
286}
287
288// CHIPLET LENGTHS
289// ================================================================================================
290
291/// Contains trace lengths of all chiplets: hash, bitwise, memory, ACE, and kernel ROM.
292#[derive(Default, Clone, Copy, Debug, PartialEq, Eq)]
293pub struct ChipletsLengths {
294    hash_chiplet_len: usize,
295    bitwise_chiplet_len: usize,
296    memory_chiplet_len: usize,
297    ace_chiplet_len: usize,
298    kernel_rom_len: usize,
299}
300
301impl ChipletsLengths {
302    pub fn new(chiplets: &Chiplets) -> Self {
303        ChipletsLengths {
304            hash_chiplet_len: chiplets.bitwise_start().into(),
305            bitwise_chiplet_len: chiplets.memory_start() - chiplets.bitwise_start(),
306            memory_chiplet_len: chiplets.ace_start() - chiplets.memory_start(),
307            ace_chiplet_len: chiplets.kernel_rom_start() - chiplets.ace_start(),
308            kernel_rom_len: chiplets.padding_start() - chiplets.kernel_rom_start(),
309        }
310    }
311
312    pub fn from_parts(
313        hash_len: usize,
314        bitwise_len: usize,
315        memory_len: usize,
316        ace_len: usize,
317        kernel_len: usize,
318    ) -> Self {
319        ChipletsLengths {
320            hash_chiplet_len: hash_len,
321            bitwise_chiplet_len: bitwise_len,
322            memory_chiplet_len: memory_len,
323            ace_chiplet_len: ace_len,
324            kernel_rom_len: kernel_len,
325        }
326    }
327
328    /// Returns the length of the hash chiplet trace.
329    pub fn hash_chiplet_len(&self) -> usize {
330        self.hash_chiplet_len
331    }
332
333    /// Returns the length of the bitwise trace.
334    pub fn bitwise_chiplet_len(&self) -> usize {
335        self.bitwise_chiplet_len
336    }
337
338    /// Returns the length of the memory trace.
339    pub fn memory_chiplet_len(&self) -> usize {
340        self.memory_chiplet_len
341    }
342
343    /// Returns the length of the ACE chiplet trace.
344    pub fn ace_chiplet_len(&self) -> usize {
345        self.ace_chiplet_len
346    }
347
348    /// Returns the length of the kernel ROM trace.
349    pub fn kernel_rom_len(&self) -> usize {
350        self.kernel_rom_len
351    }
352
353    /// Returns the length of the trace required to accommodate chiplet components and 1
354    /// mandatory padding row required for ensuring sufficient trace length for auxiliary connector
355    /// columns that rely on the memory chiplet.
356    pub fn trace_len(&self) -> usize {
357        self.hash_chiplet_len()
358            + self.bitwise_chiplet_len()
359            + self.memory_chiplet_len()
360            + self.ace_chiplet_len()
361            + self.kernel_rom_len()
362            + 1
363    }
364}
365
366// U32 HELPERS
367// ================================================================================================
368
369/// Splits an element into two 16 bit integer limbs. It assumes that the field element contains a
370/// valid 32-bit integer value.
371pub(crate) fn split_element_u32_into_u16(value: Felt) -> (Felt, Felt) {
372    let (hi, lo) = split_u32_into_u16(value.as_canonical_u64());
373    (Felt::new_unchecked(hi as u64), Felt::new_unchecked(lo as u64))
374}
375
376/// Splits a u64 integer assumed to contain a 32-bit value into two u16 integers.
377///
378/// # Errors
379/// Fails in debug mode if the provided value is not a 32-bit value.
380pub(crate) fn split_u32_into_u16(value: u64) -> (u16, u16) {
381    const U32MAX: u64 = u32::MAX as u64;
382    debug_assert!(value <= U32MAX, "not a 32-bit value");
383
384    let lo = value as u16;
385    let hi = (value >> 16) as u16;
386
387    (hi, lo)
388}
389
390// TEST HELPERS
391// ================================================================================================
392
393/// Builds a 17-op basic block payload that straddles a RESPAN batch boundary, plus the initial
394/// values its `Push` ops emit. Consumed by decoder / hasher tests that exercise multi-batch
395/// SPAN execution.
396#[cfg(test)]
397pub(super) fn build_span_with_respan_ops() -> (Vec<Operation>, Vec<Felt>) {
398    let iv = [1, 3, 5, 7, 9, 11, 13, 15, 17].to_elements();
399    let ops = alloc::vec![
400        Operation::Push(iv[0]),
401        Operation::Push(iv[1]),
402        Operation::Push(iv[2]),
403        Operation::Push(iv[3]),
404        Operation::Push(iv[4]),
405        Operation::Push(iv[5]),
406        Operation::Push(iv[6]),
407        // next batch
408        Operation::Push(iv[7]),
409        Operation::Push(iv[8]),
410        Operation::Add,
411        // drops to make sure stack overflow is empty on exit
412        Operation::Drop,
413        Operation::Drop,
414        Operation::Drop,
415        Operation::Drop,
416        Operation::Drop,
417        Operation::Drop,
418        Operation::Drop,
419        Operation::Drop,
420    ];
421    (ops, iv)
422}
miden_processor/trace/utils.rs

miden_processor/trace/
utils.rs
