feat: enable arbitrary constraint degree (#593)

Cargo.toml

@@ -1,19 +1,19 @@
 [workspace]
 members = [
-    "cli",
-    "core",
-    "derive",
-    "eval",
-    "helper",
-    "primitives",
-    "prover",
-    "recursion/circuit",
-    "recursion/compiler",
-    "recursion/core",
-    "recursion/gnark-ffi",
-    "recursion/program",
-    "sdk",
-    "zkvm/*",
+  "cli",
+  "core",
+  "derive",
+  "eval",
+  "helper",
+  "primitives",
+  "prover",
+  "recursion/circuit",
+  "recursion/compiler",
+  "recursion/core",
+  "recursion/gnark-ffi",
+  "recursion/program",
+  "sdk",
+  "zkvm/*",
 ]
 exclude = ["examples/target"]
 resolver = "2"

core/src/air/builder.rs

@@ -656,4 +656,5 @@ impl<'a, SC: StarkGenericConfig> EmptyMessageBuilder for VerifierConstraintFolde
 impl<F: Field> EmptyMessageBuilder for SymbolicAirBuilder<F> {}
 
 #[cfg(debug_assertions)]
+#[cfg(not(doctest))]
 impl<'a, F: Field> EmptyMessageBuilder for p3_uni_stark::DebugConstraintBuilder<'a, F> {}

core/src/alu/add_sub/mod.rs

@@ -180,12 +180,6 @@ where
         builder.assert_bool(local.is_add);
         builder.assert_bool(local.is_sub);
         builder.assert_bool(is_real);
-
-        // Degree 3 constraint to avoid "OodEvaluationMismatch".
-        builder.assert_zero(
-            local.operand_1[0] * local.operand_1[0] * local.operand_1[0]
-                - local.operand_1[0] * local.operand_1[0] * local.operand_1[0],
-        );
     }
 }
 

core/src/alu/bitwise/mod.rs

@@ -162,11 +162,6 @@ where
         builder.assert_bool(local.is_or);
         builder.assert_bool(local.is_and);
         builder.assert_bool(is_real);
-
-        // Degree 3 constraint to avoid "OodEvaluationMismatch".
-        builder.assert_zero(
-            local.a[0] * local.b[0] * local.c[0] - local.a[0] * local.b[0] * local.c[0],
-        );
     }
 }
 

core/src/alu/divrem/mod.rs

@@ -766,11 +766,6 @@ where
                 local.is_real,
             );
         }
-
-        // A dummy constraint to keep the degree 3.
-        builder.assert_zero(
-            local.a[0] * local.b[0] * local.c[0] - local.a[0] * local.b[0] * local.c[0],
-        )
     }
 }
 

core/src/alu/lt/mod.rs

@@ -240,11 +240,6 @@ where
 
         let is_real = local.is_slt + local.is_sltu;
 
-        // Dummy degree 3 constraint to avoid "OodEvaluationMismatch".
-        builder.assert_zero(
-            local.a[0] * local.b[0] * local.c[0] - local.a[0] * local.b[0] * local.c[0],
-        );
-
         // We can compute the signed set-less-than as follows:
         // SLT (signed) = b_s * (1 - c_s) + (b_s == c_s) * SLTU(b_<s, c_<s)
         // Source: Jolt 5.3: Set Less Than (https://people.cs.georgetown.edu/jthaler/Jolt-paper.pdf)

core/src/alu/mul/mod.rs

@@ -441,11 +441,6 @@ where
             local.shard,
             local.is_real,
         );
-
-        // A dummy constraint to keep the degree at least 3.
-        builder.assert_zero(
-            local.a[0] * local.b[0] * local.c[0] - local.a[0] * local.b[0] * local.c[0],
-        );
     }
 }
 

core/src/alu/sll/mod.rs

@@ -343,11 +343,6 @@ where
             local.shard,
             local.is_real,
         );
-
-        // A dummy constraint to keep the degree at least 3.
-        builder.assert_zero(
-            local.a[0] * local.b[0] * local.c[0] - local.a[0] * local.b[0] * local.c[0],
-        );
     }
 }
 

core/src/bytes/air.rs

@@ -72,8 +72,5 @@ impl<AB: SP1AirBuilder + PairBuilder> Air<AB> for ByteChip<AB::F> {
                 ),
             }
         }
-
-        // Dummy constraint for normalizing to degree 3.
-        builder.assert_zero(local.b * local.b * local.b - local.b * local.b * local.b);
     }
 }

core/src/memory/global.rs

@@ -130,12 +130,6 @@ where
         let local = main.row_slice(0);
         let local: &MemoryInitCols<AB::Var> = (*local).borrow();
 
-        // Dummy constraint of degree 3.
-        builder.assert_eq(
-            local.is_real * local.is_real * local.is_real,
-            local.is_real * local.is_real * local.is_real,
-        );
-
         if self.kind == MemoryChipType::Initialize {
             let mut values = vec![AB::Expr::zero(), AB::Expr::zero(), local.addr.into()];
             values.extend(local.value.map(Into::into));

core/src/operations/add.rs

@@ -107,8 +107,5 @@ impl<F: Field> AddOperation<F> {
             builder.slice_range_check_u8(&b.0, shard, is_real.clone());
             builder.slice_range_check_u8(&cols.value.0, shard, is_real);
         }
-
-        // Degree 3 constraint to avoid "OodEvaluationMismatch".
-        builder.assert_zero(a[0] * b[0] * cols.value[0] - a[0] * b[0] * cols.value[0]);
     }
 }

core/src/operations/add4.rs

@@ -149,8 +149,5 @@ impl<F: Field> Add4Operation<F> {
                 builder_is_real.assert_eq(cols.carry[i] * base, overflow.clone());
             }
         }
-
-        // Degree 3 constraint to avoid "OodEvaluationMismatch".
-        builder.assert_zero(a[0] * b[0] * cols.value[0] - a[0] * b[0] * cols.value[0]);
     }
 }

core/src/operations/add5.rs

@@ -164,8 +164,5 @@ impl<F: Field> Add5Operation<F> {
                 builder_is_real.assert_eq(cols.carry[i] * base, overflow.clone());
             }
         }
-
-        // Degree 3 constraint to avoid "OodEvaluationMismatch".
-        builder.assert_zero(is_real * is_real * is_real - is_real * is_real * is_real);
     }
 }

core/src/operations/field/field_den.rs

@@ -279,11 +279,6 @@ mod tests {
                 AB::F::one(),
                 AB::F::one(),
             );
-
-            // A dummy constraint to keep the degree 3.
-            builder.assert_zero(
-                local.a[0] * local.b[0] * local.a[0] - local.a[0] * local.b[0] * local.a[0],
-            )
         }
     }
 

core/src/operations/field/field_inner_product.rs

@@ -269,12 +269,6 @@ mod tests {
             local
                 .a_ip_b
                 .eval(builder, &local.a, &local.b, AB::F::one(), AB::F::one());
-
-            // A dummy constraint to keep the degree 3.
-            builder.assert_zero(
-                local.a[0][0] * local.b[0][0] * local.a[0][0]
-                    - local.a[0][0] * local.b[0][0] * local.a[0][0],
-            )
         }
     }
 

core/src/operations/field/field_op.rs

@@ -349,11 +349,6 @@ mod tests {
                 AB::F::one(),
                 AB::F::one(),
             );
-
-            // A dummy constraint to keep the degree 3.
-            builder.assert_zero(
-                local.a[0] * local.b[0] * local.a[0] - local.a[0] * local.b[0] * local.a[0],
-            )
         }
     }
 

core/src/operations/field/field_sqrt.rs

@@ -266,11 +266,6 @@ mod tests {
             local
                 .sqrt
                 .eval(builder, &local.a, AB::F::zero(), AB::F::one(), AB::F::one());
-
-            // A dummy constraint to keep the degree 3.
-            builder.assert_zero(
-                local.a[0] * local.a[0] * local.a[0] - local.a[0] * local.a[0] * local.a[0],
-            )
         }
     }
 

core/src/operations/is_equal_word.rs

@@ -48,11 +48,5 @@ impl<F: Field> IsEqualWordOperation<F> {
 
         // Check if the difference is 0.
         IsZeroWordOperation::<AB::F>::eval(builder, diff, cols.is_diff_zero, is_real.clone());
-
-        // Degree 3 constraint to avoid "OodEvaluationMismatch".
-        builder.assert_zero(
-            is_real.clone() * is_real.clone() * is_real.clone()
-                - is_real.clone() * is_real.clone() * is_real.clone(),
-        );
     }
 }

core/src/operations/not.rs

@@ -55,8 +55,5 @@ impl<F: Field> NotOperation<F> {
                 .when(is_real)
                 .assert_eq(cols.value[i] + a[i], AB::F::from_canonical_u8(u8::MAX));
         }
-
-        // A dummy constraint to keep the degree 3.
-        builder.assert_zero(a[0] * a[0] * a[0] - a[0] * a[0] * a[0]);
     }
 }

core/src/program/mod.rs

@@ -166,12 +166,6 @@ where
         let mult_local = main.row_slice(0);
         let mult_local: &ProgramMultiplicityCols<AB::Var> = (*mult_local).borrow();
 
-        // Dummy constraint of degree 3.
-        builder.assert_eq(
-            prep_local.pc * prep_local.pc * prep_local.pc,
-            prep_local.pc * prep_local.pc * prep_local.pc,
-        );
-
         // Contrain the interaction with CPU table
         builder.receive_program(
             prep_local.pc,

core/src/stark/chip.rs

@@ -3,14 +3,18 @@ use std::hash::Hash;
 use p3_air::{Air, BaseAir, PairBuilder};
 use p3_field::{ExtensionField, Field, PrimeField, PrimeField32};
 use p3_matrix::dense::RowMajorMatrix;
+use p3_uni_stark::{get_max_constraint_degree, SymbolicAirBuilder};
 use p3_util::log2_ceil_usize;
 
 use crate::{
     air::{MachineAir, MultiTableAirBuilder, SP1AirBuilder},
     lookup::{Interaction, InteractionBuilder, InteractionKind},
 };
 
-use super::{eval_permutation_constraints, generate_permutation_trace, permutation_trace_width};
+use super::{
+    eval_permutation_constraints, generate_permutation_trace, permutation_trace_width,
+    PROOF_MAX_NUM_PVS,
+};
 
 /// An Air that encodes lookups based on interactions.
 pub struct Chip<F: Field, A> {
@@ -55,15 +59,20 @@ where
     /// Records the interactions and constraint degree from the air and crates a new chip.
     pub fn new(air: A) -> Self
     where
-        A: MachineAir<F> + Air<InteractionBuilder<F>>,
+        A: MachineAir<F> + Air<InteractionBuilder<F>> + Air<SymbolicAirBuilder<F>>,
     {
         // Todo: correct values
         let mut builder = InteractionBuilder::new(air.preprocessed_width(), air.width());
         air.eval(&mut builder);
         let (sends, receives) = builder.interactions();
 
-        // TODO: count constraints from the air.
-        let max_constraint_degree = 3;
+        // TODO: enable different numbers of public values.
+        let mut max_constraint_degree =
+            get_max_constraint_degree(&air, air.preprocessed_width(), PROOF_MAX_NUM_PVS);
+
+        if !sends.is_empty() || !receives.is_empty() {
+            max_constraint_degree = max_constraint_degree.max(3);
+        }
         let log_quotient_degree = log2_ceil_usize(max_constraint_degree - 1);
 
         Self {
@@ -124,8 +133,7 @@ where
 
     #[inline]
     pub fn logup_batch_size(&self) -> usize {
-        // TODO: calculate by log_quotient_degree.
-        2
+        1 << self.log_quotient_degree
     }
 }
 

core/src/stark/prover.rs

@@ -15,7 +15,6 @@ use p3_field::PrimeField32;
 use p3_matrix::dense::RowMajorMatrix;
 use p3_matrix::Matrix;
 use p3_maybe_rayon::prelude::*;
-use p3_util::log2_ceil_usize;
 use p3_util::log2_strict_usize;
 use web_time::Instant;
 
@@ -249,10 +248,10 @@ where
             .map(|degree| log2_strict_usize(*degree))
             .collect::<Vec<_>>();
 
-        // TODO: read dynamically from Chip.
-        let max_constraint_degree = 3;
-        let log_quotient_degree = log2_ceil_usize(max_constraint_degree - 1);
-        let quotient_degree = 1 << log_quotient_degree;
+        let log_quotient_degrees = chips
+            .iter()
+            .map(|chip| chip.log_quotient_degree())
+            .collect::<Vec<_>>();
 
         let pcs = config.pcs();
         let trace_domains = degrees
@@ -336,8 +335,9 @@ where
 
         let quotient_domains = trace_domains
             .iter()
-            .zip(log_degrees.iter())
-            .map(|(domain, log_degree)| {
+            .zip_eq(log_degrees.iter())
+            .zip_eq(log_quotient_degrees.iter())
+            .map(|((domain, log_degree), log_quotient_degree)| {
                 domain.create_disjoint_domain(1 << (log_degree + log_quotient_degree))
             })
             .collect::<Vec<_>>();
@@ -385,16 +385,27 @@ where
         let quotient_domains_and_chunks = quotient_domains
             .into_iter()
             .zip_eq(quotient_values)
-            .flat_map(|(quotient_domain, quotient_values)| {
-                let quotient_flat = RowMajorMatrix::new_col(quotient_values).flatten_to_base();
-                let quotient_chunks = quotient_domain.split_evals(quotient_degree, quotient_flat);
-                let qc_domains = quotient_domain.split_domains(quotient_degree);
-                qc_domains.into_iter().zip_eq(quotient_chunks)
-            })
+            .zip_eq(log_quotient_degrees.iter())
+            .flat_map(
+                |((quotient_domain, quotient_values), log_quotient_degree)| {
+                    let quotient_degree = 1 << *log_quotient_degree;
+                    let quotient_flat = RowMajorMatrix::new_col(quotient_values).flatten_to_base();
+                    let quotient_chunks =
+                        quotient_domain.split_evals(quotient_degree, quotient_flat);
+                    let qc_domains = quotient_domain.split_domains(quotient_degree);
+                    qc_domains.into_iter().zip_eq(quotient_chunks)
+                },
+            )
             .collect::<Vec<_>>();
 
         let num_quotient_chunks = quotient_domains_and_chunks.len();
-        assert_eq!(num_quotient_chunks, chips.len() * quotient_degree);
+        assert_eq!(
+            num_quotient_chunks,
+            chips
+                .iter()
+                .map(|c| 1 << c.log_quotient_degree())
+                .sum::<usize>()
+        );
 
         let (quotient_commit, quotient_data) = tracing::debug_span!("commit to quotient traces")
             .in_scope(|| pcs.commit(quotient_domains_and_chunks));
@@ -465,15 +476,12 @@ where
                 AirOpenedValues { local, next }
             })
             .collect::<Vec<_>>();
-        let quotient_opened_values = quotient_values
-            .chunks_exact_mut(quotient_degree)
-            .map(|slice| {
-                slice
-                    .iter_mut()
-                    .map(|op| op.pop().unwrap())
-                    .collect::<Vec<_>>()
-            })
-            .collect::<Vec<_>>();
+        let mut quotient_opened_values = Vec::with_capacity(log_quotient_degrees.len());
+        for log_quotient_degree in log_quotient_degrees.iter() {
+            let degree = 1 << *log_quotient_degree;
+            let slice = quotient_values.drain(0..degree);
+            quotient_opened_values.push(slice.map(|mut op| op.pop().unwrap()).collect::<Vec<_>>());
+        }
 
         let opened_values = main_opened_values
             .into_iter()

core/src/syscall/precompiles/blake3/compress/g.rs

@@ -192,7 +192,5 @@ impl<F: Field> GOperation<F> {
                 builder.assert_eq(cols.result[i][j], results[i][j]);
             }
         }
-        // Degree 3 constraint to avoid "OodEvaluationMismatch".
-        builder.assert_zero(is_real * is_real * is_real - is_real * is_real * is_real);
     }
 }