Skip to content

Commit 87ab5e1

Browse files
committed
feat(prover): add compiled constraint evaluator for CPU table
Eliminates 77 virtual dispatch calls per LDE row for the CPU table by inlining all constraint computations. Uses F×E multiplication (3 base muls) instead of to_extension() + E×E multiplication (6 base muls) for the 66 native constraints that produce base-field results. Changes: - AIR trait: add evaluate_transitions_compiled() optional method - AirWithBuses: add CompiledEvaluator trait + with_compiled_evaluator() - Evaluator: try compiled path before virtual dispatch fallback - lookup: add compute_fingerprint_raw/compute_multiplicity_raw helpers - CpuCompiledEvaluator: inline all 66 native + 11 LogUp constraints
1 parent 9574292 commit 87ab5e1

6 files changed

Lines changed: 929 additions & 3 deletions

File tree

crypto/stark/src/constraints/evaluator.rs

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,33 @@ where
9898
|(transition_buf, periodic_buf, frame), (i, boundary)| {
9999
frame.fill_from_lde(lde_trace, i, offsets);
100100

101+
// Try compiled path: bypasses virtual dispatch + to_extension()
102+
// by computing the weighted sum directly with F×E multiplication.
103+
// Requires at least 2 evaluation steps (current + next row).
104+
if num_offsets >= 2 {
105+
let step0 = frame.get_evaluation_step(0);
106+
let step1 = frame.get_evaluation_step(1);
107+
if let Some(compiled_sum) = air.evaluate_transitions_compiled(
108+
&step0.data[0],
109+
&step1.data[0],
110+
&step0.aux_data[0],
111+
&step1.aux_data[0],
112+
transition_coefficients,
113+
rap_challenges,
114+
&logup_alpha_powers,
115+
logup_table_offset,
116+
) {
117+
// Compiled path already returns the weighted sum.
118+
// Apply zerofier (uniform fast path — CPU has uniform zerofier).
119+
let z = if is_uniform {
120+
zerofier_data.get_uniform(i)
121+
} else {
122+
zerofier_data.get(0, i)
123+
};
124+
return z * &compiled_sum + boundary;
125+
}
126+
}
127+
101128
for (j, col) in lde_periodic_columns.iter().enumerate() {
102129
periodic_buf[j] = col[i].clone();
103130
}
@@ -150,6 +177,30 @@ where
150177
.map(|(i, boundary)| {
151178
frame.fill_from_lde(lde_trace, i, offsets);
152179

180+
// Try compiled path: bypasses virtual dispatch + to_extension()
181+
// Requires at least 2 evaluation steps (current + next row).
182+
if num_offsets >= 2 {
183+
let step0 = frame.get_evaluation_step(0);
184+
let step1 = frame.get_evaluation_step(1);
185+
if let Some(compiled_sum) = air.evaluate_transitions_compiled(
186+
&step0.data[0],
187+
&step1.data[0],
188+
&step0.aux_data[0],
189+
&step1.aux_data[0],
190+
transition_coefficients,
191+
rap_challenges,
192+
&logup_alpha_powers,
193+
logup_table_offset,
194+
) {
195+
let z = if is_uniform {
196+
zerofier_data.get_uniform(i)
197+
} else {
198+
zerofier_data.get(0, i)
199+
};
200+
return z * &compiled_sum + boundary;
201+
}
202+
}
203+
153204
for (j, col) in lde_periodic_columns.iter().enumerate() {
154205
periodic_buf[j] = col[i].clone();
155206
}

crypto/stark/src/lookup.rs

Lines changed: 165 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -778,6 +778,31 @@ impl BusValue {
778778
}
779779
}
780780

781+
// =============================================================================
782+
// Compiled Constraint Evaluator
783+
// =============================================================================
784+
785+
/// Trait for compiled constraint evaluation, bypassing virtual dispatch.
786+
///
787+
/// Implementors compute the weighted sum `sum_i beta_i * constraint_i(row)` directly,
788+
/// using F×E multiplication (3 base muls) instead of `to_extension()` + E×E (6 base muls).
789+
pub trait CompiledEvaluator<F: IsField, E: IsField>: Send + Sync {
790+
/// Compute the weighted sum of ALL constraints (native + LogUp) for one LDE row.
791+
///
792+
/// Must return the same value as the virtual-dispatch path to maintain soundness.
793+
fn evaluate(
794+
&self,
795+
main_curr: &[FieldElement<F>],
796+
main_next: &[FieldElement<F>],
797+
aux_curr: &[FieldElement<E>],
798+
aux_next: &[FieldElement<E>],
799+
transition_coefficients: &[FieldElement<E>],
800+
rap_challenges: &[FieldElement<E>],
801+
logup_alpha_powers: &[FieldElement<E>],
802+
logup_table_offset: &FieldElement<E>,
803+
) -> FieldElement<E>;
804+
}
805+
781806
// =============================================================================
782807
// AirWithBuses
783808
// =============================================================================
@@ -804,6 +829,8 @@ pub struct AirWithBuses<
804829
/// Maximum number of bus elements across all interactions.
805830
/// Used to compute the correct number of alpha powers.
806831
max_bus_elements: usize,
832+
/// Optional compiled evaluator: bypasses virtual dispatch in constraint evaluation hot loop.
833+
compiled_evaluator: Option<Box<dyn CompiledEvaluator<F, E>>>,
807834
}
808835

809836
impl<
@@ -896,6 +923,7 @@ impl<
896923
num_precomputed_cols: None,
897924
name: None,
898925
max_bus_elements,
926+
compiled_evaluator: None,
899927
}
900928
}
901929

@@ -933,6 +961,19 @@ impl<
933961
self.name = Some(name.to_string());
934962
self
935963
}
964+
965+
/// Attach a compiled constraint evaluator that bypasses virtual dispatch.
966+
///
967+
/// When set, the evaluator's hot loop calls this instead of iterating over
968+
/// `transition_constraints` with dynamic dispatch. The compiled evaluator
969+
/// must compute the same weighted sum as the virtual-dispatch path.
970+
pub fn with_compiled_evaluator(
971+
mut self,
972+
evaluator: Box<dyn CompiledEvaluator<F, E>>,
973+
) -> Self {
974+
self.compiled_evaluator = Some(evaluator);
975+
self
976+
}
936977
}
937978

938979
impl<F, E, B, PI> crate::traits::AIR for AirWithBuses<F, E, B, PI>
@@ -996,6 +1037,31 @@ where
9961037
&self.transition_constraints
9971038
}
9981039

1040+
fn evaluate_transitions_compiled(
1041+
&self,
1042+
main_curr: &[FieldElement<Self::Field>],
1043+
main_next: &[FieldElement<Self::Field>],
1044+
aux_curr: &[FieldElement<Self::FieldExtension>],
1045+
aux_next: &[FieldElement<Self::FieldExtension>],
1046+
transition_coefficients: &[FieldElement<Self::FieldExtension>],
1047+
rap_challenges: &[FieldElement<Self::FieldExtension>],
1048+
logup_alpha_powers: &[FieldElement<Self::FieldExtension>],
1049+
logup_table_offset: &FieldElement<Self::FieldExtension>,
1050+
) -> Option<FieldElement<Self::FieldExtension>> {
1051+
self.compiled_evaluator.as_ref().map(|ce| {
1052+
ce.evaluate(
1053+
main_curr,
1054+
main_next,
1055+
aux_curr,
1056+
aux_next,
1057+
transition_coefficients,
1058+
rap_challenges,
1059+
logup_alpha_powers,
1060+
logup_table_offset,
1061+
)
1062+
})
1063+
}
1064+
9991065
fn build_auxiliary_trace(
10001066
&self,
10011067
trace: &mut TraceTable<F, E>,
@@ -1849,6 +1915,105 @@ fn compute_fingerprint_from_step<A: IsSubFieldOf<B>, B: IsField>(
18491915
z - &linear_combination
18501916
}
18511917

1918+
/// Computes multiplicity from a raw main trace slice (compiled evaluator path).
1919+
pub fn compute_multiplicity_raw<F: IsField>(
1920+
main: &[FieldElement<F>],
1921+
multiplicity: &Multiplicity,
1922+
) -> FieldElement<F> {
1923+
match multiplicity {
1924+
Multiplicity::One => FieldElement::<F>::one(),
1925+
Multiplicity::Column(col) => main[*col].clone(),
1926+
Multiplicity::Sum(col_a, col_b) => &main[*col_a] + &main[*col_b],
1927+
Multiplicity::Negated(col) => FieldElement::<F>::one() - &main[*col],
1928+
Multiplicity::Diff(col_a, col_b) => &main[*col_a] - &main[*col_b],
1929+
Multiplicity::Sum3(col_a, col_b, col_c) => &main[*col_a] + &main[*col_b] + &main[*col_c],
1930+
Multiplicity::Linear(terms) => {
1931+
let mut result = FieldElement::<F>::zero();
1932+
for term in terms {
1933+
match term {
1934+
LinearTerm::Column {
1935+
coefficient,
1936+
column,
1937+
} => {
1938+
let coeff = FieldElement::<F>::from(*coefficient);
1939+
result += &main[*column] * coeff;
1940+
}
1941+
LinearTerm::ColumnUnsigned {
1942+
coefficient,
1943+
column,
1944+
} => {
1945+
let coeff = FieldElement::<F>::from(*coefficient);
1946+
result += &main[*column] * coeff;
1947+
}
1948+
LinearTerm::Constant(value) => {
1949+
result += FieldElement::<F>::from(*value);
1950+
}
1951+
}
1952+
}
1953+
result
1954+
}
1955+
}
1956+
}
1957+
1958+
/// Computes fingerprint from a raw main trace slice (compiled evaluator path).
1959+
///
1960+
/// Returns `z - (bus_id*alpha^0 + v[0]*alpha^1 + ...)`
1961+
pub fn compute_fingerprint_raw<F: IsSubFieldOf<E>, E: IsField>(
1962+
main: &[FieldElement<F>],
1963+
interaction: &BusInteraction,
1964+
z: &FieldElement<E>,
1965+
alpha_powers: &[FieldElement<E>],
1966+
shifts: &PackingShifts<F>,
1967+
) -> FieldElement<E> {
1968+
let bus_id_f: FieldElement<F> = FieldElement::from(interaction.bus_id);
1969+
let mut linear_combination = bus_id_f * &alpha_powers[0];
1970+
let mut alpha_idx = 1;
1971+
for bv in &interaction.values {
1972+
match bv {
1973+
BusValue::Packed {
1974+
start_column,
1975+
packing,
1976+
} => {
1977+
alpha_idx += packing.accumulate_fingerprint_with(
1978+
*start_column,
1979+
|col| &main[col],
1980+
alpha_powers,
1981+
alpha_idx,
1982+
&mut linear_combination,
1983+
shifts,
1984+
);
1985+
}
1986+
BusValue::Linear(terms) => {
1987+
let mut result = FieldElement::<F>::zero();
1988+
for term in terms {
1989+
match term {
1990+
LinearTerm::Column {
1991+
coefficient,
1992+
column,
1993+
} => {
1994+
let coeff = FieldElement::<F>::from(*coefficient);
1995+
result += &main[*column] * coeff;
1996+
}
1997+
LinearTerm::ColumnUnsigned {
1998+
coefficient,
1999+
column,
2000+
} => {
2001+
let coeff = FieldElement::<F>::from(*coefficient);
2002+
result += &main[*column] * coeff;
2003+
}
2004+
LinearTerm::Constant(value) => {
2005+
result += FieldElement::<F>::from(*value);
2006+
}
2007+
}
2008+
}
2009+
linear_combination += result * &alpha_powers[alpha_idx];
2010+
alpha_idx += 1;
2011+
}
2012+
}
2013+
}
2014+
z - &linear_combination
2015+
}
2016+
18522017
/// Constraint for a batched pair of interactions sharing one aux column.
18532018
///
18542019
/// Verifies: `c = m_a/fp_a + m_b/fp_b` where signs are baked into m_a, m_b.

crypto/stark/src/traits.rs

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -306,6 +306,38 @@ pub trait AIR: Send + Sync {
306306
result
307307
}
308308

309+
/// Compiled constraint evaluation: computes the weighted sum
310+
/// `sum_i beta_i * constraint_i(row)` directly, bypassing virtual dispatch.
311+
/// Returns the pre-zerofier composition value for one LDE row.
312+
///
313+
/// The key optimization is using F*E multiplication (3 base muls) for
314+
/// base-field constraints instead of the `to_extension()` + E*E mul path
315+
/// (embed + 6 base muls).
316+
///
317+
/// Parameters:
318+
/// - `main_curr`/`main_next`: raw column values at current and next LDE rows
319+
/// - `aux_curr`/`aux_next`: extension-field column values (LogUp)
320+
/// - `transition_coefficients`: random linear combination weights (one per constraint)
321+
/// - `rap_challenges`: LogUp challenges `[z, alpha]`
322+
/// - `logup_alpha_powers`: precomputed powers of alpha for fingerprint computation
323+
/// - `logup_table_offset`: `L/N` per table for circular LogUp constraint
324+
///
325+
/// Returns `Some(weighted_sum)` if compiled path is available, `None` to fall back.
326+
#[allow(unused_variables)]
327+
fn evaluate_transitions_compiled(
328+
&self,
329+
main_curr: &[FieldElement<Self::Field>],
330+
main_next: &[FieldElement<Self::Field>],
331+
aux_curr: &[FieldElement<Self::FieldExtension>],
332+
aux_next: &[FieldElement<Self::FieldExtension>],
333+
transition_coefficients: &[FieldElement<Self::FieldExtension>],
334+
rap_challenges: &[FieldElement<Self::FieldExtension>],
335+
logup_alpha_powers: &[FieldElement<Self::FieldExtension>],
336+
logup_table_offset: &FieldElement<Self::FieldExtension>,
337+
) -> Option<FieldElement<Self::FieldExtension>> {
338+
None
339+
}
340+
309341
fn transition_constraints(
310342
&self,
311343
) -> &Vec<Box<dyn TransitionConstraint<Self::Field, Self::FieldExtension>>>;

0 commit comments

Comments
 (0)