Skip to content

Commit

Permalink
avx poly
Browse files Browse the repository at this point in the history
  • Loading branch information
spapinistarkware committed Feb 22, 2024
1 parent 47a0728 commit 0811c5f
Show file tree
Hide file tree
Showing 3 changed files with 66 additions and 1 deletion.
2 changes: 1 addition & 1 deletion benches/bit_rev.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
})
});
}
Expand Down
53 changes: 53 additions & 0 deletions src/core/backend/avx512/circle.rs
Original file line number Diff line number Diff line change
@@ -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,
))
}
}
12 changes: 12 additions & 0 deletions src/core/backend/avx512/mod.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
pub mod bit_reverse;
pub mod circle;

use std::ops::Index;

Expand Down Expand Up @@ -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 {
Expand Down

0 comments on commit 0811c5f

Please sign in to comment.