avx poly

src/core/backend/avx512/circle.rs

@@ -0,0 +1,52 @@
+use super::{as_cpu_vec, AVX512Backend};
+use crate::core::backend::CPUBackend;
+use crate::core::circle::CirclePoint;
+use crate::core::fields::m31::BaseField;
+use crate::core::fields::{Col, ExtensionOf};
+use crate::core::poly::circle::{
+    CanonicCoset, CircleDomain, CircleEvaluation, CirclePoly, PolyOps,
+};
+
+impl PolyOps<BaseField> for AVX512Backend {
+    fn new_canonical_ordered(
+        coset: CanonicCoset,
+        values: Col<Self, BaseField>,
+    ) -> CircleEvaluation<Self, BaseField> {
+        let eval = CPUBackend::new_canonical_ordered(coset, as_cpu_vec(values));
+        CircleEvaluation::new(eval.domain, Col::<AVX512Backend, _>::from_iter(eval.values))
+    }
+
+    fn interpolate(eval: CircleEvaluation<Self, BaseField>) -> CirclePoly<Self, BaseField> {
+        let cpu_eval = CircleEvaluation::<CPUBackend, _>::new(eval.domain, as_cpu_vec(eval.values));
+        let cpu_poly = cpu_eval.interpolate();
+        CirclePoly::new(Col::<AVX512Backend, _>::from_iter(cpu_poly.coeffs))
+    }
+
+    fn eval_at_point<E: ExtensionOf<BaseField>>(
+        poly: &CirclePoly<Self, BaseField>,
+        point: CirclePoint<E>,
+    ) -> E {
+        // TODO(spapini): Unnecessary allocation here.
+        let cpu_poly = CirclePoly::<CPUBackend, _>::new(as_cpu_vec(poly.coeffs.clone()));
+        cpu_poly.eval_at_point(point)
+    }
+
+    fn evaluate(
+        poly: &CirclePoly<Self, BaseField>,
+        domain: CircleDomain,
+    ) -> CircleEvaluation<Self, BaseField> {
+        let cpu_poly = CirclePoly::<CPUBackend, _>::new(as_cpu_vec(poly.coeffs.clone()));
+        let cpu_eval = cpu_poly.evaluate(domain);
+        CircleEvaluation::new(
+            cpu_eval.domain,
+            Col::<AVX512Backend, _>::from_iter(cpu_eval.values),
+        )
+    }
+
+    fn extend(poly: &CirclePoly<Self, BaseField>, log_size: u32) -> CirclePoly<Self, BaseField> {
+        let cpu_poly = CirclePoly::<CPUBackend, _>::new(as_cpu_vec(poly.coeffs.clone()));
+        CirclePoly::new(Col::<AVX512Backend, _>::from_iter(
+            cpu_poly.extend(log_size).coeffs,
+        ))
+    }
+}

src/core/backend/avx512/mod.rs

@@ -1,4 +1,5 @@
 pub mod bit_reverse;
+pub mod circle;
 pub mod cm31;
 pub mod m31;
 pub mod qm31;
@@ -78,6 +79,19 @@ impl Column<BaseField> for BaseFieldVec {
     }
 }
 
+fn as_cpu_vec(values: BaseFieldVec) -> Vec<BaseField> {
+    let capacity = values.data.capacity() * K_ELEMENTS;
+    unsafe {
+        let res = Vec::from_raw_parts(
+            values.data.as_ptr() as *mut BaseField,
+            values.length,
+            capacity,
+        );
+        std::mem::forget(values);
+        res
+    }
+}
+
 impl FromIterator<BaseField> for BaseFieldVec {
     fn from_iter<I: IntoIterator<Item = BaseField>>(iter: I) -> Self {
         let mut chunks = iter.into_iter().array_chunks();
@@ -138,4 +152,12 @@ mod tests {
             );
         }
     }
+
+    #[test]
+    fn test_as_cpu_vec() {
+        let original_vec = (1000..1100).map(BaseField::from).collect::<Vec<_>>();
+        let col = Col::<B, BaseField>::from_iter(original_vec.clone());
+        let vec = as_cpu_vec(col);
+        assert_eq!(vec, original_vec);
+    }
 }