Skip to content

Commit

Permalink
cleanup
Browse files Browse the repository at this point in the history
  • Loading branch information
@plafer
plafer committed Mar 19, 2024
1 parent 35c9413 commit 8cfa8ad
Show file tree
Hide file tree
Showing 2 changed files with 23 additions and 140 deletions.
2 changes: 1 addition & 1 deletion prover/src/constraints/evaluator/default.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -322,7 +322,7 @@ where
);
}

evaluator.evaluate_lagrange_kernel_constraints_2::<A, T>(
evaluator.evaluate_lagrange_kernel_constraints::<A, T>(
trace,
lagrange_kernel_aux_column_idx,
domain,
Expand Down
161 changes: 22 additions & 139 deletions prover/src/constraints/evaluator/lagrange.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,13 @@ impl<E: FieldElement> LagrangeKernelConstraintsBatchEvaluator<E> {
}
}

pub fn evaluate_lagrange_kernel_constraints_2<A, T>(
/// Evaluates the transition and boundary constraints. Specifically, the constraint evaluations
/// are divided by their corresponding divisors, and the resulting terms are linearly combined
/// using the composition coefficients.
///
/// Returns a buffer with the same length as the CE domain, where each element contains the
/// constraint evaluations (as explained above) at the corresponding domain point.
pub fn evaluate_lagrange_kernel_constraints<A, T>(
&self,
trace: &T,
lagrange_kernel_column_idx: usize,
Expand Down Expand Up @@ -107,146 +113,23 @@ impl<E: FieldElement> LagrangeKernelConstraintsBatchEvaluator<E> {

combined_evaluations_acc
}

/// Evaluates the transition and boundary constraints. Specifically, the constraint evaluations
/// are divided by their corresponding divisors, and the resulting terms are linearly combined
/// using the composition coefficients.
///
/// Returns a buffer with the same length as the CE domain, where each element contains the
/// constraint evaluations (as explained above) at the corresponding domain point.
pub fn evaluate_lagrange_kernel_constraints<A>(
&self,
num_trans_constraints: usize,
lagrange_kernel_column_frames: Vec<LagrangeKernelEvaluationFrame<E>>,
domain: &StarkDomain<A::BaseField>,
) -> Vec<E>
where
A: Air,
E: FieldElement<BaseField = A::BaseField>,
{
let transition_constraint_combined_evaluations = self
.evaluate_combined_transition_constraints::<A>(
num_trans_constraints,
&lagrange_kernel_column_frames,
domain,
);

let boundary_constraint_combined_evaluations: Vec<E> = self
.evaluate_combined_boundary_constraints::<A>(&lagrange_kernel_column_frames, domain);

// combine boundary and transition constraint combined evaluations
transition_constraint_combined_evaluations
.into_iter()
.zip(boundary_constraint_combined_evaluations)
.map(|(transitions_combined, boundaries_combined)| {
transitions_combined + boundaries_combined
})
.collect()
}

// HELPERS
// ---------------------------------------------------------------------------------------------

/// Evaluate the transition constraints where the divisors' evaluations are batched to reduce
/// the number of divisions performed (which has a big effect on performance).
///
/// This algorithm takes advantage of some structure in the divisors' evaluations. Recall that
/// the divisor for the i'th transition constraint is `x^(2^i) - 1`. When substituting `x` for
/// each value of the constraint evaluation domain, for constraints `i>0`, the divisor
/// evaluations "wrap-around" such that some values repeat. For example,
/// i=0: no repetitions
/// i=1: the first half of the buffer is equal to the second half
/// i=2: each 1/4th of the buffer are equal
/// i=3: each 1/8th of the buffer are equal
/// ...
/// Therefore, we only compute the non-repeating section of the buffer in each iteration, and index
/// into it accordingly.
fn evaluate_combined_transition_constraints<A>(
&self,
num_trans_constraints: usize,
lagrange_kernel_column_frames: &[LagrangeKernelEvaluationFrame<E>],
domain: &StarkDomain<A::BaseField>,
) -> Vec<E>
where
A: Air,
E: FieldElement<BaseField = A::BaseField>,
{
let mut combined_evaluations_acc = E::zeroed_vector(domain.ce_domain_size());
let mut denominators: Vec<E> = Vec::with_capacity(domain.ce_domain_size());

for trans_constraint_idx in 0..num_trans_constraints {
let num_non_repeating_denoms =
domain.ce_domain_size() / 2_usize.pow(trans_constraint_idx as u32);

for step in 0..num_non_repeating_denoms {
let domain_point = domain.get_ce_x_at(step);
let denominator = self
.lagrange_kernel_constraints
.transition
.evaluate_ith_divisor(trans_constraint_idx, domain_point);
denominators.push(denominator);
}
let denominators_inv = batch_inversion(&denominators);

for step in 0..domain.ce_domain_size() {
let numerator = self.lagrange_kernel_constraints.transition.evaluate_ith_numerator(
&lagrange_kernel_column_frames[step],
&self.rand_elements,
trans_constraint_idx,
);
combined_evaluations_acc[step] +=
numerator * denominators_inv[step % denominators_inv.len()];
}

denominators.truncate(0);
}

combined_evaluations_acc
}

/// Evaluate the boundary constraint where the divisors' evaluations are batched to reduce the
/// number of divisions performed (which has a big effect on performance).
fn evaluate_combined_boundary_constraints<A>(
&self,
lagrange_kernel_column_frames: &[LagrangeKernelEvaluationFrame<E>],
domain: &StarkDomain<A::BaseField>,
) -> Vec<E>
where
A: Air,
E: FieldElement<BaseField = A::BaseField>,
{
let mut boundary_numerator_evals = Vec::with_capacity(domain.ce_domain_size());
let mut boundary_denominator_evals = Vec::with_capacity(domain.ce_domain_size());

for (step, frame) in lagrange_kernel_column_frames.iter().enumerate() {
let domain_point = domain.get_ce_x_at(step);

{
let boundary_numerator =
self.lagrange_kernel_constraints.boundary.evaluate_numerator_at(frame);
boundary_numerator_evals.push(boundary_numerator);
}

{
let boundary_denominator = self
.lagrange_kernel_constraints
.boundary
.evaluate_denominator_at(domain_point.into());
boundary_denominator_evals.push(boundary_denominator);
}
}

let boundary_denominators_inv = batch_inversion(&boundary_denominator_evals);

boundary_numerator_evals
.into_iter()
.zip(boundary_denominators_inv)
.map(|(numerator, denom_inv)| numerator * denom_inv)
.collect()
}
}

/// Holds all the transition constraint inverse divisor evaluations over the constraint evaluation domain.
/// Holds all the transition constraint inverse divisor evaluations over the constraint evaluation
/// domain.
///
/// [`LagrangeKernelTransitionConstraintsDivisor`] takes advantage of some structure in the
/// divisors' evaluations. Recall that the divisor for the i'th transition constraint is `x^(2^i) -
/// 1`. When substituting `x` for each value of the constraint evaluation domain, for constraints
/// `i>0`, the divisor evaluations "wrap-around" such that some values repeat. For example,
///
/// i=0: no repetitions
/// i=1: the first half of the buffer is equal to the second half
/// i=2: each 1/4th of the buffer are equal
/// i=3: each 1/8th of the buffer are equal
/// ...
/// Therefore, we only compute the non-repeating section of the buffer in each iteration, and index
/// into it accordingly.
struct LagrangeKernelTransitionConstraintsDivisor<E: FieldElement> {
divisor_evals_inv: Vec<E>,

Expand Down

0 comments on commit 8cfa8ad

Please sign in to comment.