Parallelized add_mul trace gen.

Author: Alon-Ti

crates/prover/src/components/add_mul_opcode/prover.rs

@@ -1,15 +1,15 @@
 use itertools::{zip_eq, Itertools};
 use num_traits::One;
+use rayon::iter::{IndexedParallelIterator, IntoParallelRefIterator, ParallelIterator};
+use stwo_air_utils::trace::component_trace::ComponentTrace;
+use stwo_air_utils_derive::{IterMut, ParIterMut, Uninitialized};
 use stwo_prover::constraint_framework::logup::LogupTraceGenerator;
 use stwo_prover::constraint_framework::Relation;
 use stwo_prover::core::backend::simd::m31::{PackedM31, LOG_N_LANES, N_LANES};
 use stwo_prover::core::backend::simd::qm31::PackedQM31;
 use stwo_prover::core::backend::simd::SimdBackend;
-use stwo_prover::core::backend::{Col, Column};
 use stwo_prover::core::fields::m31::M31;
 use stwo_prover::core::pcs::TreeBuilder;
-use stwo_prover::core::poly::circle::{CanonicCoset, CircleEvaluation};
-use stwo_prover::core::poly::BitReversedOrder;
 use stwo_prover::core::vcs::blake2_merkle::Blake2sMerkleChannel;
 
 use super::component::{Claim, InteractionClaim, INSTRUCTION_BASE};
@@ -63,51 +63,48 @@ impl ClaimGenerator {
         tree_builder: &mut TreeBuilder<'_, '_, SimdBackend, Blake2sMerkleChannel>,
         memory_trace_generator: &mut memory::ClaimGenerator,
     ) -> (Claim, InteractionClaimGenerator) {
-        let (trace, interaction_claim_generator) =
-            write_trace_simd(&self.inputs, memory_trace_generator);
-        interaction_claim_generator.memory.iter().for_each(|c| {
+        let (trace, lookup_data) = write_trace_simd(&self.inputs, memory_trace_generator);
+        lookup_data.memory.iter().for_each(|c| {
             c.iter()
                 .for_each(|v| memory_trace_generator.add_inputs_simd(&v[0]))
         });
-        tree_builder.extend_evals(trace);
-        let claim = Claim {
-            n_rows: self.inputs.len() * N_LANES,
-        };
-        (claim, interaction_claim_generator)
+        tree_builder.extend_evals(trace.to_evals());
+        let n_rows = self.inputs.len() * N_LANES;
+        (
+            Claim { n_rows },
+            InteractionClaimGenerator {
+                n_rows,
+                lookup_data,
+            },
+        )
     }
 }
 
-pub struct InteractionClaimGenerator {
-    pub n_rows: usize,
+#[derive(Debug, Uninitialized, IterMut, ParIterMut)]
+pub struct LookupData {
     pub memory: [Vec<[PackedM31; N_MEMORY_ELEMS]>; N_MEMORY_LOOKUPS],
     pub state: [Vec<[PackedM31; STATE_SIZE]>; N_STATE_LOOKUPS],
 }
-impl InteractionClaimGenerator {
-    pub fn with_capacity(capacity: usize) -> Self {
-        Self {
-            n_rows: capacity,
-            memory: [
-                Vec::with_capacity(capacity),
-                Vec::with_capacity(capacity),
-                Vec::with_capacity(capacity),
-                Vec::with_capacity(capacity),
-            ],
-            state: [Vec::with_capacity(capacity), Vec::with_capacity(capacity)],
-        }
-    }
 
+#[derive(Debug)]
+pub struct InteractionClaimGenerator {
+    pub n_rows: usize,
+    pub lookup_data: LookupData,
+}
+
+impl InteractionClaimGenerator {
     pub fn write_interaction_trace(
         &self,
         tree_builder: &mut TreeBuilder<'_, '_, SimdBackend, Blake2sMerkleChannel>,
         memory_relation: &MemoryRelation,
         state_relation: &StateRelation,
     ) -> InteractionClaim {
-        let log_size = self.memory[0].len().ilog2() + LOG_N_LANES;
+        let log_size = std::cmp::max(self.n_rows.next_power_of_two().ilog2(), LOG_N_LANES);
         let mut logup_gen = LogupTraceGenerator::new(log_size);
 
         let mut col0 = logup_gen.new_col();
-        let state_use = &self.state[0];
-        let read_pc = &self.memory[0];
+        let state_use = &self.lookup_data.state[0];
+        let read_pc = &self.lookup_data.memory[0];
         for (i, (x, y)) in zip_eq(state_use, read_pc).enumerate() {
             let denom_x: PackedQM31 = state_relation.combine(x);
             let denom_y: PackedQM31 = memory_relation.combine(y);
@@ -117,8 +114,8 @@ impl InteractionClaimGenerator {
         col0.finalize_col();
 
         let mut col1 = logup_gen.new_col();
-        let read_dst = &self.memory[1];
-        let read_lhs = &self.memory[2];
+        let read_dst = &self.lookup_data.memory[1];
+        let read_lhs = &self.lookup_data.memory[2];
         for (i, (x, y)) in zip_eq(read_dst, read_lhs).enumerate() {
             let denom_x: PackedQM31 = memory_relation.combine(x);
             let denom_y: PackedQM31 = memory_relation.combine(y);
@@ -128,8 +125,8 @@ impl InteractionClaimGenerator {
         col1.finalize_col();
 
         let mut col2 = logup_gen.new_col();
-        let read_rhs = &self.memory[3];
-        let state_yield = &self.state[1];
+        let read_rhs = &self.lookup_data.memory[3];
+        let state_yield = &self.lookup_data.state[1];
         for (i, (x, y)) in zip_eq(read_rhs, state_yield).enumerate() {
             let denom_x: PackedQM31 = memory_relation.combine(x);
             let denom_y: PackedQM31 = state_relation.combine(y);
@@ -151,119 +148,88 @@ impl InteractionClaimGenerator {
 fn write_trace_simd(
     inputs: &[PackedVmState],
     memory_trace_generator: &memory::ClaimGenerator,
-) -> (
-    Vec<CircleEvaluation<SimdBackend, M31, BitReversedOrder>>,
-    InteractionClaimGenerator,
-) {
-    let n_trace_columns = N_TRACE_COLUMNS;
-    let mut trace_values = (0..n_trace_columns)
-        .map(|_| Col::<SimdBackend, M31>::zeros(inputs.len() * N_LANES))
-        .collect_vec();
-    let mut sub_components_inputs = InteractionClaimGenerator::with_capacity(inputs.len());
-    inputs.iter().enumerate().for_each(|(i, input)| {
-        write_trace_row(
-            &mut trace_values,
-            input,
-            i,
-            &mut sub_components_inputs,
-            memory_trace_generator,
-        );
-    });
-
-    let trace = trace_values
-        .into_iter()
-        .map(|eval| {
-            // TODO(Ohad): Support non-power of 2 inputs.
-            let domain = CanonicCoset::new(
-                eval.len()
-                    .checked_ilog2()
-                    .expect("Input is not a power of 2!"),
-            )
-            .circle_domain();
-            CircleEvaluation::<SimdBackend, M31, BitReversedOrder>::new(domain, eval)
-        })
-        .collect_vec();
-
-    (trace, sub_components_inputs)
-}
+) -> (ComponentTrace<N_TRACE_COLUMNS>, LookupData) {
+    let log_n_packed_rows = inputs.len().ilog2();
+    let log_size = log_n_packed_rows + LOG_N_LANES;
+    let (mut trace, mut lookup_data) = unsafe {
+        (
+            ComponentTrace::<N_TRACE_COLUMNS>::uninitialized(log_size),
+            LookupData::uninitialized(log_n_packed_rows),
+        )
+    };
+
+    trace
+        .par_iter_mut()
+        .zip(inputs.par_iter())
+        .zip(lookup_data.par_iter_mut())
+        .for_each(|((row, input), lookup_data)| {
+            // Initial state
+            *row[0] = input.pc;
+            *row[1] = input.ap;
+            *row[2] = input.fp;
+            *lookup_data.state[0] = [input.pc, input.ap, input.fp];
+
+            // Flags
+            let [opcode, off0, off1, off2] = memory_trace_generator
+                .deduce_output(input.pc)
+                .into_packed_m31s();
+            *lookup_data.memory[0] = [input.pc, opcode, off0, off1, off2];
+
+            let [op_type, reg0, reg1, reg2, appp] =
+                decode_opcode(INSTRUCTION_BASE, opcode, [2, 2, 2, 2, 2]);
+
+            *row[3] = op_type;
+            *row[4] = reg0;
+            *row[5] = reg1;
+            *row[6] = reg2;
+            *row[7] = appp;
+
+            // Offsets
+            *row[8] = off0;
+            *row[9] = off1;
+            *row[10] = off2;
+
+            // Addresses
+            let dst_addr = Selector::select(&reg0, [&input.ap, &input.fp]) + off0;
+            let lhs_addr = Selector::select(&reg1, [&input.ap, &input.fp]) + off1;
+            let rhs_addr = Selector::select(&reg2, [&input.ap, &input.fp]) + off2;
+
+            *row[11] = dst_addr;
+            *row[12] = lhs_addr;
+            *row[13] = rhs_addr;
+
+            // Values
+            let [dst0, dst1, dst2, dst3] = memory_trace_generator
+                .deduce_output(dst_addr)
+                .into_packed_m31s();
+            let [lhs0, lhs1, lhs2, lhs3] = memory_trace_generator
+                .deduce_output(lhs_addr)
+                .into_packed_m31s();
+            let [rhs0, rhs1, rhs2, rhs3] = memory_trace_generator
+                .deduce_output(rhs_addr)
+                .into_packed_m31s();
+
+            *row[14] = dst0;
+            *row[15] = dst1;
+            *row[16] = dst2;
+            *row[17] = dst3;
+
+            *row[18] = lhs0;
+            *row[19] = lhs1;
+            *row[20] = lhs2;
+            *row[21] = lhs3;
+
+            *row[22] = rhs0;
+            *row[23] = rhs1;
+            *row[24] = rhs2;
+            *row[25] = rhs3;
+
+            *lookup_data.memory[1] = [dst_addr, dst0, dst1, dst2, dst3];
+            *lookup_data.memory[2] = [lhs_addr, lhs0, lhs1, lhs2, lhs3];
+            *lookup_data.memory[3] = [rhs_addr, rhs0, rhs1, rhs2, rhs3];
+
+            *lookup_data.state[1] = [input.pc + PackedM31::one(), input.ap + appp, input.fp];
+        });
 
-// Add / Mul trace row:
-// | State (3) | flags (5) | offsets (3) | addrs (3) | values (3 * 4) |
-fn write_trace_row(
-    trace: &mut [Col<SimdBackend, M31>],
-    input: &PackedVmState,
-    row_index: usize,
-    interaction_claim_generator: &mut InteractionClaimGenerator,
-    memory_trace_generator: &memory::ClaimGenerator,
-) {
-    // Initial state
-    trace[0].data[row_index] = input.pc;
-    trace[1].data[row_index] = input.ap;
-    trace[2].data[row_index] = input.fp;
-    interaction_claim_generator.state[0].push([input.pc, input.ap, input.fp]);
-
-    // Flags
-    let [opcode, off0, off1, off2] = memory_trace_generator
-        .deduce_output(input.pc)
-        .into_packed_m31s();
-    interaction_claim_generator.memory[0].push([input.pc, opcode, off0, off1, off2]);
-
-    let [op_type, reg0, reg1, reg2, appp] =
-        decode_opcode(INSTRUCTION_BASE, opcode, [2, 2, 2, 2, 2]);
-
-    trace[3].data[row_index] = op_type;
-    trace[4].data[row_index] = reg0;
-    trace[5].data[row_index] = reg1;
-    trace[6].data[row_index] = reg2;
-    trace[7].data[row_index] = appp;
-
-    // Offsets
-    trace[8].data[row_index] = off0;
-    trace[9].data[row_index] = off1;
-    trace[10].data[row_index] = off2;
-
-    // Addresses
-    let dst_addr = Selector::select(&reg0, [&input.ap, &input.fp]) + off0;
-    let lhs_addr = Selector::select(&reg1, [&input.ap, &input.fp]) + off1;
-    let rhs_addr = Selector::select(&reg2, [&input.ap, &input.fp]) + off2;
-
-    trace[11].data[row_index] = dst_addr;
-    trace[12].data[row_index] = lhs_addr;
-    trace[13].data[row_index] = rhs_addr;
-
-    // Values
-    let [dst0, dst1, dst2, dst3] = memory_trace_generator
-        .deduce_output(dst_addr)
-        .into_packed_m31s();
-    let [lhs0, lhs1, lhs2, lhs3] = memory_trace_generator
-        .deduce_output(lhs_addr)
-        .into_packed_m31s();
-    let [rhs0, rhs1, rhs2, rhs3] = memory_trace_generator
-        .deduce_output(rhs_addr)
-        .into_packed_m31s();
-
-    trace[14].data[row_index] = dst0;
-    trace[15].data[row_index] = dst1;
-    trace[16].data[row_index] = dst2;
-    trace[17].data[row_index] = dst3;
-
-    trace[18].data[row_index] = lhs0;
-    trace[19].data[row_index] = lhs1;
-    trace[20].data[row_index] = lhs2;
-    trace[21].data[row_index] = lhs3;
-
-    trace[22].data[row_index] = rhs0;
-    trace[23].data[row_index] = rhs1;
-    trace[24].data[row_index] = rhs2;
-    trace[25].data[row_index] = rhs3;
-
-    interaction_claim_generator.memory[1].push([dst_addr, dst0, dst1, dst2, dst3]);
-    interaction_claim_generator.memory[2].push([lhs_addr, lhs0, lhs1, lhs2, lhs3]);
-    interaction_claim_generator.memory[3].push([rhs_addr, rhs0, rhs1, rhs2, rhs3]);
-
-    interaction_claim_generator.state[1].push([
-        input.pc + PackedM31::one(),
-        input.ap + appp,
-        input.fp,
-    ]);
+    (trace, lookup_data)
 }