Skip to content

Commit 06e4522

Browse files
diegokingstondiegokingston
committed
perf: base-field constraint evaluation for F×E accumulation
Add evaluate_prover method to TransitionConstraint trait and split the evaluator's accumulation loop into F×E (base) and E×E (extension) parts. Base-field constraints (all table-specific ones) now write to a separate base-field buffer, enabling cheaper F×E multiplication (3 base muls) instead of E×E multiplication (6 base muls) during accumulation. LogUp constraints remain in the extension field via the default implementation. Changes: - TransitionConstraint::evaluate_prover with default delegating to evaluate - AIR::num_base_transition_constraints + compute_transition_prover - AirWithBuses stores and reports num_base_constraints - Evaluator uses split accumulation (base_buf + transition_buf) - All 14 table constraint structs override evaluate_prover
1 parent 64bad0e commit 06e4522

16 files changed

Lines changed: 507 additions & 33 deletions

File tree

‎crypto/stark/src/constraints/evaluator.rs‎

Lines changed: 112 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@ where
5151
logup_table_offset: &FieldElement<FieldExtension>,
5252
) -> Vec<FieldElement<FieldExtension>> {
5353
let is_uniform = zerofier_data.is_uniform();
54+
let num_base = air.num_base_transition_constraints();
5455

5556
// Pre-compute LogUp alpha powers once for all LDE domain points.
5657
let logup_alpha_powers: Vec<FieldElement<FieldExtension>> =
@@ -93,9 +94,10 @@ where
9394
num_main_cols,
9495
num_aux_cols,
9596
),
97+
vec![FieldElement::<Field>::zero(); num_base],
9698
)
9799
},
98-
|(transition_buf, periodic_buf, frame), (i, boundary)| {
100+
|(transition_buf, periodic_buf, frame, base_buf), (i, boundary)| {
99101
frame.fill_from_lde(lde_trace, i, offsets);
100102

101103
for (j, col) in lde_periodic_columns.iter().enumerate() {
@@ -110,24 +112,66 @@ where
110112
logup_table_offset,
111113
&packing_shifts,
112114
);
113-
air.compute_transition_into(&ctx, transition_buf);
114115

115-
let acc_transition = if is_uniform {
116-
// All constraints share one zerofier: factor it out of the sum.
117-
let z = zerofier_data.get_uniform(i);
118-
let sum = transition_buf
119-
.iter()
120-
.zip(transition_coefficients)
121-
.fold(FieldElement::zero(), |acc, (eval, beta)| acc + eval * beta);
122-
z * &sum
116+
let acc_transition = if num_base > 0 {
117+
air.compute_transition_prover(&ctx, base_buf, transition_buf);
118+
119+
if is_uniform {
120+
let z = zerofier_data.get_uniform(i);
121+
// F×E accumulation for base constraints (cheaper)
122+
let mut sum = base_buf
123+
.iter()
124+
.zip(&transition_coefficients[..num_base])
125+
.fold(FieldElement::zero(), |acc, (eval, beta)| {
126+
acc + eval * beta
127+
});
128+
// E×E accumulation for extension constraints
129+
sum = transition_buf[num_base..]
130+
.iter()
131+
.zip(&transition_coefficients[num_base..])
132+
.fold(sum, |acc, (eval, beta)| acc + eval * beta);
133+
z * &sum
134+
} else {
135+
// F×E accumulation for base constraints with per-constraint zerofiers
136+
let mut sum = base_buf
137+
.iter()
138+
.enumerate()
139+
.zip(&transition_coefficients[..num_base])
140+
.fold(FieldElement::zero(), |acc, ((c_idx, eval), beta)| {
141+
acc + zerofier_data.get(c_idx, i) * eval * beta
142+
});
143+
// E×E accumulation for extension constraints
144+
sum = transition_buf[num_base..]
145+
.iter()
146+
.enumerate()
147+
.zip(&transition_coefficients[num_base..])
148+
.fold(sum, |acc, ((j, eval), beta)| {
149+
acc + zerofier_data.get(num_base + j, i) * eval * beta
150+
});
151+
sum
152+
}
123153
} else {
124-
transition_buf
125-
.iter()
126-
.enumerate()
127-
.zip(transition_coefficients)
128-
.fold(FieldElement::zero(), |acc, ((c_idx, eval), beta)| {
129-
acc + zerofier_data.get(c_idx, i) * eval * beta
130-
})
154+
// Old path (no base constraints)
155+
air.compute_transition_into(&ctx, transition_buf);
156+
157+
if is_uniform {
158+
let z = zerofier_data.get_uniform(i);
159+
let sum = transition_buf
160+
.iter()
161+
.zip(transition_coefficients)
162+
.fold(FieldElement::zero(), |acc, (eval, beta)| {
163+
acc + eval * beta
164+
});
165+
z * &sum
166+
} else {
167+
transition_buf
168+
.iter()
169+
.enumerate()
170+
.zip(transition_coefficients)
171+
.fold(FieldElement::zero(), |acc, ((c_idx, eval), beta)| {
172+
acc + zerofier_data.get(c_idx, i) * eval * beta
173+
})
174+
}
131175
};
132176

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

147192
boundary_evaluation
148193
.into_iter()
@@ -162,23 +207,57 @@ where
162207
logup_table_offset,
163208
&packing_shifts,
164209
);
165-
air.compute_transition_into(&ctx, &mut transition_buf);
166-
167-
let acc_transition = if is_uniform {
168-
let z = zerofier_data.get_uniform(i);
169-
let sum = transition_buf
170-
.iter()
171-
.zip(transition_coefficients)
172-
.fold(FieldElement::zero(), |acc, (eval, beta)| acc + eval * beta);
173-
z * &sum
210+
211+
let acc_transition = if num_base > 0 {
212+
air.compute_transition_prover(&ctx, &mut base_buf, &mut transition_buf);
213+
214+
if is_uniform {
215+
let z = zerofier_data.get_uniform(i);
216+
let mut sum = base_buf
217+
.iter()
218+
.zip(&transition_coefficients[..num_base])
219+
.fold(FieldElement::zero(), |acc, (eval, beta)| acc + eval * beta);
220+
sum = transition_buf[num_base..]
221+
.iter()
222+
.zip(&transition_coefficients[num_base..])
223+
.fold(sum, |acc, (eval, beta)| acc + eval * beta);
224+
z * &sum
225+
} else {
226+
let mut sum = base_buf
227+
.iter()
228+
.enumerate()
229+
.zip(&transition_coefficients[..num_base])
230+
.fold(FieldElement::zero(), |acc, ((c_idx, eval), beta)| {
231+
acc + zerofier_data.get(c_idx, i) * eval * beta
232+
});
233+
sum = transition_buf[num_base..]
234+
.iter()
235+
.enumerate()
236+
.zip(&transition_coefficients[num_base..])
237+
.fold(sum, |acc, ((j, eval), beta)| {
238+
acc + zerofier_data.get(num_base + j, i) * eval * beta
239+
});
240+
sum
241+
}
174242
} else {
175-
transition_buf
176-
.iter()
177-
.enumerate()
178-
.zip(transition_coefficients)
179-
.fold(FieldElement::zero(), |acc, ((c_idx, eval), beta)| {
180-
acc + zerofier_data.get(c_idx, i) * eval * beta
181-
})
243+
air.compute_transition_into(&ctx, &mut transition_buf);
244+
245+
if is_uniform {
246+
let z = zerofier_data.get_uniform(i);
247+
let sum = transition_buf
248+
.iter()
249+
.zip(transition_coefficients)
250+
.fold(FieldElement::zero(), |acc, (eval, beta)| acc + eval * beta);
251+
z * &sum
252+
} else {
253+
transition_buf
254+
.iter()
255+
.enumerate()
256+
.zip(transition_coefficients)
257+
.fold(FieldElement::zero(), |acc, ((c_idx, eval), beta)| {
258+
acc + zerofier_data.get(c_idx, i) * eval * beta
259+
})
260+
}
182261
};
183262

184263
acc_transition + boundary

‎crypto/stark/src/constraints/transition.rs‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,22 @@ where
3636
transition_evaluations: &mut [FieldElement<E>],
3737
);
3838

39+
/// Prover-only: evaluate into a base-field buffer for F×E accumulation.
40+
///
41+
/// Override this for constraints that can produce base-field results (all main
42+
/// trace constraints). The default delegates to `evaluate()` into `ext_evaluations`,
43+
/// used by LogUp constraints that inherently operate in the extension field.
44+
fn evaluate_prover(
45+
&self,
46+
evaluation_context: &TransitionEvaluationContext<F, E>,
47+
base_evaluations: &mut [FieldElement<F>],
48+
ext_evaluations: &mut [FieldElement<E>],
49+
) {
50+
// Default: write to extension buffer (for LogUp constraints)
51+
let _ = base_evaluations;
52+
self.evaluate(evaluation_context, ext_evaluations);
53+
}
54+
3955
/// The periodicity the constraint is applied over the trace.
4056
///
4157
/// Default value is 1, meaning that the constraint is applied to every

‎crypto/stark/src/lookup.rs‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -804,6 +804,9 @@ pub struct AirWithBuses<
804804
/// Maximum number of bus elements across all interactions.
805805
/// Used to compute the correct number of alpha powers.
806806
max_bus_elements: usize,
807+
/// Number of table-specific (base-field) transition constraints.
808+
/// LogUp constraints are extension-field and come after these.
809+
num_base_constraints: usize,
807810
}
808811

809812
impl<
@@ -832,6 +835,7 @@ impl<
832835
mut transition_constraints: Vec<Box<dyn TransitionConstraint<F, E>>>,
833836
) -> Self {
834837
let num_interactions = auxiliary_trace_build_data.interactions.len();
838+
let num_base_constraints = transition_constraints.len();
835839

836840
// Split interactions: committed pairs get term columns, last 1-2 are absorbed
837841
let (num_committed_pairs, absorbed_count) = split_interactions(num_interactions);
@@ -896,6 +900,7 @@ impl<
896900
num_precomputed_cols: None,
897901
name: None,
898902
max_bus_elements,
903+
num_base_constraints,
899904
}
900905
}
901906

@@ -990,6 +995,13 @@ where
990995
&self.context
991996
}
992997

998+
fn num_base_transition_constraints(&self) -> usize {
999+
// All table-specific constraints are base-field.
1000+
// LogUp constraints (batched term + accumulated) are extension-field
1001+
// and are appended after the table-specific ones.
1002+
self.num_base_constraints
1003+
}
1004+
9931005
fn transition_constraints(
9941006
&self,
9951007
) -> &Vec<Box<dyn TransitionConstraint<Self::Field, Self::FieldExtension>>> {

‎crypto/stark/src/traits.rs‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -282,6 +282,37 @@ pub trait AIR: Send + Sync {
282282
self.context().num_transition_constraints
283283
}
284284

285+
/// Number of transition constraints that produce base-field evaluations.
286+
///
287+
/// Constraints `0..num_base` write to the base-field buffer via `evaluate_prover`.
288+
/// Constraints `num_base..num_total` write to the extension buffer.
289+
/// Default: 0 (all E×E, backward compatible).
290+
fn num_base_transition_constraints(&self) -> usize {
291+
0
292+
}
293+
294+
/// Evaluate all transition constraints into split base/extension buffers.
295+
///
296+
/// Base-field constraints write to `base_evaluations` (length = `num_base_transition_constraints()`).
297+
/// Extension-field constraints write to `ext_evaluations` (length = `num_transition_constraints()`).
298+
fn compute_transition_prover(
299+
&self,
300+
evaluation_context: &TransitionEvaluationContext<Self::Field, Self::FieldExtension>,
301+
base_evaluations: &mut [FieldElement<Self::Field>],
302+
ext_evaluations: &mut [FieldElement<Self::FieldExtension>],
303+
) {
304+
let num_base = base_evaluations.len();
305+
for e in base_evaluations.iter_mut() {
306+
*e = FieldElement::zero();
307+
}
308+
for e in ext_evaluations[num_base..].iter_mut() {
309+
*e = FieldElement::zero();
310+
}
311+
self.transition_constraints()
312+
.iter()
313+
.for_each(|c| c.evaluate_prover(evaluation_context, base_evaluations, ext_evaluations));
314+
}
315+
285316
fn get_periodic_column_values(&self) -> Vec<Vec<FieldElement<Self::Field>>> {
286317
vec![]
287318
}

0 commit comments

Comments
 (0)