avx poly

src/core/backend/avx512/circle.rs

@@ -0,0 +1,51 @@
+use super::{as_cpu_vec, AVX512Backend};
+use crate::core::backend::CPUBackend;
+use crate::core::fields::m31::BaseField;
+use crate::core::fields::Col;
+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: crate::core::fields::ExtensionOf<BaseField>>(
+        poly: &CirclePoly<Self, BaseField>,
+        point: crate::core::circle::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;
 
 use std::ops::Index;
 
@@ -53,6 +54,17 @@ impl Column<BaseField> for BaseFieldVec {
     }
 }
 
+fn as_cpu_vec(values: BaseFieldVec) -> Vec<BaseField> {
+    let capacity = values.len() * 16;
+    unsafe {
+        Vec::from_raw_parts(
+            values.data.as_ptr() as *mut BaseField,
+            values.length,
+            capacity,
+        )
+    }
+}
+
 impl Index<usize> for BaseFieldVec {
     type Output = BaseField;
     fn index(&self, index: usize) -> &Self::Output {