1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153

use std::borrow::BorrowMut;

use crate::{
    air::BinomialExtensionUtils,
    memory::MemoryCols,
    runtime::{
        get_heap_size_range_check_events, instruction_is_heap_expand, ExecutionRecord, Opcode,
        RecursionProgram, D,
    },
};
use p3_field::{extension::BinomiallyExtendable, PrimeField32};
use p3_matrix::dense::RowMajorMatrix;
use p3_maybe_rayon::prelude::IndexedParallelIterator;
use p3_maybe_rayon::prelude::ParallelIterator;
use p3_maybe_rayon::prelude::ParallelSliceMut;
use sp1_core::{
    air::{BinomialExtension, MachineAir},
    utils::{next_power_of_two, par_for_each_row},
};
use tracing::instrument;

use super::{CpuChip, CpuCols, NUM_CPU_COLS};

impl<F: PrimeField32 + BinomiallyExtendable<D>, const L: usize> MachineAir<F> for CpuChip<F, L> {
    type Record = ExecutionRecord<F>;
    type Program = RecursionProgram<F>;

    fn name(&self) -> String {
        "CPU".to_string()
    }

    fn generate_dependencies(&self, _: &Self::Record, _: &mut Self::Record) {
        // There are no dependencies, since we do it all in the runtime. This is just a placeholder.
    }

    #[instrument(name = "generate cpu trace", level = "debug", skip_all, fields(rows = input.cpu_events.len()))]
    fn generate_trace(
        &self,
        input: &ExecutionRecord<F>,
        _: &mut ExecutionRecord<F>,
    ) -> RowMajorMatrix<F> {
        let nb_events = input.cpu_events.len();
        let padded_nb_rows = next_power_of_two(nb_events, self.fixed_log2_rows);
        let mut values = vec![F::zero(); padded_nb_rows * NUM_CPU_COLS];

        par_for_each_row(&mut values, NUM_CPU_COLS, |i, row| {
            if i >= nb_events {
                return;
            }
            let event = &input.cpu_events[i];
            let cols: &mut CpuCols<F> = row.borrow_mut();

            cols.clk = event.clk;
            cols.pc = event.pc;
            cols.fp = event.fp;

            // Populate the instruction related columns.
            cols.selectors.populate(&event.instruction);
            cols.instruction.populate(&event.instruction);

            // Populate the register columns.
            if let Some(record) = &event.a_record {
                cols.a.populate(record);
            }
            if let Some(record) = &event.b_record {
                cols.b.populate(record);
            } else {
                *cols.b.value_mut() = event.instruction.op_b;
            }
            if let Some(record) = &event.c_record {
                cols.c.populate(record);
            } else {
                *cols.c.value_mut() = event.instruction.op_c;
            }
            if let Some(record) = &event.memory_record {
                let memory_cols = cols.opcode_specific.memory_mut();
                memory_cols.memory.populate(record);
                memory_cols.memory_addr = record.addr;
            }

            // Populate the heap columns.
            if instruction_is_heap_expand(&event.instruction) {
                let (u16_range_check, u12_range_check) =
                    get_heap_size_range_check_events(cols.a.value()[0]);

                let heap_cols = cols.opcode_specific.heap_expand_mut();
                heap_cols.diff_16bit_limb = F::from_canonical_u16(u16_range_check.val);
                heap_cols.diff_12bit_limb = F::from_canonical_u16(u12_range_check.val);
            }

            // Populate the branch columns.
            if matches!(
                event.instruction.opcode,
                Opcode::BEQ | Opcode::BNE | Opcode::BNEINC
            ) {
                let branch_cols = cols.opcode_specific.branch_mut();
                let a_ext: BinomialExtension<F> =
                    BinomialExtensionUtils::from_block(*cols.a.value());
                let b_ext: BinomialExtension<F> =
                    BinomialExtensionUtils::from_block(*cols.b.value());

                let (comparison_diff, do_branch) = match event.instruction.opcode {
                    Opcode::BEQ => (a_ext - b_ext, a_ext == b_ext),
                    Opcode::BNE | Opcode::BNEINC => (a_ext - b_ext, a_ext != b_ext),
                    _ => unreachable!(),
                };

                branch_cols
                    .comparison_diff
                    .populate((comparison_diff).as_block());
                branch_cols.comparison_diff_val = comparison_diff;
                branch_cols.do_branch = F::from_bool(do_branch);
                branch_cols.next_pc = if do_branch {
                    event.pc + event.instruction.op_c[0]
                } else {
                    event.pc + F::one()
                };
            }

            // Populate the public values columns.
            if event.instruction.opcode == Opcode::Commit {
                let public_values_cols = cols.opcode_specific.public_values_mut();
                let idx = cols.b.prev_value()[0].as_canonical_u32() as usize;
                public_values_cols.idx_bitmap[idx] = F::one();
            }

            cols.is_real = F::one();
        });

        let mut trace = RowMajorMatrix::new(values, NUM_CPU_COLS);

        // Fill in the dummy values for the padding rows.
        let padded_rows = trace
            .values
            .par_chunks_mut(NUM_CPU_COLS)
            .enumerate()
            .skip(input.cpu_events.len());
        padded_rows.for_each(|(i, row)| {
            let cols: &mut CpuCols<F> = row.borrow_mut();
            cols.selectors.is_noop = F::one();
            cols.instruction.imm_b = F::one();
            cols.instruction.imm_c = F::one();
            cols.clk = F::from_canonical_u32(4) * F::from_canonical_usize(i);
            cols.instruction.imm_b = F::from_canonical_u32(1);
            cols.instruction.imm_c = F::from_canonical_u32(1);
        });
        trace
    }

    fn included(&self, _: &Self::Record) -> bool {
        true
    }
}