Skip to content

perf: base-field constraint evaluation for F×E accumulation #510

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
diegokingston wants to merge 2 commits into from
Closed

perf: base-field constraint evaluation for F×E accumulation #510

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
06e4522
perf: base-field constraint evaluation for F×E accumulation
diegokingston Apr 18, 2026
85845a3
perf: monolithic constraint evaluation (single function per table)
diegokingston Apr 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
145 changes: 112 additions & 33 deletions crypto/stark/src/constraints/evaluator.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,7 @@ where
logup_table_offset: &FieldElement<FieldExtension>,
) -> Vec<FieldElement<FieldExtension>> {
let is_uniform = zerofier_data.is_uniform();
let num_base = air.num_base_transition_constraints();

// Pre-compute LogUp alpha powers once for all LDE domain points.
let logup_alpha_powers: Vec<FieldElement<FieldExtension>> =
Expand Down Expand Up @@ -93,9 +94,10 @@ where
num_main_cols,
num_aux_cols,
),
vec![FieldElement::<Field>::zero(); num_base],
)
},
|(transition_buf, periodic_buf, frame), (i, boundary)| {
|(transition_buf, periodic_buf, frame, base_buf), (i, boundary)| {
frame.fill_from_lde(lde_trace, i, offsets);

for (j, col) in lde_periodic_columns.iter().enumerate() {
Expand All @@ -110,24 +112,66 @@ where
logup_table_offset,
&packing_shifts,
);
air.compute_transition_into(&ctx, transition_buf);

let acc_transition = if is_uniform {
// All constraints share one zerofier: factor it out of the sum.
let z = zerofier_data.get_uniform(i);
let sum = transition_buf
.iter()
.zip(transition_coefficients)
.fold(FieldElement::zero(), |acc, (eval, beta)| acc + eval * beta);
z * &sum
let acc_transition = if num_base > 0 {

@claude claude Bot Apr 18, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[Medium] ~90 lines of logic duplicated between parallel and non-parallel paths

The num_base > 0 branch (both is_uniform and non-uniform sub-cases) is copy-pasted nearly verbatim in the #[cfg(not(feature = "parallel"))] path below. This PR doubled the duplication surface compared to the original. Any future correctness fix will need to be applied in two places.

Consider extracting the inner accumulation into a shared free function, e.g.:

fn accumulate_transition<F, E>(
    base_buf: &[FieldElement<F>],
    transition_buf: &[FieldElement<E>],
    num_base: usize,
    transition_coefficients: &[FieldElement<E>],
    zerofier_data: &ZerofierData,
    i: usize,
    is_uniform: bool,
) -> FieldElement<E> { ... }

Both map_init and the sequential loop can then call it with no duplication.

air.compute_transition_prover(&ctx, base_buf, transition_buf);

if is_uniform {
let z = zerofier_data.get_uniform(i);
// F×E accumulation for base constraints (cheaper)
let mut sum = base_buf
.iter()
.zip(&transition_coefficients[..num_base])
.fold(FieldElement::zero(), |acc, (eval, beta)| {
acc + eval * beta
});
// E×E accumulation for extension constraints
sum = transition_buf[num_base..]
.iter()
.zip(&transition_coefficients[num_base..])
.fold(sum, |acc, (eval, beta)| acc + eval * beta);
z * &sum
} else {
// F×E accumulation for base constraints with per-constraint zerofiers
let mut sum = base_buf
.iter()
.enumerate()
.zip(&transition_coefficients[..num_base])
.fold(FieldElement::zero(), |acc, ((c_idx, eval), beta)| {
acc + zerofier_data.get(c_idx, i) * eval * beta
});
// E×E accumulation for extension constraints
sum = transition_buf[num_base..]
.iter()
.enumerate()
.zip(&transition_coefficients[num_base..])
.fold(sum, |acc, ((j, eval), beta)| {
acc + zerofier_data.get(num_base + j, i) * eval * beta
});
sum
}
} else {
transition_buf
.iter()
.enumerate()
.zip(transition_coefficients)
.fold(FieldElement::zero(), |acc, ((c_idx, eval), beta)| {
acc + zerofier_data.get(c_idx, i) * eval * beta
})
// Old path (no base constraints)
air.compute_transition_into(&ctx, transition_buf);

if is_uniform {
let z = zerofier_data.get_uniform(i);
let sum = transition_buf
.iter()
.zip(transition_coefficients)
.fold(FieldElement::zero(), |acc, (eval, beta)| {
acc + eval * beta
});
z * &sum
} else {
transition_buf
.iter()
.enumerate()
.zip(transition_coefficients)
.fold(FieldElement::zero(), |acc, ((c_idx, eval), beta)| {
acc + zerofier_data.get(c_idx, i) * eval * beta
})
}
};

acc_transition + boundary
Expand All @@ -143,6 +187,7 @@ where
let mut periodic_buf = vec![FieldElement::<Field>::zero(); num_periodic];
let mut frame =
Frame::preallocate(num_offsets, rows_per_step, num_main_cols, num_aux_cols);
let mut base_buf = vec![FieldElement::<Field>::zero(); num_base];

boundary_evaluation
.into_iter()
Expand All @@ -162,23 +207,57 @@ where
logup_table_offset,
&packing_shifts,
);
air.compute_transition_into(&ctx, &mut transition_buf);

let acc_transition = if is_uniform {
let z = zerofier_data.get_uniform(i);
let sum = transition_buf
.iter()
.zip(transition_coefficients)
.fold(FieldElement::zero(), |acc, (eval, beta)| acc + eval * beta);
z * &sum

let acc_transition = if num_base > 0 {
air.compute_transition_prover(&ctx, &mut base_buf, &mut transition_buf);

if is_uniform {
let z = zerofier_data.get_uniform(i);
let mut sum = base_buf
.iter()
.zip(&transition_coefficients[..num_base])
.fold(FieldElement::zero(), |acc, (eval, beta)| acc + eval * beta);
sum = transition_buf[num_base..]
.iter()
.zip(&transition_coefficients[num_base..])
.fold(sum, |acc, (eval, beta)| acc + eval * beta);
z * &sum
} else {
let mut sum = base_buf
.iter()
.enumerate()
.zip(&transition_coefficients[..num_base])
.fold(FieldElement::zero(), |acc, ((c_idx, eval), beta)| {
acc + zerofier_data.get(c_idx, i) * eval * beta
});
sum = transition_buf[num_base..]
.iter()
.enumerate()
.zip(&transition_coefficients[num_base..])
.fold(sum, |acc, ((j, eval), beta)| {
acc + zerofier_data.get(num_base + j, i) * eval * beta
});
sum
}
} else {
transition_buf
.iter()
.enumerate()
.zip(transition_coefficients)
.fold(FieldElement::zero(), |acc, ((c_idx, eval), beta)| {
acc + zerofier_data.get(c_idx, i) * eval * beta
})
air.compute_transition_into(&ctx, &mut transition_buf);

if is_uniform {
let z = zerofier_data.get_uniform(i);
let sum = transition_buf
.iter()
.zip(transition_coefficients)
.fold(FieldElement::zero(), |acc, (eval, beta)| acc + eval * beta);
z * &sum
} else {
transition_buf
.iter()
.enumerate()
.zip(transition_coefficients)
.fold(FieldElement::zero(), |acc, ((c_idx, eval), beta)| {
acc + zerofier_data.get(c_idx, i) * eval * beta
})
}
};

acc_transition + boundary
Expand Down
16 changes: 16 additions & 0 deletions crypto/stark/src/constraints/transition.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,22 @@ where
transition_evaluations: &mut [FieldElement<E>],
);

/// Prover-only: evaluate into a base-field buffer for F×E accumulation.
///
/// Override this for constraints that can produce base-field results (all main
/// trace constraints). The default delegates to `evaluate()` into `ext_evaluations`,
/// used by LogUp constraints that inherently operate in the extension field.
fn evaluate_prover(
&self,
evaluation_context: &TransitionEvaluationContext<F, E>,
base_evaluations: &mut [FieldElement<F>],
ext_evaluations: &mut [FieldElement<E>],
) {
// Default: write to extension buffer (for LogUp constraints)
let _ = base_evaluations;
self.evaluate(evaluation_context, ext_evaluations);
}

/// The periodicity the constraint is applied over the trace.
///
/// Default value is 1, meaning that the constraint is applied to every
Expand Down
46 changes: 46 additions & 0 deletions crypto/stark/src/lookup.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -804,6 +804,14 @@ pub struct AirWithBuses<
/// Maximum number of bus elements across all interactions.
/// Used to compute the correct number of alpha powers.
max_bus_elements: usize,
/// Number of table-specific (base-field) transition constraints.
/// LogUp constraints are extension-field and come after these.
num_base_constraints: usize,
/// Optional monolithic evaluator for all base-field constraints.
/// When set, evaluates ALL base constraints in ONE function call instead of N vtable calls.
#[allow(clippy::type_complexity)]
monolithic_eval_fn:
Option<Box<dyn Fn(&crate::table::TableView<F, E>, &mut [FieldElement<F>]) + Send + Sync>>,
}

impl<
Expand Down Expand Up @@ -832,6 +840,7 @@ impl<
mut transition_constraints: Vec<Box<dyn TransitionConstraint<F, E>>>,
) -> Self {
let num_interactions = auxiliary_trace_build_data.interactions.len();
let num_base_constraints = transition_constraints.len();

@claude claude Bot Apr 18, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[Medium] Ordering invariant is implicit and unguarded

num_base_constraints is captured here, before LogUp constraints are appended. The correctness of the whole F×E optimization depends on the invariant that all base constraints occupy indices 0..num_base and all LogUp constraints occupy num_base..total. This ordering is never enforced.

If a refactor ever inserts LogUp constraints before the table-specific ones (e.g. a prepend instead of an append), num_base_constraints would silently count LogUp constraints as base-field and the F×E path would produce wrong evaluations.

Consider adding a debug assertion after LogUp constraints are pushed — verifying that every constraint in 0..num_base is not a LogUp type — or at minimum a comment cross-referencing where the append happens.


// Split interactions: committed pairs get term columns, last 1-2 are absorbed
let (num_committed_pairs, absorbed_count) = split_interactions(num_interactions);
Expand Down Expand Up @@ -896,6 +905,8 @@ impl<
num_precomputed_cols: None,
name: None,
max_bus_elements,
num_base_constraints,
monolithic_eval_fn: None,
}
}

Expand Down Expand Up @@ -933,6 +944,21 @@ impl<
self.name = Some(name.to_string());
self
}

/// Set a monolithic evaluator for all base-field constraints.
///
/// The closure evaluates ALL base-field constraints in a single function call,
/// avoiding N vtable calls through `Box<dyn TransitionConstraint>::evaluate_prover()`.
/// The compiler can inline and optimize across all constraint expressions.
///
/// LogUp constraints (extension-field) are still dispatched individually.
pub fn with_monolithic_eval<CB>(mut self, f: CB) -> Self
where
CB: Fn(&crate::table::TableView<F, E>, &mut [FieldElement<F>]) + Send + Sync + 'static,
{
self.monolithic_eval_fn = Some(Box::new(f));
self
}
}

impl<F, E, B, PI> crate::traits::AIR for AirWithBuses<F, E, B, PI>
Expand Down Expand Up @@ -990,6 +1016,26 @@ where
&self.context
}

fn num_base_transition_constraints(&self) -> usize {
// All table-specific constraints are base-field.
// LogUp constraints (batched term + accumulated) are extension-field
// and are appended after the table-specific ones.
self.num_base_constraints
}

fn eval_all_base_constraints(
&self,
step: &crate::table::TableView<Self::Field, Self::FieldExtension>,
evals: &mut [FieldElement<Self::Field>],
) -> bool {
if let Some(ref f) = self.monolithic_eval_fn {
f(step, evals);
true
} else {
false
}
}

fn transition_constraints(
&self,
) -> &Vec<Box<dyn TransitionConstraint<Self::Field, Self::FieldExtension>>> {
Expand Down
73 changes: 73 additions & 0 deletions crypto/stark/src/traits.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -282,6 +282,64 @@ pub trait AIR: Send + Sync {
self.context().num_transition_constraints
}

/// Number of transition constraints that produce base-field evaluations.
///
/// Constraints `0..num_base` write to the base-field buffer via `evaluate_prover`.
/// Constraints `num_base..num_total` write to the extension buffer.
/// Default: 0 (all E×E, backward compatible).
fn num_base_transition_constraints(&self) -> usize {
0
}

/// Evaluate all transition constraints into split base/extension buffers.
///
/// Base-field constraints write to `base_evaluations` (length = `num_base_transition_constraints()`).
/// Extension-field constraints write to `ext_evaluations` (length = `num_transition_constraints()`).
///
/// When a monolithic evaluator is available (`eval_all_base_constraints` returns true),
/// all base-field constraints are evaluated in ONE function call instead of N vtable calls.
/// LogUp constraints (extension-field) always use per-constraint dispatch.
fn compute_transition_prover(
&self,
evaluation_context: &TransitionEvaluationContext<Self::Field, Self::FieldExtension>,
base_evaluations: &mut [FieldElement<Self::Field>],
ext_evaluations: &mut [FieldElement<Self::FieldExtension>],
) {
let num_base = base_evaluations.len();
for e in base_evaluations.iter_mut() {
*e = FieldElement::zero();
}
for e in ext_evaluations[num_base..].iter_mut() {

@claude claude Bot Apr 18, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[High] Silent correctness gap: ext_evaluations[0..num_base] is never zeroed

The docstring says ext_evaluations has length num_transition_constraints(), so slots 0..num_base exist in the buffer but are skipped by the zeroing loop here. The accumulator in evaluator.rs only reads transition_buf[num_base..], so there's no bug today — but this relies on an implicit invariant that base-field constraints only ever write to base_evaluations, never to ext_evaluations[0..num_base].

If a constraint's evaluate_prover override is ever mis-implemented (or the default fallback is called unexpectedly) and writes into ext_evaluations[c_idx] with c_idx < num_base, stale data from a previous iteration would be accumulated silently, producing an incorrect proof with no panic or error.

Consider either zeroing the full ext_evaluations slice here, or adding a debug_assert in the accumulator that transition_buf[0..num_base] are all zero after each compute_transition_prover call.

*e = FieldElement::zero();
}

// Try monolithic evaluation for base constraints (ONE call for all)
let step = match evaluation_context {
TransitionEvaluationContext::Prover { frame, .. } => Some(frame.get_evaluation_step(0)),
_ => None,
};

let used_monolithic = if let Some(step) = step {
self.eval_all_base_constraints(step, base_evaluations)
} else {
false
};

if used_monolithic {
// Monolithic handled base constraints. Only dispatch LogUp (extension):
for c in self.transition_constraints().iter() {
if c.constraint_idx() >= num_base {
c.evaluate_prover(evaluation_context, base_evaluations, ext_evaluations);
}
}
} else {
// Fallback: per-constraint dispatch for ALL constraints
self.transition_constraints().iter().for_each(|c| {
c.evaluate_prover(evaluation_context, base_evaluations, ext_evaluations)
});
}
}

fn get_periodic_column_values(&self) -> Vec<Vec<FieldElement<Self::Field>>> {
vec![]
}
Expand All @@ -306,6 +364,21 @@ pub trait AIR: Send + Sync {
result
}

/// Evaluate ALL base-field transition constraints in a single function call.
///
/// Returns `true` if this AIR has a monolithic evaluator (overrides individual dispatch).
/// When `true`, the evaluator calls this instead of per-constraint `evaluate_prover()`.
/// LogUp constraints are still dispatched individually (they're in extension field).
///
/// Default: `false` (no monolithic evaluator, use per-constraint dispatch).
fn eval_all_base_constraints(
&self,
_step: &crate::table::TableView<Self::Field, Self::FieldExtension>,
_evals: &mut [FieldElement<Self::Field>],
) -> bool {
false
}

fn transition_constraints(
&self,
) -> &Vec<Box<dyn TransitionConstraint<Self::Field, Self::FieldExtension>>>;
Expand Down
Loading
Loading