From 9eff1ff2cb16f680e00a581ea3304573adf2141e Mon Sep 17 00:00:00 2001 From: sragss Date: Wed, 26 Jun 2024 12:14:29 -0700 Subject: [PATCH 1/6] init Vec> --- jolt-core/src/r1cs/builder.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/jolt-core/src/r1cs/builder.rs b/jolt-core/src/r1cs/builder.rs index e69cdecc8..e548e358c 100644 --- a/jolt-core/src/r1cs/builder.rs +++ b/jolt-core/src/r1cs/builder.rs @@ -552,7 +552,7 @@ pub struct CombinedUniformBuilder { /// Padded to the nearest power of 2 uniform_repeat: usize, - offset_equality_constraint: OffsetEqConstraint, + offset_equality_constraints: Vec>, } #[tracing::instrument(skip_all, name = "batch_inputs")] @@ -574,13 +574,13 @@ impl CombinedUniformBuilder { pub fn construct( uniform_builder: R1CSBuilder, uniform_repeat: usize, - offset_equality_constraint: OffsetEqConstraint, + offset_equality_constraints: Vec>, ) -> Self { assert!(uniform_repeat.is_power_of_two()); Self { uniform_builder, uniform_repeat, - offset_equality_constraint, + offset_equality_constraints, } } From 9f6a31bfef45b213255a7824d6db3e1e242f7a87 Mon Sep 17 00:00:00 2001 From: sragss Date: Fri, 28 Jun 2024 13:07:56 -0700 Subject: [PATCH 2/6] key stuff --- jolt-core/src/r1cs/key.rs | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/jolt-core/src/r1cs/key.rs b/jolt-core/src/r1cs/key.rs index f4dcfb166..4da3ad438 100644 --- a/jolt-core/src/r1cs/key.rs +++ b/jolt-core/src/r1cs/key.rs @@ -67,16 +67,16 @@ pub struct UniformR1CS { pub num_rows: usize, } -/// NonUniformR1CS only supports a single additional equality constraint. 'a' holds the equality (something minus something), +/// NonUniformR1CSConstraint only supports a single additional equality constraint. 'a' holds the equality (something minus something), /// 'b' holds the condition. 'a' * 'b' == 0. Each SparseEqualityItem stores a uniform_column (pointing to a variable) and an offset /// suggesting which other step to point to. #[derive(CanonicalSerialize, CanonicalDeserialize)] -pub struct NonUniformR1CS { +pub struct NonUniformR1CSConstraint { pub eq: SparseEqualityItem, pub condition: SparseEqualityItem, } -impl NonUniformR1CS { +impl NonUniformR1CSConstraint { pub fn new(eq: SparseEqualityItem, condition: SparseEqualityItem) -> Self { Self { eq, condition } } @@ -89,6 +89,12 @@ impl NonUniformR1CS { } } +/// NonUniformR1CS stores a vector of NonUniformR1CSConstraint +#[derive(CanonicalSerialize, CanonicalDeserialize)] +pub struct NonUniformR1CS { + pub constraints: Vec>, +} + /// Represents a single constraint row where the variables are either from the current step (offset = false) /// or from the proceeding step (offset = true). #[derive(CanonicalSerialize, CanonicalDeserialize, Debug, PartialEq)] From 73962827c67888c8ec17fff796638312b426de96 Mon Sep 17 00:00:00 2001 From: sragss Date: Tue, 16 Jul 2024 15:17:24 -0700 Subject: [PATCH 3/6] init working --- jolt-core/src/r1cs/builder.rs | 233 +++++++++++++------------ jolt-core/src/r1cs/jolt_constraints.rs | 18 +- jolt-core/src/r1cs/key.rs | 202 +++++++++++---------- jolt-core/src/r1cs/test.rs | 7 +- 4 files changed, 243 insertions(+), 217 deletions(-) diff --git a/jolt-core/src/r1cs/builder.rs b/jolt-core/src/r1cs/builder.rs index e548e358c..38985bcff 100644 --- a/jolt-core/src/r1cs/builder.rs +++ b/jolt-core/src/r1cs/builder.rs @@ -16,7 +16,7 @@ use rayon::prelude::*; use std::{collections::HashMap, fmt::Debug}; use super::{ - key::{NonUniformR1CS, SparseEqualityItem}, + key::{NonUniformR1CS, NonUniformR1CSConstraint, SparseEqualityItem}, ops::{ConstraintInput, Term, Variable, LC}, special_polys::SparsePolynomial, }; @@ -684,71 +684,75 @@ impl CombinedUniformBuilder { self.uniform_builder.materialize() } + /// Converts builder::OffsetEqConstraints into key::NonUniformR1CSConstraint pub fn materialize_offset_eq(&self) -> NonUniformR1CS { // (a - b) * condition == 0 // A: a - b // B: condition // C: 0 - let mut eq = SparseEqualityItem::::empty(); - let mut condition = SparseEqualityItem::::empty(); + let mut constraints = Vec::with_capacity(self.offset_equality_constraints.len()); + for constraint in &self.offset_equality_constraints { + let mut eq = SparseEqualityItem::::empty(); + let mut condition = SparseEqualityItem::::empty(); - let constraint = &self.offset_equality_constraint; + constraint + .cond + .1 + .terms() + .iter() + .filter(|term| matches!(term.0, Variable::Input(_) | Variable::Auxiliary(_))) + .for_each(|term| { + condition.offset_vars.push(( + self.uniform_builder.variable_to_column(term.0), + constraint.cond.0, + F::from_i64(term.1), + )) + }); + if let Some(term) = constraint.cond.1.constant_term() { + condition.constant = F::from_i64(term.1); + } - constraint - .cond - .1 - .terms() - .iter() - .filter(|term| matches!(term.0, Variable::Input(_) | Variable::Auxiliary(_))) - .for_each(|term| { - condition.offset_vars.push(( - self.uniform_builder.variable_to_column(term.0), - constraint.cond.0, - F::from_i64(term.1), - )) - }); - if let Some(term) = constraint.cond.1.constant_term() { - condition.constant = F::from_i64(term.1); - } + // Can't simply combine like terms because of the offset + let lhs = constraint.a.1.clone(); + let rhs = -constraint.b.1.clone(); - // Can't simply combine like terms because of the offset - let lhs = constraint.a.1.clone(); - let rhs = -constraint.b.1.clone(); + lhs.terms() + .iter() + .filter(|term| matches!(term.0, Variable::Input(_) | Variable::Auxiliary(_))) + .for_each(|term| { + eq.offset_vars.push(( + self.uniform_builder.variable_to_column(term.0), + constraint.a.0, + F::from_i64(term.1), + )) + }); + rhs.terms() + .iter() + .filter(|term| matches!(term.0, Variable::Input(_) | Variable::Auxiliary(_))) + .for_each(|term| { + eq.offset_vars.push(( + self.uniform_builder.variable_to_column(term.0), + constraint.b.0, + F::from_i64(term.1), + )) + }); - lhs.terms() - .iter() - .filter(|term| matches!(term.0, Variable::Input(_) | Variable::Auxiliary(_))) - .for_each(|term| { - eq.offset_vars.push(( - self.uniform_builder.variable_to_column(term.0), - constraint.a.0, - F::from_i64(term.1), - )) - }); - rhs.terms() - .iter() - .filter(|term| matches!(term.0, Variable::Input(_) | Variable::Auxiliary(_))) - .for_each(|term| { - eq.offset_vars.push(( - self.uniform_builder.variable_to_column(term.0), - constraint.b.0, - F::from_i64(term.1), - )) + // Handle constants + lhs.terms().iter().for_each(|term| { + assert!( + !matches!(term.0, Variable::Constant), + "Constants only supported in RHS" + ) }); + if let Some(term) = rhs.constant_term() { + eq.constant = F::from_i64(term.1); + } - // Handle constants - lhs.terms().iter().for_each(|term| { - assert!( - !matches!(term.0, Variable::Constant), - "Constants only supported in RHS" - ) - }); - if let Some(term) = rhs.constant_term() { - eq.constant = F::from_i64(term.1); + constraints.push(NonUniformR1CSConstraint::new(eq, condition)); } - NonUniformR1CS::new(eq, condition) + NonUniformR1CS { constraints } } /// inputs should be of the format [[I::0, I::0, ...], [I::1, I::1, ...], ... [I::N, I::N]] @@ -829,53 +833,54 @@ impl CombinedUniformBuilder { let span = tracing::span!(tracing::Level::DEBUG, "offset_eq"); let _enter = span.enter(); - let constr = &self.offset_equality_constraint; - let condition_evals = constr - .cond - .1 - .evaluate_batch(&batch_inputs(&constr.cond.1), self.uniform_repeat); - let eq_a_evals = constr - .a - .1 - .evaluate_batch(&batch_inputs(&constr.a.1), self.uniform_repeat); - let eq_b_evals = constr - .b - .1 - .evaluate_batch(&batch_inputs(&constr.b.1), self.uniform_repeat); - - let dense_az_bz: Vec<(F, F)> = (0..self.uniform_repeat) - .into_par_iter() - .map(|step_index| { - // Write corresponding values, if outside the step range, only include the constant. - let a_step = step_index + if constr.a.0 { 1 } else { 0 }; - let b_step = step_index + if constr.b.0 { 1 } else { 0 }; - let a = eq_a_evals - .get(a_step) - .cloned() - .unwrap_or(constr.a.1.constant_term_field()); - let b = eq_b_evals - .get(b_step) - .cloned() - .unwrap_or(constr.b.1.constant_term_field()); - let az = a - b; - - let condition_step = step_index + if constr.cond.0 { 1 } else { 0 }; - let bz = condition_evals - .get(condition_step) - .cloned() - .unwrap_or(constr.cond.1.constant_term_field()); - (az, bz) - }) - .collect(); + for constr in &self.offset_equality_constraints { + let condition_evals = constr + .cond + .1 + .evaluate_batch(&batch_inputs(&constr.cond.1), self.uniform_repeat); + let eq_a_evals = constr + .a + .1 + .evaluate_batch(&batch_inputs(&constr.a.1), self.uniform_repeat); + let eq_b_evals = constr + .b + .1 + .evaluate_batch(&batch_inputs(&constr.b.1), self.uniform_repeat); + + let dense_az_bz: Vec<(F, F)> = (0..self.uniform_repeat) + .into_par_iter() + .map(|step_index| { + // Write corresponding values, if outside the step range, only include the constant. + let a_step = step_index + if constr.a.0 { 1 } else { 0 }; + let b_step = step_index + if constr.b.0 { 1 } else { 0 }; + let a = eq_a_evals + .get(a_step) + .cloned() + .unwrap_or(constr.a.1.constant_term_field()); + let b = eq_b_evals + .get(b_step) + .cloned() + .unwrap_or(constr.b.1.constant_term_field()); + let az = a - b; + + let condition_step = step_index + if constr.cond.0 { 1 } else { 0 }; + let bz = condition_evals + .get(condition_step) + .cloned() + .unwrap_or(constr.cond.1.constant_term_field()); + (az, bz) + }) + .collect(); - // Sparsify: take only the non-zero elements - for (local_index, (az, bz)) in dense_az_bz.iter().enumerate() { - let global_index = uniform_constraint_rows + local_index; - if !az.is_zero() { - az_sparse.push((*az, global_index)); - } - if !bz.is_zero() { - bz_sparse.push((*bz, global_index)); + // Sparsify: take only the non-zero elements + for (local_index, (az, bz)) in dense_az_bz.iter().enumerate() { + let global_index = uniform_constraint_rows + local_index; + if !az.is_zero() { + az_sparse.push((*az, global_index)); + } + if !bz.is_zero() { + bz_sparse.push((*bz, global_index)); + } } } @@ -1304,11 +1309,8 @@ mod tests { assert_eq!(uniform_builder.constraints.len(), 1); assert_eq!(uniform_builder.next_aux, 1); let num_steps = 2; - let combined_builder = CombinedUniformBuilder::construct( - uniform_builder, - num_steps, - OffsetEqConstraint::empty(), - ); + let combined_builder = + CombinedUniformBuilder::construct(uniform_builder, num_steps, vec![]); let mut inputs = vec![vec![Fr::zero(); num_steps]; TestInputs::COUNT]; inputs[TestInputs::OpFlags0 as usize][0] = Fr::from(5); @@ -1348,11 +1350,8 @@ mod tests { assert_eq!(uniform_builder.next_aux, 2); let num_steps = 2; - let combined_builder = CombinedUniformBuilder::construct( - uniform_builder, - num_steps, - OffsetEqConstraint::empty(), - ); + let combined_builder = + CombinedUniformBuilder::construct(uniform_builder, num_steps, vec![]); let mut inputs = vec![vec![Fr::zero(); num_steps]; TestInputs::COUNT]; inputs[TestInputs::OpFlags0 as usize][0] = Fr::from(5); @@ -1409,8 +1408,11 @@ mod tests { (TestInputs::OpFlags0, false), (TestInputs::OpFlags0, true), ); - let combined_builder = - CombinedUniformBuilder::construct(uniform_builder, num_steps, non_uniform_constraint); + let combined_builder = CombinedUniformBuilder::construct( + uniform_builder, + num_steps, + vec![non_uniform_constraint], + ); let mut inputs = vec![vec![Fr::zero(); num_steps]; TestInputs::COUNT]; inputs[TestInputs::OpFlags0 as usize][0] = Fr::from(5); @@ -1453,8 +1455,11 @@ mod tests { (TestInputs::OpFlags0, false), (TestInputs::OpFlags0, true), ); - let combined_builder = - CombinedUniformBuilder::construct(uniform_builder, num_steps, non_uniform_constraint); + let combined_builder = CombinedUniformBuilder::construct( + uniform_builder, + num_steps, + vec![non_uniform_constraint], + ); let offset_eq = combined_builder.materialize_offset_eq(); let mut expected_condition = SparseEqualityItem::::empty(); @@ -1466,8 +1471,8 @@ mod tests { (TestInputs::OpFlags0 as usize, true, Fr::from_i64(-1)), ]; - assert_eq!(offset_eq.condition, expected_condition); - assert_eq!(offset_eq.eq, expected_eq); + assert_eq!(offset_eq.constraints[0].condition, expected_condition); + assert_eq!(offset_eq.constraints[0].eq, expected_eq); } #[test] diff --git a/jolt-core/src/r1cs/jolt_constraints.rs b/jolt-core/src/r1cs/jolt_constraints.rs index 221d9a7c8..53d9504b0 100644 --- a/jolt-core/src/r1cs/jolt_constraints.rs +++ b/jolt-core/src/r1cs/jolt_constraints.rs @@ -22,7 +22,11 @@ pub fn construct_jolt_constraints( (4 * JoltIn::PcIn + PC_START_ADDRESS, true), ); - CombinedUniformBuilder::construct(uniform_builder, padded_trace_length, non_uniform_constraint) + CombinedUniformBuilder::construct( + uniform_builder, + padded_trace_length, + vec![non_uniform_constraint], + ) } // TODO(#377): Dedupe OpFlags / CircuitFlags @@ -282,10 +286,7 @@ impl R1CSConstraintBuilder for UniformJoltConstraints { mod tests { use super::*; - use crate::{ - jolt::vm::rv32i_vm::RV32I, - r1cs::builder::{CombinedUniformBuilder, OffsetEqConstraint}, - }; + use crate::{jolt::vm::rv32i_vm::RV32I, r1cs::builder::CombinedUniformBuilder}; use ark_bn254::Fr; use ark_std::Zero; @@ -307,11 +308,8 @@ mod tests { jolt_constraints.build_constraints(&mut uniform_builder); let num_steps = 1; - let combined_builder = CombinedUniformBuilder::construct( - uniform_builder, - num_steps, - OffsetEqConstraint::empty(), - ); + let combined_builder = + CombinedUniformBuilder::construct(uniform_builder, num_steps, vec![]); let mut inputs = vec![vec![Fr::zero(); num_steps]; JoltIn::COUNT]; // ADD instruction diff --git a/jolt-core/src/r1cs/key.rs b/jolt-core/src/r1cs/key.rs index 4da3ad438..43620b0c9 100644 --- a/jolt-core/src/r1cs/key.rs +++ b/jolt-core/src/r1cs/key.rs @@ -95,6 +95,21 @@ pub struct NonUniformR1CS { pub constraints: Vec>, } +impl NonUniformR1CS { + /// Returns a tuple of (eq_constants, condition_constants) + fn constants(&self) -> (Vec, Vec) { + let mut eq_constants = Vec::with_capacity(self.constraints.len()); + let mut condition_constants = Vec::with_capacity(self.constraints.len()); + + for constraint in &self.constraints { + eq_constants.push(constraint.eq.constant); + condition_constants.push(constraint.condition.constant); + } + + (eq_constants, condition_constants) + } +} + /// Represents a single constraint row where the variables are either from the current step (offset = false) /// or from the proceeding step (offset = true). #[derive(CanonicalSerialize, CanonicalDeserialize, Debug, PartialEq)] @@ -164,39 +179,39 @@ impl UniformSpartanKey { let eq_rx_step = EqPolynomial::evals(r_step); let eq_rx_constr = EqPolynomial::evals(r_constr); - let non_uniform_row = self.uniform_r1cs.num_rows; + let first_non_uniform_row = self.uniform_r1cs.num_rows; + let constant_column = self.uniform_r1cs.num_vars; // Computation strategy // 1. Compute the RLC of the repeated terms in A, B, C, and the constant column // 2. Expand this RLC to the full column y by multiplying by eq(rx_step, step_index) for each step // 3. Add the non uniform constraint rows - let compute_repeated = |constraints: &SparseConstraints, - non_uni_constant: Option| - -> Vec { - // +1 for constant - let mut evals = unsafe_allocate_zero_vec(self.uniform_r1cs.num_vars + 1); - for (row, col, val) in constraints.vars.iter() { - evals[*col] += mul_0_1_optimized(val, &eq_rx_constr[*row]); - } + let compute_repeated = + |constraints: &SparseConstraints, non_uni_constants: Option>| -> Vec { + // +1 for constant + let mut evals = unsafe_allocate_zero_vec(self.uniform_r1cs.num_vars + 1); + for (row, col, val) in constraints.vars.iter() { + evals[*col] += mul_0_1_optimized(val, &eq_rx_constr[*row]); + } - for (row, val) in constraints.consts.iter() { - evals[self.uniform_r1cs.num_vars] += mul_0_1_optimized(val, &eq_rx_constr[*row]); - } + for (row, val) in constraints.consts.iter() { + evals[constant_column] += mul_0_1_optimized(val, &eq_rx_constr[*row]); + } - if let Some(non_uni_constant) = non_uni_constant { - evals[self.uniform_r1cs.num_vars] += - eq_rx_constr[non_uniform_row] * non_uni_constant; - } + if let Some(non_uni_constants) = non_uni_constants { + for (i, non_uni_constant) in non_uni_constants.iter().enumerate() { + evals[constant_column] += + eq_rx_constr[first_non_uniform_row + i] * non_uni_constant; + } + } - evals - }; + evals + }; - let sm_a_r = compute_repeated(&self.uniform_r1cs.a, Some(self.offset_eq_r1cs.eq.constant)); - let sm_b_r = compute_repeated( - &self.uniform_r1cs.b, - Some(self.offset_eq_r1cs.condition.constant), - ); + let (eq_constants, condition_constants) = self.offset_eq_r1cs.constants(); + let sm_a_r = compute_repeated(&self.uniform_r1cs.a, Some(eq_constants)); + let sm_b_r = compute_repeated(&self.uniform_r1cs.b, Some(condition_constants)); let sm_c_r = compute_repeated(&self.uniform_r1cs.c, None); let r_rlc_sq = r_rlc.square(); @@ -227,7 +242,10 @@ impl UniformSpartanKey { rlc[self.num_vars_total()] = sm_rlc[self.uniform_r1cs.num_vars]; // constant // Handle non-uniform constraints - let update_non_uni = |rlc: &mut Vec, offset: &SparseEqualityItem, r: F| { + let update_non_uni = |rlc: &mut Vec, + offset: &SparseEqualityItem, + non_uni_constraint_index: usize, + r: F| { for (col, is_offset, coeff) in offset.offset_vars.iter() { let offset = if *is_offset { 1 } else { 0 }; @@ -241,7 +259,7 @@ impl UniformSpartanKey { .for_each(|(rlc_col, step_index)| { *rlc_col += mul_0_1_optimized(&r, coeff) * eq_rx_step[step_index] - * eq_rx_constr[non_uniform_row]; + * eq_rx_constr[first_non_uniform_row + non_uni_constraint_index]; }); } }; @@ -249,8 +267,10 @@ impl UniformSpartanKey { { let span = tracing::span!(tracing::Level::INFO, "update_non_uniform"); let _guard = span.enter(); - update_non_uni(&mut rlc, &self.offset_eq_r1cs.eq, F::one()); - update_non_uni(&mut rlc, &self.offset_eq_r1cs.condition, r_rlc); + for (i, constraint) in self.offset_eq_r1cs.constraints.iter().enumerate() { + update_non_uni(&mut rlc, &constraint.eq, i, F::one()); + update_non_uni(&mut rlc, &constraint.condition, i, r_rlc); + } } rlc @@ -305,63 +325,69 @@ impl UniformSpartanKey { let constant_column = index_to_field_bitvector(self.num_cols_total() / 2, total_cols_bits); let col_eq_constant = EqPolynomial::new(r_col.to_vec()).evaluate(&constant_column); - let non_uni_constraint_index = - index_to_field_bitvector(self.uniform_r1cs.num_rows, constraint_rows_bits); - let row_constr_eq_non_uni = - EqPolynomial::new(r_row_constr.to_vec()).evaluate(&non_uni_constraint_index); - assert_eq!( - row_constr_eq_non_uni, - eq_rx_constr[self.uniform_r1cs.num_rows] - ); + let compute_uniform_matrix_mle = |constraints: &SparseConstraints| -> F { + let mut full_mle_evaluation: F = constraints + .vars + .iter() + .map(|(row, col, coeff)| *coeff * eq_rx_constr[*row] * eq_ry_var[*col]) + .sum::() + * eq_rx_ry_step; + + full_mle_evaluation += constraints + .consts + .iter() + .map(|(constraint_row, constant_coeff)| { + *constant_coeff * eq_rx_constr[*constraint_row] + }) + .sum::() + * col_eq_constant; + + full_mle_evaluation + }; - let compute_matrix_mle = - |constraints: &SparseConstraints, non_uni: Option<&SparseEqualityItem>| -> F { - let mut full_mle_evaluation: F = constraints - .vars - .iter() - .map(|(row, col, coeff)| *coeff * eq_rx_constr[*row] * eq_ry_var[*col]) - .sum::() - * eq_rx_ry_step; - - full_mle_evaluation += constraints - .consts - .iter() - .map(|(constraint_row, constant_coeff)| { - *constant_coeff * eq_rx_constr[*constraint_row] - }) - .sum::() - * col_eq_constant; - - // Non uniform - let mut non_uni_mle = F::zero(); - if let Some(non_uni) = non_uni { - let eq_step_offset_1 = eq_plus_one(r_row_step, r_col_step, steps_bits); - - non_uni_mle = non_uni - .offset_vars - .iter() - .map(|(col, offset, coeff)| { - if !offset { - *coeff * eq_ry_var[*col] * eq_rx_ry_step - } else { - *coeff * eq_ry_var[*col] * eq_step_offset_1 - } - }) - .sum::(); - - non_uni_mle += non_uni.constant * col_eq_constant; - } + let mut a_mle = compute_uniform_matrix_mle(&self.uniform_r1cs.a); + let mut b_mle = compute_uniform_matrix_mle(&self.uniform_r1cs.b); + let c_mle = compute_uniform_matrix_mle(&self.uniform_r1cs.c); + + // Non-uniform constraints + let eq_step_offset_1 = eq_plus_one(r_row_step, r_col_step, steps_bits); + let compute_non_uniform = |non_uni: &SparseEqualityItem| -> F { + let mut non_uni_mle = non_uni + .offset_vars + .iter() + .map(|(col, offset, coeff)| { + if !offset { + *coeff * eq_ry_var[*col] * eq_rx_ry_step + } else { + *coeff * eq_ry_var[*col] * eq_step_offset_1 + } + }) + .sum::(); - full_mle_evaluation += non_uni_mle * row_constr_eq_non_uni; + non_uni_mle += non_uni.constant * col_eq_constant; - full_mle_evaluation - }; + non_uni_mle + }; - ( - compute_matrix_mle(&self.uniform_r1cs.a, Some(&self.offset_eq_r1cs.eq)), - compute_matrix_mle(&self.uniform_r1cs.b, Some(&self.offset_eq_r1cs.condition)), - compute_matrix_mle(&self.uniform_r1cs.c, None), - ) + for (i, constraint) in self.offset_eq_r1cs.constraints.iter().enumerate() { + let non_uni_a = compute_non_uniform(&constraint.eq); + let non_uni_b = compute_non_uniform(&constraint.condition); + + let non_uni_constraint_index = + index_to_field_bitvector(self.uniform_r1cs.num_rows + i, constraint_rows_bits); + let row_constr_eq_non_uni = + EqPolynomial::new(r_row_constr.to_vec()).evaluate(&non_uni_constraint_index); + + debug_assert_eq!( + row_constr_eq_non_uni, + eq_rx_constr[self.uniform_r1cs.num_rows + i] + ); + + a_mle += non_uni_a * row_constr_eq_non_uni; + b_mle += non_uni_b * row_constr_eq_non_uni; + } + + (a_mle, b_mle, c_mle) } /// Returns the digest of the r1cs shape @@ -408,7 +434,7 @@ mod test { use crate::{ poly::dense_mlpoly::DensePolynomial, r1cs::{ - builder::{OffsetEqConstraint, R1CSBuilder, R1CSConstraintBuilder}, + builder::{R1CSBuilder, R1CSConstraintBuilder}, test::{ materialize_full_uniform, simp_test_big_matrices, simp_test_builder_key, TestInputs, }, @@ -433,11 +459,8 @@ mod test { constraints.build_constraints(&mut uniform_builder); let _num_steps: usize = 3; let num_steps_pad = 4; - let combined_builder = CombinedUniformBuilder::construct( - uniform_builder, - num_steps_pad, - OffsetEqConstraint::empty(), - ); + let combined_builder = + CombinedUniformBuilder::construct(uniform_builder, num_steps_pad, vec![]); let key = UniformSpartanKey::from_builder(&combined_builder); let materialized_a = materialize_full_uniform(&key, &key.uniform_r1cs.a); @@ -552,11 +575,8 @@ mod test { let constraints = TestConstraints(); constraints.build_constraints(&mut uniform_builder); let num_steps_pad = 4; - let combined_builder = CombinedUniformBuilder::construct( - uniform_builder, - num_steps_pad, - OffsetEqConstraint::empty(), - ); + let combined_builder = + CombinedUniformBuilder::construct(uniform_builder, num_steps_pad, vec![]); let mut inputs = vec![vec![Fr::zero(); num_steps_pad]; TestInputs::COUNT]; inputs[TestInputs::OpFlags0 as usize][0] = Fr::from(1); diff --git a/jolt-core/src/r1cs/test.rs b/jolt-core/src/r1cs/test.rs index 086abd7ae..a6b1c09b7 100644 --- a/jolt-core/src/r1cs/test.rs +++ b/jolt-core/src/r1cs/test.rs @@ -135,8 +135,11 @@ pub fn simp_test_builder_key( let _num_steps: usize = 3; let num_steps_pad = 4; - let combined_builder = - CombinedUniformBuilder::construct(uniform_builder, num_steps_pad, offset_eq_constraint); + let combined_builder = CombinedUniformBuilder::construct( + uniform_builder, + num_steps_pad, + vec![offset_eq_constraint], + ); let key = UniformSpartanKey::from_builder(&combined_builder); (combined_builder, key) From 2ac0d7768acce9c8989788235d9b0863539a6d0c Mon Sep 17 00:00:00 2001 From: sragss Date: Tue, 16 Jul 2024 15:47:07 -0700 Subject: [PATCH 4/6] fix missing sparse index update --- jolt-core/src/r1cs/builder.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/jolt-core/src/r1cs/builder.rs b/jolt-core/src/r1cs/builder.rs index 38985bcff..0010b5337 100644 --- a/jolt-core/src/r1cs/builder.rs +++ b/jolt-core/src/r1cs/builder.rs @@ -833,7 +833,7 @@ impl CombinedUniformBuilder { let span = tracing::span!(tracing::Level::DEBUG, "offset_eq"); let _enter = span.enter(); - for constr in &self.offset_equality_constraints { + for (constr_i, constr) in self.offset_equality_constraints.iter().enumerate() { let condition_evals = constr .cond .1 @@ -874,7 +874,8 @@ impl CombinedUniformBuilder { // Sparsify: take only the non-zero elements for (local_index, (az, bz)) in dense_az_bz.iter().enumerate() { - let global_index = uniform_constraint_rows + local_index; + let global_index = + uniform_constraint_rows + constr_i * self.uniform_repeat + local_index; if !az.is_zero() { az_sparse.push((*az, global_index)); } From 26b3117a39aee32b329fc47eb00615ab3de255da Mon Sep 17 00:00:00 2001 From: sragss Date: Tue, 23 Jul 2024 08:11:06 -0700 Subject: [PATCH 5/6] fix compute_spartan_Az_Bz_Cz allocation size --- jolt-core/src/r1cs/builder.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/jolt-core/src/r1cs/builder.rs b/jolt-core/src/r1cs/builder.rs index 0010b5337..a162ae29b 100644 --- a/jolt-core/src/r1cs/builder.rs +++ b/jolt-core/src/r1cs/builder.rs @@ -667,7 +667,7 @@ impl CombinedUniformBuilder { } pub(super) fn offset_eq_constraint_rows(&self) -> usize { - self.uniform_repeat + self.uniform_repeat * self.offset_equality_constraints.len() } /// Total number of rows used across all repeated constraints. Not padded to nearest power of two. @@ -824,7 +824,7 @@ impl CombinedUniformBuilder { let (mut az_sparse, mut bz_sparse, cz_sparse) = par_flatten_triple( uni_constraint_evals, unsafe_allocate_sparse_zero_vec, - self.uniform_repeat, // Capacity overhead for offset_eq constraints. + self.offset_eq_constraint_rows(), // Capacity overhead for offset_eq constraints. ); // offset_equality_constraints: Xz[uniform_constraint_rows..uniform_constraint_rows + 1] From b46d72a4a6dcb6c7353e75f395686fc990c13a0e Mon Sep 17 00:00:00 2001 From: sragss Date: Tue, 23 Jul 2024 08:13:15 -0700 Subject: [PATCH 6/6] debug_assert -> assert --- jolt-core/src/r1cs/key.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/jolt-core/src/r1cs/key.rs b/jolt-core/src/r1cs/key.rs index 43620b0c9..cf4ab26ab 100644 --- a/jolt-core/src/r1cs/key.rs +++ b/jolt-core/src/r1cs/key.rs @@ -378,7 +378,7 @@ impl UniformSpartanKey { let row_constr_eq_non_uni = EqPolynomial::new(r_row_constr.to_vec()).evaluate(&non_uni_constraint_index); - debug_assert_eq!( + assert_eq!( row_constr_eq_non_uni, eq_rx_constr[self.uniform_r1cs.num_rows + i] );