Precompute alphas for quotients. (#558)

src/core/backend/avx512/quotients.rs

@@ -1,4 +1,4 @@
-use itertools::zip_eq;
+use itertools::{izip, zip_eq};
 
 use super::qm31::PackedQM31;
 use super::{AVX512Backend, VECS_LOG_SIZE};
@@ -41,7 +41,6 @@ impl QuotientOps for AVX512Backend {
                 columns,
                 &quotient_constants,
                 vec_row,
-                random_coeff,
                 (domain_points_x, domain_points_y),
             );
             values.set_packed(vec_row, row_accumulator);
@@ -55,15 +54,16 @@ pub fn accumulate_row_quotients(
     columns: &[&CircleEvaluation<AVX512Backend, BaseField, BitReversedOrder>],
     quotient_constants: &QuotientConstants,
     vec_row: usize,
-    random_coeff: SecureField,
     domain_point_vec: (PackedBaseField, PackedBaseField),
 ) -> PackedQM31 {
     let mut row_accumulator = PackedQM31::zero();
-    for (sample_batch, sample_constants) in zip_eq(sample_batches, &quotient_constants.line_coeffs)
-    {
+    for (sample_batch, line_coeffs, batch_coeff) in izip!(
+        sample_batches,
+        &quotient_constants.line_coeffs,
+        &quotient_constants.batch_random_coeffs
+    ) {
         let mut numerator = PackedQM31::zero();
-        for ((column_index, _), (a, b, c)) in
-            zip_eq(&sample_batch.columns_and_values, sample_constants)
+        for ((column_index, _), (a, b, c)) in zip_eq(&sample_batch.columns_and_values, line_coeffs)
         {
             let column = &columns[*column_index];
             let value = PackedQM31::broadcast(*c) * column.data[vec_row];
@@ -85,10 +85,7 @@ pub fn accumulate_row_quotients(
             domain_point_vec,
         );
 
-        row_accumulator = row_accumulator
-            * PackedQM31::broadcast(
-                random_coeff.pow(sample_batch.columns_and_values.len() as u128),
-            )
+        row_accumulator = row_accumulator * PackedQM31::broadcast(*batch_coeff)
             + numerator * denominator.inverse();
     }
     row_accumulator

src/core/backend/cpu/quotients.rs

@@ -1,4 +1,4 @@
-use itertools::zip_eq;
+use itertools::{izip, zip_eq};
 use num_traits::{One, Zero};
 
 use super::CPUBackend;
@@ -31,7 +31,6 @@ impl QuotientOps for CPUBackend {
                 columns,
                 &quotient_constants,
                 row,
-                random_coeff,
                 domain_point,
             );
             values.set(row, row_value);
@@ -45,15 +44,16 @@ pub fn accumulate_row_quotients(
     columns: &[&CircleEvaluation<CPUBackend, BaseField, BitReversedOrder>],
     quotient_constants: &QuotientConstants,
     row: usize,
-    random_coeff: SecureField,
     domain_point: CirclePoint<BaseField>,
 ) -> SecureField {
     let mut row_accumulator = SecureField::zero();
-    for (sample_batch, sample_constants) in zip_eq(sample_batches, &quotient_constants.line_coeffs)
-    {
+    for (sample_batch, line_coeffs, batch_coeff) in izip!(
+        sample_batches,
+        &quotient_constants.line_coeffs,
+        &quotient_constants.batch_random_coeffs
+    ) {
         let mut numerator = SecureField::zero();
-        for ((column_index, _), (a, b, c)) in
-            zip_eq(&sample_batch.columns_and_values, sample_constants)
+        for ((column_index, _), (a, b, c)) in zip_eq(&sample_batch.columns_and_values, line_coeffs)
         {
             let column = &columns[*column_index];
             let value = column[row] * *c;
@@ -67,9 +67,7 @@ pub fn accumulate_row_quotients(
             domain_point.into_ef(),
         );
 
-        row_accumulator = row_accumulator
-            * random_coeff.pow(sample_batch.columns_and_values.len() as u128)
-            + numerator / denominator;
+        row_accumulator = row_accumulator * *batch_coeff + numerator / denominator;
     }
     row_accumulator
 }
@@ -102,19 +100,39 @@ pub fn column_line_coeffs(
         .collect()
 }
 
+/// Precompute the random coefficients used to linearly combine the batched quotients.
+/// Specifically, for each sample batch we compute random_coeff^(number of columns in the batch),
+/// which is used to linearly combine the batch with the next one.
+pub fn batch_random_coeffs(
+    sample_batches: &[ColumnSampleBatch],
+    random_coeff: SecureField,
+) -> Vec<SecureField> {
+    sample_batches
+        .iter()
+        .map(|sb| random_coeff.pow(sb.columns_and_values.len() as u128))
+        .collect()
+}
+
 pub fn quotient_constants(
     sample_batches: &[ColumnSampleBatch],
     random_coeff: SecureField,
 ) -> QuotientConstants {
     let line_coeffs = column_line_coeffs(sample_batches, random_coeff);
-    QuotientConstants { line_coeffs }
+    let batch_random_coeffs = batch_random_coeffs(sample_batches, random_coeff);
+    QuotientConstants {
+        line_coeffs,
+        batch_random_coeffs,
+    }
 }
 
 /// Holds the precomputed constant values used in each quotient evaluation.
 pub struct QuotientConstants {
     /// The line coefficients for each quotient numerator term. For more details see
     /// [self::column_line_coeffs].
     pub line_coeffs: Vec<Vec<(SecureField, SecureField, SecureField)>>,
+    /// The random coefficients used to linearly combine the batched quotients For more details see
+    /// [self::batch_random_coeffs].
+    pub batch_random_coeffs: Vec<SecureField>,
 }
 
 #[cfg(test)]

src/core/commitment_scheme/quotients.rs

@@ -157,7 +157,6 @@ pub fn fri_answers_for_log_size(
                 &column_evals.iter().collect_vec(),
                 &quotient_constants,
                 row,
-                random_coeff,
                 domain_point,
             );
             values.push(value);