Integrate buffer abstraction to Zook

benches/expand_from_coeff.rs

@@ -1,5 +1,8 @@
 use divan::{black_box, AllocProfiler, Bencher};
-use whir::algebra::{fields::Field64, ntt, random_vector};
+use whir::{
+    algebra::{fields::Field64, ntt, random_vector},
+    buffer::{Buffer, BufferOps},
+};
 
 #[global_allocator]
 static ALLOC: AllocProfiler = AllocProfiler::system();
@@ -28,15 +31,17 @@ fn interleaved_rs_encode(bencher: Bencher, case: &(usize, usize, usize)) {
             let message_length = 1 << (exp - coset_sz);
             let num_messages = 1 << coset_sz;
             let mut rng = ark_std::rand::thread_rng();
-            let coeffs: Vec<Vec<Field64>> = (0..num_messages)
-                .map(|_| random_vector(&mut rng, message_length))
+            let coeffs: Vec<Buffer<Field64>> = (0..num_messages)
+                .map(|_| Buffer::from(random_vector(&mut rng, message_length)))
                 .collect();
-            (coeffs, expansion, coset_sz)
+            (coeffs, expansion)
         })
-        .bench_values(|(coeffs, expansion, _coset_sz)| {
-            let coeffs_refs = coeffs.iter().map(|v| v.as_slice()).collect::<Vec<_>>();
+        .bench_values(|(coeffs, expansion)| {
+            let coefficient_refs = coeffs.iter().collect::<Vec<_>>();
+            let segments = [ntt::PolynomialSegment::from_rows(&coefficient_refs, 1)];
+            let polynomials = ntt::Polynomials::from_segments(&segments);
             black_box(ntt::interleaved_rs_encode(
-                &coeffs_refs,
+                polynomials,
                 coeffs[0].len() * expansion,
             ))
         });

src/algebra/linear_form/univariate_evaluation.rs

@@ -1,9 +1,7 @@
 use ark_ff::Field;
 
 use super::LinearForm;
-use crate::algebra::{
-    embedding::Embedding, geometric_accumulate, linear_form::Evaluate, mixed_univariate_evaluate,
-};
+use crate::algebra::{embedding::Embedding, linear_form::Evaluate, mixed_univariate_evaluate};
 
 /// Linear form to represent univariate polynomial evaluation.
 ///
@@ -21,21 +19,6 @@ impl<F: Field> UnivariateEvaluation<F> {
     pub const fn new(point: F, size: usize) -> Self {
         Self { size, point }
     }
-
-    /// Batched version of [`LinearForm::accumulate`] for many [`UnivariateEvaluation`]s.
-    pub fn accumulate_many(evaluators: &[Self], accumulator: &mut [F], scalars: &[F]) {
-        assert_eq!(evaluators.len(), scalars.len());
-        let Some(size) = evaluators.first().map(|e| e.size) else {
-            return;
-        };
-        assert_eq!(accumulator.len(), size);
-        for evaluator in evaluators {
-            assert_eq!(evaluator.size, size);
-        }
-        let points = evaluators.iter().map(|e| e.point).collect::<Vec<F>>();
-        let scalars = scalars.to_vec();
-        geometric_accumulate(accumulator, scalars, &points);
-    }
 }
 
 impl<F: Field> LinearForm<F> for UnivariateEvaluation<F> {
@@ -55,7 +38,6 @@ impl<F: Field> LinearForm<F> for UnivariateEvaluation<F> {
         result
     }
 
-    /// See also [`Self::accumulate_many`] for a more efficient batched version.
     fn accumulate(&self, accumulator: &mut [F], scalar: F) {
         assert_eq!(accumulator.len(), self.size);
         let mut power = scalar;

src/algebra/ntt/cooley_tukey.rs

@@ -14,11 +14,14 @@ use {crate::utils::workload_size, rayon::prelude::*, std::cmp::max};
 use super::{
     transpose,
     utils::{lcm, sqrt_factor},
-    ReedSolomon,
+    PolynomialSegment, Polynomials, ReedSolomon,
 };
 #[cfg(not(feature = "rs_in_order"))]
 use crate::algebra::ntt::transpose::transpose_permute;
-use crate::{algebra::ntt::utils::divisors, buffer::Buffer};
+use crate::{
+    algebra::ntt::utils::divisors,
+    buffer::{Buffer, BufferOps},
+};
 
 // Supported primes
 const PRIMES: [usize; 2] = [2, 3];
@@ -45,6 +48,43 @@ pub struct NttEngine<F: Field> {
     roots: RwLock<Vec<F>>,
 }
 
+struct SegmentRows<'segment, 'buffer, F> {
+    segment: &'segment PolynomialSegment<'buffer, F>,
+    buffer_index: usize,
+    row_index: usize,
+    source: Option<&'buffer [F]>,
+}
+
+impl<'segment, 'buffer, F: Copy> SegmentRows<'segment, 'buffer, F> {
+    const fn new(segment: &'segment PolynomialSegment<'buffer, F>) -> Self {
+        Self {
+            segment,
+            buffer_index: 0,
+            row_index: 0,
+            source: None,
+        }
+    }
+
+    fn next_row(&mut self) -> &'buffer [F] {
+        if self.source.is_none() {
+            self.source = Some(self.segment.buffer(self.buffer_index).to_slice());
+        }
+
+        let source = self.source.expect("Polynomial segment has no buffer.");
+        let row_width = self.segment.row_width();
+        let start = self.row_index * row_width;
+        let row = &source[start..start + row_width];
+
+        self.row_index += 1;
+        if self.row_index == self.segment.rows_per_buffer() {
+            self.row_index = 0;
+            self.buffer_index += 1;
+            self.source = None;
+        }
+        row
+    }
+}
+
 impl<F: FftField> NttEngine<F> {
     /// Construct a new engine from the field's `FftField` trait.
     pub fn new_from_fftfield() -> Self {
@@ -398,20 +438,23 @@ impl<F: Field> ReedSolomon<F> for NttEngine<F> {
         result
     }
 
-    #[cfg_attr(feature = "tracing", instrument(skip(self, polys), fields(
-        num_polys = polys.len(),
-        poly_length = polys.first().map(|p| p.len()),
+    #[cfg_attr(feature = "tracing", instrument(skip(self, polynomials), fields(
+        num_polys = polynomials.len(),
+        poly_length = polynomials.polynomial_length(),
         codeword_length = codeword_length,
     )))]
-    fn interleaved_encode(&self, polys: &[&[F]], codeword_length: usize) -> Buffer<F> {
+    fn interleaved_encode(
+        &self,
+        polynomials: Polynomials<'_, F>,
+        codeword_length: usize,
+    ) -> Buffer<F> {
         assert!(self.order.is_multiple_of(codeword_length));
-        if polys.is_empty() {
+        let num_polys = polynomials.len();
+        let poly_length = polynomials.polynomial_length();
+        assert!(poly_length <= codeword_length);
+        if num_polys == 0 {
             return Buffer::from(Vec::new());
         }
-        let num_polys = polys.len();
-        let poly_length = polys[0].len();
-        assert!(polys.iter().all(|p| p.len() == poly_length));
-        assert!(poly_length <= codeword_length);
 
         // Coset-NTT: instead of doing one codeword-length NTT on mostly zeros,
         // do `num_cosets` many `coset_size`-point NTTs on twisted coefficient
@@ -433,15 +476,57 @@ impl<F: Field> ReedSolomon<F> for NttEngine<F> {
 
         // Lay out twisted coefficients in contiguous coset blocks of length
         // `coset_size`, zero-padding each block as needed.
-        let mut result = Vec::with_capacity(num_polys * codeword_length);
-        for poly in polys {
-            // FFT[a 0 0 0] = [a a a a], so just replicate input in coset dimension.
-            for _ in 0..num_cosets {
-                result.extend_from_slice(poly);
-                result.resize(result.len() + coset_padding, F::ZERO);
+        let output_length = num_polys
+            .checked_mul(codeword_length)
+            .expect("Encoded polynomial length overflow.");
+        let mut result = Vec::with_capacity(output_length);
+        match polynomials.segments() {
+            [] => {
+                for _ in 0..num_polys * num_cosets {
+                    result.resize(result.len() + coset_size, F::ZERO);
+                }
+            }
+            [segment] => {
+                let mut rows = SegmentRows::new(segment);
+                for _ in 0..num_polys {
+                    let polynomial = rows.next_row();
+                    // FFT[a 0 0 0] = [a a a a], so replicate the input
+                    // in the coset dimension.
+                    for _ in 0..num_cosets {
+                        result.extend_from_slice(polynomial);
+                        result.resize(result.len() + coset_padding, F::ZERO);
+                    }
+                }
+            }
+            [first, second] => {
+                let mut first_rows = SegmentRows::new(first);
+                let mut second_rows = SegmentRows::new(second);
+                for _ in 0..num_polys {
+                    let first = first_rows.next_row();
+                    let second = second_rows.next_row();
+                    for _ in 0..num_cosets {
+                        result.extend_from_slice(first);
+                        result.extend_from_slice(second);
+                        result.resize(result.len() + coset_padding, F::ZERO);
+                    }
+                }
+            }
+            segments => {
+                let mut segments = segments.iter().map(SegmentRows::new).collect::<Vec<_>>();
+                let mut rows = Vec::with_capacity(segments.len());
+                for _ in 0..num_polys {
+                    rows.clear();
+                    rows.extend(segments.iter_mut().map(SegmentRows::next_row));
+                    for _ in 0..num_cosets {
+                        for row in &rows {
+                            result.extend_from_slice(row);
+                        }
+                        result.resize(result.len() + coset_padding, F::ZERO);
+                    }
+                }
             }
         }
-        assert_eq!(result.len(), num_polys * codeword_length);
+        assert_eq!(result.len(), output_length);
 
         // NTT each coset block, then transpose each codeword block from
         // coset-major `(num_cosets × coset_size)` layout into standard codeword