From 8af70b44608c4b177005cc9855cdd6bb0877acde Mon Sep 17 00:00:00 2001
From: jay white <JayWhite2357@users.noreply.github.com>
Date: Thu, 24 Oct 2024 20:23:45 -0400
Subject: [PATCH] perf: make cpu commitment computation more efficient

---
 .../dynamic_dory_commitment_helper_cpu.rs     | 57 +++++++++++++------
 .../dory/dynamic_dory_structure.rs            |  2 +-
 2 files changed, 42 insertions(+), 17 deletions(-)

diff --git a/crates/proof-of-sql/src/proof_primitive/dory/dynamic_dory_commitment_helper_cpu.rs b/crates/proof-of-sql/src/proof_primitive/dory/dynamic_dory_commitment_helper_cpu.rs
index 8ed2ddbb5..b36e1177d 100644
--- a/crates/proof-of-sql/src/proof_primitive/dory/dynamic_dory_commitment_helper_cpu.rs
+++ b/crates/proof-of-sql/src/proof_primitive/dory/dynamic_dory_commitment_helper_cpu.rs
@@ -1,10 +1,14 @@
 use super::{
-    dynamic_dory_structure::row_and_column_from_index, pairings, DoryScalar, DynamicDoryCommitment,
-    G1Affine, G1Projective, ProverSetup, GT,
+    dynamic_dory_structure::{full_width_of_row, row_and_column_from_index, row_start_index},
+    pairings, DoryScalar, DynamicDoryCommitment, G1Projective, ProverSetup, GT,
 };
-use crate::base::commitment::CommittableColumn;
-use alloc::{vec, vec::Vec};
+use crate::base::{commitment::CommittableColumn, if_rayon, slice_ops::slice_cast};
+use alloc::vec::Vec;
+use ark_ec::VariableBaseMSM;
+use bytemuck::TransparentWrapper;
 use num_traits::Zero;
+#[cfg(feature = "rayon")]
+use rayon::iter::{IntoParallelIterator, IntoParallelRefIterator, ParallelIterator};
 
 #[tracing::instrument(name = "compute_dory_commitment_impl (cpu)", level = "debug", skip_all)]
 /// # Panics
@@ -13,6 +17,7 @@ use num_traits::Zero;
 /// - `setup.Gamma_1.last()` returns `None`, indicating that `Gamma_1` is empty.
 /// - `setup.Gamma_2.last()` returns `None`, indicating that `Gamma_2` is empty.
 /// - The indexing for `Gamma_2` with `first_row..=last_row` goes out of bounds.
+#[allow(clippy::range_plus_one)]
 fn compute_dory_commitment_impl<'a, T>(
     column: &'a [T],
     offset: usize,
@@ -22,18 +27,39 @@ where
     &'a T: Into<DoryScalar>,
     T: Sync,
 {
+    if column.is_empty() {
+        return DynamicDoryCommitment::default();
+    }
     let Gamma_1 = setup.Gamma_1.last().unwrap();
     let Gamma_2 = setup.Gamma_2.last().unwrap();
-    let (first_row, _) = row_and_column_from_index(offset);
-    let (last_row, _) = row_and_column_from_index(offset + column.len() - 1);
-    let row_commits = column.iter().enumerate().fold(
-        vec![G1Projective::from(G1Affine::identity()); last_row - first_row + 1],
-        |mut row_commits, (i, v)| {
-            let (row, col) = row_and_column_from_index(i + offset);
-            row_commits[row - first_row] += Gamma_1[col] * v.into().0;
-            row_commits
-        },
-    );
+    let (first_row, first_col) = row_and_column_from_index(offset);
+    let (last_row, last_col) = row_and_column_from_index(offset + column.len() - 1);
+
+    let row_commits: Vec<_> = if_rayon!(
+        (first_row..=last_row).into_par_iter(),
+        (first_row..=last_row)
+    )
+    .map(|row| {
+        let width = full_width_of_row(row);
+        let row_start = row_start_index(row);
+        let (gamma_range, column_range) = if first_row == last_row {
+            (first_col..last_col + 1, 0..column.len())
+        } else if row == 1 {
+            (1..2, (1 - offset)..(2 - offset))
+        } else if row == first_row {
+            (first_col..width, 0..width - first_col)
+        } else if row == last_row {
+            (0..last_col + 1, column.len() - last_col - 1..column.len())
+        } else {
+            (0..width, row_start - offset..width + row_start - offset)
+        };
+        G1Projective::msm_unchecked(
+            &Gamma_1[gamma_range],
+            TransparentWrapper::peel_slice(&slice_cast::<_, DoryScalar>(&column[column_range])),
+        )
+    })
+    .collect();
+
     DynamicDoryCommitment(pairings::multi_pairing(
         row_commits,
         &Gamma_2[first_row..=last_row],
@@ -70,8 +96,7 @@ pub(super) fn compute_dynamic_dory_commitments(
     offset: usize,
     setup: &ProverSetup,
 ) -> Vec<DynamicDoryCommitment> {
-    committable_columns
-        .iter()
+    if_rayon!(committable_columns.par_iter(), committable_columns.iter())
         .map(|column| {
             column
                 .is_empty()
diff --git a/crates/proof-of-sql/src/proof_primitive/dory/dynamic_dory_structure.rs b/crates/proof-of-sql/src/proof_primitive/dory/dynamic_dory_structure.rs
index 2598db988..03f6ffaa3 100644
--- a/crates/proof-of-sql/src/proof_primitive/dory/dynamic_dory_structure.rs
+++ b/crates/proof-of-sql/src/proof_primitive/dory/dynamic_dory_structure.rs
@@ -40,7 +40,7 @@ pub(crate) const fn full_width_of_row(row: usize) -> usize {
 /// Returns the index that belongs in the first column in a particular row.
 ///
 /// Note: when row = 1, this correctly returns 0, even though no data belongs there.
-#[cfg(test)]
+#[cfg(any(test, not(feature = "blitzar")))]
 pub(crate) const fn row_start_index(row: usize) -> usize {
     let width_of_row = full_width_of_row(row);
     width_of_row * (row - width_of_row / 2)