Skip to content

Commit

Permalink
Revert "removed computing line coeffs (#643)" (#684)
Browse files Browse the repository at this point in the history
* Revert "removed computing line coeffs (#643)"

This reverts commit a150b8f.

* Revert "integrated point vanishing in fri quotients (#619)"

This reverts commit a6eabdd.
  • Loading branch information
spapinistarkware committed Jul 2, 2024
1 parent 436541e commit 60027bd
Show file tree
Hide file tree
Showing 4 changed files with 215 additions and 96 deletions.
87 changes: 63 additions & 24 deletions crates/prover/src/core/backend/cpu/quotients.rs
Original file line number Diff line number Diff line change
@@ -1,17 +1,18 @@
use itertools::izip;
use num_traits::Zero;
use itertools::{izip, zip_eq};
use num_traits::{One, Zero};

use super::CpuBackend;
use crate::core::backend::{Backend, Col};
use crate::core::constraints::point_vanishing_fraction;
use crate::core::circle::CirclePoint;
use crate::core::constraints::{complex_conjugate_line_coeffs, pair_vanishing};
use crate::core::fields::m31::BaseField;
use crate::core::fields::qm31::SecureField;
use crate::core::fields::secure_column::SecureColumn;
use crate::core::fields::FieldExpOps;
use crate::core::pcs::quotients::{ColumnSampleBatch, QuotientOps};
use crate::core::fields::{ComplexConjugate, FieldExpOps};
use crate::core::pcs::quotients::{ColumnSampleBatch, PointSample, QuotientOps};
use crate::core::poly::circle::{CircleDomain, CircleEvaluation, SecureEvaluation};
use crate::core::poly::BitReversedOrder;
use crate::core::utils::bit_reverse;
use crate::core::utils::{bit_reverse, bit_reverse_index};

impl QuotientOps for CpuBackend {
fn accumulate_quotients(
Expand All @@ -26,40 +27,76 @@ impl QuotientOps for CpuBackend {
// TODO(spapini): bit reverse iterator.
for row in 0..domain.size() {
// TODO(alonh): Make an efficient bit reverse domain iterator, possibly for AVX backend.
let row_value =
accumulate_row_quotients(sample_batches, columns, &quotient_constants, row);
let domain_point = domain.at(bit_reverse_index(row, domain.log_size()));
let row_value = accumulate_row_quotients(
sample_batches,
columns,
&quotient_constants,
row,
domain_point,
);
values.set(row, row_value);
}
SecureEvaluation { domain, values }
}
}

// TODO(Ohad): no longer using pair_vanishing, remove domain_point_vec and line_coeffs, or write a
// function that deals with quotients over pair_vanishing polynomials.
pub fn accumulate_row_quotients(
sample_batches: &[ColumnSampleBatch],
columns: &[&CircleEvaluation<CpuBackend, BaseField, BitReversedOrder>],
quotient_constants: &QuotientConstants<CpuBackend>,
row: usize,
domain_point: CirclePoint<BaseField>,
) -> SecureField {
let mut row_accumulator = SecureField::zero();
for (sample_batch, batch_coeff, denominator_inverses) in izip!(
for (sample_batch, line_coeffs, batch_coeff, denominator_inverses) in izip!(
sample_batches,
&quotient_constants.line_coeffs,
&quotient_constants.batch_random_coeffs,
&quotient_constants.denominator_inverses
) {
let mut numerator = SecureField::zero();
for (column_index, sampled_value) in sample_batch.columns_and_values.iter() {
for ((column_index, _), (a, b, c)) in zip_eq(&sample_batch.columns_and_values, line_coeffs)
{
let column = &columns[*column_index];
let value = column[row];
numerator += value - *sampled_value;
let value = column[row] * *c;
let linear_term = *a * domain_point.y + *b;
numerator += value - linear_term;
}

row_accumulator = row_accumulator * *batch_coeff + numerator * denominator_inverses[row];
}
row_accumulator
}

/// Precompute the complex conjugate line coefficients for each column in each sample batch.
/// Specifically, for the i-th (in a sample batch) column's numerator term
/// `alpha^i * (c * F(p) - (a * p.y + b))`, we precompute and return the constants:
/// (`alpha^i * a`, `alpha^i * b`, `alpha^i * c`).
pub fn column_line_coeffs(
sample_batches: &[ColumnSampleBatch],
random_coeff: SecureField,
) -> Vec<Vec<(SecureField, SecureField, SecureField)>> {
sample_batches
.iter()
.map(|sample_batch| {
let mut alpha = SecureField::one();
sample_batch
.columns_and_values
.iter()
.map(|(_, sampled_value)| {
alpha *= random_coeff;
let sample = PointSample {
point: sample_batch.point,
value: *sampled_value,
};
complex_conjugate_line_coeffs(&sample, alpha)
})
.collect()
})
.collect()
}

/// Precompute the random coefficients used to linearly combine the batched quotients.
/// Specifically, for each sample batch we compute random_coeff^(number of columns in the batch),
/// which is used to linearly combine the batch with the next one.
Expand All @@ -77,24 +114,21 @@ fn denominator_inverses(
sample_batches: &[ColumnSampleBatch],
domain: CircleDomain,
) -> Vec<Col<CpuBackend, SecureField>> {
let n_fracions = sample_batches.len() * domain.size();
let mut flat_denominators = Vec::with_capacity(n_fracions);
let mut numerator_terms = Vec::with_capacity(n_fracions);
let mut flat_denominators = Vec::with_capacity(sample_batches.len() * domain.size());
for sample_batch in sample_batches {
for row in 0..domain.size() {
let domain_point = domain.at(row);
let (num, denom) = point_vanishing_fraction(sample_batch.point, domain_point);
flat_denominators.push(num);
numerator_terms.push(denom);
let denominator = pair_vanishing(
sample_batch.point,
sample_batch.point.complex_conjugate(),
domain_point.into_ef(),
);
flat_denominators.push(denominator);
}
}

let mut flat_denominator_inverses = vec![SecureField::zero(); flat_denominators.len()];
SecureField::batch_inverse(&flat_denominators, &mut flat_denominator_inverses);
flat_denominator_inverses
.iter_mut()
.zip(&numerator_terms)
.for_each(|(inv, num_term)| *inv *= *num_term);

flat_denominator_inverses
.chunks_mut(domain.size())
Expand All @@ -110,16 +144,21 @@ pub fn quotient_constants(
random_coeff: SecureField,
domain: CircleDomain,
) -> QuotientConstants<CpuBackend> {
let line_coeffs = column_line_coeffs(sample_batches, random_coeff);
let batch_random_coeffs = batch_random_coeffs(sample_batches, random_coeff);
let denominator_inverses = denominator_inverses(sample_batches, domain);
QuotientConstants {
line_coeffs,
batch_random_coeffs,
denominator_inverses,
}
}

/// Holds the precomputed constant values used in each quotient evaluation.
pub struct QuotientConstants<B: Backend> {
/// The line coefficients for each quotient numerator term. For more details see
/// [self::column_line_coeffs].
pub line_coeffs: Vec<Vec<(SecureField, SecureField, SecureField)>>,
/// The random coefficients used to linearly combine the batched quotients For more details see
/// [self::batch_random_coeffs].
pub batch_random_coeffs: Vec<SecureField>,
Expand Down
121 changes: 61 additions & 60 deletions crates/prover/src/core/backend/simd/quotients.rs
Original file line number Diff line number Diff line change
@@ -1,19 +1,19 @@
use std::iter::zip;

use itertools::izip;
use num_traits::{One, Zero};
use itertools::{izip, zip_eq, Itertools};
use num_traits::Zero;

use super::column::SecureFieldVec;
use super::m31::{PackedBaseField, LOG_N_LANES, N_LANES};
use super::qm31::PackedSecureField;
use super::SimdBackend;
use crate::core::backend::cpu::quotients::{batch_random_coeffs, QuotientConstants};
use crate::core::backend::cpu::quotients::{
batch_random_coeffs, column_line_coeffs, QuotientConstants,
};
use crate::core::backend::{Col, Column};
use crate::core::circle::CirclePoint;
use crate::core::fields::m31::BaseField;
use crate::core::fields::qm31::SecureField;
use crate::core::fields::secure_column::SecureColumn;
use crate::core::fields::FieldOps;
use crate::core::fields::{ComplexConjugate, FieldOps};
use crate::core::pcs::quotients::{ColumnSampleBatch, QuotientOps};
use crate::core::poly::circle::{CircleDomain, CircleEvaluation, SecureEvaluation};
use crate::core::poly::BitReversedOrder;
Expand Down Expand Up @@ -59,19 +59,30 @@ pub fn accumulate_row_quotients(
columns: &[&CircleEvaluation<SimdBackend, BaseField, BitReversedOrder>],
quotient_constants: &QuotientConstants<SimdBackend>,
vec_row: usize,
_domain_point_vec: (PackedBaseField, PackedBaseField),
domain_point_vec: (PackedBaseField, PackedBaseField),
) -> PackedSecureField {
let mut row_accumulator = PackedSecureField::zero();
for (sample_batch, batch_coeff, denominator_inverses) in izip!(
for (sample_batch, line_coeffs, batch_coeff, denominator_inverses) in izip!(
sample_batches,
&quotient_constants.line_coeffs,
&quotient_constants.batch_random_coeffs,
&quotient_constants.denominator_inverses
) {
let mut numerator = PackedSecureField::zero();
for (column_index, sampled_value) in sample_batch.columns_and_values.iter() {
for ((column_index, _), (a, b, c)) in zip_eq(&sample_batch.columns_and_values, line_coeffs)
{
let column = &columns[*column_index];
let value = column.data[vec_row];
numerator += PackedSecureField::broadcast(-*sampled_value) + value;
let value = PackedSecureField::broadcast(*c) * column.data[vec_row];
// The numerator is a line equation passing through
// (sample_point.y, sample_value), (conj(sample_point), conj(sample_value))
// evaluated at (domain_point.y, value).
// When substituting a polynomial in this line equation, we get a polynomial with a root
// at sample_point and conj(sample_point) if the original polynomial had the values
// sample_value and conj(sample_value) at these points.
// TODO(AlonH): Use single point vanishing to save a multiplication.
let linear_term = PackedSecureField::broadcast(*a) * domain_point_vec.1
+ PackedSecureField::broadcast(*b);
numerator += value - linear_term;
}

row_accumulator = row_accumulator * PackedSecureField::broadcast(*batch_coeff)
Expand All @@ -80,65 +91,53 @@ pub fn accumulate_row_quotients(
row_accumulator
}

/// Point vanishing for the packed representation of the points. skips the division.
/// See [crate::core::constraints::point_vanishing_fraction] for more details.
fn packed_point_vanishing_fraction(
excluded: CirclePoint<SecureField>,
p: (PackedBaseField, PackedBaseField),
) -> (PackedSecureField, PackedSecureField) {
let e_conjugate = excluded.conjugate();
let h_x = PackedSecureField::broadcast(e_conjugate.x) * p.0
- PackedSecureField::broadcast(e_conjugate.y) * p.1;
let h_y = PackedSecureField::broadcast(e_conjugate.y) * p.0
+ PackedSecureField::broadcast(e_conjugate.x) * p.1;
(h_y, PackedSecureField::one() + h_x)
/// Pair vanishing for the packed representation of the points. See
/// [crate::core::constraints::pair_vanishing] for more details.
fn packed_pair_vanishing(
excluded0: CirclePoint<SecureField>,
excluded1: CirclePoint<SecureField>,
packed_p: (PackedBaseField, PackedBaseField),
) -> PackedSecureField {
PackedSecureField::broadcast(excluded0.y - excluded1.y) * packed_p.0
+ PackedSecureField::broadcast(excluded1.x - excluded0.x) * packed_p.1
+ PackedSecureField::broadcast(excluded0.x * excluded1.y - excluded0.y * excluded1.x)
}

fn denominator_inverses(
sample_batches: &[ColumnSampleBatch],
domain: CircleDomain,
) -> Vec<Col<SimdBackend, SecureField>> {
let mut numerators = Vec::new();
let mut denominators = Vec::new();

for sample_batch in sample_batches {
for vec_row in 0..1 << (domain.log_size() - LOG_N_LANES) {
// TODO(spapini): Optimize this, for the small number of columns case.
let points = std::array::from_fn(|i| {
domain.at(bit_reverse_index(
(vec_row << LOG_N_LANES) + i,
domain.log_size(),
))
});
let domain_points_x = PackedBaseField::from_array(points.map(|p| p.x));
let domain_points_y = PackedBaseField::from_array(points.map(|p| p.y));
let domain_point_vec = (domain_points_x, domain_points_y);
let (denominator, numerator) =
packed_point_vanishing_fraction(sample_batch.point, domain_point_vec);
denominators.push(denominator);
numerators.push(numerator);
}
}

let denominators = SecureFieldVec {
length: denominators.len() * N_LANES,
data: denominators,
};

let numerators = SecureFieldVec {
length: numerators.len() * N_LANES,
data: numerators,
};

let mut flat_denominator_inverses = SecureFieldVec::zeros(denominators.len());
let flat_denominators: SecureFieldVec = sample_batches
.iter()
.flat_map(|sample_batch| {
(0..(1 << (domain.log_size() - LOG_N_LANES)))
.map(|vec_row| {
// TODO(spapini): Optimize this, for the small number of columns case.
let points = std::array::from_fn(|i| {
domain.at(bit_reverse_index(
(vec_row << LOG_N_LANES) + i,
domain.log_size(),
))
});
let domain_points_x = PackedBaseField::from_array(points.map(|p| p.x));
let domain_points_y = PackedBaseField::from_array(points.map(|p| p.y));
let domain_point_vec = (domain_points_x, domain_points_y);
packed_pair_vanishing(
sample_batch.point,
sample_batch.point.complex_conjugate(),
domain_point_vec,
)
})
.collect_vec()
})
.collect();

let mut flat_denominator_inverses = SecureFieldVec::zeros(flat_denominators.len());
<SimdBackend as FieldOps<SecureField>>::batch_inverse(
&denominators,
&flat_denominators,
&mut flat_denominator_inverses,
);

zip(&mut flat_denominator_inverses.data, &numerators.data)
.for_each(|(inv, denom_denom)| *inv *= *denom_denom);

flat_denominator_inverses
.data
.chunks(domain.size() / N_LANES)
Expand All @@ -151,9 +150,11 @@ fn quotient_constants(
random_coeff: SecureField,
domain: CircleDomain,
) -> QuotientConstants<SimdBackend> {
let line_coeffs = column_line_coeffs(sample_batches, random_coeff);
let batch_random_coeffs = batch_random_coeffs(sample_batches, random_coeff);
let denominator_inverses = denominator_inverses(sample_batches, domain);
QuotientConstants {
line_coeffs,
batch_random_coeffs,
denominator_inverses,
}
Expand Down
Loading

0 comments on commit 60027bd

Please sign in to comment.