Skip to content

Commit

Permalink
Backend
Browse files Browse the repository at this point in the history
  • Loading branch information
spapinistarkware committed Feb 21, 2024
1 parent 5dc08ff commit d9e5379
Show file tree
Hide file tree
Showing 15 changed files with 410 additions and 230 deletions.
42 changes: 23 additions & 19 deletions src/core/air/evaluation.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
//! Given a random alpha, the combined polynomial is defined as
//! f(p) = sum_i alpha^{N-1-i} u_i (P).
use crate::core::backend::{Backend, CPUBackend, Column};
use crate::core::fields::m31::BaseField;
use crate::core::fields::qm31::SecureField;
use crate::core::fields::Field;
Expand Down Expand Up @@ -56,20 +57,18 @@ impl PointEvaluationAccumulator {
}
}

type Column = Vec<SecureField>;

/// Accumulates evaluations of u_i(P), each at an evaluation domain of the size of that polynomial.
/// Computes the coefficients of f(P).
pub struct DomainEvaluationAccumulator {
pub struct DomainEvaluationAccumulator<B: Backend> {
random_coeff: SecureField,
// Accumulated evaluations for each log_size.
// Each `sub_accumulation` holds `sum_{i=0}^{n-1} evaluation_i * alpha^(n-1-i)`,
// where `n` is the number of accumulated evaluations for this log_size.
sub_accumulations: Vec<Column>,
sub_accumulations: Vec<Column<B, SecureField>>,
// Number of accumulated evaluations for each log_size.
n_cols_per_size: Vec<usize>,
}
impl DomainEvaluationAccumulator {
impl<B: Backend> DomainEvaluationAccumulator<B> {
/// Creates a new accumulator.
/// `random_coeff` should be a secure random field element, drawn from the channel.
/// `max_log_size` is the maximum log_size of the accumulated evaluations.
Expand All @@ -78,7 +77,7 @@ impl DomainEvaluationAccumulator {
Self {
random_coeff,
sub_accumulations: (0..(max_log_size + 1))
.map(|n| vec![SecureField::default(); 1 << n])
.map(|n| Column::<B, SecureField>::from_iter(vec![SecureField::default(); 1 << n]))
.collect(),
n_cols_per_size: vec![0; max_log_size + 1],
}
Expand All @@ -92,7 +91,7 @@ impl DomainEvaluationAccumulator {
pub fn columns<const N: usize>(
&mut self,
n_cols_per_size: [(u32, usize); N],
) -> [ColumnAccumulator<'_>; N] {
) -> [ColumnAccumulator<'_, B>; N] {
n_cols_per_size.iter().for_each(|(log_size, n_col)| {
self.n_cols_per_size[*log_size as usize] += n_col;
});
Expand All @@ -105,8 +104,15 @@ impl DomainEvaluationAccumulator {
})
}

/// Returns the log_size of the resulting polynomial.
pub fn log_size(&self) -> u32 {
(self.sub_accumulations.len() - 1) as u32
}
}

impl DomainEvaluationAccumulator<CPUBackend> {
/// Computes f(P) as coefficients.
pub fn finalize(self) -> CirclePoly<SecureField> {
pub fn finalize(self) -> CirclePoly<CPUBackend, SecureField> {
let mut res_coeffs = vec![SecureField::default(); 1 << self.log_size()];
let res_log_size = self.log_size();
for (coeffs, n_cols) in self
Expand All @@ -117,7 +123,7 @@ impl DomainEvaluationAccumulator {
if log_size == 0 {
return values;
}
CircleEvaluation::new(
CircleEvaluation::<CPUBackend, SecureField>::new(
CircleDomain::constraint_evaluation_domain(log_size as u32),
values,
)
Expand All @@ -139,19 +145,14 @@ impl DomainEvaluationAccumulator {

CirclePoly::new(res_coeffs)
}

/// Returns the log_size of the resulting polynomial.
pub fn log_size(&self) -> u32 {
(self.sub_accumulations.len() - 1) as u32
}
}

/// An domain accumulator for polynomials of a single size.
pub struct ColumnAccumulator<'a> {
pub struct ColumnAccumulator<'a, B: Backend> {
random_coeff: SecureField,
col: &'a mut Column,
col: &'a mut Column<B, SecureField>,
}
impl<'a> ColumnAccumulator<'a> {
impl<'a> ColumnAccumulator<'a, CPUBackend> {
pub fn accumulate(&mut self, index: usize, evaluation: BaseField) {
let accum = &mut self.col[index];
*accum = *accum * self.random_coeff + evaluation;
Expand All @@ -166,10 +167,13 @@ mod tests {
use rand::{Rng, SeedableRng};

use super::*;
use crate::core::backend::CPUBackend;
use crate::core::circle::CirclePoint;
use crate::core::fields::m31::{M31, P};
use crate::qm31;

type B = CPUBackend;

#[test]
fn test_point_evaluation_accumulator() {
// Generate a vector of random sizes with a constant seed.
Expand Down Expand Up @@ -229,7 +233,7 @@ mod tests {
let alpha = qm31!(2, 3, 4, 5);

// Use accumulator.
let mut accumulator = DomainEvaluationAccumulator::new(alpha, LOG_SIZE_BOUND);
let mut accumulator = DomainEvaluationAccumulator::<B>::new(alpha, LOG_SIZE_BOUND);
let n_cols_per_size: [(u32, usize); (LOG_SIZE_BOUND - LOG_SIZE_MIN) as usize] =
array::from_fn(|i| {
let current_log_size = LOG_SIZE_MIN + i as u32;
Expand Down Expand Up @@ -260,7 +264,7 @@ mod tests {
let mut res = SecureField::default();
for (log_size, values) in pairs.into_iter() {
res = res * alpha
+ CircleEvaluation::new(
+ CircleEvaluation::<B, _>::new(
CircleDomain::constraint_evaluation_domain(log_size),
values,
)
Expand Down
25 changes: 13 additions & 12 deletions src/core/air/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ use std::iter::zip;
use std::ops::Deref;

use self::evaluation::{DomainEvaluationAccumulator, PointEvaluationAccumulator};
use super::backend::Backend;
use super::circle::CirclePoint;
use super::fields::m31::BaseField;
use super::fields::qm31::SecureField;
Expand All @@ -15,11 +16,11 @@ pub mod evaluation;
/// For instance, all interaction elements are assumed to be present in it.
/// Therefore, an AIR is generated only after the initial trace commitment phase.
// TODO(spapini): consider renaming this struct.
pub trait Air {
fn visit_components<V: ComponentVisitor>(&self, v: &mut V);
pub trait Air<B: Backend> {
fn visit_components<V: ComponentVisitor<B>>(&self, v: &mut V);
}
pub trait ComponentVisitor {
fn visit<C: Component>(&mut self, component: &C);
pub trait ComponentVisitor<B: Backend> {
fn visit<C: Component<B>>(&mut self, component: &C);
}

/// Holds the mask offsets at each column.
Expand Down Expand Up @@ -54,7 +55,7 @@ impl Deref for Mask {

/// A component is a set of trace columns of various sizes along with a set of
/// constraints on them.
pub trait Component {
pub trait Component<B: Backend> {
fn max_constraint_log_degree_bound(&self) -> u32;

/// Returns the degree bounds of each trace column.
Expand All @@ -66,8 +67,8 @@ pub trait Component {
// Note: This will be computed using a MaterializedGraph.
fn evaluate_constraint_quotients_on_domain(
&self,
trace: &ComponentTrace<'_>,
evaluation_accumulator: &mut DomainEvaluationAccumulator,
trace: &ComponentTrace<'_, B>,
evaluation_accumulator: &mut DomainEvaluationAccumulator<B>,
);

fn mask(&self) -> Mask;
Expand All @@ -79,7 +80,7 @@ pub trait Component {
fn mask_points_and_values(
&self,
point: CirclePoint<SecureField>,
trace: &ComponentTrace<'_>,
trace: &ComponentTrace<'_, B>,
) -> (Vec<Vec<CirclePoint<SecureField>>>, Vec<Vec<SecureField>>) {
let domains = trace
.columns
Expand Down Expand Up @@ -110,12 +111,12 @@ pub trait Component {
// TODO(spapini): Extra functions for FRI and decommitment.
}

pub struct ComponentTrace<'a> {
pub columns: Vec<&'a CirclePoly<BaseField>>,
pub struct ComponentTrace<'a, B: Backend> {
pub columns: Vec<&'a CirclePoly<B, BaseField>>,
}

impl<'a> ComponentTrace<'a> {
pub fn new(columns: Vec<&'a CirclePoly<BaseField>>) -> Self {
impl<'a, B: Backend> ComponentTrace<'a, B> {
pub fn new(columns: Vec<&'a CirclePoly<B, BaseField>>) -> Self {
Self { columns }
}
}
20 changes: 20 additions & 0 deletions src/core/backend/cpu/mod.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
use std::fmt::Debug;

use super::{Backend, ColumnTrait, FieldOps};
use crate::core::fields::Field;

mod poly;

#[derive(Copy, Clone, Debug)]
pub struct CPUBackend;
impl Backend for CPUBackend {}

impl<F: Field> FieldOps<F> for CPUBackend {
type Column = Vec<F>;
}

impl<F: Clone + Debug> ColumnTrait<F> for Vec<F> {
fn len(&self) -> usize {
self.len()
}
}
137 changes: 137 additions & 0 deletions src/core/backend/cpu/poly.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,137 @@
use super::CPUBackend;
use crate::core::backend::Column;
use crate::core::circle::CirclePoint;
use crate::core::fft::{butterfly, ibutterfly};
use crate::core::fields::m31::BaseField;
use crate::core::fields::{ExtensionOf, Field};
use crate::core::poly::circle::{
CanonicCoset, CircleDomain, CircleEvaluation, CirclePoly, PolyOps,
};
use crate::core::poly::utils::fold;
use crate::core::poly::{BitReversedOrder, NaturalOrder};
use crate::core::utils::bit_reverse;

impl<F: ExtensionOf<BaseField>> PolyOps<F> for CPUBackend {
fn new_canonical_ordered(
coset: CanonicCoset,
values: Column<Self, F>,
) -> CircleEvaluation<Self, F> {
let domain = coset.circle_domain();
assert_eq!(values.len(), domain.size());
let mut new_values = Vec::with_capacity(values.len());
let half_len = 1 << (coset.log_size() - 1);
for i in 0..half_len {
new_values.push(values[i << 1]);
}
for i in 0..half_len {
new_values.push(values[domain.size() - 1 - (i << 1)]);
}
CircleEvaluation::new(domain, new_values)
}

fn interpolate(eval: CircleEvaluation<Self, F>) -> CirclePoly<Self, F> {
// Use CFFT to interpolate.
let mut coset = eval.domain.half_coset;
let mut values = eval.values;
let (l, r) = values.split_at_mut(coset.size());
for (i, p) in coset.iter().enumerate() {
ibutterfly(&mut l[i], &mut r[i], p.y.inverse());
}
while coset.size() > 1 {
for chunk in values.chunks_exact_mut(coset.size()) {
let (l, r) = chunk.split_at_mut(coset.size() / 2);
for (i, p) in coset.iter().take(coset.size() / 2).enumerate() {
ibutterfly(&mut l[i], &mut r[i], p.x.inverse());
}
}
coset = coset.double();
}

// Divide all values by 2^log_size.
let inv = BaseField::from_u32_unchecked(eval.domain.size() as u32).inverse();
for val in &mut values {
*val *= inv;
}

CirclePoly::new(values)
}

// TODO(spapini): Remove
fn bit_reverse_natural(
eval: CircleEvaluation<Self, F, NaturalOrder>,
) -> CircleEvaluation<Self, F, BitReversedOrder> {
CircleEvaluation::new(eval.domain, bit_reverse(eval.values))
}

fn bit_reverse_reversed(
eval: CircleEvaluation<Self, F, BitReversedOrder>,
) -> CircleEvaluation<Self, F, NaturalOrder> {
CircleEvaluation::new(eval.domain, bit_reverse(eval.values))
}

fn eval_at_point<E: ExtensionOf<F>>(poly: &CirclePoly<Self, F>, point: CirclePoint<E>) -> E {
// TODO(Andrew): Allocation here expensive for small polynomials.
let mut mappings = vec![point.y, point.x];
let mut x = point.x;
for _ in 2..poly.log_size() {
x = CirclePoint::double_x(x);
mappings.push(x);
}
fold(&poly.coeffs, &mappings)
}

fn extend(poly: &CirclePoly<Self, F>, log_size: u32) -> CirclePoly<Self, F> {
assert!(log_size >= poly.log_size());
let mut coeffs = vec![F::zero(); 1 << log_size];
let log_jump = log_size - poly.log_size();
for (i, val) in poly.coeffs.iter().enumerate() {
coeffs[i << log_jump] = *val;
}
CirclePoly::new(coeffs)
}

fn evaluate(poly: &CirclePoly<Self, F>, domain: CircleDomain) -> CircleEvaluation<Self, F> {
// Use CFFT to evaluate.
let mut coset = domain.half_coset;
let mut cosets = vec![];

// TODO(spapini): extend better.
assert!(domain.log_size() >= poly.log_size());
let mut values = poly.clone().extend(domain.log_size()).coeffs;

while coset.size() > 1 {
cosets.push(coset);
coset = coset.double();
}
for coset in cosets.iter().rev() {
for chunk in values.chunks_exact_mut(coset.size()) {
let (l, r) = chunk.split_at_mut(coset.size() / 2);
for (i, p) in coset.iter().take(coset.size() / 2).enumerate() {
butterfly(&mut l[i], &mut r[i], p.x);
}
}
}
let coset = domain.half_coset;
let (l, r) = values.split_at_mut(coset.size());
for (i, p) in coset.iter().enumerate() {
butterfly(&mut l[i], &mut r[i], p.y);
}
CircleEvaluation::new(domain, values)
}
}

impl<F: ExtensionOf<BaseField>, EvalOrder> IntoIterator
for CircleEvaluation<CPUBackend, F, EvalOrder>
{
type Item = F;
type IntoIter = std::vec::IntoIter<F>;

/// Creates a consuming iterator over the evaluations.
///
/// Evaluations are returned in the same order as elements of the domain.
fn into_iter(self) -> Self::IntoIter {
self.values.into_iter()
}
}

// impl<F: ExtensionOf<BaseField>> PolyOps<CPUBackend, F>
35 changes: 35 additions & 0 deletions src/core/backend/mod.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
use std::fmt::Debug;
use std::ops::Index;

pub use cpu::CPUBackend;

use super::fields::m31::BaseField;
use super::fields::qm31::SecureField;
use super::fields::Field;
use super::poly::circle::PolyOps;

pub mod cpu;

pub trait Backend:
Copy
+ Clone
+ Debug
+ FieldOps<BaseField>
+ FieldOps<SecureField>
+ PolyOps<BaseField>
+ PolyOps<SecureField>
{
}

pub trait FieldOps<F: Field> {
type Column: ColumnTrait<F>;
}

pub type Column<B, F> = <B as FieldOps<F>>::Column;

pub trait ColumnTrait<F>: Clone + Debug + Index<usize, Output = F> + FromIterator<F> {
fn len(&self) -> usize;
fn is_empty(&self) -> bool {
self.len() == 0
}
}
Loading

0 comments on commit d9e5379

Please sign in to comment.