Skip to main content

miden_processor/trace/
utils.rs

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