Skip to main content

miden_processor/trace/chiplets/
mod.rs

1use alloc::vec::Vec;
2
3use miden_air::trace::{
4    CHIPLETS_CLK_COL, CHIPLETS_WIDTH,
5    chiplets::{
6        KERNEL_ROM_TRACE_WIDTH,
7        ace::ACE_CHIPLET_NUM_COLS,
8        bitwise::TRACE_WIDTH as BITWISE_WIDTH,
9        hasher::{HasherState, TRACE_WIDTH as HASHER_WIDTH},
10        memory::TRACE_WIDTH as MEMORY_WIDTH,
11    },
12    eidos_compression::NUM_EIDOS_COMPRESSION_COLS,
13};
14use miden_core::{field::PrimeCharacteristicRing, mast::OpBatch, program::KernelDescriptor};
15
16use crate::{
17    Felt, ONE, Word, ZERO,
18    crypto::merkle::MerklePath,
19    trace::{ChipletTraceFragment, RowIndex, range::RangeChecker},
20};
21
22mod bitwise;
23use bitwise::AEAD_STREAM_FRAGMENT_WIDTH;
24pub(crate) use bitwise::{AEAD_STREAM_CYCLE_LEN, Bitwise};
25
26const BITWISE_COL_START: usize = 2;
27const BITWISE_FRAGMENT_WIDTH: usize = if BITWISE_WIDTH > AEAD_STREAM_FRAGMENT_WIDTH {
28    BITWISE_WIDTH
29} else {
30    AEAD_STREAM_FRAGMENT_WIDTH
31};
32const _: () = assert!(BITWISE_COL_START + BITWISE_FRAGMENT_WIDTH == CHIPLETS_CLK_COL);
33
34mod hasher;
35pub(crate) use hasher::Hasher;
36pub use hasher::{
37    build_external_eidos_compression_traces, build_ordered_external_eidos_compression_traces,
38};
39
40pub(crate) fn build_and8_lookup_trace(and8_counts: &[u64]) -> Vec<Felt> {
41    hasher::build_and8_lookup_trace(and8_counts)
42}
43
44mod memory;
45pub(crate) use memory::Memory;
46
47mod ace;
48pub use ace::{
49    Ace, CircuitEvaluation, MAX_EVAL_CIRCUIT_INVOCATIONS, MAX_EVAL_CIRCUIT_WIRES,
50    MAX_NUM_ACE_WIRES, PTR_OFFSET_ELEM, PTR_OFFSET_WORD,
51};
52
53mod kernel_rom;
54pub(crate) use kernel_rom::KernelRom;
55
56#[cfg(test)]
57mod tests;
58
59// TRACE
60// ================================================================================================
61
62pub struct ChipletsTrace {
63    pub(crate) trace: Vec<Felt>,
64}
65
66pub struct EidosCompressionTrace {
67    pub(crate) trace: Vec<Felt>,
68}
69
70// CHIPLETS MODULE OF HASHER, BITWISE, MEMORY, ACE, AND KERNEL ROM CHIPLETS
71// ================================================================================================
72
73/// This module manages the VM's hasher, bitwise, memory, arithmetic circuit evaluation (ACE)
74/// and kernel ROM chiplets and is responsible for building a final execution trace from their
75/// stacked execution traces and chiplet selectors.
76///
77/// The module's trace can be thought of as 5 stacked segments in the following form.
78///
79/// The chiplet system uses `s_ctrl = column 0` to select the hasher controller.
80/// Columns 1-4 (`s1..s4`) subdivide the remaining region. The final shared data column is
81/// row-local: bitwise rows use it as normal/AEAD stream mode, while controller rows use it as
82/// the Merkle/padding flag.
83///
84/// * Hasher segment: fills the first rows of the trace up to the hasher `trace_len`.
85///   - column 0 (s_ctrl): ONE
86///   - columns 1-19: hasher-controller trace
87///   - column 22: Merkle/padding discriminator
88///
89/// * Bitwise segment: begins at the end of the hasher segment.
90///   - column 0 (s_ctrl): ZERO
91///   - column 1 (s1): ZERO
92///   - columns 2-14: execution trace of bitwise chiplet
93///   - columns 15-21: unused columns padded with ZERO
94///   - column 22: ZERO for normal bitwise rows
95///
96/// * Memory segment: begins at the end of the bitwise segment.
97///   - column 0 (s_ctrl): ZERO
98///   - column 1 (s1): ONE
99///   - column 2 (s2): ZERO
100///   - columns 3-19: execution trace of memory chiplet
101///   - columns 20-21: unused columns padded with ZERO
102///   - column 22: ZERO
103///
104/// * ACE segment: begins at the end of the memory segment.
105///   - column 0 (s_ctrl): ZERO
106///   - column 1-2 (s1, s2): ONE
107///   - column 3 (s3): ZERO
108///   - columns 4-20: execution trace of ACE chiplet
109///   - column 22: ZERO
110///
111/// * Kernel ROM segment: begins at the end of the ACE segment.
112///   - column 0 (s_ctrl): ZERO
113///   - columns 1-3 (s1, s2, s3): ONE
114///   - column 4 (s4): ZERO
115///   - columns 5-9: execution trace of kernel ROM chiplet
116///   - columns 10-21: unused columns padded with ZERO
117///   - column 22: ZERO
118///
119/// * Padding segment: fills the rest of the trace.
120///   - column 0 (s_ctrl): ZERO
121///   - columns 1-4 (s1..s4): ONE
122///   - columns 5-22: ZERO
123///
124/// Column ranges are stable for existing chiplets: controller uses columns 1..20,
125/// normal bitwise uses 2..15, memory uses 3..18, ACE uses 4..20, and kernel ROM
126/// uses 5..10. AEAD stream rows reuse the bitwise region and use columns 2..21 as
127/// payload, with column 22 as the AEAD stream selector.
128#[derive(Debug)]
129pub struct Chiplets {
130    pub hasher: Hasher,
131    pub bitwise: Bitwise,
132    pub memory: Memory,
133    pub ace: Ace,
134    pub kernel_rom: KernelRom,
135}
136
137impl Chiplets {
138    // CONSTRUCTOR
139    // --------------------------------------------------------------------------------------------
140
141    /// Returns a new chiplet set instantiated with the provided kernel descriptor.
142    pub fn new(kernel: KernelDescriptor) -> Self {
143        Self {
144            hasher: Hasher::default(),
145            bitwise: Bitwise::default(),
146            memory: Memory::default(),
147            ace: Ace::default(),
148            kernel_rom: KernelRom::new(kernel),
149        }
150    }
151
152    // PUBLIC ACCESSORS
153    // --------------------------------------------------------------------------------------------
154
155    /// Returns the chiplets trace length, including the mandatory padding row used by auxiliary
156    /// connector columns that read the memory chiplet.
157    pub fn trace_len(&self) -> usize {
158        self.hasher.trace_len()
159            + self.bitwise.trace_len()
160            + self.memory.trace_len()
161            + self.ace.trace_len()
162            + self.kernel_rom.trace_len()
163            + 1
164    }
165
166    /// Returns the unpadded trace length of the standalone Eidos compression AIR.
167    pub fn eidos_compression_trace_len(&self) -> usize {
168        self.hasher.eidos_compression_trace_len()
169    }
170
171    /// Returns the index of the first row of `Bitwise` execution trace.
172    pub fn bitwise_start(&self) -> RowIndex {
173        self.hasher.trace_len().into()
174    }
175
176    /// Returns the index of the first row of the `Memory` execution trace.
177    pub fn memory_start(&self) -> RowIndex {
178        self.bitwise_start() + self.bitwise.trace_len()
179    }
180
181    /// Returns the index of the first row of `ACE` execution trace.
182    pub fn ace_start(&self) -> RowIndex {
183        self.memory_start() + self.memory.trace_len()
184    }
185
186    /// Returns the index of the first row of `KernelRom` execution trace.
187    pub fn kernel_rom_start(&self) -> RowIndex {
188        self.ace_start() + self.ace.trace_len()
189    }
190
191    /// Returns the index of the first row of the padding section of the execution trace.
192    pub fn padding_start(&self) -> RowIndex {
193        self.kernel_rom_start() + self.kernel_rom.trace_len()
194    }
195
196    // EXECUTION TRACE
197    // --------------------------------------------------------------------------------------------
198
199    /// Adds all range checks required by the memory chiplet to the provided `RangeChecker`
200    /// instance.
201    pub fn append_range_checks(&self, range_checker: &mut RangeChecker) {
202        self.memory.append_range_checks(range_checker);
203    }
204
205    /// Adds range checks emitted by the standalone Eidos compression AIR.
206    pub fn append_eidos_compression_range_checks(
207        &self,
208        eidos_compression_height: usize,
209        range_checker: &mut RangeChecker,
210    ) {
211        self.hasher
212            .append_eidos_compression_range_checks(eidos_compression_height, range_checker);
213    }
214
215    /// Returns execution traces for `ChipletsAir` and `EidosCompressionAir`.
216    pub fn into_traces(
217        self,
218        trace_len: usize,
219        eidos_compression_trace_len: usize,
220    ) -> (ChipletsTrace, EidosCompressionTrace, Vec<u64>) {
221        assert!(self.trace_len() <= trace_len, "target trace length too small");
222        assert!(
223            self.eidos_compression_trace_len() <= eidos_compression_trace_len,
224            "target Eidos compression trace length too small"
225        );
226
227        let mut trace = Felt::zero_vec(CHIPLETS_WIDTH * trace_len);
228        let mut eidos_compression_trace =
229            Felt::zero_vec(NUM_EIDOS_COMPRESSION_COLS * eidos_compression_trace_len);
230        let and8_counts = self.fill_trace(&mut trace, trace_len, &mut eidos_compression_trace);
231
232        (
233            ChipletsTrace { trace },
234            EidosCompressionTrace { trace: eidos_compression_trace },
235            and8_counts,
236        )
237    }
238
239    // HELPER METHODS
240    // --------------------------------------------------------------------------------------------
241
242    /// Fills the provided trace for the chiplets module with the stacked execution traces of the
243    /// Hasher, Bitwise, Memory, ACE, and kernel ROM chiplets along with selector columns
244    /// to identify each individual chiplet trace in addition to padding to fill the rest of
245    /// the trace.
246    fn fill_trace(
247        self,
248        trace: &mut [Felt],
249        trace_len: usize,
250        eidos_compression_trace: &mut [Felt],
251    ) -> Vec<u64> {
252        const W: usize = CHIPLETS_WIDTH;
253        debug_assert_eq!(trace.len(), W * trace_len);
254
255        let memory_start: usize = self.memory_start().into();
256        let ace_start: usize = self.ace_start().into();
257        let kernel_rom_start: usize = self.kernel_rom_start().into();
258        let padding_start: usize = self.padding_start().into();
259
260        let Chiplets { hasher, bitwise, memory, ace, kernel_rom } = self;
261
262        // Per-chiplet row counts. Chiplets are stacked vertically, so each one's region is a
263        // contiguous band of rows: hasher [0, h), bitwise [h, h+b), and so on.
264        let hasher_len = hasher.trace_len();
265        let bitwise_len = bitwise.trace_len();
266        let memory_len = memory.trace_len();
267        let ace_len = ace.trace_len();
268        let kernel_rom_len = kernel_rom.trace_len();
269
270        // Chiplets are nested as hasher > bitwise > memory > ace > kernel_rom and begin at columns
271        // 1, 2, 3, 4, 5. Each chiplet's `copy_rows_from` writes its prefix selector ONEs and
272        // `chip_clk` along with its data; `s_ctrl` (col 0) is hasher-set per row, and
273        // padding rows are filled directly below.
274
275        // Carve `trace` into the per-chiplet contiguous row bands.
276        let (hasher_band, rest) = trace.split_at_mut(hasher_len * W);
277        let (bitwise_band, rest) = rest.split_at_mut(bitwise_len * W);
278        let (memory_band, rest) = rest.split_at_mut(memory_len * W);
279        let (ace_band, rest) = rest.split_at_mut(ace_len * W);
280        let (kernel_band, padding_band) = rest.split_at_mut(kernel_rom_len * W);
281
282        let mut hasher_fragment =
283            ChipletTraceFragment::with_overheads(hasher_band, W, 1, HASHER_WIDTH, 0, &[]);
284        let mut bitwise_fragment = ChipletTraceFragment::with_overheads(
285            bitwise_band,
286            W,
287            BITWISE_COL_START,
288            BITWISE_FRAGMENT_WIDTH,
289            hasher_len,
290            &[],
291        );
292        let mut memory_fragment = ChipletTraceFragment::with_overheads(
293            memory_band,
294            W,
295            3,
296            MEMORY_WIDTH,
297            memory_start,
298            &[1],
299        );
300        let mut ace_fragment = ChipletTraceFragment::with_overheads(
301            ace_band,
302            W,
303            4,
304            ACE_CHIPLET_NUM_COLS,
305            ace_start,
306            &[1, 2],
307        );
308        let mut kernel_rom_fragment = ChipletTraceFragment::with_overheads(
309            kernel_band,
310            W,
311            5,
312            KERNEL_ROM_TRACE_WIDTH,
313            kernel_rom_start,
314            &[1, 2, 3],
315        );
316
317        let mut and8_counts = Vec::new();
318        let mut bitwise_and8_counts = Vec::new();
319        rayon::scope(|s| {
320            let and8_counts = &mut and8_counts;
321            s.spawn(move |_| {
322                *and8_counts = hasher.fill_trace(&mut hasher_fragment, eidos_compression_trace);
323            });
324            let bitwise_and8_counts = &mut bitwise_and8_counts;
325            s.spawn(move |_| {
326                *bitwise_and8_counts = bitwise.fill_trace(&mut bitwise_fragment);
327            });
328            s.spawn(move |_| {
329                memory.fill_trace(&mut memory_fragment);
330            });
331            s.spawn(move |_| {
332                ace.fill_trace(&mut ace_fragment);
333            });
334            s.spawn(move |_| {
335                kernel_rom.fill_trace(&mut kernel_rom_fragment);
336            });
337            s.spawn(move |_| {
338                fill_padding_rows(padding_band, padding_start);
339            });
340        });
341
342        for (count, bitwise_count) in and8_counts.iter_mut().zip(bitwise_and8_counts) {
343            *count += bitwise_count;
344        }
345
346        and8_counts
347    }
348}
349
350/// Fills padding rows after the kernel ROM region: cols 1..=4 = ONE, chip_clk = row + 1.
351fn fill_padding_rows(band: &mut [Felt], row_offset: usize) {
352    const W: usize = CHIPLETS_WIDTH;
353    let (rows, _) = band.as_chunks_mut::<W>();
354    for (i, row) in rows.iter_mut().enumerate() {
355        row[1] = ONE;
356        row[2] = ONE;
357        row[3] = ONE;
358        row[4] = ONE;
359        row[W - 1] = Felt::from_u32((row_offset + i + 1) as u32);
360    }
361}
362
363// HELPER STRUCTS
364// ================================================================================================
365
366/// Result of a Merkle tree node update.
367///
368/// Contains the old root, the new root, and the trace row where the computation started.
369#[derive(Debug, Copy, Clone)]
370pub struct MerkleRootUpdate {
371    address: Felt,
372    old_root: Word,
373    new_root: Word,
374}
375
376impl MerkleRootUpdate {
377    pub fn get_address(&self) -> Felt {
378        self.address
379    }
380    pub fn get_old_root(&self) -> Word {
381        self.old_root
382    }
383    pub fn get_new_root(&self) -> Word {
384        self.new_root
385    }
386}