Skip to main content

miden_processor/trace/chiplets/ace/
trace.rs

1use alloc::vec::Vec;
2use core::borrow::BorrowMut;
3
4use miden_air::{
5    AceCols, QuadFeltExpr,
6    trace::{RowIndex, chiplets::ace::ACE_CHIPLET_NUM_COLS},
7};
8use miden_core::{
9    Felt, Word,
10    field::{BasedVectorSpace, QuadFelt},
11    serde::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable},
12};
13
14use super::{
15    MAX_NUM_ACE_WIRES,
16    instruction::{Op, decode_instruction},
17};
18use crate::{ContextId, errors::AceError};
19
20/// One row of the ACE chiplet trace in `READ` mode: two memory-loaded wires per row, plus the
21/// pointer of the word that was loaded.
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23struct ReadNode {
24    ptr: Felt,
25    id_0: Felt,
26    v_0: QuadFelt,
27    id_1: Felt,
28    v_1: QuadFelt,
29}
30
31/// One row of the ACE chiplet trace in `EVAL` mode: a single arithmetic gate `(id_0, v_0)` with
32/// two inputs `(id_1, v_1)` (left) and `(id_2, v_2)` (right), the instruction pointer that
33/// produced it, and the gate's `eval_op` selector.
34#[derive(Debug, Clone, Copy, PartialEq, Eq)]
35struct EvalNode {
36    ptr: Felt,
37    eval_op: Felt,
38    id_0: Felt,
39    v_0: QuadFelt,
40    id_1: Felt,
41    v_1: QuadFelt,
42    id_2: Felt,
43    v_2: QuadFelt,
44}
45
46/// Contains the variable and evaluation nodes resulting from the evaluation of a circuit.
47/// The output value is checked to be equal to 0.
48///
49/// The set of nodes is used to fill the ACE chiplet trace.
50#[derive(Debug, Clone, PartialEq, Eq)]
51pub struct CircuitEvaluation {
52    ctx: ContextId,
53    clk: RowIndex,
54    wire_bus: WireBus,
55    read_nodes: Vec<ReadNode>,
56    eval_nodes: Vec<EvalNode>,
57}
58
59impl CircuitEvaluation {
60    /// Generates the nodes in the graph generated by evaluating the inputs and circuit
61    /// located in a contiguous memory region.
62    ///
63    /// # Panics:
64    /// This function panics if the number of rows for each section leads to more than
65    /// [`MAX_NUM_ACE_WIRES`] wires.
66    pub fn new(ctx: ContextId, clk: RowIndex, num_read_rows: u32, num_eval_rows: u32) -> Self {
67        let num_wires = 2 * (num_read_rows as u64) + (num_eval_rows as u64);
68        assert!(num_wires <= MAX_NUM_ACE_WIRES as u64, "too many wires");
69
70        Self {
71            ctx,
72            clk,
73            wire_bus: WireBus::new(num_wires as u32),
74            read_nodes: Vec::with_capacity(num_read_rows as usize),
75            eval_nodes: Vec::with_capacity(num_eval_rows as usize),
76        }
77    }
78
79    pub fn num_rows(&self) -> usize {
80        self.read_nodes.len() + self.eval_nodes.len()
81    }
82
83    pub fn clk(&self) -> u32 {
84        self.clk.into()
85    }
86
87    pub fn ctx(&self) -> u32 {
88        self.ctx.into()
89    }
90
91    pub fn num_read_rows(&self) -> u32 {
92        self.read_nodes.len() as u32
93    }
94
95    pub fn num_eval_rows(&self) -> u32 {
96        self.eval_nodes.len() as u32
97    }
98
99    /// Reads the word from memory at `ptr`, interpreting it as `[v_00, v_01, v_10, v_11]`, and
100    /// adds wires with values `v_0 = QuadFelt(v_00, v_01)` and `v_1 = QuadFelt(v_10, v_11)`.
101    pub fn do_read(&mut self, ptr: Felt, word: Word) {
102        let v_0 = QuadFelt::from_basis_coefficients_fn(|i: usize| [word[0], word[1]][i]);
103        let id_0 = self.wire_bus.insert(v_0);
104
105        let v_1 = QuadFelt::from_basis_coefficients_fn(|i: usize| [word[2], word[3]][i]);
106        let id_1 = self.wire_bus.insert(v_1);
107
108        self.read_nodes.push(ReadNode { ptr, id_0, v_0, id_1, v_1 });
109    }
110
111    /// Reads the next instruction at `ptr`, requests the inputs from the wire bus
112    /// and inserts a new wire with the result.
113    pub fn do_eval(&mut self, ptr: Felt, instruction: Felt) -> Result<(), AceError> {
114        let (id_l, id_r, op) = decode_instruction(instruction)
115            .ok_or(AceError("failed to decode instruction".into()))?;
116
117        let v_l = self
118            .wire_bus
119            .read_value(id_l)
120            .ok_or(AceError("failed to read from the wiring bus".into()))?;
121        let id_1 = Felt::from_u32(id_l);
122
123        let v_r = self
124            .wire_bus
125            .read_value(id_r)
126            .ok_or(AceError("failed to read from the wiring bus".into()))?;
127        let id_2 = Felt::from_u32(id_r);
128
129        let v_0 = match op {
130            Op::Sub => v_l - v_r,
131            Op::Mul => v_l * v_r,
132            Op::Add => v_l + v_r,
133        };
134        let id_0 = self.wire_bus.insert(v_0);
135
136        let eval_op = match op {
137            Op::Sub => -Felt::ONE,
138            Op::Mul => Felt::ZERO,
139            Op::Add => Felt::ONE,
140        };
141
142        self.eval_nodes.push(EvalNode {
143            ptr,
144            eval_op,
145            id_0,
146            v_0,
147            id_1,
148            v_1: v_l,
149            id_2,
150            v_2: v_r,
151        });
152        Ok(())
153    }
154
155    /// Writes this circuit evaluation's rows into the row-major buffer `out`
156    /// (`ACE_CHIPLET_NUM_COLS` contiguous cells per row), starting at row `offset`. `out`
157    /// is assumed zero-initialized, so columns that are zero on a row are left untouched.
158    pub fn fill(&self, offset: usize, out: &mut [Felt]) {
159        const W: usize = ACE_CHIPLET_NUM_COLS;
160        let (out_rows, _) = out.as_chunks_mut::<W>();
161        let num_read_rows = self.read_nodes.len();
162        let num_eval_rows = self.eval_nodes.len();
163
164        let ctx_felt: Felt = self.ctx.into();
165        let clk_felt: Felt = self.clk.into();
166        let eval_section_first_idx = Felt::from_u32(num_eval_rows as u32 - 1);
167        let mut multiplicities_iter = self.wire_bus.wires.iter().map(|(_v, m)| Felt::from_u32(*m));
168
169        // READ rows.
170        for (i, node) in self.read_nodes.iter().enumerate() {
171            let cols: &mut AceCols<Felt> = out_rows[offset + i].as_mut_slice().borrow_mut();
172            cols.s_start = if i == 0 { Felt::ONE } else { Felt::ZERO };
173            cols.s_block = Felt::ZERO;
174            cols.ctx = ctx_felt;
175            cols.clk = clk_felt;
176            cols.ptr = node.ptr;
177            cols.id_0 = node.id_0;
178            cols.v_0 = quad_to_expr(node.v_0);
179            cols.id_1 = node.id_1;
180            cols.v_1 = quad_to_expr(node.v_1);
181
182            let m_0 = multiplicities_iter
183                .next()
184                .expect("the m0 multiplicities were not constructed properly");
185            let m_1 = multiplicities_iter
186                .next()
187                .expect("the m1 multiplicities were not constructed properly");
188
189            let read = cols.read_mut();
190            read.num_eval = eval_section_first_idx;
191            read.m_0 = m_0;
192            read.m_1 = m_1;
193        }
194
195        // EVAL rows.
196        for (i, node) in self.eval_nodes.iter().enumerate() {
197            let cols: &mut AceCols<Felt> =
198                out_rows[offset + num_read_rows + i].as_mut_slice().borrow_mut();
199            cols.s_start = Felt::ZERO;
200            cols.s_block = Felt::ONE;
201            cols.ctx = ctx_felt;
202            cols.clk = clk_felt;
203            cols.ptr = node.ptr;
204            cols.eval_op = node.eval_op;
205            cols.id_0 = node.id_0;
206            cols.v_0 = quad_to_expr(node.v_0);
207            cols.id_1 = node.id_1;
208            cols.v_1 = quad_to_expr(node.v_1);
209
210            let m_0 = multiplicities_iter
211                .next()
212                .expect("the m0 multiplicities were not constructed properly");
213
214            let eval = cols.eval_mut();
215            eval.id_2 = node.id_2;
216            eval.v_2 = quad_to_expr(node.v_2);
217            eval.m_0 = m_0;
218        }
219
220        let next = multiplicities_iter.next();
221        debug_assert!(next.is_none());
222    }
223
224    /// Returns the output value, if the circuit has finished evaluating.
225    pub fn output_value(&self) -> Option<QuadFelt> {
226        if !self.wire_bus.is_finalized() {
227            return None;
228        }
229        self.wire_bus.wires.last().map(|(v, _m)| *v)
230    }
231}
232
233/// Lifts a `QuadFelt` value into the [`QuadFeltExpr<Felt>`] basis-coefficient pair expected by the
234/// chiplet column structs.
235fn quad_to_expr(v: QuadFelt) -> QuadFeltExpr<Felt> {
236    let c = v.as_basis_coefficients_slice();
237    QuadFeltExpr(c[0], c[1])
238}
239
240/// Processor-local state used to construct the circuit witness sequentially.
241///
242/// Unlike the ACE AIR's order-independent wiring relation, this resolves only wires already
243/// inserted by the processor.
244///
245/// Gates are fan-in 2 but can have fan-out up to the field characteristic which, given the bounds
246/// on the execution trace length, means practically arbitrary fan-out.
247/// The main idea, with some slight variations between the `READ` and `EVAL` sections, is, for each
248/// gate, to "receive" the values of the input wires from the bus and to "send" the value of
249/// the value of the output wire back with multiplicity equal to the fan-out of the respective gate.
250/// Note that the messages include extra data in order to avoid collisions.
251#[derive(Debug, Clone, PartialEq, Eq)]
252struct WireBus {
253    // Circuit ID as Felt of the next wire to be inserted
254    id_next: Felt,
255    // Pairs of values and multiplicities
256    // The wire with index `id` is stored at `num_wires - 1 - id`
257    wires: Vec<(QuadFelt, u32)>,
258    // Total expected number of wires to be inserted.
259    num_wires: u32,
260}
261
262impl WireBus {
263    fn new(num_wires: u32) -> Self {
264        Self {
265            wires: Vec::with_capacity(num_wires as usize),
266            num_wires,
267            id_next: Felt::from_u32(num_wires - 1),
268        }
269    }
270
271    /// Inserts a new value into the bus, and returns its expected id as `Felt`
272    fn insert(&mut self, value: QuadFelt) -> Felt {
273        debug_assert!(!self.is_finalized());
274        self.wires.push((value, 0));
275        let id = self.id_next;
276        self.id_next -= Felt::ONE;
277        id
278    }
279
280    /// Reads the value of a wire with given `id`, incrementing its multiplicity.
281    /// Returns `None` if the sequential witness construction has not resolved the wire yet.
282    fn read_value(&mut self, id: u32) -> Option<QuadFelt> {
283        // Ensures subtracting the id from num_wires results in a valid wire index
284        let (v, m) = self
285            .num_wires
286            .checked_sub(id + 1)
287            .and_then(|id| self.wires.get_mut(id as usize))?;
288        *m += 1;
289        Some(*v)
290    }
291
292    /// Return true if the expected number of wires have been inserted.
293    fn is_finalized(&self) -> bool {
294        self.wires.len() == self.num_wires as usize
295    }
296}
297
298// SERIALIZATION
299// ================================================================================================
300
301impl Serializable for ReadNode {
302    fn write_into<W: ByteWriter>(&self, target: &mut W) {
303        self.ptr.write_into(target);
304        self.id_0.write_into(target);
305        self.v_0.write_into(target);
306        self.id_1.write_into(target);
307        self.v_1.write_into(target);
308    }
309}
310
311impl Deserializable for ReadNode {
312    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
313        Ok(Self {
314            ptr: Felt::read_from(source)?,
315            id_0: Felt::read_from(source)?,
316            v_0: QuadFelt::read_from(source)?,
317            id_1: Felt::read_from(source)?,
318            v_1: QuadFelt::read_from(source)?,
319        })
320    }
321
322    fn min_serialized_size() -> usize {
323        Felt::min_serialized_size() * 3 + QuadFelt::min_serialized_size() * 2
324    }
325}
326
327impl Serializable for EvalNode {
328    fn write_into<W: ByteWriter>(&self, target: &mut W) {
329        self.ptr.write_into(target);
330        self.eval_op.write_into(target);
331        self.id_0.write_into(target);
332        self.v_0.write_into(target);
333        self.id_1.write_into(target);
334        self.v_1.write_into(target);
335        self.id_2.write_into(target);
336        self.v_2.write_into(target);
337    }
338}
339
340impl Deserializable for EvalNode {
341    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
342        Ok(Self {
343            ptr: Felt::read_from(source)?,
344            eval_op: Felt::read_from(source)?,
345            id_0: Felt::read_from(source)?,
346            v_0: QuadFelt::read_from(source)?,
347            id_1: Felt::read_from(source)?,
348            v_1: QuadFelt::read_from(source)?,
349            id_2: Felt::read_from(source)?,
350            v_2: QuadFelt::read_from(source)?,
351        })
352    }
353
354    fn min_serialized_size() -> usize {
355        Felt::min_serialized_size() * 5 + QuadFelt::min_serialized_size() * 3
356    }
357}
358
359impl Serializable for WireBus {
360    fn write_into<W: ByteWriter>(&self, target: &mut W) {
361        self.id_next.write_into(target);
362        self.wires.write_into(target);
363        self.num_wires.write_into(target);
364    }
365}
366
367impl Deserializable for WireBus {
368    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
369        let id_next = Felt::read_from(source)?;
370        let wires = Vec::<(QuadFelt, u32)>::read_from(source)?;
371        let wire_count = wires.len();
372        let num_wires = u32::read_from(source)?;
373        if num_wires == 0 {
374            return Err(DeserializationError::InvalidValue(
375                "ACE wire bus must contain at least one wire".into(),
376            ));
377        }
378        if num_wires > MAX_NUM_ACE_WIRES {
379            return Err(DeserializationError::InvalidValue(format!(
380                "ACE declared wire count {num_wires} exceeds maximum {MAX_NUM_ACE_WIRES}"
381            )));
382        }
383        if wire_count != num_wires as usize {
384            return Err(DeserializationError::InvalidValue(format!(
385                "ACE wire count {wire_count} does not match declared wire count {num_wires}"
386            )));
387        }
388        Ok(Self { id_next, wires, num_wires })
389    }
390
391    fn min_serialized_size() -> usize {
392        Felt::min_serialized_size() + Vec::<u8>::min_serialized_size() + u32::min_serialized_size()
393    }
394}
395
396impl Serializable for CircuitEvaluation {
397    fn write_into<W: ByteWriter>(&self, target: &mut W) {
398        self.ctx.write_into(target);
399        self.clk.write_into(target);
400        self.wire_bus.write_into(target);
401        self.read_nodes.write_into(target);
402        self.eval_nodes.write_into(target);
403    }
404}
405
406impl Deserializable for CircuitEvaluation {
407    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
408        let evaluation = Self {
409            ctx: ContextId::read_from(source)?,
410            clk: RowIndex::read_from(source)?,
411            wire_bus: WireBus::read_from(source)?,
412            read_nodes: Vec::<ReadNode>::read_from(source)?,
413            eval_nodes: Vec::<EvalNode>::read_from(source)?,
414        };
415        evaluation.validate_wire_count()?;
416        Ok(evaluation)
417    }
418
419    fn min_serialized_size() -> usize {
420        ContextId::min_serialized_size()
421            + RowIndex::min_serialized_size()
422            + WireBus::min_serialized_size()
423            + Vec::<ReadNode>::min_serialized_size()
424            + Vec::<EvalNode>::min_serialized_size()
425    }
426}
427
428impl CircuitEvaluation {
429    fn validate_wire_count(&self) -> Result<(), DeserializationError> {
430        if self.eval_nodes.is_empty() {
431            return Err(DeserializationError::InvalidValue(
432                "ACE circuit evaluation must contain at least one eval node".into(),
433            ));
434        }
435        let read_wires = self.read_nodes.len().checked_mul(2).ok_or_else(|| {
436            DeserializationError::InvalidValue("ACE read-node wire count overflow".into())
437        })?;
438        let expected_wires = read_wires.checked_add(self.eval_nodes.len()).ok_or_else(|| {
439            DeserializationError::InvalidValue("ACE total wire count overflow".into())
440        })?;
441
442        if expected_wires == 0 {
443            return Err(DeserializationError::InvalidValue(
444                "ACE circuit evaluation must contain at least one wire".into(),
445            ));
446        }
447        if expected_wires > MAX_NUM_ACE_WIRES as usize {
448            return Err(DeserializationError::InvalidValue(format!(
449                "ACE circuit evaluation wire count {expected_wires} exceeds maximum {MAX_NUM_ACE_WIRES}"
450            )));
451        }
452        if self.wire_bus.num_wires as usize != expected_wires {
453            return Err(DeserializationError::InvalidValue(format!(
454                "ACE wire bus count {} does not match read/eval node wire count {expected_wires}",
455                self.wire_bus.num_wires
456            )));
457        }
458
459        Ok(())
460    }
461}
462
463#[cfg(test)]
464mod serialization_tests {
465    use alloc::vec;
466
467    use super::*;
468
469    fn sample_read_node(value: QuadFelt) -> ReadNode {
470        ReadNode {
471            ptr: Felt::ZERO,
472            id_0: Felt::ZERO,
473            v_0: value,
474            id_1: Felt::ONE,
475            v_1: value,
476        }
477    }
478
479    fn sample_eval_node(value: QuadFelt) -> EvalNode {
480        EvalNode {
481            ptr: Felt::ZERO,
482            eval_op: Felt::ZERO,
483            id_0: Felt::ZERO,
484            v_0: value,
485            id_1: Felt::ZERO,
486            v_1: value,
487            id_2: Felt::ONE,
488            v_2: value,
489        }
490    }
491
492    #[test]
493    fn circuit_evaluation_read_rejects_mismatched_wire_bus_count() {
494        let value = QuadFelt::new([Felt::ONE, Felt::ZERO]);
495        let evaluation = CircuitEvaluation {
496            ctx: ContextId::from(0),
497            clk: RowIndex::from(0_u32),
498            wire_bus: WireBus {
499                id_next: Felt::ZERO,
500                wires: vec![(value, 0), (value, 0)],
501                num_wires: 2,
502            },
503            read_nodes: vec![sample_read_node(value)],
504            eval_nodes: vec![sample_eval_node(value)],
505        };
506
507        let err = CircuitEvaluation::read_from_bytes(&evaluation.to_bytes()).unwrap_err();
508        let DeserializationError::InvalidValue(message) = err else {
509            panic!("expected invalid ACE wire count error");
510        };
511        assert!(message.contains("does not match read/eval node wire count"));
512    }
513
514    #[test]
515    fn circuit_evaluation_read_rejects_empty_eval_section() {
516        let value = QuadFelt::new([Felt::ONE, Felt::ZERO]);
517        let evaluation = CircuitEvaluation {
518            ctx: ContextId::from(0),
519            clk: RowIndex::from(0_u32),
520            wire_bus: WireBus {
521                id_next: Felt::ZERO,
522                wires: vec![(value, 0), (value, 0)],
523                num_wires: 2,
524            },
525            read_nodes: vec![sample_read_node(value)],
526            eval_nodes: Vec::new(),
527        };
528
529        let err = CircuitEvaluation::read_from_bytes(&evaluation.to_bytes()).unwrap_err();
530        let DeserializationError::InvalidValue(message) = err else {
531            panic!("expected invalid ACE eval section error");
532        };
533        assert!(message.contains("at least one eval node"));
534    }
535}