AVX backend

Cargo.toml

@@ -13,6 +13,7 @@ itertools = "0.12.0"
 num-traits = "0.2.17"
 thiserror = "1.0.56"
 merging-iterator = "1.3.0"
+bytemuck = { version = "1.14.3", features = ["derive"] }
 
 [dev-dependencies]
 criterion = { version = "0.5.1", features = ["html_reports"] }

benches/bit_rev.rs

@@ -13,7 +13,7 @@ pub fn cpu_bit_rev(c: &mut criterion::Criterion) {
 
     c.bench_function("cpu bit_rev", |b| {
         b.iter(|| {
-            data = stwo::core::utils::bit_reverse(std::mem::take(&mut data));
+            stwo::core::utils::bit_reverse(&mut data);
         })
     });
 }

src/core/backend/avx512/bit_reverse.rs

@@ -1,16 +1,15 @@
 use std::arch::x86_64::{__m512i, _mm512_permutex2var_epi32};
 
-use crate::core::fields::m31::BaseField;
+use super::PackedBaseField;
 use crate::core::utils::bit_reverse_index;
 
 const VEC_BITS: u32 = 4;
 const W_BITS: u32 = 3;
-const MIN_LOG_SIZE: u32 = 2 * W_BITS + VEC_BITS;
+pub const MIN_LOG_SIZE: u32 = 2 * W_BITS + VEC_BITS;
 
-// TODO(spapini): Use  PackedBaseField type.
 /// Bit reverses packed M31 values.
-/// Given an array A[0..2^n), computes B[i] = A[bit_reverse(i)].
-pub fn bit_reverse_m31(data: &mut [[BaseField; 16]]) {
+/// Given an array `A[0..2^n)`, computes `B[i] = A[bit_reverse(i)]`.
+pub fn bit_reverse_m31(data: &mut [PackedBaseField]) {
     assert!(data.len().is_power_of_two());
     assert!(data.len().ilog2() >= MIN_LOG_SIZE);
 
@@ -74,7 +73,7 @@ pub fn bit_reverse_m31(data: &mut [[BaseField; 16]]) {
 }
 
 /// Bit reverses 16 packed M31 values.
-fn bit_reverse16(data: [[BaseField; 16]; 16]) -> [[BaseField; 16]; 16] {
+fn bit_reverse16(data: [PackedBaseField; 16]) -> [PackedBaseField; 16] {
     let mut data: [__m512i; 16] = unsafe { std::mem::transmute(data) };
     // L is an input to _mm512_permutex2var_epi32, and it is used to
     // interleave the first half of a with the first half of b.
@@ -159,7 +158,8 @@ mod tests {
         let data: Vec<_> = (0..SIZE as u32)
             .map(BaseField::from_u32_unchecked)
             .collect();
-        let expected = bit_reverse(data.clone());
+        let mut expected = data.clone();
+        bit_reverse(&mut expected);
         let mut data: Vec<_> = data.into_iter().array_chunks::<16>().collect();
         let expected: Vec<_> = expected.into_iter().array_chunks::<16>().collect();
 

src/core/backend/avx512/mod.rs

@@ -1 +1,124 @@
 pub mod bit_reverse;
+
+use std::ops::Index;
+
+use bytemuck::checked::cast_slice_mut;
+use num_traits::Zero;
+
+use self::bit_reverse::bit_reverse_m31;
+use crate::core::fields::m31::BaseField;
+use crate::core::fields::{Column, FieldOps};
+use crate::core::utils;
+
+#[derive(Copy, Clone, Debug)]
+pub struct AVX512Backend;
+
+// BaseField.
+// TODO(spapini): Unite with the M31AVX512 type.
+pub const K_ELEMENTS: usize = 16;
+type PackedBaseField = [BaseField; K_ELEMENTS];
+#[derive(Clone, Debug)]
+pub struct BaseFieldVec {
+    data: Vec<PackedBaseField>,
+    length: usize,
+}
+impl FieldOps<BaseField> for AVX512Backend {
+    type Column = BaseFieldVec;
+
+    fn bit_reverse_column(column: &mut Self::Column) {
+        // Fallback to cpu bit_reverse.
+        if column.data.len().ilog2() < bit_reverse::MIN_LOG_SIZE {
+            let data: &mut [BaseField] = cast_slice_mut(&mut column.data[..]);
+            utils::bit_reverse(&mut data[..column.length]);
+            return;
+        }
+        bit_reverse_m31(&mut column.data);
+    }
+}
+
+impl Column<BaseField> for BaseFieldVec {
+    fn zeros(len: usize) -> Self {
+        Self {
+            data: vec![PackedBaseField::default(); len.div_ceil(K_ELEMENTS)],
+            length: len,
+        }
+    }
+    fn to_vec(&self) -> Vec<BaseField> {
+        self.data
+            .iter()
+            .flatten()
+            .copied()
+            .take(self.length)
+            .collect()
+    }
+    fn len(&self) -> usize {
+        self.length
+    }
+}
+
+impl Index<usize> for BaseFieldVec {
+    type Output = BaseField;
+    fn index(&self, index: usize) -> &Self::Output {
+        &self.data[index / K_ELEMENTS][index % K_ELEMENTS]
+    }
+}
+
+impl FromIterator<BaseField> for BaseFieldVec {
+    fn from_iter<I: IntoIterator<Item = BaseField>>(iter: I) -> Self {
+        let mut chunks = iter.into_iter().array_chunks();
+        let mut res: Vec<_> = (&mut chunks).collect();
+        let mut length = res.len() * K_ELEMENTS;
+
+        if let Some(remainder) = chunks.into_remainder() {
+            if !remainder.is_empty() {
+                length += remainder.len();
+                let pad_len = 16 - remainder.len();
+                let last: PackedBaseField = remainder
+                    .chain(std::iter::repeat(BaseField::zero()).take(pad_len))
+                    .collect::<Vec<_>>()
+                    .try_into()
+                    .unwrap();
+                res.push(last);
+            }
+        }
+
+        Self { data: res, length }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::core::fields::{Col, Column};
+
+    type B = AVX512Backend;
+
+    #[test]
+    fn test_column() {
+        for i in 0..100 {
+            let col = Col::<B, BaseField>::from_iter((0..i).map(BaseField::from));
+            assert_eq!(
+                col.to_vec(),
+                (0..i).map(BaseField::from).collect::<Vec<_>>()
+            );
+            for j in 0..i {
+                assert_eq!(col[j], BaseField::from(j));
+            }
+        }
+    }
+
+    #[test]
+    fn test_bit_reverse() {
+        for i in 1..16 {
+            let len = 1 << i;
+            let mut col = Col::<B, BaseField>::from_iter((0..len).map(BaseField::from));
+            B::bit_reverse_column(&mut col);
+            assert_eq!(
+                col.to_vec(),
+                (0..len)
+                    .map(|x| BaseField::from(utils::bit_reverse_index(x, i as u32)))
+                    .collect::<Vec<_>>()
+            );
+        }
+    }
+}

src/core/backend/cpu/mod.rs

@@ -18,7 +18,7 @@ impl Backend for CPUBackend {}
 impl<F: Field> FieldOps<F> for CPUBackend {
     type Column = Vec<F>;
 
-    fn bit_reverse_column(column: Self::Column) -> Self::Column {
+    fn bit_reverse_column(column: &mut Self::Column) {
         bit_reverse(column)
     }
 }

src/core/fields/m31.rs

@@ -3,14 +3,17 @@ use std::ops::{
     Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Rem, RemAssign, Sub, SubAssign,
 };
 
+use bytemuck::{Pod, Zeroable};
+
 use super::ComplexConjugate;
 use crate::impl_field;
 
 pub const MODULUS_BITS: u32 = 31;
 pub const N_BYTES_FELT: usize = 4;
 pub const P: u32 = 2147483647; // 2 ** 31 - 1
 
-#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
+#[repr(transparent)]
+#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash, Pod, Zeroable)]
 pub struct M31(u32);
 pub type BaseField = M31;
 

src/core/fields/mod.rs

@@ -11,7 +11,7 @@ pub mod qm31;
 
 pub trait FieldOps<F: Field> {
     type Column: Column<F>;
-    fn bit_reverse_column(column: Self::Column) -> Self::Column;
+    fn bit_reverse_column(column: &mut Self::Column);
 }
 
 pub type Col<B, F> = <B as FieldOps<F>>::Column;

src/core/poly/circle.rs

@@ -254,8 +254,9 @@ impl<F: ExtensionOf<BaseField>, B: PolyOps<F>> CircleEvaluation<B, F> {
         self.values[self.domain.find(point_index).expect("Not in domain")]
     }
 
-    pub fn bit_reverse(self) -> CircleEvaluation<B, F, BitReversedOrder> {
-        CircleEvaluation::new(self.domain, B::bit_reverse_column(self.values))
+    pub fn bit_reverse(mut self) -> CircleEvaluation<B, F, BitReversedOrder> {
+        B::bit_reverse_column(&mut self.values);
+        CircleEvaluation::new(self.domain, self.values)
     }
 }
 
@@ -281,8 +282,9 @@ impl<F: ExtensionOf<BaseField>> CPUCircleEvaluation<F> {
 }
 
 impl<B: PolyOps<F>, F: ExtensionOf<BaseField>> CircleEvaluation<B, F, BitReversedOrder> {
-    pub fn bit_reverse(self) -> CircleEvaluation<B, F, NaturalOrder> {
-        CircleEvaluation::new(self.domain, B::bit_reverse_column(self.values))
+    pub fn bit_reverse(mut self) -> CircleEvaluation<B, F, NaturalOrder> {
+        B::bit_reverse_column(&mut self.values);
+        CircleEvaluation::new(self.domain, self.values)
     }
 
     pub fn get_at(&self, point_index: CirclePointIndex) -> F {
@@ -389,7 +391,8 @@ impl<F: ExtensionOf<BaseField>, B: PolyOps<F>> CirclePoly<B, F> {
 #[cfg(test)]
 impl<F: ExtensionOf<BaseField>> crate::core::backend::cpu::CPUCirclePoly<F> {
     pub fn is_in_fft_space(&self, log_fft_size: u32) -> bool {
-        let mut coeffs = crate::core::utils::bit_reverse(self.coeffs.clone());
+        let mut coeffs = self.coeffs.clone();
+        crate::core::utils::bit_reverse(&mut coeffs);
         while coeffs.last() == Some(&F::zero()) {
             coeffs.pop();
         }

src/core/poly/line.rs

@@ -156,17 +156,19 @@ impl<F: ExtensionOf<BaseField>> LinePoly<F> {
     }
 
     /// Returns the polynomial's coefficients in their natural order.
-    pub fn into_ordered_coefficients(self) -> Vec<F> {
-        bit_reverse(self.coeffs)
+    pub fn into_ordered_coefficients(mut self) -> Vec<F> {
+        bit_reverse(&mut self.coeffs);
+        self.coeffs
     }
 
     /// Creates a new line polynomial from coefficients in their natural order.
     ///
     /// # Panics
     ///
     /// Panics if the number of coefficients is not a power of two.
-    pub fn from_ordered_coefficients(coeffs: Vec<F>) -> Self {
-        Self::new(bit_reverse(coeffs))
+    pub fn from_ordered_coefficients(mut coeffs: Vec<F>) -> Self {
+        bit_reverse(&mut coeffs);
+        Self::new(coeffs)
     }
 }
 
@@ -242,19 +244,21 @@ impl<F: ExtensionOf<BaseField>> CPULineEvaluation<F> {
 }
 
 impl<B: FieldOps<F>, F: Field> LineEvaluation<B, F> {
-    pub fn bit_reverse(self) -> LineEvaluation<B, F, BitReversedOrder> {
+    pub fn bit_reverse(mut self) -> LineEvaluation<B, F, BitReversedOrder> {
+        B::bit_reverse_column(&mut self.values);
         LineEvaluation {
-            values: B::bit_reverse_column(self.values),
+            values: self.values,
             domain: self.domain,
             _eval_order: PhantomData,
         }
     }
 }
 
 impl<B: FieldOps<F>, F: Field> LineEvaluation<B, F, BitReversedOrder> {
-    pub fn bit_reverse(self) -> LineEvaluation<B, F, NaturalOrder> {
+    pub fn bit_reverse(mut self) -> LineEvaluation<B, F, NaturalOrder> {
+        B::bit_reverse_column(&mut self.values);
         LineEvaluation {
-            values: B::bit_reverse_column(self.values),
+            values: self.values,
             domain: self.domain,
             _eval_order: PhantomData,
         }

src/core/queries.rs

@@ -158,13 +158,13 @@ mod tests {
     pub fn test_folded_queries() {
         let log_domain_size = 7;
         let domain = CanonicCoset::new(log_domain_size).circle_domain();
-        let values = domain.iter().collect::<Vec<_>>();
-        let values = bit_reverse(values.clone());
+        let mut values = domain.iter().collect::<Vec<_>>();
+        bit_reverse(&mut values);
 
         let log_folded_domain_size = 5;
         let folded_domain = CanonicCoset::new(log_folded_domain_size).circle_domain();
-        let folded_values = folded_domain.iter().collect::<Vec<_>>();
-        let folded_values = bit_reverse(folded_values.clone());
+        let mut folded_values = folded_domain.iter().collect::<Vec<_>>();
+        bit_reverse(&mut folded_values);
 
         // Generate all possible queries.
         let queries = Queries {
@@ -192,8 +192,8 @@ mod tests {
         let channel = &mut Blake2sChannel::new(Blake2sHash::default());
         let log_domain_size = 7;
         let domain = CanonicCoset::new(log_domain_size).circle_domain();
-        let values = domain.iter().collect::<Vec<_>>();
-        let values = bit_reverse(values.clone());
+        let mut values = domain.iter().collect::<Vec<_>>();
+        bit_reverse(&mut values);
 
         // Test random queries one by one because the conjugate queries are sorted.
         for _ in 0..100 {