adf

core/src/cpu/trace.rs

@@ -15,6 +15,8 @@ use crate::runtime::{ExecutionRecord, Opcode};
 use crate::utils::env;
 use p3_field::PrimeField;
 use p3_matrix::dense::RowMajorMatrix;
+use p3_maybe_rayon::prelude::IntoParallelRefIterator;
+use p3_maybe_rayon::prelude::ParallelIterator;
 use std::borrow::BorrowMut;
 use std::collections::HashMap;
 use tracing::instrument;
@@ -46,32 +48,33 @@ impl<F: PrimeField> MachineAir<F> for CpuChip {
 
     #[instrument(name = "generate CPU trace", skip_all)]
     fn generate_trace(&self, record: &mut ExecutionRecord) -> RowMajorMatrix<F> {
-        let mut new_blu_events = Vec::new();
         let mut new_alu_events = HashMap::new();
-        let mut new_field_events = Vec::new();
+        let mut new_blu_events = Vec::new();
+        let mut new_field_events: Vec<FieldEvent> = Vec::new();
 
         // Generate the trace rows for each event.
-        let rows = record
+        let rows_with_events = record
             .cpu_events
-            .iter()
-            .map(|op| {
-                self.event_to_row(
-                    *op,
-                    &mut new_alu_events,
-                    &mut new_blu_events,
-                    &mut new_field_events,
-                )
-            })
+            .par_iter()
+            .map(|op: &CpuEvent| self.event_to_row::<F>(*op))
             .collect::<Vec<_>>();
 
+        let mut rows = Vec::<F>::new();
+        rows_with_events.into_iter().for_each(|row_with_events| {
+            let (row, alu_events, blu_events, field_events) = row_with_events;
+            rows.extend(row);
+            new_alu_events.extend(alu_events);
+            new_blu_events.extend(blu_events);
+            new_field_events.extend(field_events);
+        });
+
         // Add the dependency events to the shard.
         record.add_alu_events(new_alu_events);
         record.add_byte_lookup_events(new_blu_events);
         record.field_events.extend(new_field_events);
 
         // Convert the trace to a row major matrix.
-        let mut trace =
-            RowMajorMatrix::new(rows.into_iter().flatten().collect::<Vec<_>>(), NUM_CPU_COLS);
+        let mut trace = RowMajorMatrix::new(rows, NUM_CPU_COLS);
 
         // Pad the trace to a power of two.
         Self::pad_to_power_of_two::<F>(&mut trace.values);
@@ -85,10 +88,16 @@ impl CpuChip {
     fn event_to_row<F: PrimeField>(
         &self,
         event: CpuEvent,
-        new_alu_events: &mut HashMap<Opcode, Vec<alu::AluEvent>>,
-        new_blu_events: &mut Vec<ByteLookupEvent>,
-        new_field_events: &mut Vec<FieldEvent>,
-    ) -> [F; NUM_CPU_COLS] {
+    ) -> (
+        [F; NUM_CPU_COLS],
+        HashMap<Opcode, Vec<alu::AluEvent>>,
+        Vec<ByteLookupEvent>,
+        Vec<FieldEvent>,
+    ) {
+        let mut new_alu_events = HashMap::new();
+        let mut new_blu_events = Vec::new();
+        let mut new_field_events = Vec::new();
+
         let mut row = [F::zero(); NUM_CPU_COLS];
         let cols: &mut CpuCols<F> = row.as_mut_slice().borrow_mut();
 
@@ -104,13 +113,13 @@ impl CpuChip {
 
         // Populate memory accesses for a, b, and c.
         if let Some(record) = event.a_record {
-            cols.op_a_access.populate(record, new_field_events)
+            cols.op_a_access.populate(record, &mut new_field_events)
         }
         if let Some(MemoryRecordEnum::Read(record)) = event.b_record {
-            cols.op_b_access.populate(record, new_field_events)
+            cols.op_b_access.populate(record, &mut new_field_events)
         }
         if let Some(MemoryRecordEnum::Read(record)) = event.c_record {
-            cols.op_c_access.populate(record, new_field_events)
+            cols.op_c_access.populate(record, &mut new_field_events)
         }
 
         // Populate memory accesses for reading from memory.
@@ -120,19 +129,19 @@ impl CpuChip {
         if let Some(record) = event.memory_record {
             memory_columns
                 .memory_access
-                .populate(record, new_field_events)
+                .populate(record, &mut new_field_events)
         }
 
         // Populate memory, branch, jump, and auipc specific fields.
-        self.populate_memory(cols, event, new_alu_events, new_blu_events);
-        self.populate_branch(cols, event, new_alu_events);
-        self.populate_jump(cols, event, new_alu_events);
-        self.populate_auipc(cols, event, new_alu_events);
+        self.populate_memory(cols, event, &mut new_alu_events, &mut new_blu_events);
+        self.populate_branch(cols, event, &mut new_alu_events);
+        self.populate_jump(cols, event, &mut new_alu_events);
+        self.populate_auipc(cols, event, &mut new_alu_events);
 
         // Assert that the instruction is not a no-op.
         cols.is_real = F::one();
 
-        row
+        (row, new_alu_events, new_blu_events, new_field_events)
     }
 
     /// Populates columns related to memory.

core/src/stark/prover.rs

@@ -17,7 +17,6 @@ use serde::de::DeserializeOwned;
 use serde::Serialize;
 use std::cmp::max;
 use std::marker::PhantomData;
-use std::mem;
 
 use super::folder::ProverConstraintFolder;
 use super::util::decompose_and_flatten;