From 0811c5fa5602f8ea30e15d35bf582dd52381038c Mon Sep 17 00:00:00 2001
From: Shahar Papini <spapini@starkware.co>
Date: Sat, 17 Feb 2024 20:20:28 +0200
Subject: [PATCH] avx poly

---
 benches/bit_rev.rs                |  2 +-
 src/core/backend/avx512/circle.rs | 53 +++++++++++++++++++++++++++++++
 src/core/backend/avx512/mod.rs    | 12 +++++++
 3 files changed, 66 insertions(+), 1 deletion(-)
 create mode 100644 src/core/backend/avx512/circle.rs

diff --git a/benches/bit_rev.rs b/benches/bit_rev.rs
index 3ee16c0ca..9dfa74b2d 100644
--- a/benches/bit_rev.rs
+++ b/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 = prover_research::core::utils::bit_reverse(std::mem::take(&mut data));
+            prover_research::core::utils::bit_reverse(&mut data);
         })
     });
 }
diff --git a/src/core/backend/avx512/circle.rs b/src/core/backend/avx512/circle.rs
new file mode 100644
index 000000000..9abd0a36f
--- /dev/null
+++ b/src/core/backend/avx512/circle.rs
@@ -0,0 +1,53 @@
+use super::{as_cpu_vec, AVX512Backend};
+use crate::core::backend::{CPUBackend, Column, ColumnTrait};
+use crate::core::fields::m31::BaseField;
+use crate::core::poly::circle::{
+    CanonicCoset, CircleDomain, CircleEvaluation, CirclePoly, PolyOps,
+};
+
+impl PolyOps<BaseField> for AVX512Backend {
+    fn new_canonical_ordered(
+        coset: CanonicCoset,
+        values: Column<Self, BaseField>,
+    ) -> CircleEvaluation<Self, BaseField> {
+        let eval = CPUBackend::new_canonical_ordered(coset, as_cpu_vec(values));
+        CircleEvaluation::new(
+            eval.domain,
+            Column::<AVX512Backend, _>::from_vec(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(Column::<AVX512Backend, _>::from_vec(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,
+            Column::<AVX512Backend, _>::from_vec(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(Column::<AVX512Backend, _>::from_vec(
+            cpu_poly.extend(log_size).coeffs,
+        ))
+    }
+}
diff --git a/src/core/backend/avx512/mod.rs b/src/core/backend/avx512/mod.rs
index 553323149..ebb309189 100644
--- a/src/core/backend/avx512/mod.rs
+++ b/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 {