openvm_sha2_circuit/sha2_chips/main_chip/
trace.rs

1use openvm_circuit::{
2    arch::*,
3    system::memory::{
4        offline_checker::{MemoryReadAuxRecord, MemoryWriteBytesAuxRecord},
5        MemoryAuxColsFactory,
6    },
7    utils::next_power_of_two_or_zero,
8};
9use openvm_circuit_primitives::Chip;
10use openvm_cpu_backend::CpuBackend;
11use openvm_instructions::riscv::{RV32_CELL_BITS, RV32_REGISTER_NUM_LIMBS};
12use openvm_sha2_air::set_arrayview_from_u8_slice;
13use openvm_stark_backend::{
14    p3_field::{PrimeCharacteristicRing, PrimeField32},
15    p3_matrix::{dense::RowMajorMatrix, Matrix},
16    p3_maybe_rayon::prelude::*,
17    prover::AirProvingContext,
18    StarkProtocolConfig, Val,
19};
20
21use crate::{
22    Sha2ColsRefMut, Sha2Config, Sha2MainChip, Sha2Metadata, Sha2RecordLayout, Sha2RecordMut,
23    Sha2SharedRecords, SHA2_WRITE_SIZE,
24};
25
26// We will allocate a new trace matrix instead of using the record arena directly,
27// because we want to pass the record arena to Sha2BlockHasherChip when we are done.
28impl<RA, SC: StarkProtocolConfig, C: Sha2Config> Chip<RA, CpuBackend<SC>>
29    for Sha2MainChip<Val<SC>, C>
30where
31    Val<SC>: PrimeField32,
32    SC: StarkProtocolConfig,
33    RA: RowMajorMatrixArena<Val<SC>> + Send + Sync,
34{
35    fn generate_proving_ctx(&self, arena: RA) -> AirProvingContext<CpuBackend<SC>> {
36        // Since Sha2Metadata::get_num_rows() = 1, the number of rows used is equal to the number of
37        // SHA-2 instructions executed.
38        let rows_used = arena.trace_offset() / arena.width();
39
40        // We will fill the trace into a separate buffer, because we want to pass the arena to the
41        // Sha2BlockHasherChip when we are done.
42        // Sha2MainChip uses 1 row per instruction, we allocate rows_used * arena.width() space for
43        // the trace.
44        let height = next_power_of_two_or_zero(rows_used);
45        let trace = Val::<SC>::zero_vec(height * arena.width());
46        let mut trace_matrix = RowMajorMatrix::new(trace, arena.width());
47        let mem_helper = self.mem_helper.as_borrowed();
48
49        let mut records = arena.into_matrix();
50
51        self.fill_trace(&mem_helper, &mut trace_matrix, rows_used, &mut records);
52
53        // Pass the records to Sha2BlockHasherChip
54        *self.records.lock().unwrap() = Some(Sha2SharedRecords {
55            num_records: rows_used,
56            matrix: records,
57        });
58
59        AirProvingContext::simple_no_pis(trace_matrix)
60    }
61}
62
63// Note: we would like to just impl TraceFiller here, but we can't because we need to pass the
64// records and row_idx to the tracegen functions.
65impl<F: PrimeField32, C: Sha2Config> Sha2MainChip<F, C> {
66    // Preconditions:
67    // - trace should be a matrix with width = Sha2MainAir::width() and height = rows_used
68    // - trace should be filled with all zeros
69    // - records should be a matrix with height = rows_used, where each row stores a record
70    pub fn fill_trace(
71        &self,
72        mem_helper: &MemoryAuxColsFactory<F>,
73        trace: &mut RowMajorMatrix<F>,
74        rows_used: usize,
75        records: &mut RowMajorMatrix<F>,
76    ) {
77        let width = trace.width();
78        trace.values[..rows_used * width]
79            .par_chunks_exact_mut(width)
80            .zip(records.par_rows_mut())
81            .enumerate()
82            .for_each(|(row_idx, (row_slice, record))| {
83                self.fill_trace_row_with_row_idx(mem_helper, row_slice, row_idx, record);
84            });
85    }
86
87    // Same as TraceFiller::fill_trace_row, except we also take the row index as a parameter.
88    //
89    // Note: the only reason the record parameter is mutable is that get_record_from_slice
90    // requires a &mut &mut [F] slice. This parameter type is useful in other places where
91    // get_record_from_slice is used, to circumvent the borrow checker. Here, we don't actually need
92    // this workaround (we could duplicate get_record_from_slice and modify it to take a &mut
93    // [F] slice), but we just use the existing function for simplicity.
94    fn fill_trace_row_with_row_idx(
95        &self,
96        mem_helper: &MemoryAuxColsFactory<F>,
97        row_slice: &mut [F],
98        row_idx: usize,
99        mut record: &mut [F],
100    ) where
101        F: Clone,
102    {
103        // SAFETY:
104        // - caller ensures `record` contains a valid record representation that was previously
105        //   written by the executor
106        // - record contains a valid Sha2RecordMut with the exact layout specified
107        // - get_record_from_slice will correctly split the buffer into header and other components
108        //   based on this layout.
109        let record: Sha2RecordMut = unsafe {
110            get_record_from_slice(
111                &mut record,
112                Sha2RecordLayout::new(Sha2Metadata {
113                    variant: C::VARIANT,
114                }),
115            )
116        };
117
118        // save all the components of the record on the stack so that we don't overwrite them when
119        // filling in the trace matrix.
120        let vm_record = record.inner.clone();
121
122        let mut message_bytes = Vec::with_capacity(C::BLOCK_BYTES);
123        message_bytes.extend_from_slice(record.message_bytes);
124
125        let mut prev_state = Vec::with_capacity(C::STATE_BYTES);
126        prev_state.extend_from_slice(record.prev_state);
127
128        let mut new_state = Vec::with_capacity(C::STATE_BYTES);
129        new_state.extend_from_slice(record.new_state);
130
131        let mut input_reads_aux =
132            Vec::with_capacity(C::BLOCK_READS * size_of::<MemoryReadAuxRecord>());
133        input_reads_aux.extend_from_slice(record.input_reads_aux);
134
135        let mut state_reads_aux =
136            Vec::with_capacity(C::STATE_READS * size_of::<MemoryReadAuxRecord>());
137        state_reads_aux.extend_from_slice(record.state_reads_aux);
138
139        let mut write_aux = Vec::with_capacity(
140            C::STATE_WRITES * size_of::<MemoryWriteBytesAuxRecord<SHA2_WRITE_SIZE>>(),
141        );
142        write_aux.extend_from_slice(record.write_aux);
143
144        let mut cols = Sha2ColsRefMut::from::<C>(row_slice);
145
146        *cols.block.request_id = F::from_usize(row_idx);
147        set_arrayview_from_u8_slice(&mut cols.block.message_bytes, message_bytes);
148        set_arrayview_from_u8_slice(&mut cols.block.prev_state, prev_state);
149        set_arrayview_from_u8_slice(&mut cols.block.new_state, new_state);
150
151        *cols.instruction.is_enabled = F::ONE;
152        cols.instruction.from_state.timestamp = F::from_u32(vm_record.timestamp);
153        cols.instruction.from_state.pc = F::from_u32(vm_record.from_pc);
154        *cols.instruction.dst_reg_ptr = F::from_u32(vm_record.dst_reg_ptr);
155        *cols.instruction.state_reg_ptr = F::from_u32(vm_record.state_reg_ptr);
156        *cols.instruction.input_reg_ptr = F::from_u32(vm_record.input_reg_ptr);
157
158        let dst_ptr_limbs = vm_record.dst_ptr.to_le_bytes();
159        let state_ptr_limbs = vm_record.state_ptr.to_le_bytes();
160        let input_ptr_limbs = vm_record.input_ptr.to_le_bytes();
161        set_arrayview_from_u8_slice(&mut cols.instruction.dst_ptr_limbs, dst_ptr_limbs);
162        set_arrayview_from_u8_slice(&mut cols.instruction.state_ptr_limbs, state_ptr_limbs);
163        set_arrayview_from_u8_slice(&mut cols.instruction.input_ptr_limbs, input_ptr_limbs);
164        let needs_range_check = [
165            dst_ptr_limbs[RV32_REGISTER_NUM_LIMBS - 1],
166            state_ptr_limbs[RV32_REGISTER_NUM_LIMBS - 1],
167            input_ptr_limbs[RV32_REGISTER_NUM_LIMBS - 1],
168            input_ptr_limbs[RV32_REGISTER_NUM_LIMBS - 1],
169        ];
170        let shift: u32 = 1 << (RV32_REGISTER_NUM_LIMBS * RV32_CELL_BITS - self.pointer_max_bits);
171        for pair in needs_range_check.chunks_exact(2) {
172            self.bitwise_lookup_chip
173                .request_range(pair[0] as u32 * shift, pair[1] as u32 * shift);
174        }
175
176        // fill in the register reads aux
177        let mut timestamp = vm_record.timestamp;
178        for (cols, vm_record) in cols
179            .mem
180            .register_aux
181            .iter_mut()
182            .zip(vm_record.register_reads_aux.iter())
183        {
184            mem_helper.fill(vm_record.prev_timestamp, timestamp, cols.as_mut());
185            timestamp += 1;
186        }
187
188        input_reads_aux.iter().zip(cols.mem.input_reads).for_each(
189            |(read_aux_record, read_aux_cols)| {
190                mem_helper.fill(
191                    read_aux_record.prev_timestamp,
192                    timestamp,
193                    read_aux_cols.as_mut(),
194                );
195                timestamp += 1;
196            },
197        );
198
199        state_reads_aux.iter().zip(cols.mem.state_reads).for_each(
200            |(state_aux_record, state_aux_cols)| {
201                mem_helper.fill(
202                    state_aux_record.prev_timestamp,
203                    timestamp,
204                    state_aux_cols.as_mut(),
205                );
206                timestamp += 1;
207            },
208        );
209
210        write_aux
211            .iter()
212            .zip(cols.mem.write_aux)
213            .for_each(|(write_aux_record, write_aux_cols)| {
214                write_aux_cols.set_prev_data(write_aux_record.prev_data.map(F::from_u8));
215                mem_helper.fill(
216                    write_aux_record.prev_timestamp,
217                    timestamp,
218                    write_aux_cols.as_mut(),
219                );
220                timestamp += 1;
221            });
222    }
223}
openvm_sha2_circuit/sha2_chips/main_chip/trace.rs

openvm_sha2_circuit/sha2_chips/main_chip/
trace.rs
