Use BaseColumnPool in CommitmentSchemeProver::commit to avoid allocation during polynomial evaluation.

Add BaseColumnPool for Col<B, BaseField> buffers, add evaluate_into to PolyOps trait
(CPU and SIMD backends), and thread the pool through commit → CommitmentTreeProver::new
→ evaluate_polynomials. Buffers are pre-taken before parallel iteration.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: ilyalesokhin-starkware

crates/stwo/src/prover/backend/cpu/circle.rs

@@ -227,6 +227,59 @@ impl PolyOps for CpuBackend {
         CircleEvaluation::new(domain, values)
     }
 
+    fn evaluate_into(
+        poly: &CircleCoefficients<Self>,
+        domain: CircleDomain,
+        twiddles: &TwiddleTree<Self>,
+        mut buffer: Col<Self, BaseField>,
+    ) -> CircleEvaluation<Self, BaseField, BitReversedOrder> {
+        assert!(domain.half_coset.is_doubling_of(twiddles.root_coset));
+        assert_eq!(buffer.len(), domain.size());
+
+        // Copy extended coefficients into the buffer.
+        let poly_len = poly.coeffs.len();
+        buffer[..poly_len].copy_from_slice(&poly.coeffs);
+        for v in &mut buffer[poly_len..] {
+            *v = BaseField::zero();
+        }
+
+        if domain.log_size() == 1 {
+            let (mut v0, mut v1) = (buffer[0], buffer[1]);
+            butterfly(&mut v0, &mut v1, domain.half_coset.initial.y);
+            buffer[0] = v0;
+            buffer[1] = v1;
+            return CircleEvaluation::new(domain, buffer);
+        }
+
+        if domain.log_size() == 2 {
+            let (mut v0, mut v1, mut v2, mut v3) = (buffer[0], buffer[1], buffer[2], buffer[3]);
+            let CirclePoint { x, y } = domain.half_coset.initial;
+            butterfly(&mut v0, &mut v2, x);
+            butterfly(&mut v1, &mut v3, x);
+            butterfly(&mut v0, &mut v1, y);
+            butterfly(&mut v2, &mut v3, -y);
+            buffer[0] = v0;
+            buffer[1] = v1;
+            buffer[2] = v2;
+            buffer[3] = v3;
+            return CircleEvaluation::new(domain, buffer);
+        }
+
+        let line_twiddles = domain_line_twiddles_from_tree(domain, &twiddles.twiddles);
+        let circle_twiddles = circle_twiddles_from_line_twiddles(line_twiddles[0]);
+
+        for (layer, layer_twiddles) in line_twiddles.iter().enumerate().rev() {
+            for (h, &t) in layer_twiddles.iter().enumerate() {
+                fft_layer_loop(&mut buffer, layer + 1, h, t, butterfly);
+            }
+        }
+        for (h, t) in circle_twiddles.enumerate() {
+            fft_layer_loop(&mut buffer, 0, h, t, butterfly);
+        }
+
+        CircleEvaluation::new(domain, buffer)
+    }
+
     fn precompute_twiddles(coset: Coset) -> TwiddleTree<Self> {
         const CHUNK_LOG_SIZE: usize = 12;
         const CHUNK_SIZE: usize = 1 << CHUNK_LOG_SIZE;

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

@@ -442,6 +442,63 @@ impl PolyOps for SimdBackend {
         )
     }
 
+    fn evaluate_into(
+        poly: &CircleCoefficients<Self>,
+        domain: CircleDomain,
+        twiddles: &TwiddleTree<Self>,
+        mut buffer: Col<Self, BaseField>,
+    ) -> CircleEvaluation<Self, BaseField, BitReversedOrder> {
+        let _span = span!(Level::TRACE, "", class = "rFFT").entered();
+        let log_size = domain.log_size();
+        let fft_log_size = poly.log_size();
+        assert!(
+            log_size >= fft_log_size,
+            "Can only evaluate on larger domains"
+        );
+        assert_eq!(buffer.len(), domain.size());
+
+        if fft_log_size < MIN_FFT_LOG_SIZE {
+            let cpu_poly: CircleCoefficients<CpuBackend> =
+                CircleCoefficients::new(poly.coeffs.to_cpu());
+            let cpu_eval = cpu_poly.evaluate(domain);
+            return CircleEvaluation::new(
+                cpu_eval.domain,
+                Col::<SimdBackend, BaseField>::from_iter(cpu_eval.values),
+            );
+        }
+
+        let twiddles = domain_line_twiddles_from_tree(domain, &twiddles.twiddles);
+
+        // Evaluate on big domains by evaluating on several subdomains.
+        let log_subdomains = log_size - fft_log_size;
+
+        for i in 0..(1 << log_subdomains) {
+            // The subdomain twiddles are a slice of the large domain twiddles.
+            let subdomain_twiddles = (0..(fft_log_size - 1))
+                .map(|layer_i| {
+                    &twiddles[layer_i as usize]
+                        [i << (fft_log_size - 2 - layer_i)..(i + 1) << (fft_log_size - 2 - layer_i)]
+                })
+                .collect::<Vec<_>>();
+
+            // FFT from the coefficients buffer directly into the provided buffer.
+            unsafe {
+                rfft::fft(
+                    transmute::<*const PackedBaseField, *const u32>(poly.coeffs.data.as_ptr()),
+                    transmute::<*mut PackedBaseField, *mut u32>(
+                        buffer.data[i << (fft_log_size - LOG_N_LANES)
+                            ..(i + 1) << (fft_log_size - LOG_N_LANES)]
+                            .as_mut_ptr(),
+                    ),
+                    &subdomain_twiddles,
+                    fft_log_size as usize,
+                );
+            }
+        }
+
+        CircleEvaluation::new(domain, buffer)
+    }
+
     /// Precomputes the (doubled) twiddles for a given coset tower.
     /// The twiddles are the x values of each coset in bit-reversed order.
     /// Note: the coset point are symmetrical over the x-axis so only the first half of the coset is

crates/stwo/src/prover/mempool.rs

@@ -7,7 +7,7 @@
 use std::collections::HashMap;
 
 use crate::core::fields::m31::BaseField;
-use crate::prover::backend::ColumnOps;
+use crate::prover::backend::{Col, Column, ColumnOps};
 use crate::prover::secure_column::SecureColumnByCoords;
 
 /// A pool of pre-allocated [`SecureColumnByCoords`] buffers, organized by log_size.
@@ -87,6 +87,76 @@ impl<B: ColumnOps<BaseField>> Default for ColumnPool<B> {
     }
 }
 
+/// A pool of pre-allocated [`Col<B, BaseField>`] buffers, organized by log_size.
+///
+/// Used to avoid repeated allocation of evaluation buffers during polynomial commitment.
+pub struct BaseColumnPool<B: ColumnOps<BaseField>> {
+    /// Map from log_size -> stack of available buffers.
+    pools: HashMap<u32, Vec<Col<B, BaseField>>>,
+}
+
+impl<B: ColumnOps<BaseField>> BaseColumnPool<B> {
+    /// Creates a new empty base column pool.
+    pub fn new() -> Self {
+        Self {
+            pools: HashMap::new(),
+        }
+    }
+
+    /// Pre-allocates `count` zero-initialized buffers of size `1 << log_size`.
+    pub fn reserve(&mut self, log_size: u32, count: usize) {
+        let pool = self.pools.entry(log_size).or_default();
+        for _ in 0..count {
+            pool.push(Col::<B, BaseField>::zeros(1 << log_size));
+        }
+    }
+
+    /// Takes a buffer from the pool for the given `log_size`.
+    ///
+    /// # Panics
+    ///
+    /// Panics if no buffer of the requested size is available.
+    pub fn take(&mut self, log_size: u32) -> Col<B, BaseField> {
+        self.pools
+            .get_mut(&log_size)
+            .and_then(|pool| pool.pop())
+            .unwrap_or_else(|| {
+                panic!("BaseColumnPool: no buffer available for log_size={log_size}")
+            })
+    }
+
+    /// Takes a buffer from the pool, or allocates a new zero-initialized one if none is available.
+    pub fn take_or_alloc(&mut self, log_size: u32) -> Col<B, BaseField> {
+        self.pools
+            .get_mut(&log_size)
+            .and_then(|pool| pool.pop())
+            .unwrap_or_else(|| Col::<B, BaseField>::zeros(1 << log_size))
+    }
+
+    /// Returns a buffer to the pool. The caller is responsible for ensuring the buffer's log_size
+    /// matches.
+    pub fn give_back(&mut self, log_size: u32, buf: Col<B, BaseField>) {
+        debug_assert_eq!(buf.len(), 1 << log_size);
+        self.pools.entry(log_size).or_default().push(buf);
+    }
+
+    /// Returns the number of available buffers for a given log_size.
+    pub fn available(&self, log_size: u32) -> usize {
+        self.pools.get(&log_size).map_or(0, |pool| pool.len())
+    }
+
+    /// Returns the total number of buffers across all sizes.
+    pub fn total_available(&self) -> usize {
+        self.pools.values().map(|pool| pool.len()).sum()
+    }
+}
+
+impl<B: ColumnOps<BaseField>> Default for BaseColumnPool<B> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
 /// Zeroes out all columns in a [`SecureColumnByCoords`].
 fn zero_secure_column<B: ColumnOps<BaseField>>(col: &mut SecureColumnByCoords<B>) {
     let len = col.len();
@@ -101,13 +171,16 @@ fn zero_secure_column<B: ColumnOps<BaseField>>(col: &mut SecureColumnByCoords<B>
 pub struct ProverMemPool<B: ColumnOps<BaseField>> {
     /// Pool of reusable [`SecureColumnByCoords`] buffers.
     pub column_pool: ColumnPool<B>,
+    /// Pool of reusable base field column buffers.
+    pub base_column_pool: BaseColumnPool<B>,
 }
 
 impl<B: ColumnOps<BaseField>> ProverMemPool<B> {
     /// Creates a new workspace with an empty column pool.
     pub fn new() -> Self {
         Self {
             column_pool: ColumnPool::new(),
+            base_column_pool: BaseColumnPool::new(),
         }
     }
 
@@ -203,4 +276,49 @@ mod tests {
     }
 
     use num_traits::Zero;
+
+    #[test]
+    fn test_base_column_pool_reserve_and_take() {
+        let mut pool = BaseColumnPool::<CpuBackend>::new();
+        pool.reserve(4, 3);
+        assert_eq!(pool.available(4), 3);
+
+        let buf = pool.take(4);
+        assert_eq!(buf.len(), 1 << 4);
+        assert_eq!(pool.available(4), 2);
+    }
+
+    #[test]
+    fn test_base_column_pool_give_back() {
+        let mut pool = BaseColumnPool::<CpuBackend>::new();
+        pool.reserve(5, 1);
+        let buf = pool.take(5);
+        assert_eq!(pool.available(5), 0);
+
+        pool.give_back(5, buf);
+        assert_eq!(pool.available(5), 1);
+    }
+
+    #[test]
+    fn test_base_column_pool_take_or_alloc() {
+        let mut pool = BaseColumnPool::<CpuBackend>::new();
+
+        // No pre-allocated buffer, should allocate.
+        let buf = pool.take_or_alloc(3);
+        assert_eq!(buf.len(), 1 << 3);
+        assert_eq!(pool.available(3), 0);
+
+        // Return and take again.
+        pool.give_back(3, buf);
+        assert_eq!(pool.available(3), 1);
+        let _buf = pool.take_or_alloc(3);
+        assert_eq!(pool.available(3), 0);
+    }
+
+    #[test]
+    #[should_panic(expected = "no buffer available")]
+    fn test_base_column_pool_take_panics_when_empty() {
+        let mut pool = BaseColumnPool::<CpuBackend>::new();
+        pool.take(4);
+    }
 }

crates/stwo/src/prover/pcs/mod.rs

@@ -20,6 +20,7 @@ use crate::core::ColumnVec;
 use crate::prover::air::component_prover::{Poly, Trace, WeightsHashMap};
 use crate::prover::backend::{BackendForChannel, Col};
 use crate::prover::fri::{FriDecommitResult, FriProver};
+use crate::prover::mempool::{BaseColumnPool, ColumnPool};
 use crate::prover::pcs::quotient_ops::compute_fri_quotients;
 use crate::prover::poly::circle::{CircleCoefficients, CircleEvaluation};
 use crate::prover::poly::twiddles::TwiddleTree;
@@ -34,6 +35,10 @@ pub struct CommitmentSchemeProver<'a, B: BackendForChannel<MC>, MC: MerkleChanne
     pub config: PcsConfig,
     twiddles: &'a TwiddleTree<B>,
     pub store_polynomials_coefficients: bool,
+    /// Pre-allocated memory pool for the proving pipeline.
+    pub mempool: ColumnPool<B>,
+    /// Pre-allocated base field column pool for polynomial evaluation during commit.
+    pub base_column_pool: BaseColumnPool<B>,
 }
 impl<'a, B: BackendForChannel<MC>, MC: MerkleChannel> CommitmentSchemeProver<'a, B, MC> {
     /// Creates a new empty commitment scheme prover with the given configuration and twiddles. The
@@ -44,6 +49,8 @@ impl<'a, B: BackendForChannel<MC>, MC: MerkleChannel> CommitmentSchemeProver<'a,
             config,
             twiddles,
             store_polynomials_coefficients: false,
+            mempool: ColumnPool::new(),
+            base_column_pool: BaseColumnPool::new(),
         }
     }
 
@@ -62,6 +69,7 @@ impl<'a, B: BackendForChannel<MC>, MC: MerkleChannel> CommitmentSchemeProver<'a,
             self.twiddles,
             self.store_polynomials_coefficients,
             self.config.lifting_log_size,
+            &mut self.base_column_pool,
         );
         self.trees.push(tree);
     }
@@ -319,13 +327,15 @@ impl<B: BackendForChannel<MC>, MC: MerkleChannel> CommitmentTreeProver<B, MC> {
         twiddles: &TwiddleTree<B>,
         store_polynomials_coefficients: bool,
         lifting_log_size: Option<u32>,
+        base_column_pool: &mut BaseColumnPool<B>,
     ) -> Self {
         let span = span!(Level::INFO, "Extension").entered();
         let polynomials = B::evaluate_polynomials(
             polynomials,
             log_blowup_factor,
             twiddles,
             store_polynomials_coefficients,
+            base_column_pool,
         );
         span.exit();
 

crates/stwo/src/prover/poly/circle/ops.rs

@@ -9,6 +9,7 @@ use crate::core::poly::circle::{CanonicCoset, CircleDomain};
 use crate::core::ColumnVec;
 use crate::prover::air::component_prover::Poly;
 use crate::prover::backend::{Col, ColumnOps};
+use crate::prover::mempool::BaseColumnPool;
 use crate::prover::poly::twiddles::TwiddleTree;
 use crate::prover::poly::BitReversedOrder;
 
@@ -82,25 +83,43 @@ pub trait PolyOps: ColumnOps<BaseField> + ColumnOps<SecureField> + Sized {
         twiddles: &TwiddleTree<Self>,
     ) -> CircleEvaluation<Self, BaseField, BitReversedOrder>;
 
+    /// Evaluates the polynomial at all points in the domain, writing results into the provided
+    /// buffer instead of allocating a new one. The buffer must have size `domain.size()`.
+    fn evaluate_into(
+        poly: &CircleCoefficients<Self>,
+        domain: CircleDomain,
+        twiddles: &TwiddleTree<Self>,
+        buffer: Col<Self, BaseField>,
+    ) -> CircleEvaluation<Self, BaseField, BitReversedOrder>;
+
     fn evaluate_polynomials(
         polynomials: ColumnVec<CircleCoefficients<Self>>,
         log_blowup_factor: u32,
         twiddles: &TwiddleTree<Self>,
         store_polynomials_coefficients: bool,
+        pool: &mut BaseColumnPool<Self>,
     ) -> Vec<Poly<Self>>
     where
         Self: crate::prover::backend::Backend,
     {
+        // Pre-take all buffers from the pool before the parallel section.
+        let buffers: Vec<_> = polynomials
+            .iter()
+            .map(|poly_coeffs| {
+                let log_eval_size = poly_coeffs.log_size() + log_blowup_factor;
+                pool.take_or_alloc(log_eval_size)
+            })
+            .collect();
+
         #[cfg(feature = "parallel")]
-        let iter = polynomials.into_par_iter();
+        let iter = polynomials.into_par_iter().zip(buffers.into_par_iter());
         #[cfg(not(feature = "parallel"))]
-        let iter = polynomials.into_iter();
+        let iter = polynomials.into_iter().zip(buffers.into_iter());
 
-        iter.map(|poly_coeffs| {
-            let evals = poly_coeffs.evaluate_with_twiddles(
-                CanonicCoset::new(poly_coeffs.log_size() + log_blowup_factor).circle_domain(),
-                twiddles,
-            );
+        iter.map(|(poly_coeffs, buffer)| {
+            let domain =
+                CanonicCoset::new(poly_coeffs.log_size() + log_blowup_factor).circle_domain();
+            let evals = Self::evaluate_into(&poly_coeffs, domain, twiddles, buffer);
             Poly::new(store_polynomials_coefficients.then_some(poly_coeffs), evals)
         })
         .collect()