don't hardcode WHIR_N_VARS in recursion program

Author: TomWambsgans

crates/lean_compiler/snark_lib.py

@@ -122,4 +122,7 @@ def match_range(value: int, *args):
     raise AssertionError(f"Value {value} not in any range")
 
 def hint_decompose_bits_xmss(*args):
-    _ = args
+    _ = args
+
+def hint_log2_ceil(n):
+    return log2_ceil(n)

crates/lean_compiler/src/a_simplify_lang.rs

@@ -2182,7 +2182,7 @@ fn simplify_lines(
                                     "Custom hint {function_name} should not return values, at {location}"
                                 ));
                             }
-                            if !hint.n_args_range().contains(&args.len()) {
+                            if args.len() != hint.n_args() {
                                 return Err(format!(
                                     "Custom hint {function_name}: invalid number of arguments, at {location}"
                                 ));

crates/lean_compiler/tests/test_compiler.rs

@@ -156,17 +156,3 @@ def main():
    "#;
     compile_and_run(&ProgramSource::Raw(program.to_string()), (&[], &[]), false);
 }
-
-#[test]
-fn bug() {
-    let program = r#"
-def main():
-    three = double(1) + 1
-    assert three == 3
-    return
-
-def double(a: Const):
-    return a + a
-   "#;
-    compile_and_run(&ProgramSource::Raw(program.to_string()), (&[], &[]), false);
-}

crates/lean_prover/src/prove_execution.rs

@@ -15,7 +15,7 @@ pub struct ExecutionProof {
     pub proof: Vec<F>,
     pub proof_size_fe: usize,
     pub exec_summary: String,
-    pub first_whir_n_vars: usize,
+    pub whir_n_vars: usize,
 }
 
 pub fn prove_execution(
@@ -110,7 +110,7 @@ pub fn prove_execution(
         &bytecode_acc,
         &traces,
     );
-    let first_whir_n_vars = stacked_pcs_witness.global_polynomial.by_ref().n_vars();
+    let whir_n_vars = stacked_pcs_witness.global_polynomial.by_ref().n_vars();
 
     // logup (GKR)
     let logup_c = prover_state.sample();
@@ -207,7 +207,7 @@ pub fn prove_execution(
         proof: prover_state.raw_proof(),
         proof_size_fe,
         exec_summary,
-        first_whir_n_vars,
+        whir_n_vars,
     }
 }
 

crates/lean_vm/src/core/constants.rs

@@ -7,17 +7,17 @@ pub const DIGEST_LEN: usize = 8;
 
 /// Minimum and maximum memory size (as powers of two)
 pub const MIN_LOG_MEMORY_SIZE: usize = 16;
-pub const MAX_LOG_MEMORY_SIZE: usize = 29;
+pub const MAX_LOG_MEMORY_SIZE: usize = 25;
 
 /// Maximum memory size for VM runner (specific to this implementation)
 pub const MAX_RUNNER_MEMORY_SIZE: usize = 1 << 24;
 
 /// Minimum and maximum number of rows per table (as powers of two), both inclusive
 pub const MIN_LOG_N_ROWS_PER_TABLE: usize = 8; // Zero padding will be added to each at least, if this minimum is not reached, (ensuring AIR / GKR work fine, with SIMD, without too much edge cases). Long term, we should find a more elegant solution.
 pub const MAX_LOG_N_ROWS_PER_TABLE: [(Table, usize); 3] = [
-    (Table::execution(), 25),
-    (Table::dot_product(), 22),
-    (Table::poseidon16(), 21),
+    (Table::execution(), 24),
+    (Table::dot_product(), 20),
+    (Table::poseidon16(), 20),
 ];
 
 /// Starting program counter
@@ -63,7 +63,7 @@ mod tests {
     use multilinear_toolkit::prelude::PrimeField64;
     use p3_util::log2_ceil_u64;
 
-    use crate::{DIMENSION, F, MAX_LOG_N_ROWS_PER_TABLE, Table, TableT};
+    use crate::{DIMENSION, F, MAX_LOG_MEMORY_SIZE, MAX_LOG_N_ROWS_PER_TABLE, Table, TableT};
 
     /// CRITICAL FOUR SOUNDNESS: TODO tripple check
     #[test]
@@ -94,10 +94,10 @@ mod tests {
 
     #[test]
     fn ensure_not_too_big_commitment_surface() {
-        let mut max_surface: u64 = 0;
+        let mut max_surface: u64 = 2 * (1 << MAX_LOG_MEMORY_SIZE) as u64; // memory and acc_memory
         for (table, max_log_n_rows) in MAX_LOG_N_ROWS_PER_TABLE {
             max_surface += (table.n_committed_columns() as u64) << (max_log_n_rows as u64);
         }
-        assert!(max_surface < F::ORDER_U64.next_power_of_two() / 2);
+        assert!(max_surface <= 1 << 29); // Maximum data we can commit via WHIR using an initial folding factor of 7
     }
 }

crates/lean_vm/src/isa/hint.rs

@@ -6,7 +6,6 @@ use multilinear_toolkit::prelude::*;
 use std::fmt::Debug;
 use std::fmt::{Display, Formatter};
 use std::hash::Hash;
-use std::ops::Range;
 use strum::IntoEnumIterator;
 use utils::{ToUsize, pretty_integer, to_big_endian_in_field, to_little_endian_in_field};
 
@@ -77,6 +76,7 @@ pub enum CustomHint {
     DecomposeBitsXMSS,
     DecomposeBits,
     LessThan,
+    Log2Ceil,
 }
 
 impl CustomHint {
@@ -85,14 +85,16 @@ impl CustomHint {
             Self::DecomposeBitsXMSS => "hint_decompose_bits_xmss",
             Self::DecomposeBits => "hint_decompose_bits",
             Self::LessThan => "hint_less_than",
+            Self::Log2Ceil => "hint_log2_ceil",
         }
     }
 
-    pub fn n_args_range(&self) -> Range<usize> {
+    pub fn n_args(&self) -> usize {
         match self {
-            Self::DecomposeBitsXMSS => 5..6,
-            Self::DecomposeBits => 4..5,
-            Self::LessThan => 3..4,
+            Self::DecomposeBitsXMSS => 5,
+            Self::DecomposeBits => 4,
+            Self::LessThan => 3,
+            Self::Log2Ceil => 2,
         }
     }
 
@@ -145,6 +147,11 @@ impl CustomHint {
                 let result = if a.to_usize() < b.to_usize() { F::ONE } else { F::ZERO };
                 ctx.memory.set(res_ptr, result)?;
             }
+            Self::Log2Ceil => {
+                let n = args[0].read_value(ctx.memory, ctx.fp)?.to_usize();
+                let res_ptr = args[1].memory_address(ctx.fp)?;
+                ctx.memory.set(res_ptr, F::from_usize(log2_ceil_usize(n)))?;
+            }
         }
         Ok(())
     }

crates/rec_aggregation/fiat_shamir.py

@@ -24,7 +24,8 @@ def fs_grinding(fs, bits):
     new_fs[8] = transcript_ptr + 8
 
     sampled = new_fs[0]
-    _, sampled_low_bits_value = checked_decompose_bits(sampled, bits)
+    _, partial_sums_24 = checked_decompose_bits(sampled)
+    sampled_low_bits_value = partial_sums_24[bits - 1]
     assert sampled_low_bits_value == 0
 
     return new_fs
@@ -107,22 +108,22 @@ def fs_print_state(fs_state):
     return
 
 
-def sample_bits_const(fs, n_samples: Const, K):
+def sample_bits_const(fs, n_samples: Const):
     # return the updated fiat-shamir, and a pointer to n pointers, each pointing to 31 (boolean) field elements,
     sampled_bits = Array(n_samples)
     n_chunks = div_ceil(n_samples, 8)
     new_fs, sampled = fs_sample_chunks(fs, n_chunks)
     for i in unroll(0, n_samples):
-        bits, _ = checked_decompose_bits(sampled[i], K)
+        bits, _ = checked_decompose_bits(sampled[i])
         sampled_bits[i] = bits
     return new_fs, sampled_bits
 
 
-def sample_bits_dynamic(fs_state, n_samples, K):
+def sample_bits_dynamic(fs_state, n_samples):
     new_fs_state: Imu
     sampled_bits: Imu
     for r in unroll(0, WHIR_N_ROUNDS + 1):
         if n_samples == WHIR_NUM_QUERIES[r]:
-            new_fs_state, sampled_bits = sample_bits_const(fs_state, WHIR_NUM_QUERIES[r], K)
+            new_fs_state, sampled_bits = sample_bits_const(fs_state, WHIR_NUM_QUERIES[r])
             return new_fs_state, sampled_bits
     assert False, "sample_bits_dynamic called with unsupported n_samples"

crates/rec_aggregation/recursion.py

@@ -67,19 +67,26 @@ def recursion(inner_public_memory_log_size, inner_public_memory, proof_transcrip
         assert mem_and_table_dims[i] == 0
     log_memory = mem_and_table_dims[0]
 
-    table_dims = mem_and_table_dims + 1
-    log_cycles = table_dims[EXECUTION_TABLE_INDEX]
-    assert log_cycles <= log_memory
+    table_log_heights = mem_and_table_dims + 1
+    log_n_cycles = table_log_heights[EXECUTION_TABLE_INDEX]
+    assert log_n_cycles <= log_memory
 
+    log_bytecode = log2_ceil(GUEST_BYTECODE_LEN)
+    log_bytecode_padded = maximum(log_bytecode, log_n_cycles)
+
+    table_heights = Array(N_TABLES)
     for i in unroll(0, N_TABLES):
-        n_vars_for_table = table_dims[i]
-        assert n_vars_for_table <= log_cycles
-        assert MIN_LOG_N_ROWS_PER_TABLE <= n_vars_for_table
-        assert n_vars_for_table <= MAX_LOG_N_ROWS_PER_TABLE[i]
+        table_log_height = table_log_heights[i]
+        table_heights[i] = powers_of_two(table_log_height)
+        assert table_log_height <= log_n_cycles
+        assert MIN_LOG_N_ROWS_PER_TABLE <= table_log_height
+        assert table_log_height <= MAX_LOG_N_ROWS_PER_TABLE[i]
     assert MIN_LOG_MEMORY_SIZE <= log_memory
     assert log_memory <= MAX_LOG_MEMORY_SIZE
     assert log_memory <= GUEST_BYTECODE_LEN
 
+    stacked_n_vars = compute_stacked_n_vars(log_memory, log_bytecode_padded, table_heights)
+
     fs, whir_base_root, whir_base_ood_points, whir_base_ood_evals = parse_whir_commitment_const(fs, WHIR_NUM_OOD_COMMIT)
 
     fs, logup_c = fs_sample_ef(fs)
@@ -107,9 +114,6 @@ def recursion(inner_public_memory_log_size, inner_public_memory, proof_transcrip
 
     offset: Mut = powers_of_two(log_memory)
 
-    log_bytecode = log2_ceil(GUEST_BYTECODE_LEN)
-    log_n_cycles = table_dims[EXECUTION_TABLE_INDEX]
-    log_bytecode_padded = maximum(log_bytecode, log_n_cycles)
     bytecode_and_acc_point = point_gkr + (N_VARS_LOGUP_GKR - log_bytecode) * DIM
     bytecode_multilinear_location_prefix = multilinear_location_prefix(
         offset / 2 ** log2_ceil(GUEST_BYTECODE_LEN), N_VARS_LOGUP_GKR - log_bytecode, point_gkr
@@ -185,8 +189,8 @@ def recursion(inner_public_memory_log_size, inner_public_memory, proof_transcrip
     for table_index in unroll(0, N_TABLES):
         # I] Bus (data flow between tables)
 
-        log_n_rows = table_dims[table_index]
-        n_rows = powers_of_two(log_n_rows)
+        log_n_rows = table_log_heights[table_index]
+        n_rows = table_heights[table_index]
         inner_point = point_gkr + (N_VARS_LOGUP_GKR - log_n_rows) * DIM
         pcs_points[table_index].push(inner_point)
 
@@ -309,7 +313,7 @@ def recursion(inner_public_memory_log_size, inner_public_memory, proof_transcrip
     air_alpha_powers = powers_const(air_alpha, MAX_NUM_AIR_CONSTRAINTS + 1)
 
     for table_index in unroll(0, N_TABLES):
-        log_n_rows = table_dims[table_index]
+        log_n_rows = table_log_heights[table_index]
         bus_numerator_value = bus_numerators_values[table_index]
         bus_denominator_value = bus_denominators_values[table_index]
         total_num_cols = NUM_COLS_F_AIR[table_index] + DIM * NUM_COLS_EF_AIR[table_index]
@@ -479,6 +483,7 @@ def recursion(inner_public_memory_log_size, inner_public_memory, proof_transcrip
     end_sum: Mut
     fs, folding_randomness_global, s, final_value, end_sum = whir_open(
         fs,
+        stacked_n_vars,
         whir_base_root,
         whir_base_ood_points,
         combination_randomness_powers,
@@ -488,31 +493,31 @@ def recursion(inner_public_memory_log_size, inner_public_memory, proof_transcrip
     curr_randomness = combination_randomness_powers + WHIR_NUM_OOD_COMMIT * DIM
 
     eq_memory_and_acc_point = eq_mle_extension(
-        folding_randomness_global + (WHIR_N_VARS - log_memory) * DIM,
+        folding_randomness_global + (stacked_n_vars - log_memory) * DIM,
         memory_and_acc_point,
         log_memory,
     )
-    prefix_memory = multilinear_location_prefix(0, WHIR_N_VARS - log_memory, folding_randomness_global)
+    prefix_memory = multilinear_location_prefix(0, stacked_n_vars - log_memory, folding_randomness_global)
     s = add_extension_ret(
         s,
         mul_extension_ret(mul_extension_ret(curr_randomness, prefix_memory), eq_memory_and_acc_point),
     )
     curr_randomness += DIM
 
-    prefix_acc_memory = multilinear_location_prefix(1, WHIR_N_VARS - log_memory, folding_randomness_global)
+    prefix_acc_memory = multilinear_location_prefix(1, stacked_n_vars - log_memory, folding_randomness_global)
     s = add_extension_ret(
         s,
         mul_extension_ret(mul_extension_ret(curr_randomness, prefix_acc_memory), eq_memory_and_acc_point),
     )
     curr_randomness += DIM
 
     eq_pub_mem = eq_mle_extension(
-        folding_randomness_global + (WHIR_N_VARS - inner_public_memory_log_size) * DIM,
+        folding_randomness_global + (stacked_n_vars - inner_public_memory_log_size) * DIM,
         public_memory_random_point,
         inner_public_memory_log_size,
     )
     prefix_pub_mem = multilinear_location_prefix(
-        0, WHIR_N_VARS - inner_public_memory_log_size, folding_randomness_global
+        0, stacked_n_vars - inner_public_memory_log_size, folding_randomness_global
     )
     s = add_extension_ret(
         s,
@@ -523,13 +528,13 @@ def recursion(inner_public_memory_log_size, inner_public_memory, proof_transcrip
     offset = powers_of_two(log_memory) * 2  # memory and acc_memory
 
     eq_bytecode_acc = eq_mle_extension(
-        folding_randomness_global + (WHIR_N_VARS - log2_ceil(GUEST_BYTECODE_LEN)) * DIM,
+        folding_randomness_global + (stacked_n_vars - log2_ceil(GUEST_BYTECODE_LEN)) * DIM,
         bytecode_and_acc_point,
         log2_ceil(GUEST_BYTECODE_LEN),
     )
     prefix_bytecode_acc = multilinear_location_prefix(
         offset / 2 ** log2_ceil(GUEST_BYTECODE_LEN),
-        WHIR_N_VARS - log2_ceil(GUEST_BYTECODE_LEN),
+        stacked_n_vars - log2_ceil(GUEST_BYTECODE_LEN),
         folding_randomness_global,
     )
     s = add_extension_ret(
@@ -541,36 +546,36 @@ def recursion(inner_public_memory_log_size, inner_public_memory, proof_transcrip
 
     prefix_pc_start = multilinear_location_prefix(
         offset + COL_PC * powers_of_two(log_n_cycles),
-        WHIR_N_VARS,
+        stacked_n_vars,
         folding_randomness_global,
     )
     s = add_extension_ret(s, mul_extension_ret(curr_randomness, prefix_pc_start))
     curr_randomness += DIM
 
     prefix_pc_end = multilinear_location_prefix(
         offset + (COL_PC + 1) * powers_of_two(log_n_cycles) - 1,
-        WHIR_N_VARS,
+        stacked_n_vars,
         folding_randomness_global,
     )
     s = add_extension_ret(s, mul_extension_ret(curr_randomness, prefix_pc_end))
     curr_randomness += DIM
 
     for table_index in unroll(0, N_TABLES):
-        log_n_rows = table_dims[table_index]
-        n_rows = powers_of_two(log_n_rows)
+        log_n_rows = table_log_heights[table_index]
+        n_rows = table_heights[table_index]
         total_num_cols = NUM_COLS_F_AIR[table_index] + DIM * NUM_COLS_EF_AIR[table_index]
         for i in unroll(0, len(pcs_points[table_index])):
             point = pcs_points[table_index][i]
             eq_factor = eq_mle_extension(
                 point,
-                folding_randomness_global + (WHIR_N_VARS - log_n_rows) * DIM,
+                folding_randomness_global + (stacked_n_vars - log_n_rows) * DIM,
                 log_n_rows,
             )
             for j in unroll(0, total_num_cols):
                 if len(pcs_values[table_index][i][j]) == 1:
                     prefix = multilinear_location_prefix(
                         offset / n_rows + j,
-                        WHIR_N_VARS - log_n_rows,
+                        stacked_n_vars - log_n_rows,
                         folding_randomness_global,
                     )
                     s = add_extension_ret(
@@ -675,6 +680,17 @@ def verify_gkr_quotient_step(fs: Mut, n_vars, point, claim_num, claim_den):
     return fs, postponed_point, new_claim_num, new_claim_den
 
 
+@inline
+def compute_stacked_n_vars(log_memory, log_bytecode_padded, tables_heights):
+    total: Mut = powers_of_two(log_memory + 1) # memory + acc_memory
+    total += powers_of_two(log_bytecode_padded)
+    for table_index in unroll(0, N_TABLES):
+        n_rows = tables_heights[table_index]
+        total += n_rows * (NUM_COLS_F_AIR[table_index] + DIM * NUM_COLS_EF_AIR[table_index])
+    debug_assert(30 - 24 < MIN_LOG_N_ROWS_PER_TABLE) # cf log2_ceil
+    return MIN_LOG_N_ROWS_PER_TABLE + log2_ceil_runtime(total / 2**MIN_LOG_N_ROWS_PER_TABLE)
+
+
 def evaluate_air_constraints(table_index, inner_evals, air_alpha_powers, bus_beta, logup_alphas_eq_poly):
     res: Imu
     debug_assert(table_index < 3)

crates/rec_aggregation/src/recursion.rs

@@ -69,7 +69,7 @@ def main():
     )
     .unwrap();
 
-    let outer_whir_config = WhirConfig::<EF>::new(&inner_whir_config, proof_to_prove.first_whir_n_vars);
+    let outer_whir_config = WhirConfig::<EF>::new(&inner_whir_config, proof_to_prove.whir_n_vars);
 
     // let guest_program_commitment = {
     //     let mut prover_state = build_prover_state();
@@ -400,7 +400,6 @@ pub(crate) fn whir_recursion_placeholder_replacements(whir_config: &WhirConfig<E
         format!("WHIR_FOLDING_FACTORS{}", end),
         format!("[{}]", folding_factors.join(", ")),
     );
-    replacements.insert(format!("WHIR_N_VARS{}", end), whir_config.num_variables.to_string());
     replacements.insert(
         format!("WHIR_LOG_INV_RATE{}", end),
         whir_config.starting_log_inv_rate.to_string(),