Uninitialized column (#763)

.gitignore

@@ -6,3 +6,6 @@ build
 dist
 target*
 */.vscode/*
+flamegraph.svg
+perf.data
+perf.data.old

crates/prover/src/constraint_framework/simd_domain.rs

@@ -63,7 +63,14 @@ impl<'a> EvalAtRow for SimdDomainEvaluator<'a> {
         offsets.map(|off| {
             // If the offset is 0, we can just return the value directly from this row.
             if off == 0 {
-                return self.trace_eval[interaction][col_index].data[self.vec_row];
+                return unsafe {
+                    *self
+                        .trace_eval
+                        .get_unchecked(interaction)
+                        .get_unchecked(col_index)
+                        .data
+                        .get_unchecked(self.vec_row)
+                };
             }
             // Otherwise, we need to look up the value at the offset.
             // Since the domain is bit-reversed circle domain ordered, we need to look up the value

crates/prover/src/core/backend/cpu/fri.rs

@@ -27,7 +27,7 @@ impl FriOps for CpuBackend {
 
     fn decompose(eval: &SecureEvaluation<Self>) -> (SecureEvaluation<Self>, SecureField) {
         let lambda = Self::decomposition_coefficient(eval);
-        let mut g_values = SecureColumnByCoords::<Self>::zeros(eval.len());
+        let mut g_values = unsafe { SecureColumnByCoords::<Self>::uninitialized(eval.len()) };
 
         let domain_size = eval.len();
         let half_domain_size = domain_size / 2;

crates/prover/src/core/backend/cpu/mod.rs

@@ -40,6 +40,12 @@ impl<T: Debug + Clone + Default> Column<T> for Vec<T> {
     fn zeros(len: usize) -> Self {
         vec![T::default(); len]
     }
+    #[allow(clippy::uninit_vec)]
+    unsafe fn uninitialized(length: usize) -> Self {
+        let mut data = Vec::with_capacity(length);
+        data.set_len(length);
+        data
+    }
     fn to_cpu(&self) -> Vec<T> {
         self.clone()
     }

crates/prover/src/core/backend/cpu/quotients.rs

@@ -21,7 +21,7 @@ impl QuotientOps for CpuBackend {
         random_coeff: SecureField,
         sample_batches: &[ColumnSampleBatch],
     ) -> SecureEvaluation<Self> {
-        let mut values = SecureColumnByCoords::zeros(domain.size());
+        let mut values = unsafe { SecureColumnByCoords::uninitialized(domain.size()) };
         let quotient_constants = quotient_constants(sample_batches, random_coeff, domain);
 
         for row in 0..domain.size() {

crates/prover/src/core/backend/mod.rs

@@ -37,6 +37,10 @@ pub type Col<B, T> = <B as ColumnOps<T>>::Column;
 pub trait Column<T>: Clone + Debug + FromIterator<T> {
     /// Creates a new column of zeros with the given length.
     fn zeros(len: usize) -> Self;
+    /// Creates a new column of uninitialized values with the given length.
+    /// # Safety
+    /// The caller must ensure that the column is populated before being used.
+    unsafe fn uninitialized(len: usize) -> Self;
     /// Returns a cpu vector of the column.
     fn to_cpu(&self) -> Vec<T>;
     /// Returns the length of the column.

crates/prover/src/core/backend/simd/circle.rs

@@ -235,7 +235,7 @@ impl PolyOps for SimdBackend {
         // Evaluate on a big domains by evaluating on several subdomains.
         let log_subdomains = log_size - fft_log_size;
 
-        // Alllocate the destination buffer without initializing.
+        // Allocate the destination buffer without initializing.
         let mut values = Vec::with_capacity(domain.size() >> LOG_N_LANES);
         #[allow(clippy::uninit_vec)]
         unsafe {

crates/prover/src/core/backend/simd/column.rs

@@ -71,6 +71,13 @@ impl Column<BaseField> for BaseColumn {
         Self { data, length }
     }
 
+    #[allow(clippy::uninit_vec)]
+    unsafe fn uninitialized(length: usize) -> Self {
+        let mut data = Vec::with_capacity(length.div_ceil(N_LANES));
+        data.set_len(length.div_ceil(N_LANES));
+        Self { data, length }
+    }
+
     fn to_cpu(&self) -> Vec<BaseField> {
         self.as_slice().to_vec()
     }
@@ -124,6 +131,13 @@ impl Column<CM31> for CM31Column {
         }
     }
 
+    #[allow(clippy::uninit_vec)]
+    unsafe fn uninitialized(length: usize) -> Self {
+        let mut data = Vec::with_capacity(length.div_ceil(N_LANES));
+        data.set_len(length.div_ceil(N_LANES));
+        Self { data, length }
+    }
+
     fn to_cpu(&self) -> Vec<CM31> {
         self.data
             .iter()
@@ -207,6 +221,13 @@ impl Column<SecureField> for SecureColumn {
         }
     }
 
+    #[allow(clippy::uninit_vec)]
+    unsafe fn uninitialized(length: usize) -> Self {
+        let mut data = Vec::with_capacity(length.div_ceil(N_LANES));
+        data.set_len(length.div_ceil(N_LANES));
+        Self { data, length }
+    }
+
     fn to_cpu(&self) -> Vec<SecureField> {
         self.data
             .iter()

crates/prover/src/core/backend/simd/quotients.rs

@@ -93,7 +93,8 @@ fn accumulate_quotients_on_subdomain(
     [crate::core::poly::circle::CirclePoly<SimdBackend>; 4],
 ) {
     assert!(subdomain.log_size() >= LOG_N_LANES + 2);
-    let mut values = SecureColumnByCoords::<SimdBackend>::zeros(subdomain.size());
+    let mut values =
+        unsafe { SecureColumnByCoords::<SimdBackend>::uninitialized(subdomain.size()) };
     let quotient_constants = quotient_constants(sample_batches, random_coeff, subdomain);
 
     let span = span!(Level::INFO, "Quotient accumulation").entered();
@@ -121,7 +122,8 @@ fn accumulate_quotients_on_subdomain(
     let span = span!(Level::INFO, "Quotient extension").entered();
 
     // Extend the evaluation to the full domain.
-    let extended_eval = SecureColumnByCoords::<SimdBackend>::zeros(domain.size());
+    let extended_eval =
+        unsafe { SecureColumnByCoords::<SimdBackend>::uninitialized(domain.size()) };
 
     let mut i = 0;
     let values = values.columns;
@@ -213,7 +215,8 @@ fn denominator_inverses(
         })
         .collect();
 
-    let mut flat_denominator_inverses = CM31Column::zeros(flat_denominators.len());
+    let mut flat_denominator_inverses =
+        unsafe { CM31Column::uninitialized(flat_denominators.len()) };
     FieldExpOps::batch_inverse(
         &flat_denominators.data,
         &mut flat_denominator_inverses.data[..],

crates/prover/src/core/fields/secure_column.rs

@@ -29,6 +29,13 @@ impl<B: FieldOps<BaseField>> SecureColumnByCoords<B> {
         }
     }
 
+    /// # Safety
+    pub unsafe fn uninitialized(len: usize) -> Self {
+        Self {
+            columns: std::array::from_fn(|_| Col::<B, BaseField>::uninitialized(len)),
+        }
+    }
+
     pub fn len(&self) -> usize {
         self.columns[0].len()
     }