diff --git a/provekit/r1cs-compiler/src/lib.rs b/provekit/r1cs-compiler/src/lib.rs index 9d902681c..2ffa41285 100644 --- a/provekit/r1cs-compiler/src/lib.rs +++ b/provekit/r1cs-compiler/src/lib.rs @@ -8,6 +8,7 @@ pub mod noir_to_r1cs; mod optimize; mod poseidon2; pub mod range_check; +mod row_scaling; mod sha256_compression; mod spread; mod uints; @@ -18,5 +19,6 @@ pub use { noir_proof_scheme::{MavrosCompiler, NoirCompiler}, noir_to_r1cs::{noir_to_r1cs, noir_to_r1cs_with_breakdown, R1CSBreakdown}, optimize::{optimize_r1cs, OptimizationStats}, + row_scaling::{minimize_coefficient_triplets, RowScalingStats}, whir_r1cs::MavrosSchemeBuilder, }; diff --git a/provekit/r1cs-compiler/src/noir_proof_scheme.rs b/provekit/r1cs-compiler/src/noir_proof_scheme.rs index dda07774f..5e5ed5095 100644 --- a/provekit/r1cs-compiler/src/noir_proof_scheme.rs +++ b/provekit/r1cs-compiler/src/noir_proof_scheme.rs @@ -92,6 +92,20 @@ impl NoirCompiler { program: ProgramArtifact, hash_config: provekit_common::HashConfig, witness_mode: provekit_common::WitnessCommitmentMode, + ) -> Result { + Self::from_program_with_ge(program, hash_config, witness_mode, true) + } + + /// Build a proof scheme, optionally running Gaussian elimination on the + /// lowered R1CS before scaling rows and preparing witness layers. Row + /// scaling and other preparation steps always run; `from_program` enables + /// GE by default. + #[instrument(skip_all)] + pub fn from_program_with_ge( + program: ProgramArtifact, + hash_config: provekit_common::HashConfig, + witness_mode: provekit_common::WitnessCommitmentMode, + ge_optimization: bool, ) -> Result { provekit_backend_bn254::register(); @@ -119,21 +133,29 @@ impl NoirCompiler { main.public_inputs().indices().iter().cloned().collect(); r1cs.num_public_inputs = acir_public_inputs_indices_set.len(); - // Gaussian elimination optimization pass - let opt_stats = crate::optimize_r1cs( - &mut r1cs, - &mut witness_builders, - &mut witness_map, - &acir_public_inputs_indices_set, - )?; - info!( - "After GE optimization: {} constraints, {} witnesses ({} eliminated, {:.1}% \ - constraint reduction)", - opt_stats.constraints_after, - opt_stats.witnesses_after, - opt_stats.eliminated, - opt_stats.constraint_reduction_percent() - ); + if ge_optimization { + let opt_stats = crate::optimize_r1cs( + &mut r1cs, + &mut witness_builders, + &mut witness_map, + &acir_public_inputs_indices_set, + )?; + info!( + "After GE optimization: {} constraints, {} witnesses ({} eliminated, {:.1}% \ + constraint reduction)", + opt_stats.constraints_after, + opt_stats.witnesses_after, + opt_stats.eliminated, + opt_stats.constraint_reduction_percent() + ); + } else { + info!("Skipping GE optimization"); + } + + // Scale the actual R1CS after GE, before witness remapping and scheme + // construction. Extraction, WHIR preprocessing and serialized keys + // must all see these same coefficients. + crate::minimize_coefficient_triplets(&mut r1cs); let has_public_inputs = !acir_public_inputs_indices_set.is_empty(); let (split_witness_builders, remapped_r1cs, remapped_witness_map, challenge_offsets) = @@ -257,7 +279,7 @@ impl MavrosCompiler { #[cfg(test)] mod tests { use { - crate::NoirCompiler, + crate::{noir_to_r1cs, NoirCompiler}, ark_std::One, provekit_backend_bn254::{ witness::{ConstantTerm, DigitalDecompositionWitnesses, SumTerm, WitnessBuilder}, @@ -281,6 +303,35 @@ mod tests { assert_eq!(value, &deserialized); } + #[test] + fn ge_can_be_disabled_when_building_a_scheme() { + let path = PathBuf::from("../../tooling/provekit-bench/benches/poseidon_rounds.json"); + let program: noirc_artifacts::program::ProgramArtifact = + serde_json::from_reader(std::fs::File::open(path).unwrap()).unwrap(); + let (raw_r1cs, ..) = noir_to_r1cs(&program.bytecode.functions[0]).unwrap(); + let unoptimized = NoirCompiler::from_program_with_ge( + program.clone(), + provekit_common::HashConfig::default(), + provekit_common::WitnessCommitmentMode::default(), + false, + ) + .unwrap(); + let optimized = NoirCompiler::from_program( + program, + provekit_common::HashConfig::default(), + provekit_common::WitnessCommitmentMode::default(), + ) + .unwrap(); + + // Witness preparation still runs, but skipping GE retains the raw + // constraint rows. The default constructor still eliminates rows. + assert_eq!( + unoptimized.r1cs().num_constraints(), + raw_r1cs.num_constraints() + ); + assert!(optimized.r1cs().num_constraints() < unoptimized.r1cs().num_constraints()); + } + #[test] fn test_noir_proof_scheme_serde() { let path = PathBuf::from("../../tooling/provekit-bench/benches/poseidon_rounds.json"); diff --git a/provekit/r1cs-compiler/src/row_scaling.rs b/provekit/r1cs-compiler/src/row_scaling.rs new file mode 100644 index 000000000..ed8ab2b72 --- /dev/null +++ b/provekit/r1cs-compiler/src/row_scaling.rs @@ -0,0 +1,625 @@ +//! Reduce the coefficient alphabet before extracting coefficient-triplet sets. +//! +//! Row scaling preserves the constraint because +//! `(alpha A w)(beta B w) - alpha beta C w = alpha beta (Aw Bw - Cw)`. +//! No witness or witness-builder changes are required. Run after GE and before +//! constructing any scheme/preprocessing data from the R1CS. + +use { + ark_ff::Field, + provekit_common::R1CS, + std::collections::{HashMap, HashSet}, + tracing::info, +}; + +/// Bound candidate generation independently of the size of a row. +const MAX_PIVOTS: usize = 3; +const MAX_SWEEPS: usize = 8; + +#[derive(Debug, Default)] +pub struct RowScalingStats { + pub triplets_before: usize, + pub triplets_after: usize, + pub rows_scaled: usize, + pub sweeps: usize, +} + +#[derive(Clone)] +struct Term { + value: [F; 3], + /// Number of occupied coordinates with this triplet across the group. + occurrences: usize, +} + +struct Group { + /// Distinct triplets, shared by every row in the group. + terms: Vec>, + /// Cumulative multipliers for A, B and C. + scales: [F; 3], +} + +/// Greedily reduce the global number of distinct nonzero `(a, b, c)` triplets. +/// +/// Rows with the same triplet palette are scaled together. A bounded candidate +/// search normalizes frequent A/B, A/C or B/C coefficients to +/-1. Accepted +/// moves strictly improve `(distinct triplets, non-ternary matrix entries)`; +/// this is a heuristic, not a minimum guarantee. Equal scores retain the +/// current representation. Search stops at a fixed point or after eight sweeps. +/// +/// With N occupied coordinates, P distinct row palettes, and K <= 160 scale +/// candidates per palette, each sweep takes expected O(K N) field/hash +/// operations, plus cached field inversions. Initial grouping additionally +/// sorts each row's triplet IDs. Storage is O(N + P + number of rows). +/// +/// Only matrix coefficients change. Dimensions, sparsity, witness indices, +/// public inputs and virtual witnesses are preserved. +pub fn minimize_coefficient_triplets(r1cs: &mut R1CS) -> RowScalingStats { + let (mut groups, row_groups) = collect_groups(r1cs); + // Count groups containing each triplet, not entries: removing an entire + // repeated-row group can then release a triplet in a single move. + let mut counts = HashMap::<[F; 3], usize>::new(); + for group in &groups { + for term in &group.terms { + *counts.entry(term.value).or_default() += 1; + } + } + let mut stats = RowScalingStats { + triplets_before: counts.len(), + ..RowScalingStats::default() + }; + let mut inverses = HashMap::from([(F::ONE, F::ONE), (-F::ONE, -F::ONE)]); + + for sweep in 1..=MAX_SWEEPS { + let mut accepted = 0; + for group in &mut groups { + accepted += usize::from(optimize_group(group, &mut counts, &mut inverses)); + } + stats.sweeps = sweep; + info!( + "R1CS row scaling sweep {sweep}: {} unique triplets, {accepted} palette updates", + counts.len() + ); + if accepted == 0 { + break; + } + } + stats.triplets_after = counts.len(); + assert!(stats.triplets_after <= stats.triplets_before); + + // Apply only the final composed scales; rejected candidates never enter + // the R1CS interner and no intermediate R1CS copies are needed. + let mut entries = Vec::new(); + for (row, &group) in row_groups.iter().enumerate() { + let scales = groups[group].scales; + if scales == [F::ONE; 3] { + continue; + } + stats.rows_scaled += 1; + for (matrix, scale) in [&mut r1cs.a, &mut r1cs.b, &mut r1cs.c] + .into_iter() + .zip(scales) + { + if scale == F::ONE { + continue; + } + entries.clear(); + entries.extend(matrix.iter_row(row).map(|(column, coefficient)| { + let value = r1cs + .interner + .get(coefficient) + .expect("interned coefficient"); + (column, r1cs.interner.intern(value * scale)) + })); + matrix.replace_row(row, &entries); + } + } + info!( + "R1CS row scaling: {} -> {} unique triplets, {} of {} rows scaled", + stats.triplets_before, + stats.triplets_after, + stats.rows_scaled, + r1cs.num_constraints() + ); + stats +} + +/// Merge the sorted A/B/C rows, then group by their sets of triplet IDs. +/// Column positions and multiplicities do not affect the primary objective; +/// multiplicities are retained for the exact non-ternary-entry tie-break. +fn collect_groups(r1cs: &R1CS) -> (Vec>, Vec) { + let matrices = [r1cs.a(), r1cs.b(), r1cs.c()]; + let mut values = Vec::new(); + let mut triplet_ids = HashMap::new(); + let mut palettes = HashMap::, usize>::new(); + let mut groups: Vec> = Vec::new(); + let mut row_groups = Vec::with_capacity(r1cs.num_constraints()); + let mut ids = Vec::new(); + let mut terms = Vec::<(usize, usize)>::new(); + + for row in 0..r1cs.num_constraints() { + ids.clear(); + let mut streams = matrices + .each_ref() + .map(|matrix| matrix.iter_row(row).peekable()); + while let Some(column) = streams + .iter_mut() + .filter_map(|stream| stream.peek().map(|&(column, _)| column)) + .min() + { + let value = streams.each_mut().map(|stream| { + stream + .next_if(|&(index, _)| index == column) + .map_or(F::ZERO, |(_, value)| value) + }); + if value == [F::ZERO; 3] { + continue; + } + let id = *triplet_ids.entry(value).or_insert_with(|| { + values.push(value); + values.len() - 1 + }); + ids.push(id); + } + ids.sort_unstable(); + terms.clear(); + for &id in &ids { + if let Some((last_id, count)) = terms.last_mut().filter(|(last, _)| *last == id) { + debug_assert_eq!(*last_id, id); + *count += 1; + } else { + terms.push((id, 1)); + } + } + ids.dedup(); + let group = if let Some(&group) = palettes.get(&ids) { + for (term, &(_, count)) in groups[group].terms.iter_mut().zip(&terms) { + term.occurrences += count; + } + group + } else { + let group = groups.len(); + palettes.insert(ids.clone(), group); + groups.push(Group { + terms: terms + .iter() + .map(|&(id, occurrences)| Term { + value: values[id], + occurrences, + }) + .collect(), + scales: [F::ONE; 3], + }); + group + }; + row_groups.push(group); + } + (groups, row_groups) +} + +fn nonternary(value: [F; 3]) -> usize { + value + .into_iter() + .filter(|&x| x != F::ZERO && x != F::ONE && x != -F::ONE) + .count() +} + +fn scale_triplet(value: [F; 3], scales: [F; 3]) -> [F; 3] { + std::array::from_fn(|i| value[i] * scales[i]) +} + +fn optimize_group( + group: &mut Group, + counts: &mut HashMap<[F; 3], usize>, + inverses: &mut HashMap, +) -> bool { + let exclusive = group + .terms + .iter() + .filter(|term| counts[&term.value] == 1) + .count(); + let nonternary = group + .terms + .iter() + .map(|term| nonternary(term.value) * term.occurrences) + .sum(); + let mut best_score = (exclusive, nonternary); + if best_score == (0, 0) { + return false; + } + + // Temporarily leave only the other groups in the dictionary. + for term in &group.terms { + let count = counts.get_mut(&term.value).expect("counted triplet"); + *count -= 1; + if *count == 0 { + counts.remove(&term.value); + } + } + let mut best = [F::ONE; 3]; + for scales in candidates(&group.terms, inverses) { + if let Some(score) = score_candidate(&group.terms, scales, counts, best_score) { + best_score = score; + best = scales; + if score == (0, 0) { + break; + } + } + } + for term in &mut group.terms { + term.value = scale_triplet(term.value, best); + *counts.entry(term.value).or_default() += 1; + } + group.scales = scale_triplet(group.scales, best); + best != [F::ONE; 3] +} + +/// Scaling by nonzero elements is injective on triplets, so candidates do not +/// need their own deduplication set. Reject as soon as this monotone partial +/// score cannot beat the incumbent. +fn score_candidate( + terms: &[Term], + scales: [F; 3], + others: &HashMap<[F; 3], usize>, + incumbent: (usize, usize), +) -> Option<(usize, usize)> { + let mut score = (0, 0); + for term in terms { + let value = scale_triplet(term.value, scales); + score.0 += usize::from(!others.contains_key(&value)); + score.1 += nonternary(value) * term.occurrences; + if score >= incumbent { + return None; + } + } + (score < incumbent).then_some(score) +} + +fn inverse(value: F, cache: &mut HashMap) -> F { + *cache + .entry(value) + .or_insert_with(|| value.inverse().expect("nonzero pivot")) +} + +/// Deterministic frequency order, with field order breaking ties. +fn pivots(terms: &[Term], component: usize) -> Vec { + let mut frequencies = HashMap::::new(); + for term in terms { + let value = term.value[component]; + if value != F::ZERO { + *frequencies.entry(value).or_default() += term.occurrences; + } + } + let mut values: Vec<_> = frequencies.into_iter().collect(); + let compare = |(a, ca): &(F, usize), (b, cb): &(F, usize)| cb.cmp(ca).then(a.cmp(b)); + if values.len() > MAX_PIVOTS { + values.select_nth_unstable_by(MAX_PIVOTS, compare); + values.truncate(MAX_PIVOTS); + } + values.sort_unstable_by(compare); + values.into_iter().map(|(value, _)| value).collect() +} + +fn candidates(terms: &[Term], inverses: &mut HashMap) -> Vec<[F; 3]> { + let mut factors = [vec![F::ONE, -F::ONE], vec![F::ONE, -F::ONE]]; + for (component, factors) in factors.iter_mut().enumerate() { + for value in pivots(terms, component) { + let reciprocal = inverse(value, inverses); + factors.extend([reciprocal, -reciprocal]); + } + factors.sort_unstable(); + factors.dedup(); + } + let mut result = Vec::new(); + let mut seen = HashSet::new(); + let mut add = |alpha, beta| { + if [alpha, beta] != [F::ONE; 2] && seen.insert([alpha, beta]) { + result.push([alpha, beta, alpha * beta]); + } + }; + // A/B anchors. In particular this includes independent sign flips. + for &alpha in &factors[0] { + for &beta in &factors[1] { + add(alpha, beta); + } + } + // A/C and B/C anchors also handle purely linear C rows. + for value in pivots(terms, 2) { + let reciprocal = inverse(value, inverses); + for gamma in [reciprocal, -reciprocal] { + for &alpha in &factors[0] { + add(alpha, gamma * inverse(alpha, inverses)); + } + for &beta in &factors[1] { + add(gamma * inverse(beta, inverses), beta); + } + } + } + result +} + +#[cfg(test)] +mod tests { + use {super::*, ark_ff::AdditiveGroup, provekit_backend_bn254::FieldElement as F}; + + fn f(value: i64) -> F { + if value < 0 { + -F::from(value.unsigned_abs()) + } else { + F::from(value as u64) + } + } + + fn add_row(r1cs: &mut R1CS, entries: &[[i64; 3]]) { + let components: [Vec<_>; 3] = std::array::from_fn(|component| { + entries + .iter() + .enumerate() + .filter_map(|(column, entry)| { + (entry[component] != 0).then_some((f(entry[component]), column)) + }) + .collect() + }); + r1cs.add_constraint(&components[0], &components[1], &components[2]); + } + + fn triplets(r1cs: &R1CS) -> HashSet<[F; 3]> { + let (groups, _) = collect_groups(r1cs); + groups + .into_iter() + .flat_map(|group| group.terms.into_iter().map(|t| t.value)) + .collect() + } + + fn matrix_row(r1cs: &R1CS, row: usize, component: usize) -> Vec<(usize, F)> { + [r1cs.a(), r1cs.b(), r1cs.c()][component] + .iter_row(row) + .collect() + } + + /// Verify the whole row transformation, not just satisfiability on one + /// witness: A and B scales must multiply to the C scale. Empty factors + /// leave the corresponding free scaling unconstrained. + fn assert_equivalent(before: &R1CS, after: &R1CS) { + assert_eq!(before.num_constraints(), after.num_constraints()); + assert_eq!(before.num_witnesses(), after.num_witnesses()); + assert_eq!(before.num_public_inputs, after.num_public_inputs); + assert_eq!(before.num_virtual, after.num_virtual); + for row in 0..before.num_constraints() { + let mut scales = [None; 3]; + for (component, scale) in scales.iter_mut().enumerate() { + let a = matrix_row(before, row, component); + let b = matrix_row(after, row, component); + assert_eq!(a.len(), b.len()); + for (&(i, x), &(j, y)) in a.iter().zip(&b) { + assert_eq!(i, j); + assert_ne!(y, F::ZERO); + let ratio = y / x; + assert_eq!(*scale.get_or_insert(ratio), ratio); + } + } + if let [Some(a), Some(b), Some(c)] = scales { + assert_eq!(a * b, c); + } + // Also check the residual identity for several witnesses, + // including invalid ones, using a free nonzero factor if needed. + let residual_scale = scales[2] + .unwrap_or_else(|| scales[0].unwrap_or(F::ONE) * scales[1].unwrap_or(F::ONE)); + for seed in 0..8 { + let witness: Vec<_> = (0..before.num_witnesses()) + .map(|column| F::from((seed * 17 + column * 3) as u64)) + .collect(); + let residual = |r1cs: &R1CS| { + let products: [F; 3] = std::array::from_fn(|component| { + matrix_row(r1cs, row, component) + .into_iter() + .map(|(column, value)| value * witness[column]) + .sum() + }); + products[0] * products[1] - products[2] + }; + assert_eq!(residual(after), residual_scale * residual(before)); + } + } + } + + #[test] + fn merges_scaled_rows_and_preserves_the_constraints() { + let mut r1cs = R1CS::new(); + r1cs.add_witnesses(4); + r1cs.num_public_inputs = 2; + r1cs.num_virtual = 3; + add_row(&mut r1cs, &[[1, 0, 0], [0, 1, 0], [0, 0, 1], [1, 1, 1]]); + add_row(&mut r1cs, &[[2, 0, 0], [0, 3, 0], [0, 0, 6], [2, 3, 6]]); + let before = r1cs.clone(); + let stats = minimize_coefficient_triplets(&mut r1cs); + assert_eq!((stats.triplets_before, stats.triplets_after), (8, 4)); + assert!(stats.rows_scaled > 0); + assert_eq!(triplets(&r1cs).len(), stats.triplets_after); + assert_equivalent(&before, &r1cs); + } + + #[test] + fn scales_repeated_palettes_together_even_at_different_columns() { + let mut r1cs = R1CS::new(); + r1cs.add_witnesses(3); + add_row(&mut r1cs, &[[1, 1, 1]]); + for _ in 0..20 { + add_row(&mut r1cs, &[[0, 0, 0], [2, 3, 6], [2, 3, 6]]); + } + let before = r1cs.clone(); + let (groups, rows) = collect_groups(&r1cs); + assert_eq!(groups.len(), 2); + assert_eq!(groups[rows[1]].terms[0].occurrences, 40); + let stats = minimize_coefficient_triplets(&mut r1cs); + assert_eq!((stats.triplets_before, stats.triplets_after), (2, 1)); + assert_eq!(stats.rows_scaled, 20); + assert_equivalent(&before, &r1cs); + } + + #[test] + fn ternary_tie_break_seeds_a_shared_palette() { + let mut r1cs = R1CS::new(); + r1cs.add_witnesses(3); + add_row(&mut r1cs, &[[2, 0, 0], [0, 3, 0], [0, 0, 6]]); + add_row(&mut r1cs, &[[5, 0, 0], [0, 7, 0], [0, 0, 35]]); + let before = r1cs.clone(); + let stats = minimize_coefficient_triplets(&mut r1cs); + assert_eq!((stats.triplets_before, stats.triplets_after), (6, 3)); + assert!(triplets(&r1cs).into_iter().all(|t| nonternary(t) == 0)); + assert_equivalent(&before, &r1cs); + } + + #[test] + fn preserves_the_c_over_ab_invariant() { + let mut r1cs = R1CS::new(); + r1cs.add_witnesses(1); + add_row(&mut r1cs, &[[1, 1, 1]]); + add_row(&mut r1cs, &[[2, 3, 12]]); + let before = r1cs.clone(); + let stats = minimize_coefficient_triplets(&mut r1cs); + assert_eq!(stats.triplets_after, 2); + assert_equivalent(&before, &r1cs); + } + + #[test] + fn handles_empty_and_linear_rows() { + let mut r1cs = R1CS::new(); + r1cs.add_witnesses(2); + add_row(&mut r1cs, &[]); + add_row(&mut r1cs, &[[0, 0, 1]]); + add_row(&mut r1cs, &[[0, 0, -9], [0, 0, -9]]); + add_row(&mut r1cs, &[[7, 0, 0], [0, 0, 14]]); + add_row(&mut r1cs, &[[0, 5, 0], [0, 0, 15]]); + let before = r1cs.clone(); + let stats = minimize_coefficient_triplets(&mut r1cs); + assert!(stats.triplets_after < stats.triplets_before); + assert_equivalent(&before, &r1cs); + let mut empty = R1CS::::new(); + let stats = minimize_coefficient_triplets(&mut empty); + assert_eq!(stats.triplets_after, 0); + assert_eq!(stats.rows_scaled, 0); + } + + #[test] + fn rejected_candidates_do_not_pollute_the_interner() { + let mut r1cs = R1CS::new(); + r1cs.add_witnesses(3); + // All 26 nonzero ternary triples in one row: any non-sign scale + // worsens the secondary objective, and sign flips only tie. + let mut entries = Vec::new(); + for a in -1..=1 { + for b in -1..=1 { + for c in -1..=1 { + if [a, b, c] != [0; 3] { + entries.push([a, b, c]); + } + } + } + } + r1cs.grow_matrices(0, entries.len()); + add_row(&mut r1cs, &entries); + let before = r1cs.clone(); + let stats = minimize_coefficient_triplets(&mut r1cs); + assert_eq!(stats.rows_scaled, 0); + assert_eq!(before, r1cs); + } + + #[test] + fn candidate_scoring_matches_materialized_union() { + let terms = vec![ + Term { + value: [f(2), f(0), f(0)], + occurrences: 10, + }, + Term { + value: [f(0), f(3), f(0)], + occurrences: 2, + }, + Term { + value: [f(2), f(3), f(6)], + occurrences: 4, + }, + ]; + let others = HashMap::from([ + ([F::ONE, F::ZERO, F::ZERO], 2), + ([F::ZERO, F::ONE, F::ZERO], 1), + ([f(2), f(3), f(6)], 3), + ]); + let mut inverses = HashMap::new(); + let choices = candidates(&terms, &mut inverses); + assert!(choices.len() <= 160); + assert_eq!( + choices.iter().copied().collect::>().len(), + choices.len() + ); + for scales in choices { + assert_ne!(scales[0], F::ZERO); + assert_ne!(scales[1], F::ZERO); + assert_eq!(scales[2], scales[0] * scales[1]); + let mut union: HashSet<_> = others.keys().copied().collect(); + let mut complexity = 0; + for term in &terms { + let value = scale_triplet(term.value, scales); + union.insert(value); + complexity += nonternary(value) * term.occurrences; + } + let expected = (union.len() - others.len(), complexity); + for missing in 0..=4 { + for cost in [0, 1, 2, 10, 50] { + let incumbent = (missing, cost); + assert_eq!( + score_candidate(&terms, scales, &others, incumbent), + (expected < incumbent).then_some(expected) + ); + } + } + } + } + + #[test] + fn candidates_cover_c_normalization_and_stay_bounded() { + let terms: Vec<_> = (1..=30) + .map(|x| Term { + value: [f(x), f(x + 30), f(x + 60)], + occurrences: x as usize, + }) + .collect(); + let mut cache = HashMap::new(); + let choices = candidates(&terms, &mut cache); + assert!(choices.len() <= 160); + assert!(choices + .iter() + .any(|&[a, b, c]| { a * f(30) == F::ONE && c * f(90) == F::ONE && b == c / a })); + } + + #[test] + fn random_systems_preserve_residuals_support_and_serialization() { + let mut seed = 123456789_u64; + for _ in 0..32 { + let mut r1cs = R1CS::new(); + r1cs.add_witnesses(7); + for _ in 0..12 { + let entries: Vec<_> = (0..7) + .map(|_| { + std::array::from_fn(|_| { + seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1); + ((seed >> 32) % 9) as i64 - 4 + }) + }) + .collect(); + add_row(&mut r1cs, &entries); + } + let before = r1cs.clone(); + let mut repeat = before.clone(); + let stats = minimize_coefficient_triplets(&mut r1cs); + minimize_coefficient_triplets(&mut repeat); + assert_eq!(r1cs, repeat, "hash iteration must not affect output"); + assert_eq!(triplets(&r1cs).len(), stats.triplets_after); + assert!(stats.triplets_after <= stats.triplets_before); + assert_equivalent(&before, &r1cs); + let bytes = postcard::to_stdvec(&r1cs).unwrap(); + let restored: R1CS = postcard::from_bytes(&bytes).unwrap(); + assert_eq!(restored, r1cs); + } + } +} diff --git a/tooling/cli/src/cmd/extract.rs b/tooling/cli/src/cmd/extract.rs new file mode 100644 index 000000000..f755a9b48 --- /dev/null +++ b/tooling/cli/src/cmd/extract.rs @@ -0,0 +1,741 @@ +use { + ark_ff::Field, + provekit_common::{InternedFieldElement, Interner, SparseMatrix, R1CS}, + std::collections::HashMap, + tracing::info, +}; + +mod matrix; +mod minimize; +mod rank; +mod rescale; +mod ternary; + +pub(super) use matrix::Core; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) enum MinimizeStrategy { + Rank, + Ternary, +} + +impl argh::FromArgValue for MinimizeStrategy { + fn from_arg_value(value: &str) -> Result { + match value { + "rank" => Ok(Self::Rank), + "ternary" => Ok(Self::Ternary), + other => Err(format!( + "Unknown MBP strategy: {other}. Use \"rank\" or \"ternary\"." + )), + } + } +} + +/// Two matrices per variable in r_0, c_0, r_1, c_1, ... order. +/// +/// Their product M_0[z_0] ... M_(2K-1)[z_(2K-1)] is the 3-by-1 +/// coefficient triple at (row, column). Before minimization, internal +/// dimensions contain one value per BDD node and one shared constant-one value +/// at coordinate zero. Minimization may change this basis and remove redundant +/// coordinates. +#[derive(Debug)] +pub(super) struct MatrixBranchingProgram { + pub(super) index_bits: usize, + pub(super) layers: Vec>, + /// With no variables, the 3-by-1 product is this constant column instead. + pub(super) constant: Option<[F; 3]>, +} + +/// Each branch independently chooses sparse or dense storage. +#[derive(Debug)] +pub(super) struct MbpLayer { + pub(super) matrices: [Core; 2], +} + +pub(super) fn extract( + r1cs: &R1CS, + strategy: MinimizeStrategy, +) -> MatrixBranchingProgram { + let mut program = extract_cqbdd(r1cs).to_mbp(); + info!( + "Extracted an MBP for {} index bits with {} layers (constant: {}).", + program.index_bits, + program.layers.len(), + program.constant.is_some() + ); + program.log_stats("MBP"); + match strategy { + MinimizeStrategy::Rank => { + program.minimize(); + program.log_stats("Minimized MBP"); + program.rescale(); + program.log_stats("Rescaled MBP"); + } + MinimizeStrategy::Ternary => { + program.minimize_ternary(); + program.log_stats("Ternary MBP"); + // Diagonal rescaling can introduce nonternary internal entries. + } + } + let dense_cores = program + .layers + .iter() + .flat_map(|layer| &layer.matrices) + .filter(|core| matches!(core, Core::Dense(_))) + .count(); + info!( + "Optimized MBP ({strategy:?}) storage: {} sparse, {} dense cores.", + 2 * program.layers.len() - dense_cores, + dense_cores + ); + program +} + +impl MatrixBranchingProgram { + fn log_stats(&self, label: &str) { + if !tracing::enabled!(tracing::Level::INFO) { + return; + } + let dimensions = if self.constant.is_some() { + vec![3, 1] + } else { + std::iter::once(3) + .chain(self.layers.iter().map(|layer| layer.matrices[0].num_cols())) + .collect::>() + }; + info!("{} dimensions: {:?}", label, dimensions); + let mut total = [0; 3]; + for (level, layer) in self.layers.iter().enumerate() { + let counts = layer.matrices.each_ref().map(Core::value_counts); + for branch in counts { + for (total, count) in total.iter_mut().zip(branch) { + *total += count; + } + } + info!( + "{} level {} ({} × {}) nonzero values [1, -1, other]: branch 0 {:?}, branch 1 {:?}", + label, + level, + layer.matrices[0].num_rows(), + layer.matrices[0].num_cols(), + counts[0], + counts[1] + ); + } + if let Some(constant) = &self.constant { + total = matrix::value_counts(constant.iter().copied()); + } + info!("{} total nonzero values [1, -1, other]: {:?}", label, total); + } +} + +/// A shared, fully layered BDD in interleaved LSB order. +/// +/// Terminal ID 1 represents true, and -1 represents false. Nonterminal ID `v` +/// indexes `nodes[v - 2]`. A negative reference complements the entire node. +#[derive(Debug)] +pub(super) struct Cqbdd { + pub(super) index_bits: usize, + pub(super) nodes: Vec, + pub(super) roots: HashMap<(F, F, F), isize>, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub(super) struct CqbddNode { + /// Zero-based variable level: r_0, c_0, r_1, c_1, ... + pub(super) level: usize, + pub(super) low_complemented: bool, + pub(super) low: usize, + pub(super) high: usize, +} + +/// Extract one CQBDD root for each nonzero coefficient triple. +fn extract_cqbdd(r1cs: &R1CS) -> Cqbdd { + let entries = collect_entries(r1cs); + + let total_nonzero = entries.values().map(|v| v.len()).sum::(); + info!( + "Collected {} unique triplets among {}.", + entries.len(), + total_nonzero + ); + + // Construct a CQBDD for all index sets. Use a square power-of-two domain + // covering both matrix dimensions, including any unused rows or columns. + let dimension = r1cs.num_constraints().max(r1cs.num_witnesses()); + let index_bits = (usize::BITS - dimension.saturating_sub(1).leading_zeros()) as usize; + let mut builder = CqbddBuilder::new(index_bits); + let roots = entries + .into_iter() + .map(|(coefficients, mut indices)| (coefficients, builder.build(0, &mut indices))) + .collect(); + let graph = Cqbdd { + index_bits, + nodes: builder.nodes, + roots, + }; + info!( + "Extracted {} CQBDD nodes across {} layers for {} roots.", + graph.nodes.len(), + 2 * graph.index_bits, + graph.roots.len() + ); + let widths = (0..2 * graph.index_bits) + .map(|level| { + graph + .nodes + .iter() + .filter(|node| node.level == level) + .count() + }) + .collect::>(); + info!("Layer widths: {:?}", widths); + graph +} + +impl Cqbdd { + fn to_mbp(&self) -> MatrixBranchingProgram { + let num_variables = 2 * self.index_bits; + let mut program = MatrixBranchingProgram { + index_bits: self.index_bits, + layers: Vec::with_capacity(num_variables), + constant: None, + }; + if num_variables == 0 { + let mut values = [F::ZERO; 3]; + for (&(a, b, c), &root) in &self.roots { + if root == 1 { + for (value, coefficient) in values.iter_mut().zip([a, b, c]) { + *value += coefficient; + } + } + } + program.constant = Some(values); + return program; + } + + let mut levels = vec![Vec::new(); num_variables]; + let mut columns = vec![0; self.nodes.len()]; + for (offset, node) in self.nodes.iter().enumerate() { + // Coordinate zero holds the one additional value for complements. + columns[offset] = levels[node.level].len() + 1; + levels[node.level].push(offset); + } + let child_column = |child: usize| { + if child == 1 { + 0 + } else { + columns[child - 2] + } + }; + for level in 0..num_variables { + let mut interner = Interner::new(); + let one = interner.intern(F::ONE); + let minus_one = interner.intern(-F::ONE); + // At the terminal, the node value and the extra constant coincide. + let num_cols = levels.get(level + 1).map_or(1, |nodes| nodes.len() + 1); + let matrices = std::array::from_fn(|bit| { + let mut matrix = SparseMatrix::new(0, num_cols); + if level == 0 { + // Absorb the coefficient-weighted signed roots into this + // first pair, so the initial dimension is exactly three. + let mut values = vec![[F::ZERO; 3]; num_cols]; + for (&(a, b, c), &root) in &self.roots { + let node = &self.nodes[root.unsigned_abs() - 2]; + let child = if bit == 0 { node.low } else { node.high }; + let column = child_column(child); + let complemented = (root < 0) ^ (bit == 0 && node.low_complemented); + for (component, coefficient) in [a, b, c].into_iter().enumerate() { + if complemented { + values[0][component] += coefficient; + values[column][component] -= coefficient; + } else { + values[column][component] += coefficient; + } + } + } + for component in 0..3 { + matrix.push_row(values.iter().enumerate().filter_map( + |(column, values)| { + let value = values[component]; + (!value.is_zero()).then(|| { + ( + u32::try_from(column) + .expect("MBP column index exceeds u32::MAX"), + interner.intern(value), + ) + }) + }, + )); + } + } else { + // Preserve the constant-one coordinate in both branches. + matrix.push_row([(0, one)].into_iter()); + for &offset in &levels[level] { + let node = &self.nodes[offset]; + let child = if bit == 0 { node.low } else { node.high }; + let column = child_column(child); + let complemented = bit == 0 && node.low_complemented; + append_branch_row(&mut matrix, column, complemented, one, minus_one); + } + } + matrix + }); + program + .layers + .push(MbpLayer::from_sparse(matrices, interner)); + } + program + } +} + +fn append_branch_row( + matrix: &mut SparseMatrix, + column: usize, + complemented: bool, + one: InternedFieldElement, + minus_one: InternedFieldElement, +) { + let column = u32::try_from(column).expect("MBP column index exceeds u32::MAX"); + match (column, complemented) { + // The complement of terminal one is zero. + (0, true) => matrix.push_row(std::iter::empty()), + (_, false) => matrix.push_row([(column, one)].into_iter()), + (_, true) => matrix.push_row([(0, one), (column, minus_one)].into_iter()), + } +} + +struct CqbddBuilder { + num_variables: usize, + nodes: Vec, + unique_tables: Vec>>, + ones: Vec>, +} + +impl CqbddBuilder { + fn new(index_bits: usize) -> Self { + assert!(index_bits <= usize::BITS as usize); + let num_variables = 2 * index_bits; + let mut ones = vec![None; num_variables + 1]; + ones[num_variables] = Some(1); + Self { + num_variables, + nodes: Vec::new(), + unique_tables: (0..num_variables).map(|_| None).collect(), + ones, + } + } + + fn make(&mut self, level: usize, low: isize, high: isize) -> isize { + // Normalize the high edge to positive, moving its complement to the + // returned reference and toggling the low complement when necessary. + let node = CqbddNode { + level, + low_complemented: (low < 0) ^ (high < 0), + low: low.unsigned_abs(), + high: high.unsigned_abs(), + }; + let table = self.unique_tables[level].get_or_insert_with(HashMap::new); + let id = *table.entry(node).or_insert_with(|| { + let id = + isize::try_from(self.nodes.len() + 2).expect("CQBDD node IDs exceed isize::MAX"); + // Keep equal-child nodes: every edge must advance exactly one level. + self.nodes.push(node); + id + }); + if high < 0 { + -id + } else { + id + } + } + + fn one(&mut self, level: usize) -> isize { + if let Some(id) = self.ones[level] { + return id; + } + let child = self.one(level + 1); + let id = self.make(level, child, child); + self.ones[level] = Some(id); + id + } + + fn build(&mut self, level: usize, indices: &mut [(usize, usize)]) -> isize { + if indices.is_empty() { + return -self.one(level); + } + if level == self.num_variables { + return 1; + } + + // Partition in place by r_(level/2) or c_(level/2), LSB first. + let mut split = 0; + for index in 0..indices.len() { + let (row, column) = indices[index]; + let coordinate = if level % 2 == 0 { row } else { column }; + if (coordinate >> (level / 2)) & 1 == 0 { + indices.swap(split, index); + split += 1; + } + } + let (low_indices, high_indices) = indices.split_at_mut(split); + let low = self.build(level + 1, low_indices); + let high = self.build(level + 1, high_indices); + self.make(level, low, high) + } +} + +fn collect_entries(r1cs: &R1CS) -> HashMap<(F, F, F), Vec<(usize, usize)>> { + let matrices = [r1cs.a(), r1cs.b(), r1cs.c()]; + let mut entries = matrices.each_ref().map(|matrix| matrix.iter().peekable()); + let mut triples = HashMap::<_, Vec<_>>::new(); + + // Merge the three sorted streams, visiting each occupied position once. + loop { + let Some(index) = entries + .iter_mut() + .filter_map(|entries| entries.peek().map(|&(index, _)| index)) + .min() + else { + break; + }; + + let [a, b, c] = entries.each_mut().map(|entries| { + entries + .next_if(|&(entry_index, _)| entry_index == index) + .map_or(F::ZERO, |(_, value)| value) + }); + + if !a.is_zero() || !b.is_zero() || !c.is_zero() { + triples.entry((a, b, c)).or_default().push(index); + } + } + + triples +} + +#[cfg(test)] +mod tests { + use {super::*, provekit_backend_bn254::FieldElement, std::collections::HashSet}; + + fn evaluate_mbp( + program: &MatrixBranchingProgram, + row: usize, + column: usize, + ) -> [FieldElement; 3] { + if let Some(values) = program.constant { + assert!(program.layers.is_empty()); + return values; + } + let mut values = vec![FieldElement::from(1u64)]; + for (level, layer) in program.layers.iter().enumerate().rev() { + let coordinate = if level % 2 == 0 { row } else { column }; + let matrix = &layer.matrices[(coordinate >> (level / 2)) & 1]; + assert_eq!(matrix.num_cols(), values.len()); + let mut next = vec![FieldElement::from(0u64); matrix.num_rows()]; + for ((row, column), coefficient) in matrix.iter() { + next[row] += coefficient * values[column]; + } + values = next; + } + values.try_into().unwrap() + } + + fn assert_mbp_dimensions(graph: &Cqbdd, program: &MatrixBranchingProgram) { + assert_eq!(program.index_bits, graph.index_bits); + assert_eq!(program.layers.len(), 2 * graph.index_bits); + let mut num_rows = 3; + for (level, layer) in program.layers.iter().enumerate() { + let num_cols = if level + 1 == program.layers.len() { + 1 + } else { + graph + .nodes + .iter() + .filter(|node| node.level == level + 1) + .count() + + 1 + }; + for matrix in &layer.matrices { + assert_eq!((matrix.num_rows(), matrix.num_cols()), (num_rows, num_cols)); + for ((..), value) in matrix.iter() { + assert!(!value.is_zero()); + } + if level != 0 { + assert_eq!( + matrix + .iter() + .filter_map( + |((row, column), value)| (row == 0).then_some((column, value)) + ) + .collect::>(), + [(0, F::ONE)] + ); + for row in 1..matrix.num_rows() { + assert!( + matrix + .iter() + .filter(|((entry_row, _), _)| *entry_row == row) + .count() + <= 2 + ); + } + } + } + num_rows = num_cols; + } + } + + fn evaluate(nodes: &[CqbddNode], mut root: isize, row: usize, column: usize) -> bool { + let mut complemented = false; + let mut bits = [row, column]; + let mut level = 0; + while root.unsigned_abs() != 1 { + complemented ^= root < 0; + let node = &nodes[root.unsigned_abs() - 2]; + assert_eq!(node.level, level); + let high = bits[level % 2] & 1 != 0; + bits[level % 2] >>= 1; + root = if high { + node.high as isize + } else if node.low_complemented { + -(node.low as isize) + } else { + node.low as isize + }; + level += 1; + } + (root > 0) ^ complemented + } + + fn assert_layered(nodes: &[CqbddNode], num_variables: usize) { + let mut unique = HashSet::new(); + for (offset, node) in nodes.iter().enumerate() { + assert!(unique.insert(*node), "duplicate node: {node:?}"); + assert!(node.level < num_variables); + for child in [node.low, node.high] { + assert!((1..offset + 2).contains(&child)); + if node.level + 1 == num_variables { + assert_eq!(child, 1); + } else { + assert!(child > 1); + assert_eq!(nodes[child - 2].level, node.level + 1); + } + } + } + } + + #[test] + fn exhaustive_four_by_four_index_sets() { + let mut builder = CqbddBuilder::new(2); + let mut roots = Vec::new(); + // Exercise every Boolean function on the four interleaved variables. + for mask in 0..=u16::MAX { + let mut indices: Vec<_> = (0..16) + .filter(|bit| mask & (1 << bit) != 0) + .map(|bit| (bit / 4, bit % 4)) + .collect(); + roots.push(builder.build(0, &mut indices)); + } + + for mask in 0..=u16::MAX { + let root = roots[mask as usize]; + assert_eq!(root, -roots[(!mask) as usize]); + for row in 0..4 { + for column in 0..4 { + assert_eq!( + evaluate(&builder.nodes, root, row, column), + mask & (1 << (4 * row + column)) != 0, + "mask {mask:#06x}, position ({row}, {column})" + ); + } + } + } + assert_layered(&builder.nodes, 4); + } + + #[test] + fn constants_keep_every_layer_and_share_nodes() { + let mut builder = CqbddBuilder::new(3); + let empty = builder.build(0, &mut []); + let mut full: Vec<_> = (0..8).flat_map(|r| (0..8).map(move |c| (r, c))).collect(); + let full = builder.build(0, &mut full); + assert_eq!(full, -empty); + assert_eq!(builder.nodes.len(), 6); + for node in &builder.nodes { + assert!(!node.low_complemented); + assert_eq!(node.low, node.high); + } + assert_layered(&builder.nodes, 6); + } + + #[test] + fn input_order_does_not_change_the_root() { + let mut builder = CqbddBuilder::new(3); + let mut indices = [(7, 0), (0, 5), (3, 2), (1, 7), (4, 4)]; + let root = builder.build(0, &mut indices); + let node_count = builder.nodes.len(); + indices.reverse(); + assert_eq!(root, builder.build(0, &mut indices)); + assert_eq!(builder.nodes.len(), node_count); + } + + #[test] + fn zero_bit_and_full_word_domains() { + let mut builder = CqbddBuilder::new(0); + assert_eq!(builder.build(0, &mut []), -1); + assert_eq!(builder.build(0, &mut [(0, 0)]), 1); + assert!(builder.nodes.is_empty()); + + let mut builder = CqbddBuilder::new(usize::BITS as usize); + let root = builder.build(0, &mut [(usize::MAX, 1)]); + assert!(evaluate(&builder.nodes, root, usize::MAX, 1)); + assert!(!evaluate(&builder.nodes, root, usize::MAX >> 1, 1)); + assert!(!evaluate(&builder.nodes, root, usize::MAX, 0)); + assert_layered(&builder.nodes, 2 * usize::BITS as usize); + + let coefficients = ( + FieldElement::from(2u64), + FieldElement::from(0u64), + -FieldElement::from(3u64), + ); + let graph = Cqbdd { + index_bits: usize::BITS as usize, + nodes: builder.nodes, + roots: HashMap::from([(coefficients, root)]), + }; + let program = graph.to_mbp(); + assert_mbp_dimensions(&graph, &program); + assert_eq!(evaluate_mbp(&program, usize::MAX, 1), [ + coefficients.0, + coefficients.1, + coefficients.2 + ]); + assert_eq!( + evaluate_mbp(&program, usize::MAX >> 1, 1), + [FieldElement::from(0u64); 3] + ); + } + + #[test] + fn mbp_combines_complementary_roots_with_one_extra_coordinate() { + let coefficients = [ + ( + FieldElement::from(0u64), + FieldElement::from(5u64), + -FieldElement::from(7u64), + ), + ( + -FieldElement::from(3u64), + FieldElement::from(0u64), + FieldElement::from(13u64), + ), + ]; + for mask in 0u8..16 { + let mut builder = CqbddBuilder::new(1); + let mut indices: Vec<_> = (0..4) + .filter(|bit| mask & (1 << bit) != 0) + .map(|bit| (bit / 2, bit % 2)) + .collect(); + let root = builder.build(0, &mut indices); + let graph = Cqbdd { + index_bits: 1, + nodes: builder.nodes, + roots: HashMap::from([(coefficients[0], root), (coefficients[1], -root)]), + }; + let program = graph.to_mbp(); + assert_mbp_dimensions(&graph, &program); + for row in 0..2 { + for column in 0..2 { + let index = usize::from(mask & (1 << (2 * row + column)) == 0); + let (a, b, c) = coefficients[index]; + assert_eq!( + evaluate_mbp(&program, row, column), + [a, b, c], + "mask {mask}" + ); + } + } + } + } + + #[test] + fn mbp_handles_constant_and_empty_systems() { + let zero = FieldElement::from(0u64); + let one = FieldElement::from(1u64); + let two = FieldElement::from(2u64); + let mut r1cs = R1CS::new(); + let program = extract(&r1cs, MinimizeStrategy::Rank); + assert_eq!(program.constant, Some([zero; 3])); + assert_eq!(evaluate_mbp(&program, 0, 0), [zero; 3]); + + r1cs.add_witnesses(1); + r1cs.add_constraint(&[(one, 0)], &[(two, 0)], &[(-one, 0)]); + let program = extract(&r1cs, MinimizeStrategy::Rank); + assert_eq!(program.index_bits, 0); + assert_eq!(program.constant, Some([one, two, -one])); + assert_eq!(evaluate_mbp(&program, 0, 0), [one, two, -one]); + + let mut r1cs = R1CS::::new(); + r1cs.add_witnesses(5); + let graph = extract_cqbdd(&r1cs); + let program = graph.to_mbp(); + assert_mbp_dimensions(&graph, &program); + for row in 0..8 { + for column in 0..8 { + assert_eq!(evaluate_mbp(&program, row, column), [zero; 3]); + } + } + } + + #[test] + fn extraction_preserves_coefficients_and_pads_the_domain() { + let one = FieldElement::from(1u64); + let two = FieldElement::from(2u64); + let zero = FieldElement::from(0u64); + let mut r1cs = R1CS::new(); + r1cs.add_witnesses(5); + r1cs.add_constraint(&[(one, 0), (two, 2)], &[(one, 0)], &[(one, 4)]); + r1cs.add_constraint(&[(one, 1)], &[(one, 1), (two, 3)], &[(two, 3)]); + r1cs.add_constraint(&[], &[], &[]); + + let graph = extract_cqbdd(&r1cs); + let expected = [ + ((one, one, zero), vec![(0, 0), (1, 1)]), + ((two, zero, zero), vec![(0, 2)]), + ((zero, zero, one), vec![(0, 4)]), + ((zero, two, two), vec![(1, 3)]), + ]; + assert_eq!(graph.index_bits, 3); + assert_eq!(graph.roots.len(), expected.len()); + for (coefficients, indices) in &expected { + let root = graph.roots[coefficients]; + for row in 0..8 { + for column in 0..8 { + assert_eq!( + evaluate(&graph.nodes, root, row, column), + indices.contains(&(row, column)) + ); + } + } + } + assert_layered(&graph.nodes, 6); + + let program = graph.to_mbp(); + assert_mbp_dimensions(&graph, &program); + for row in 0..8 { + for column in 0..8 { + let &(a, b, c) = expected + .iter() + .find(|(_, indices)| indices.contains(&(row, column))) + .map_or(&(zero, zero, zero), |(coefficients, _)| coefficients); + assert_eq!(evaluate_mbp(&program, row, column), [a, b, c]); + } + } + + let graph = extract_cqbdd(&R1CS::::new()); + assert_eq!(graph.index_bits, 0); + assert!(graph.roots.is_empty()); + assert!(graph.nodes.is_empty()); + } +} diff --git a/tooling/cli/src/cmd/extract/matrix.rs b/tooling/cli/src/cmd/extract/matrix.rs new file mode 100644 index 000000000..bf62a612c --- /dev/null +++ b/tooling/cli/src/cmd/extract/matrix.rs @@ -0,0 +1,1091 @@ +use { + super::MbpLayer, + ark_ff::Field, + provekit_common::{InternedFieldElement, Interner, SparseMatrix}, + std::{ + collections::HashSet, + ops::{Index, IndexMut}, + sync::Arc, + }, +}; + +/// Each core chooses its own storage when it is produced. Sparse cores own +/// their coefficients; the two original BDD branches can share an interner. +#[derive(Debug)] +pub(in crate::cmd) enum Core { + Sparse { + matrix: SparseMatrix, + interner: Arc>, + }, + Direct(DirectSparseMatrix), + Dense(DenseMatrix), +} + +/// CSR with direct field values. No reverse dictionary or value-index array. +#[derive(Debug)] +pub(in crate::cmd) struct DirectSparseMatrix { + pub(super) num_cols: usize, + pub(super) offsets: Vec, + pub(super) columns: Vec, + pub(super) values: Vec, +} + +impl DirectSparseMatrix { + pub(super) fn new(num_cols: usize) -> Self { + Self { + num_cols, + offsets: vec![0], + columns: Vec::new(), + values: Vec::new(), + } + } + + pub(super) fn num_rows(&self) -> usize { + self.offsets.len() - 1 + } + + pub(super) fn push_row(&mut self, entries: impl Iterator) { + for (column, value) in entries { + self.columns + .push(u32::try_from(column).expect("MBP column index exceeds u32::MAX")); + self.values.push(value); + } + self.offsets.push(self.values.len()); + } + + pub(super) fn iter_row(&self, row: usize) -> impl Iterator + '_ { + let range = self.offsets[row]..self.offsets[row + 1]; + self.columns[range.clone()] + .iter() + .copied() + .zip(self.values[range].iter().copied()) + .map(|(column, value)| (column as usize, value)) + } + + fn iter(&self) -> impl Iterator + '_ { + (0..self.num_rows()).flat_map(|row| { + self.iter_row(row) + .map(move |(column, value)| ((row, column), value)) + }) + } +} + +impl Core { + pub(super) fn from_direct(mut direct: DirectSparseMatrix) -> Self { + let num_rows = direct.num_rows(); + let nnz = direct.values.len(); + let dense_bytes = num_rows + .saturating_mul(direct.num_cols) + .saturating_mul(size_of::()); + let direct_bytes = direct + .offsets + .capacity() + .saturating_mul(size_of::()) + .saturating_add(direct.columns.capacity().saturating_mul(size_of::())) + .saturating_add(direct.values.capacity().saturating_mul(size_of::())); + if Self::prefer_sparse( + num_rows, + nnz, + dense_bytes.min(direct_bytes), + direct.values.iter().copied(), + ) { + let mut interner = Interner::new(); + let values = direct + .values + .into_iter() + .map(|value| interner.intern(value)) + .collect(); + direct.offsets.pop(); + let offsets = direct + .offsets + .into_iter() + .map(|offset| u32::try_from(offset).expect("too many sparse entries")) + .collect(); + return Self::Sparse { + matrix: SparseMatrix::from_raw_parts( + num_rows, + direct.num_cols, + offsets, + direct.columns, + values, + ), + interner: Arc::new(interner), + }; + } + if direct_bytes < dense_bytes { + return Self::Direct(direct); + } + if nnz == num_rows.saturating_mul(direct.num_cols) { + return Self::Dense(DenseMatrix { + num_rows, + num_cols: direct.num_cols, + values: direct.values, + }); + } + let mut dense = DenseMatrix::zeros(num_rows, direct.num_cols); + for ((row, column), value) in direct.iter() { + dense[(row, column)] = value; + } + Self::Dense(dense) + } + + pub(super) fn num_rows(&self) -> usize { + match self { + Self::Sparse { matrix, .. } => matrix.num_rows, + Self::Direct(matrix) => matrix.num_rows(), + Self::Dense(matrix) => matrix.num_rows, + } + } + + pub(super) fn num_cols(&self) -> usize { + match self { + Self::Sparse { matrix, .. } => matrix.num_cols, + Self::Direct(matrix) => matrix.num_cols, + Self::Dense(matrix) => matrix.num_cols, + } + } + + pub(super) fn nnz(&self) -> usize { + match self { + Self::Sparse { matrix, .. } => matrix.num_entries(), + Self::Direct(matrix) => matrix.values.len(), + Self::Dense(matrix) => matrix + .values + .iter() + .filter(|value| !value.is_zero()) + .count(), + } + } + + /// Counts nonzero entries in the order [one, minus one, other]. + pub(super) fn value_counts(&self) -> [usize; 3] { + match self { + Self::Sparse { matrix, interner } => value_counts( + matrix + .iter() + .map(|(_, value)| interner.get(value).expect("invalid core coefficient")), + ), + Self::Direct(matrix) => value_counts(matrix.values.iter().copied()), + Self::Dense(matrix) => value_counts(matrix.values.iter().copied()), + } + } + + pub(super) fn is_sparse(&self) -> bool { + matches!(self, Self::Sparse { .. } | Self::Direct(_)) + } + + /// Multiply columns (or rows) by nonzero diagonal factors, preserving + /// support. + pub(super) fn scale_axis(&mut self, factors: &[F], columns: bool) { + assert_eq!( + factors.len(), + if columns { + self.num_cols() + } else { + self.num_rows() + } + ); + assert!(factors.iter().all(|factor| !factor.is_zero())); + if factors.iter().all(|factor| *factor == F::ONE) { + return; + } + match self { + Self::Sparse { matrix, interner } => { + let mut direct = DirectSparseMatrix::new(matrix.num_cols); + direct.columns.reserve(matrix.num_entries()); + direct.values.reserve(matrix.num_entries()); + direct.offsets.reserve(matrix.num_rows); + for row in 0..matrix.num_rows { + direct.push_row(matrix.iter_row(row).map(|(column, value)| { + let value = interner.get(value).expect("invalid core coefficient"); + let factor = factors[if columns { column } else { row }]; + ( + column, + if factor == F::ONE { + value + } else { + value * factor + }, + ) + })); + } + *self = Self::from_direct(direct); + } + Self::Direct(matrix) => { + if columns { + for (&column, value) in matrix.columns.iter().zip(&mut matrix.values) { + let factor = factors[column as usize]; + if factor != F::ONE { + *value *= factor; + } + } + } else { + for (row, &factor) in factors.iter().enumerate() { + if factor != F::ONE { + for value in + &mut matrix.values[matrix.offsets[row]..matrix.offsets[row + 1]] + { + *value *= factor; + } + } + } + } + } + Self::Dense(matrix) => { + for row in 0..matrix.num_rows { + for (column, value) in matrix.values + [row * matrix.num_cols..(row + 1) * matrix.num_cols] + .iter_mut() + .enumerate() + { + let factor = factors[if columns { column } else { row }]; + if factor != F::ONE && !value.is_zero() { + *value *= factor; + } + } + } + } + } + } + + pub(super) fn iter_row(&self, row: usize) -> impl Iterator + '_ { + let (sparse, direct, dense) = match self { + Self::Sparse { matrix, interner } => (Some((matrix, interner)), None, None), + Self::Direct(matrix) => (None, Some(matrix), None), + Self::Dense(matrix) => (None, None, Some(matrix)), + }; + sparse + .into_iter() + .flat_map(move |(matrix, interner)| { + matrix.iter_row(row).map(move |(column, value)| { + (column, interner.get(value).expect("invalid coefficient")) + }) + }) + .chain( + direct + .into_iter() + .flat_map(move |matrix| matrix.iter_row(row)), + ) + .chain(dense.into_iter().flat_map(move |matrix| { + matrix.values[row * matrix.num_cols..(row + 1) * matrix.num_cols] + .iter() + .copied() + .enumerate() + .filter(|(_, value)| !value.is_zero()) + })) + } + + /// Nonzero entries, independent of the chosen representation. + pub(super) fn iter(&self) -> impl Iterator + '_ { + let (sparse, direct, dense) = match self { + Self::Sparse { matrix, interner } => (Some((matrix, interner)), None, None), + Self::Direct(matrix) => (None, Some(matrix), None), + Self::Dense(matrix) => (None, None, Some(matrix)), + }; + sparse + .into_iter() + .flat_map(|(matrix, interner)| { + matrix.iter().map(move |(index, value)| { + ( + index, + interner.get(value).expect("invalid core coefficient"), + ) + }) + }) + .chain(direct.into_iter().flat_map(|matrix| matrix.iter())) + .chain(dense.into_iter().flat_map(|matrix| { + matrix + .values + .iter() + .enumerate() + .filter_map(move |(index, &value)| { + (!value.is_zero()) + .then(|| ((index / matrix.num_cols, index % matrix.num_cols), value)) + }) + })) + } + + /// Select storage by estimated heap bytes, including coefficient diversity. + /// The reverse-map allowance is conservative, not an exact allocator model. + /// Dense inputs retain their buffer. No existing core is converted in + /// place. + pub(super) fn from_dense(dense: DenseMatrix) -> Self { + let nnz = dense.values.iter().filter(|value| !value.is_zero()).count(); + let dense_bytes = dense.values.capacity().saturating_mul(size_of::()); + let direct_bytes = Self::direct_bytes(dense.num_rows, nnz); + if Self::prefer_sparse( + dense.num_rows, + nnz, + dense_bytes.min(direct_bytes), + dense.values.iter().copied(), + ) { + Self::encode_sparse(dense, nnz) + } else if direct_bytes < dense_bytes { + let mut direct = DirectSparseMatrix::new(dense.num_cols); + direct.columns.reserve(nnz); + direct.values.reserve(nnz); + for row in 0..dense.num_rows { + direct.push_row((0..dense.num_cols).filter_map(|column| { + let value = dense[(row, column)]; + (!value.is_zero()).then_some((column, value)) + })); + } + Self::Direct(direct) + } else { + Self::Dense(dense) + } + } + + fn prefer_sparse( + num_rows: usize, + nnz: usize, + dense_bytes: usize, + values: impl Iterator, + ) -> bool { + let sparse_base = num_rows + .saturating_mul(size_of::()) + .saturating_add( + nnz.saturating_mul(size_of::() + size_of::()), + ) + .saturating_add(size_of::>() + 2 * size_of::()); + if sparse_base >= dense_bytes { + return false; + } + // One stored field value plus a reverse-map entry with slack for hash + // table capacity and control bytes. Repeated +/-1 values cost very little. + let per_distinct = size_of::() + 2 * size_of::<(F, usize)>(); + // Bound dictionary probing even when the logical dense matrix is huge. + // High-diversity factors use direct CSR instead of multi-GB hash tables. + let max_distinct = ((dense_bytes - sparse_base) / per_distinct).min(4096); + let mut distinct = HashSet::new(); + for value in values { + if !value.is_zero() && distinct.insert(value) && distinct.len() > max_distinct { + return false; + } + } + true + } + + fn direct_bytes(num_rows: usize, nnz: usize) -> usize { + (num_rows + 1) + .saturating_mul(size_of::()) + .saturating_add(nnz.saturating_mul(size_of::() + size_of::())) + } + + /// Sorted, duplicate-free nonzero rows can stay sparse throughout assembly. + pub(super) fn from_rows(num_cols: usize, rows: Vec>) -> Self { + let nnz = rows.iter().map(Vec::len).sum(); + let dense_bytes = rows + .len() + .saturating_mul(num_cols) + .saturating_mul(size_of::()); + let direct_bytes = Self::direct_bytes(rows.len(), nnz); + if !Self::prefer_sparse( + rows.len(), + nnz, + dense_bytes.min(direct_bytes), + rows.iter().flatten().map(|&(_, value)| value), + ) { + if direct_bytes < dense_bytes { + let mut direct = DirectSparseMatrix::new(num_cols); + direct.columns.reserve(nnz); + direct.values.reserve(nnz); + direct.offsets.reserve(rows.len()); + for row in rows { + direct.push_row(row.into_iter()); + } + return Self::Direct(direct); + } + let mut dense = DenseMatrix::zeros(rows.len(), num_cols); + for (row, entries) in rows.into_iter().enumerate() { + for (column, value) in entries { + dense[(row, column)] = value; + } + } + return Self::Dense(dense); + } + let mut interner = Interner::new(); + let mut matrix = SparseMatrix::new(0, num_cols); + matrix.reserve(rows.len(), nnz); + for entries in rows { + matrix.push_row(entries.into_iter().map(|(column, value)| { + ( + u32::try_from(column).expect("MBP column index exceeds u32::MAX"), + interner.intern(value), + ) + })); + } + Self::Sparse { + matrix, + interner: Arc::new(interner), + } + } + + pub(super) fn block( + &self, + row: usize, + num_rows: usize, + column: usize, + num_cols: usize, + ) -> Self { + match self { + Self::Dense(matrix) => Self::from_dense(matrix.block(row, num_rows, column, num_cols)), + Self::Direct(matrix) => { + let mut block = DirectSparseMatrix::new(num_cols); + for row in row..row + num_rows { + block.push_row( + matrix + .iter_row(row) + .filter(|&(col, _)| (column..column + num_cols).contains(&col)) + .map(|(col, value)| (col - column, value)), + ); + } + Self::Direct(block) + } + Self::Sparse { matrix, interner } => { + let mut block = SparseMatrix::new(0, num_cols); + for row in row..row + num_rows { + block.push_row( + matrix + .iter_row(row) + .filter(|&(col, _)| (column..column + num_cols).contains(&col)) + .map(|(col, value)| ((col - column) as u32, value)), + ); + } + Self::Sparse { + matrix: block, + interner: Arc::clone(interner), + } + } + } + } + + pub(super) fn multiply(&self, rhs: &Self) -> Self { + assert_eq!(self.num_cols(), rhs.num_rows()); + match (self, rhs) { + (lhs, rhs) if lhs.is_sparse() && rhs.is_sparse() => { + // Gustavson row accumulation: only one output row is dense. + let num_cols = rhs.num_cols(); + let mut values = vec![F::ZERO; num_cols]; + let mut seen = vec![false; num_cols]; + let mut touched = Vec::new(); + let mut rows = Vec::with_capacity(lhs.num_rows()); + for row in 0..self.num_rows() { + for (inner, coefficient) in lhs.iter_row(row) { + for (column, value) in rhs.iter_row(inner) { + if !seen[column] { + seen[column] = true; + touched.push(column); + } + values[column] += if coefficient == F::ONE { + value + } else if coefficient == -F::ONE { + -value + } else if value == F::ONE { + coefficient + } else if value == -F::ONE { + -coefficient + } else { + coefficient * value + }; + } + } + touched.sort_unstable(); + let mut entries = Vec::with_capacity(touched.len()); + for column in touched.drain(..) { + let value = std::mem::replace(&mut values[column], F::ZERO); + seen[column] = false; + if !value.is_zero() { + entries.push((column, value)); + } + } + rows.push(entries); + } + Self::from_rows(num_cols, rows) + } + (Self::Dense(lhs), _) => Self::from_dense(lhs.multiply_core(rhs)), + (_, Self::Dense(rhs)) => Self::from_dense(DenseMatrix::core_multiply(self, rhs)), + _ => unreachable!("all sparse combinations handled above"), + } + } + + fn encode_sparse(dense: DenseMatrix, nnz: usize) -> Self { + let mut interner = Interner::new(); + let mut matrix = SparseMatrix::new(0, dense.num_cols); + matrix.reserve(dense.num_rows, nnz); + for row in 0..dense.num_rows { + matrix.push_row((0..dense.num_cols).filter_map(|column| { + let value = dense[(row, column)]; + (!value.is_zero()).then(|| { + ( + u32::try_from(column).expect("MBP column index exceeds u32::MAX"), + interner.intern(value), + ) + }) + })); + } + Self::Sparse { + matrix, + interner: Arc::new(interner), + } + } + + #[cfg(test)] + pub(super) fn force_sparse(dense: DenseMatrix) -> Self { + let nnz = dense.values.iter().filter(|value| !value.is_zero()).count(); + Self::encode_sparse(dense, nnz) + } + + #[cfg(test)] + pub(super) fn force_direct(dense: DenseMatrix) -> Self { + let mut direct = DirectSparseMatrix::new(dense.num_cols); + for row in 0..dense.num_rows { + direct.push_row((0..dense.num_cols).filter_map(|column| { + let value = dense[(row, column)]; + (!value.is_zero()).then_some((column, value)) + })); + } + Self::Direct(direct) + } +} + +/// Counts nonzero entries in the order [one, minus one, other]. +pub(super) fn value_counts(values: impl Iterator) -> [usize; 3] { + let mut counts = [0; 3]; + let minus_one = -F::ONE; + for value in values { + if !value.is_zero() { + let bucket = if value == F::ONE { + 0 + } else if value == minus_one { + 1 + } else { + 2 + }; + counts[bucket] += 1; + } + } + counts +} + +impl MbpLayer { + #[cfg(test)] + pub(super) fn from_dense(matrices: [DenseMatrix; 2]) -> Self { + Self { + matrices: matrices.map(Core::from_dense), + } + } + + pub(super) fn from_sparse(matrices: [SparseMatrix; 2], interner: Interner) -> Self { + let interner = Arc::new(interner); + Self { + matrices: matrices.map(|matrix| Core::Sparse { + matrix, + interner: Arc::clone(&interner), + }), + } + } +} + +/// Dense workspaces and products. Stored cores independently use dense values, +/// direct CSR, or interned sparse entries. +#[derive(Clone, Debug, PartialEq, Eq)] +pub(in crate::cmd) struct DenseMatrix { + pub(super) num_rows: usize, + pub(super) num_cols: usize, + pub(super) values: Vec, +} + +impl DenseMatrix { + pub(super) fn zeros(num_rows: usize, num_cols: usize) -> Self { + Self { + num_rows, + num_cols, + values: vec![F::ZERO; num_rows * num_cols], + } + } + + #[cfg(test)] + pub(super) fn from_sparse(matrix: &SparseMatrix, interner: &Interner) -> Self { + let mut dense = Self::zeros(matrix.num_rows, matrix.num_cols); + for ((row, column), value) in matrix.hydrate(interner).iter() { + dense[(row, column)] = value; + } + dense + } + + pub(super) fn stack(layer: &MbpLayer, vertically: bool) -> Self { + let num_rows = layer.matrices[0].num_rows(); + let num_cols = layer.matrices[0].num_cols(); + let mut dense = if vertically { + Self::zeros(2 * num_rows, num_cols) + } else { + Self::zeros(num_rows, 2 * num_cols) + }; + for (bit, matrix) in layer.matrices.iter().enumerate() { + assert_eq!((matrix.num_rows(), matrix.num_cols()), (num_rows, num_cols)); + for ((row, column), value) in matrix.iter() { + let row = row + if vertically { bit * num_rows } else { 0 }; + let column = column + if vertically { 0 } else { bit * num_cols }; + dense[(row, column)] = value; + } + } + dense + } + + pub(super) fn multiply_core(&self, rhs: &Core) -> Self { + match rhs { + Core::Sparse { matrix, interner } => self.multiply_sparse(matrix, interner), + Core::Direct(matrix) => { + assert_eq!(self.num_cols, matrix.num_rows()); + let mut product = Self::zeros(self.num_rows, matrix.num_cols); + let minus_one = -F::ONE; + for row in 0..self.num_rows { + for inner in 0..self.num_cols { + let coefficient = self[(row, inner)]; + if coefficient.is_zero() { + continue; + } + for (column, value) in matrix.iter_row(inner) { + product[(row, column)] += if value == F::ONE { + coefficient + } else if value == minus_one { + -coefficient + } else { + coefficient * value + }; + } + } + } + product + } + Core::Dense(matrix) => self.multiply(matrix), + } + } + + pub(super) fn core_multiply(lhs: &Core, rhs: &Self) -> Self { + match lhs { + Core::Sparse { matrix, interner } => Self::sparse_multiply(matrix, interner, rhs), + Core::Direct(matrix) => { + assert_eq!(matrix.num_cols, rhs.num_rows); + let mut product = Self::zeros(matrix.num_rows(), rhs.num_cols); + let minus_one = -F::ONE; + for ((row, inner), coefficient) in matrix.iter() { + for column in 0..rhs.num_cols { + let value = rhs[(inner, column)]; + product[(row, column)] += if coefficient == F::ONE { + value + } else if coefficient == minus_one { + -value + } else { + coefficient * value + }; + } + } + product + } + Core::Dense(matrix) => matrix.multiply(rhs), + } + } + + pub(super) fn block( + &self, + row: usize, + num_rows: usize, + column: usize, + num_cols: usize, + ) -> Self { + let mut block = Self::zeros(num_rows, num_cols); + for i in 0..num_rows { + for j in 0..num_cols { + block[(i, j)] = self[(row + i, column + j)]; + } + } + block + } + + /// Dense R times a sparse core: O(rows(R) * nnz(core)) field operations, + /// allocating only the reduced output, not a dense copy of the sparse core. + pub(super) fn multiply_sparse(&self, rhs: &SparseMatrix, interner: &Interner) -> Self { + assert_eq!(self.num_cols, rhs.num_rows); + let mut product = Self::zeros(self.num_rows, rhs.num_cols); + if self.num_rows == 0 || rhs.num_cols == 0 || rhs.num_entries() == 0 { + return product; + } + let rhs = rhs.hydrate(interner); + let minus_one = -F::ONE; + for row in 0..self.num_rows { + for inner in 0..self.num_cols { + let coefficient = self[(row, inner)]; + if coefficient.is_zero() { + continue; + } + for (column, value) in rhs.iter_row(inner) { + // Original BDD transitions use only +/-1. Avoid general + // field multiplication for these common coefficients. + if value == F::ONE { + product[(row, column)] += coefficient; + } else if value == minus_one { + product[(row, column)] -= coefficient; + } else { + product[(row, column)] += coefficient * value; + } + } + } + } + product + } + + /// Sparse core times dense L: O(nnz(core) * cols(L)) field operations. + pub(super) fn sparse_multiply(lhs: &SparseMatrix, interner: &Interner, rhs: &Self) -> Self { + assert_eq!(lhs.num_cols, rhs.num_rows); + let mut product = Self::zeros(lhs.num_rows, rhs.num_cols); + if rhs.num_cols == 0 { + return product; + } + let minus_one = -F::ONE; + for ((row, inner), coefficient) in lhs.hydrate(interner).iter() { + if coefficient.is_zero() { + continue; + } + let target = &mut product.values[row * rhs.num_cols..(row + 1) * rhs.num_cols]; + let source = &rhs.values[inner * rhs.num_cols..(inner + 1) * rhs.num_cols]; + if coefficient == F::ONE { + for (target, source) in target.iter_mut().zip(source) { + *target += source; + } + } else if coefficient == minus_one { + for (target, source) in target.iter_mut().zip(source) { + *target -= source; + } + } else { + for (target, source) in target.iter_mut().zip(source) { + *target += coefficient * source; + } + } + } + product + } + + pub(super) fn multiply(&self, rhs: &Self) -> Self { + assert_eq!(self.num_cols, rhs.num_rows); + let mut product = Self::zeros(self.num_rows, rhs.num_cols); + for row in 0..self.num_rows { + for inner in 0..self.num_cols { + let coefficient = self[(row, inner)]; + if coefficient.is_zero() { + continue; + } + for column in 0..rhs.num_cols { + product[(row, column)] += coefficient * rhs[(inner, column)]; + } + } + } + product + } + + /// A = L R: take the original pivot columns as L and the nonzero RREF + /// rows as R. This also handles rank zero and zero-sized dimensions. + pub(super) fn rank_factor(self) -> (Self, Self) { + let mut reduced = self.clone(); + let mut pivots = Vec::new(); + for column in 0..self.num_cols { + let rank = pivots.len(); + let Some(pivot) = (rank..self.num_rows).find(|&row| !reduced[(row, column)].is_zero()) + else { + continue; + }; + for j in 0..self.num_cols { + reduced + .values + .swap(rank * self.num_cols + j, pivot * self.num_cols + j); + } + let inverse = reduced[(rank, column)].inverse().expect("nonzero pivot"); + for j in column..self.num_cols { + reduced[(rank, j)] *= inverse; + } + for row in 0..self.num_rows { + if row == rank { + continue; + } + let coefficient = reduced[(row, column)]; + if coefficient.is_zero() { + continue; + } + for j in column..self.num_cols { + let pivot_value = reduced[(rank, j)]; + reduced[(row, j)] -= coefficient * pivot_value; + } + } + pivots.push(column); + } + let mut left = Self::zeros(self.num_rows, pivots.len()); + for row in 0..self.num_rows { + for (j, &column) in pivots.iter().enumerate() { + left[(row, j)] = self[(row, column)]; + } + } + let right = reduced.block(0, pivots.len(), 0, self.num_cols); + (left, right) + } +} + +impl Index<(usize, usize)> for DenseMatrix { + type Output = F; + + fn index(&self, (row, column): (usize, usize)) -> &F { + &self.values[row * self.num_cols + column] + } +} + +impl IndexMut<(usize, usize)> for DenseMatrix { + fn index_mut(&mut self, (row, column): (usize, usize)) -> &mut F { + &mut self.values[row * self.num_cols + column] + } +} + +#[cfg(test)] +mod tests { + use {super::*, provekit_backend_bn254::FieldElement as F}; + + fn sparse_values() -> DenseMatrix { + let mut dense = DenseMatrix::zeros(64, 128); + for row in 0..64 { + dense[(row, row)] = F::from(1u64); + dense[(row, row + 64)] = F::from(-1i64); + } + dense + } + + fn unique_values() -> DenseMatrix { + let mut dense = DenseMatrix::zeros(64, 128); + for (index, value) in dense.values.iter_mut().enumerate() { + *value = F::from((index + 1) as u64); + } + dense + } + + fn assert_values(core: &Core, expected: &DenseMatrix) { + assert_eq!( + (core.num_rows(), core.num_cols()), + (expected.num_rows, expected.num_cols) + ); + let mut actual = DenseMatrix::zeros(core.num_rows(), core.num_cols()); + for ((row, column), value) in core.iter() { + actual[(row, column)] = value; + } + assert_eq!(&actual, expected); + } + + #[test] + fn sparse_and_diverse_dense_cores_choose_different_storage() { + let sparse = sparse_values(); + let core = Core::from_dense(sparse.clone()); + assert!(matches!(core, Core::Sparse { .. })); + assert_values(&core, &sparse); + + let dense = unique_values(); + let expected = dense.clone(); + let pointer = dense.values.as_ptr(); + let core = Core::from_dense(dense); + let Core::Dense(ref matrix) = core else { + panic!("expected dense storage") + }; + assert_eq!( + matrix.values.as_ptr(), + pointer, + "retain the original dense buffer" + ); + assert_values(&core, &expected); + } + + #[test] + fn diverse_sparse_values_use_direct_csr_without_a_dictionary() { + let rows: Vec<_> = (0..128) + .map(|row| { + (0..64) + .map(|column| (column * 32, F::from((row * 64 + column + 1) as u64))) + .collect() + }) + .collect(); + let core = Core::from_rows(4096, rows); + assert!(matches!(core, Core::Direct(_))); + assert_eq!(core.nnz(), 8192); + assert_eq!(core.iter_row(127).count(), 64); + let block = core.block(120, 8, 32, 128); + assert_eq!( + (block.num_rows(), block.num_cols(), block.nnz()), + (8, 128, 32) + ); + let mut expected = DenseMatrix::zeros(8, 128); + for row in 0..8 { + for column in 0..4 { + expected[(row, column * 32)] = F::from(((row + 120) * 64 + column + 2) as u64); + } + } + assert_values(&block, &expected); + } + + #[test] + fn coefficient_repetition_changes_the_choice_at_equal_occupancy() { + let dense = unique_values(); + let mut repeated = dense.clone(); + repeated.values.fill(F::from(1u64)); + assert!(matches!(Core::from_dense(dense), Core::Dense(_))); + let core = Core::from_dense(repeated.clone()); + assert!(matches!(core, Core::Sparse { .. })); + assert_values(&core, &repeated); + } + + #[test] + fn mixed_layers_stack_correctly_in_both_directions() { + let sparse = sparse_values(); + let dense = unique_values(); + let layer = MbpLayer::from_dense([sparse.clone(), dense.clone()]); + assert!(matches!(layer.matrices[0], Core::Sparse { .. })); + assert!(matches!(layer.matrices[1], Core::Dense(_))); + for vertically in [false, true] { + let stacked = DenseMatrix::stack(&layer, vertically); + for row in 0..64 { + for column in 0..128 { + assert_eq!(stacked[(row, column)], sparse[(row, column)]); + let index = if vertically { + (row + 64, column) + } else { + (row, column + 128) + }; + assert_eq!(stacked[index], dense[(row, column)]); + } + } + } + } + + #[test] + fn zero_shapes_and_all_zero_values_preserve_their_dimensions() { + for rows in [0, 1, 64] { + for columns in [0, 1, 128] { + let dense = DenseMatrix::::zeros(rows, columns); + let core = Core::from_dense(dense.clone()); + assert_values(&core, &dense); + assert!(core.iter().next().is_none()); + } + } + } + + #[test] + fn sparse_assembly_splitting_and_products_avoid_dense_storage() { + let rows: Vec<_> = (0..64) + .map(|row| vec![(row, F::from(1u64)), (row + 64, F::from(-1i64))]) + .collect(); + let core = Core::from_rows(128, rows); + let Core::Sparse { interner, .. } = &core else { + panic!("expected sparse assembly") + }; + let block = core.block(0, 64, 0, 64); + let Core::Sparse { + interner: block_interner, + .. + } = &block + else { + panic!("expected sparse block") + }; + assert!(Arc::ptr_eq(interner, block_interner)); + assert_eq!(block.nnz(), 64); + + let rows: Vec<_> = (0..128) + .map(|row| vec![(row, F::from(1u64)), (row + 128, F::from(1u64))]) + .collect(); + let rhs = Core::from_rows(4096, rows); + let product = core.multiply(&rhs); + assert!(matches!(product, Core::Sparse { .. })); + assert_eq!( + (product.num_rows(), product.num_cols(), product.nnz()), + (64, 4096, 256) + ); + for ((row, column), value) in product.iter() { + assert_eq!(column % 64, row); + assert_eq!( + value, + F::from(if (column / 64) % 2 == 0 { 1i64 } else { -1i64 }) + ); + } + } + + #[test] + fn sparse_product_removes_exact_cancellations() { + let lhs = Core::from_rows(3, vec![vec![ + (0, F::from(1u64)), + (1, F::from(1u64)), + (2, F::from(1u64)), + ]]); + let rhs = Core::from_rows(2, vec![ + vec![(0, F::from(1u64)), (1, F::from(2u64))], + vec![(0, F::from(-1i64)), (1, F::from(3u64))], + vec![(0, F::from(1u64)), (1, F::from(-5i64))], + ]); + // Force sparse storage even for this deliberately tiny example. + let to_sparse = |core: Core| { + let mut dense = DenseMatrix::zeros(core.num_rows(), core.num_cols()); + for ((row, column), value) in core.iter() { + dense[(row, column)] = value; + } + Core::force_sparse(dense) + }; + let product = to_sparse(lhs).multiply(&to_sparse(rhs)); + assert_eq!(product.iter().collect::>(), [( + (0, 0), + F::from(1u64) + )]); + } + + #[test] + #[ignore = "manual comparison of adaptive and forced core representations"] + fn benchmark_adaptive_storage() { + use std::{hint::black_box, time::Instant}; + + let mut left = DenseMatrix::zeros(4, 64); + let mut right = DenseMatrix::zeros(128, 4); + left.values.fill(F::from(3u64)); + right.values.fill(F::from(5u64)); + for (name, dense) in [ + ("sparse +/-1", sparse_values()), + ("dense distinct", unique_values()), + ] { + let start = Instant::now(); + let adaptive = Core::from_dense(dense.clone()); + let selection = start.elapsed(); + let sparse = Core::force_sparse(dense.clone()); + let dense = Core::Dense(dense); + for (representation, core) in [ + ("adaptive", &adaptive), + ("forced sparse", &sparse), + ("forced dense", &dense), + ] { + let start = Instant::now(); + let mut product = DenseMatrix::zeros(0, 0); + for _ in 0..100 { + product = black_box(&left).multiply_core(black_box(core)); + } + let forward = start.elapsed(); + assert_eq!(product, left.multiply_core(&dense)); + let start = Instant::now(); + for _ in 0..100 { + product = DenseMatrix::core_multiply(black_box(core), black_box(&right)); + } + let backward = start.elapsed(); + assert_eq!(product, DenseMatrix::core_multiply(&dense, &right)); + eprintln!( + "{name}, {representation}: 100 forward {forward:?}, 100 backward {backward:?}" + ); + } + eprintln!( + "{name}: selection {selection:?}, chose {}", + if matches!(adaptive, Core::Dense(_)) { + "dense" + } else { + "sparse" + } + ); + } + } +} diff --git a/tooling/cli/src/cmd/extract/minimize.rs b/tooling/cli/src/cmd/extract/minimize.rs new file mode 100644 index 000000000..2f7ce726c --- /dev/null +++ b/tooling/cli/src/cmd/extract/minimize.rs @@ -0,0 +1,469 @@ +use { + super::{MatrixBranchingProgram, MbpLayer}, + ark_ff::Field, + tracing::{info, instrument}, +}; + +impl MatrixBranchingProgram { + /// Minimize every internal width by exact forward and backward rank sweeps. + /// Boundary dimensions are preserved, including the three coefficient rows. + /// Uses O(k w^3) field operations and O(w^2) auxiliary field elements. + #[instrument(skip_all)] + pub(super) fn minimize(&mut self) { + // Establish a full-column-rank prefix at every cut. + for level in 0..self.layers.len().saturating_sub(1) { + info!("Forwards level {}/{}", level, self.layers.len()); + let before = tracing::enabled!(tracing::Level::INFO).then(|| { + [ + layer_stats(&self.layers[level]), + layer_stats(&self.layers[level + 1]), + ] + }); + let old_rank = self.layers[level].matrices[0].num_cols(); + let num_rows = self.layers[level].matrices[0].num_rows(); + let (left, right) = self.layers[level].rank_factor(true); + let rank = right.num_rows(); + info!("Rank {} → {}", old_rank, rank); + let current = std::array::from_fn(|bit| left.block(bit * num_rows, num_rows, 0, rank)); + let next = self.layers[level + 1] + .matrices + .each_ref() + .map(|matrix| right.multiply(matrix)); + self.layers[level] = MbpLayer { matrices: current }; + self.layers[level + 1] = MbpLayer { matrices: next }; + if let Some([current, next]) = before { + log_density(level, current, layer_stats(&self.layers[level])); + log_density(level + 1, next, layer_stats(&self.layers[level + 1])); + } + } + + // Remove directions invisible to the suffix at every cut. The forward + // sweep makes these suffix ranks equal to the full unfolding ranks. + for level in (1..self.layers.len()).rev() { + info!("Backwards level {}/{}", level, self.layers.len()); + let before = tracing::enabled!(tracing::Level::INFO).then(|| { + [ + layer_stats(&self.layers[level]), + layer_stats(&self.layers[level - 1]), + ] + }); + let old_rank = self.layers[level].matrices[0].num_rows(); + let num_cols = self.layers[level].matrices[0].num_cols(); + let (left, right) = self.layers[level].rank_factor(false); + let rank = left.num_cols(); + info!("Rank {} → {}", old_rank, rank); + let current = std::array::from_fn(|bit| right.block(0, rank, bit * num_cols, num_cols)); + let previous = self.layers[level - 1] + .matrices + .each_ref() + .map(|matrix| matrix.multiply(&left)); + self.layers[level] = MbpLayer { matrices: current }; + self.layers[level - 1] = MbpLayer { matrices: previous }; + if let Some([current, previous]) = before { + log_density(level, current, layer_stats(&self.layers[level])); + log_density(level - 1, previous, layer_stats(&self.layers[level - 1])); + } + } + } +} + +/// (Nonzero count, occupancy percentage) for branches zero and one. +fn layer_stats(layer: &MbpLayer) -> [(usize, f64); 2] { + layer.matrices.each_ref().map(|core| { + let nnz = core.nnz(); + let entries = core.num_rows() as f64 * core.num_cols() as f64; + let density = if entries == 0.0 { + 0.0 + } else { + 100.0 * nnz as f64 / entries + }; + (nnz, density) + }) +} + +fn log_density(level: usize, before: [(usize, f64); 2], after: [(usize, f64); 2]) { + info!( + "Level {} nnz ({}, {}) → ({}, {}); density ({:.2}%, {:.2}%) → ({:.2}%, {:.2}%)", + level, + before[0].0, + before[1].0, + after[0].0, + after[1].0, + before[0].1, + before[1].1, + after[0].1, + after[1].1, + ); +} + +#[cfg(test)] +mod tests { + use { + super::{ + super::{extract, extract_cqbdd, matrix::DenseMatrix, Core, MinimizeStrategy}, + *, + }, + provekit_backend_bn254::FieldElement as F, + provekit_common::{Interner, SparseMatrix, R1CS}, + tracing::instrument, + }; + + fn matrix(num_rows: usize, num_cols: usize, values: &[i64]) -> DenseMatrix { + assert_eq!(values.len(), num_rows * num_cols); + DenseMatrix { + num_rows, + num_cols, + values: values.iter().copied().map(F::from).collect(), + } + } + + #[test] + fn sparse_products_match_dense_products() { + for rows in [0, 1, 3, 5] { + for inner in [0, 1, 3, 5] { + for columns in [0, 1, 3, 5] { + for all_zero in [false, true] { + let mut lhs = DenseMatrix::zeros(rows, inner); + let mut rhs = DenseMatrix::zeros(inner, columns); + for (index, value) in lhs.values.iter_mut().enumerate() { + *value = F::from((index * 7 % 11) as i64 - 5); + } + if !all_zero { + for (index, value) in rhs.values.iter_mut().enumerate() { + // Covers zeros, +/-1, and general coefficients. + *value = F::from((index % 7) as i64 - 3); + } + } + let expected = lhs.multiply(&rhs); + let sparse_rhs = Core::force_sparse(rhs.clone()); + assert_eq!( + lhs.multiply_core(&sparse_rhs), + expected, + "dense x sparse: {rows} x {inner} x {columns}, zero={all_zero}" + ); + assert_eq!(lhs.multiply_core(&Core::Dense(rhs.clone())), expected); + let sparse_lhs = Core::force_sparse(lhs.clone()); + assert_eq!( + DenseMatrix::core_multiply(&sparse_lhs, &rhs), + expected, + "sparse x dense: {rows} x {inner} x {columns}, zero={all_zero}" + ); + assert_eq!( + DenseMatrix::core_multiply(&Core::Dense(lhs.clone()), &rhs), + expected + ); + let product = sparse_lhs.multiply(&sparse_rhs); + let mut actual = DenseMatrix::zeros(product.num_rows(), product.num_cols()); + for ((row, column), value) in product.iter() { + actual[(row, column)] = value; + } + assert_eq!( + actual, expected, + "sparse x sparse: {rows} x {inner} x {columns}" + ); + let direct_lhs = Core::force_direct(lhs.clone()); + let direct_rhs = Core::force_direct(rhs.clone()); + assert_eq!(lhs.multiply_core(&direct_rhs), expected); + assert_eq!(DenseMatrix::core_multiply(&direct_lhs, &rhs), expected); + for product in [ + direct_lhs.multiply(&sparse_rhs), + sparse_lhs.multiply(&direct_rhs), + direct_lhs.multiply(&direct_rhs), + ] { + let mut actual = + DenseMatrix::zeros(product.num_rows(), product.num_cols()); + for ((row, column), value) in product.iter() { + actual[(row, column)] = value; + } + assert_eq!(actual, expected, "direct CSR multiplication"); + } + } + } + } + } + } + + #[test] + #[ignore = "manual performance comparison with a 64 MiB dense baseline"] + fn benchmark_sparse_products() { + use std::{hint::black_box, time::Instant}; + + let mut interner = Interner::new(); + let one = interner.intern(F::from(1u64)); + let minus_one = interner.intern(F::from(-1i64)); + let mut core = SparseMatrix::new(0, 4096); + for row in 0..512 { + core.push_row( + [ + (0, one), + (row + 1, if row % 2 == 0 { one } else { minus_one }), + ] + .into_iter(), + ); + } + let mut left = DenseMatrix::zeros(4, core.num_rows); + let mut right = DenseMatrix::zeros(core.num_cols, 4); + for (index, value) in left.values.iter_mut().chain(&mut right.values).enumerate() { + *value = F::from((index % 17 + 2) as u64); + } + + let start = Instant::now(); + let sparse_product = black_box(&left).multiply_sparse(black_box(&core), &interner); + let sparse_time = start.elapsed(); + let start = Instant::now(); + let dense_product = left.multiply(&DenseMatrix::from_sparse(&core, &interner)); + let dense_time = start.elapsed(); + assert_eq!(sparse_product, dense_product); + eprintln!("forward: dense including conversion {dense_time:?}, sparse {sparse_time:?}"); + + let start = Instant::now(); + let sparse_product = + DenseMatrix::sparse_multiply(black_box(&core), &interner, black_box(&right)); + let sparse_time = start.elapsed(); + let start = Instant::now(); + let dense_product = DenseMatrix::from_sparse(&core, &interner).multiply(&right); + let dense_time = start.elapsed(); + assert_eq!(sparse_product, dense_product); + eprintln!("backward: dense including conversion {dense_time:?}, sparse {sparse_time:?}"); + eprintln!( + "avoided dense core: {} bytes; forward output: {} bytes; backward output: {} bytes", + core.num_rows * core.num_cols * size_of::(), + left.num_rows * core.num_cols * size_of::(), + core.num_rows * right.num_cols * size_of::(), + ); + } + + // Independent forward elimination, without row normalization or inverses. + #[instrument(skip_all)] + fn rank(mut matrix: DenseMatrix) -> usize { + let mut rank = 0; + for column in 0..matrix.num_cols { + let Some(pivot) = + (rank..matrix.num_rows).find(|&row| matrix[(row, column)] != F::from(0u64)) + else { + continue; + }; + for j in 0..matrix.num_cols { + matrix + .values + .swap(rank * matrix.num_cols + j, pivot * matrix.num_cols + j); + } + let pivot = matrix[(rank, column)]; + for row in rank + 1..matrix.num_rows { + let coefficient = matrix[(row, column)]; + if coefficient == F::from(0u64) { + continue; + } + for j in column..matrix.num_cols { + matrix[(row, j)] = matrix[(row, j)] * pivot - coefficient * matrix[(rank, j)]; + } + } + rank += 1; + } + rank + } + + fn evaluate(program: &MatrixBranchingProgram, assignment: usize) -> Vec { + if let Some(values) = program.constant { + return values.to_vec(); + } + let mut values = matrix(1, 1, &[1]); + for (level, layer) in program.layers.iter().enumerate().rev() { + let core = &layer.matrices[(assignment >> level) & 1]; + let mut dense = DenseMatrix::zeros(core.num_rows(), core.num_cols()); + for ((row, column), value) in core.iter() { + dense[(row, column)] = value; + } + values = dense.multiply(&values); + } + values.values + } + + #[instrument(skip_all)] + fn assert_minimal_and_equivalent(program: &mut MatrixBranchingProgram) { + let num_variables = program.layers.len(); + let table: Vec<_> = (0..1 << num_variables) + .map(|assignment| evaluate(program, assignment)) + .collect(); + let old_widths: Vec<_> = program + .layers + .iter() + .map(|layer| layer.matrices[0].num_cols()) + .collect(); + let constant = program.constant; + program.minimize(); + assert_eq!(program.constant, constant); + assert_eq!(program.layers.len(), num_variables); + for (assignment, expected) in table.iter().enumerate() { + assert_eq!( + &evaluate(program, assignment), + expected, + "assignment {assignment}" + ); + } + for cut in 1..num_variables { + let num_prefixes = 1 << cut; + let num_suffixes = 1 << (num_variables - cut); + let mut unfolding = DenseMatrix::zeros(table[0].len() * num_prefixes, num_suffixes); + for prefix in 0..num_prefixes { + for suffix in 0..num_suffixes { + let values = &table[prefix | (suffix << cut)]; + for (component, &value) in values.iter().enumerate() { + unfolding[(component * num_prefixes + prefix, suffix)] = value; + } + } + } + let rank = rank(unfolding); + assert!(rank <= old_widths[cut - 1]); + for bit in 0..2 { + assert_eq!( + program.layers[cut - 1].matrices[bit].num_cols(), + rank, + "cut {cut}" + ); + assert_eq!( + program.layers[cut].matrices[bit].num_rows(), + rank, + "cut {cut}" + ); + } + } + if let Some(first) = program.layers.first() { + for core in &first.matrices { + assert_eq!(core.num_rows(), table[0].len()); + } + for core in &program.layers.last().unwrap().matrices { + assert_eq!(core.num_cols(), 1); + } + } + } + + #[test] + fn rank_factorizations_are_exact_including_degenerate_shapes() { + let cases = [ + (matrix(3, 4, &[0, 2, 4, 0, 0, 0, 0, 3, 0, 4, 8, 6]), 2), + (matrix(3, 2, &[0, 2, 3, 4, 6, 8]), 2), + (matrix(2, 3, &[2, 4, -6, -3, -6, 9]), 1), + (matrix(3, 3, &[0, 0, 2, 0, 3, 0, 4, 0, 0]), 3), + (matrix(2, 3, &[0; 6]), 0), + (matrix(0, 3, &[]), 0), + (matrix(3, 0, &[]), 0), + (matrix(0, 0, &[]), 0), + ]; + for (original, expected_rank) in cases { + let (left, right) = original.clone().rank_factor(); + assert_eq!( + (left.num_rows, left.num_cols), + (original.num_rows, expected_rank) + ); + assert_eq!( + (right.num_rows, right.num_cols), + (expected_rank, original.num_cols) + ); + assert_eq!(left.multiply(&right), original); + } + } + + #[test] + fn both_scalar_and_vector_boundaries_attain_unfolding_ranks() { + for widths in [[1, 6, 5, 7, 1], [3, 8, 7, 5, 1]] { + let mut seed = 17u64; + let layers = widths + .windows(2) + .map(|width| { + let matrices = std::array::from_fn(|_| { + let mut dense = DenseMatrix::zeros(width[0], width[1]); + for value in &mut dense.values { + seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1); + *value = F::from(((seed >> 32) % 11) as i64 - 5); + } + dense + }); + MbpLayer::from_dense(matrices) + }) + .collect(); + let mut program = MatrixBranchingProgram { + index_bits: 2, + layers, + constant: None, + }; + assert_minimal_and_equivalent(&mut program); + } + } + + #[test] + fn backward_sweep_removes_cancelling_directions() { + let mut program = MatrixBranchingProgram { + index_bits: 1, + layers: vec![ + MbpLayer::from_dense(std::array::from_fn(|_| matrix(3, 2, &[1, 1, 2, 2, 3, 3]))), + MbpLayer::from_dense(std::array::from_fn(|_| matrix(2, 1, &[1, -1]))), + ], + constant: None, + }; + assert_minimal_and_equivalent(&mut program); + assert_eq!(program.layers[0].matrices[0].num_cols(), 0); + assert_eq!(program.layers[1].matrices[0].num_rows(), 0); + for layer in program.layers { + for core in layer.matrices { + assert!(core.iter().next().is_none()); + } + } + } + + #[test] + fn extraction_minimizes_and_preserves_every_coefficient() { + let mut r1cs = R1CS::::new(); + r1cs.add_witnesses(7); + for row in 0..5 { + let a: Vec<_> = (0..7) + .map(|column| (F::from((row + 1) * (column + 1)), column as usize)) + .collect(); + let b = [(F::from(-3i64), row as usize)]; + let c: Vec<_> = (0..7) + .map(|column| (F::from((row ^ column) + 1), column as usize)) + .collect(); + r1cs.add_constraint(&a, &b, &c); + } + let mut program = extract_cqbdd(&r1cs).to_mbp(); + assert_minimal_and_equivalent(&mut program); + let extracted = extract(&r1cs, MinimizeStrategy::Rank); + for assignment in 0..64 { + let mut row = 0; + let mut column = 0; + for bit in 0..3 { + row |= ((assignment >> (2 * bit)) & 1) << bit; + column |= ((assignment >> (2 * bit + 1)) & 1) << bit; + } + let expected = [&r1cs.a, &r1cs.b, &r1cs.c].map(|matrix| { + if row < matrix.num_rows && column < matrix.num_cols { + matrix + .get(row, column) + .map_or(F::from(0u64), |value| r1cs.interner.get(value).unwrap()) + } else { + F::from(0u64) + } + }); + assert_eq!(evaluate(&extracted, assignment), expected); + } + } + + #[test] + fn constants_and_programs_without_internal_cuts_are_unchanged() { + let mut constant = MatrixBranchingProgram:: { + index_bits: 0, + layers: Vec::new(), + constant: Some([F::from(7u64); 3]), + }; + assert_minimal_and_equivalent(&mut constant); + let mut single = MatrixBranchingProgram { + index_bits: 0, + layers: vec![MbpLayer::from_dense([ + matrix(3, 1, &[1, 2, 3]), + matrix(3, 1, &[4, 5, 6]), + ])], + constant: None, + }; + assert_minimal_and_equivalent(&mut single); + } +} diff --git a/tooling/cli/src/cmd/extract/rank.rs b/tooling/cli/src/cmd/extract/rank.rs new file mode 100644 index 000000000..a866c75d4 --- /dev/null +++ b/tooling/cli/src/cmd/extract/rank.rs @@ -0,0 +1,674 @@ +use { + super::{ + matrix::{DenseMatrix, DirectSparseMatrix}, + Core, MbpLayer, + }, + ark_ff::Field, + tracing::{info, instrument}, +}; + +#[cfg(test)] +type Row = Vec<(usize, F)>; + +impl MbpLayer { + /// Exact rank factorization of a vertical or horizontal stack. Sparse + /// inputs are eliminated as sparse rows, with no dense input allocation. + #[instrument(skip_all)] + pub(super) fn rank_factor(&self, vertically: bool) -> (Core, Core) { + if self.matrices.iter().all(Core::is_sparse) { + self.sparse_rank_factor(vertically) + } else { + self.dense_rank_factor(vertically) + } + } + + fn dense_rank_factor(&self, vertically: bool) -> (Core, Core) { + let (left, right) = DenseMatrix::stack(self, vertically).rank_factor(); + (Core::from_dense(left), Core::from_dense(right)) + } + + fn sparse_rank_factor(&self, vertically: bool) -> (Core, Core) { + let height = self.matrices[0].num_rows(); + let width = self.matrices[0].num_cols(); + assert_eq!( + (self.matrices[1].num_rows(), self.matrices[1].num_cols()), + (height, width) + ); + let num_rows = if vertically { 2 * height } else { height }; + let num_cols = if vertically { width } else { 2 * width }; + let dense_bytes = num_rows + .saturating_mul(num_cols) + .saturating_mul(size_of::()); + let mut degrees = vec![0usize; num_cols]; + let mut order: Vec<_> = (0..num_rows) + .map(|row| { + let mut nnz = 0; + self.visit_stack_row(row, vertically, |column, _| { + degrees[column] += 1; + nnz += 1; + }); + (nnz, row) + }) + .collect(); + // Sparse rows first; preserve original row indices in the left factor. + order.sort_unstable(); + let mut owners: Vec> = vec![None; num_cols]; + let mut pivot_columns: Vec = Vec::new(); + let mut basis = DirectSparseMatrix::::new(num_cols); + let mut left = DirectSparseMatrix::::new(0); + let mut workspace = Accumulator::::new(num_cols, num_rows.min(num_cols)); + let mut eliminations = 0usize; + let mut updates = 0usize; + for &(_, original_row) in &order { + self.visit_stack_row(original_row, vertically, |column, value| { + workspace.set(column, value, &owners) + }); + let mut coefficients = Vec::new(); + // Basis row j is zero at all earlier pivot columns. Processing + // active pivots in creation order therefore never revisits a pivot. + while let Some(pivot) = workspace.pivots.first() { + let column = pivot_columns[pivot]; + let coefficient = workspace.values[column]; + workspace.set(column, F::ZERO, &owners); + let positive = coefficient == F::ONE; + let negative = coefficient == -F::ONE; + for (column, value) in basis.iter_row(pivot) { + if column == pivot_columns[pivot] { + continue; + } + let scaled = if positive { + value + } else if negative { + -value + } else { + coefficient * value + }; + workspace.set(column, workspace.values[column] - scaled, &owners); + updates += 1; + } + coefficients.push((pivot, coefficient)); + eliminations += 1; + } + let new_column = workspace.columns.iter().min_by_key(|&column| { + ( + degrees[column], + workspace.values[column] != F::ONE && workspace.values[column] != -F::ONE, + column, + ) + }); + if let Some(column) = new_column { + let coefficient = workspace.values[column]; + let pivot = basis.num_rows(); + owners[column] = Some(pivot); + pivot_columns.push(column); + coefficients.push((pivot, coefficient)); + let inverse = if coefficient == F::ONE { + F::ONE + } else if coefficient == -F::ONE { + -F::ONE + } else { + coefficient.inverse().expect("nonzero sparse pivot") + }; + let positive = inverse == F::ONE; + let negative = inverse == -F::ONE; + basis.push_row(workspace.columns.iter().map(|column| { + let value = workspace.values[column]; + ( + column, + if positive { + value + } else if negative { + -value + } else { + value * inverse + }, + ) + })); + workspace.clear(); + } + coefficients.sort_unstable_by_key(|&(pivot, _)| pivot); + left.push_row(coefficients.into_iter()); + let sparse_bytes = basis + .values + .capacity() + .saturating_add(left.values.capacity()) + .saturating_mul(size_of::()) + .saturating_add( + basis + .columns + .capacity() + .saturating_add(left.columns.capacity()) + .saturating_mul(size_of::()), + ); + if sparse_bytes > dense_bytes { + info!( + "Sparse rank factorization switching to dense after {} rows.", + left.num_rows() + ); + drop(left); + drop(basis); + drop(workspace); + drop(owners); + return self.dense_rank_factor(vertically); + } + } + let rank = basis.num_rows(); + left.num_cols = rank; + if order + .iter() + .enumerate() + .any(|(position, &(_, row))| position != row) + { + let mut inverse = vec![0; num_rows]; + for (position, &(_, row)) in order.iter().enumerate() { + inverse[row] = position; + } + let mut restored = DirectSparseMatrix::new(rank); + restored.values.reserve(left.values.len()); + restored.columns.reserve(left.columns.len()); + for position in inverse { + restored.push_row(left.iter_row(position)); + } + left = restored; + } + info!( + "Sparse rank factorization {}x{}: rank {}, factor nnz ({}, {}), {} eliminations, {} \ + coefficient updates.", + num_rows, + num_cols, + rank, + left.values.len(), + basis.values.len(), + eliminations, + updates + ); + (Core::from_direct(left), Core::from_direct(basis)) + } + + fn visit_stack_row(&self, row: usize, vertically: bool, mut visit: impl FnMut(usize, F)) { + let height = self.matrices[0].num_rows(); + let width = self.matrices[0].num_cols(); + if vertically { + let bit = usize::from(row >= height); + for (column, value) in self.matrices[bit].iter_row(row - bit * height) { + visit(column, value); + } + } else { + for (bit, core) in self.matrices.iter().enumerate() { + for (column, value) in core.iter_row(row) { + visit(column + bit * width, value); + } + } + } + } + + #[cfg(test)] + fn sparse_stack_row(&self, row: usize, vertically: bool) -> Row { + let mut entries = Vec::new(); + self.visit_stack_row(row, vertically, |column, value| { + entries.push((column, value)) + }); + entries + } + #[cfg(test)] + fn reference_rank_factor(&self, vertically: bool) -> (Core, Core) { + let height = self.matrices[0].num_rows(); + let width = self.matrices[0].num_cols(); + assert_eq!( + (self.matrices[1].num_rows(), self.matrices[1].num_cols()), + (height, width) + ); + let num_rows = if vertically { 2 * height } else { height }; + let num_cols = if vertically { width } else { 2 * width }; + let dense_bytes = num_rows + .saturating_mul(num_cols) + .saturating_mul(size_of::()); + let mut pivots = vec![None; num_cols]; + let mut basis: Vec> = Vec::new(); + let mut left = Vec::with_capacity(num_rows); + let mut scratch = Vec::new(); + let mut basis_nnz = 0usize; + let mut left_nnz = 0usize; + + for row_index in 0..num_rows { + let mut row = self.sparse_stack_row(row_index, vertically); + let mut coefficients = Vec::new(); + while let Some(&(column, coefficient)) = row.first() { + if let Some(pivot) = pivots[column] { + // Normalized basis rows have leading coefficient one. + subtract_scaled(&mut row, &basis[pivot], coefficient, &mut scratch); + coefficients.push((pivot, coefficient)); + } else { + let pivot = basis.len(); + pivots[column] = Some(pivot); + if coefficient != F::ONE { + let inverse = coefficient.inverse().expect("nonzero sparse pivot"); + for (_, value) in &mut row { + *value *= inverse; + } + } + basis_nnz += row.len(); + basis.push(row); + coefficients.push((pivot, coefficient)); + break; + } + } + // Pivot columns increase during reduction, but basis IDs follow + // creation order. Sorting yields canonical sparse output rows. + coefficients.sort_unstable_by_key(|&(pivot, _)| pivot); + left_nnz += coefficients.len(); + left.push(coefficients); + + // Fill-in can erase the memory advantage of sparse elimination. + // Restart the existing dense algorithm only once sparse factors + // exceed the payload of the full dense input. Drop work first. + let sparse_bytes = basis_nnz + .saturating_add(left_nnz) + .saturating_mul(size_of::<(usize, F)>()); + if sparse_bytes > dense_bytes { + info!( + "Sparse rank factorization switching to dense after {} rows ({} factor nnz).", + row_index + 1, + basis_nnz + left_nnz + ); + drop(left); + drop(basis); + drop(pivots); + drop(scratch); + return self.dense_rank_factor(vertically); + } + } + + // Every input row was expressed in the independent normalized basis. + // Thus A = L R, and each newly introduced basis row supplies a nonzero + // triangular pivot in L. Full RREF/back substitution is unnecessary. + let rank = basis.len(); + info!( + "Sparse rank factorization {}x{}: rank {}, factor nnz ({}, {}).", + num_rows, num_cols, rank, left_nnz, basis_nnz + ); + ( + Core::from_rows(rank, left), + Core::from_rows(num_cols, basis), + ) + } +} + +/// Two-level bitset: the summary finds active words without scanning the +/// whole column range at each pivot. It stores only O(width) bits. +struct ActiveSet { + words: Vec, + summary: Vec, +} + +impl ActiveSet { + fn new(len: usize) -> Self { + let words = len.div_ceil(64); + Self { + words: vec![0; words], + summary: vec![0; words.div_ceil(64)], + } + } + fn set(&mut self, index: usize, active: bool) { + let word = index / 64; + let bit = 1u64 << (index % 64); + if active { + self.words[word] |= bit; + } else { + self.words[word] &= !bit; + } + let summary_bit = 1u64 << (word % 64); + if self.words[word] == 0 { + self.summary[word / 64] &= !summary_bit; + } else { + self.summary[word / 64] |= summary_bit; + } + } + fn first(&self) -> Option { + self.iter().next() + } + fn iter(&self) -> impl Iterator + '_ { + self.summary + .iter() + .copied() + .enumerate() + .flat_map(move |(group, mut active)| { + std::iter::from_fn(move || { + if active == 0 { + return None; + } + let word = group * 64 + active.trailing_zeros() as usize; + active &= active - 1; + Some(word) + }) + .flat_map(move |word| { + let mut bits = self.words[word]; + std::iter::from_fn(move || { + if bits == 0 { + return None; + } + let index = word * 64 + bits.trailing_zeros() as usize; + bits &= bits - 1; + Some(index) + }) + }) + }) + } +} + +/// One reusable field accumulator. Elimination touches only the pivot support; +/// it never copies the unaffected remainder of the current sparse row. +struct Accumulator { + values: Vec, + columns: ActiveSet, + pivots: ActiveSet, +} +impl Accumulator { + fn new(columns: usize, pivots: usize) -> Self { + Self { + values: vec![F::ZERO; columns], + columns: ActiveSet::new(columns), + pivots: ActiveSet::new(pivots), + } + } + fn set(&mut self, column: usize, value: F, owners: &[Option]) { + let was_active = !self.values[column].is_zero(); + self.values[column] = value; + let active = !value.is_zero(); + if active != was_active { + self.columns.set(column, active); + if let Some(pivot) = owners[column] { + self.pivots.set(pivot, active); + } + } + } + fn clear(&mut self) { + for column in self.columns.iter() { + self.values[column] = F::ZERO; + } + self.columns.words.fill(0); + self.columns.summary.fill(0); + debug_assert!(self.pivots.first().is_none()); + } +} + +#[cfg(test)] +/// Merge sorted supports, dropping exact cancellations and reusing scratch. +fn subtract_scaled( + row: &mut Row, + basis: &Row, + coefficient: F, + scratch: &mut Row, +) { + scratch.clear(); + let mut i = 0; + let mut j = 0; + let scale = |value: F| { + if coefficient == F::ONE { + value + } else if coefficient == -F::ONE { + -value + } else { + coefficient * value + } + }; + while i < row.len() || j < basis.len() { + if j == basis.len() || (i < row.len() && row[i].0 < basis[j].0) { + scratch.push(row[i]); + i += 1; + } else if i == row.len() || basis[j].0 < row[i].0 { + scratch.push((basis[j].0, -scale(basis[j].1))); + j += 1; + } else { + let value = row[i].1 - scale(basis[j].1); + if !value.is_zero() { + scratch.push((row[i].0, value)); + } + i += 1; + j += 1; + } + } + std::mem::swap(row, scratch); +} + +#[cfg(test)] +mod tests { + use {super::*, provekit_backend_bn254::FieldElement as F}; + + fn dense(core: &Core) -> DenseMatrix { + let mut matrix = DenseMatrix::zeros(core.num_rows(), core.num_cols()); + for ((row, column), value) in core.iter() { + matrix[(row, column)] = value; + } + matrix + } + + fn assert_factorization(layer: &MbpLayer, vertically: bool) { + let original = DenseMatrix::stack(layer, vertically); + let (_, reference) = original.clone().rank_factor(); + let (left, right) = layer.rank_factor(vertically); + assert_eq!( + (left.num_rows(), left.num_cols()), + (original.num_rows, reference.num_rows) + ); + assert_eq!( + (right.num_rows(), right.num_cols()), + (reference.num_rows, original.num_cols) + ); + assert_eq!(dense(&left).multiply(&dense(&right)), original); + } + + #[test] + fn active_sets_handle_word_boundaries_and_cancellation() { + let indices = [0, 1, 63, 64, 65, 4095, 4096, 4999]; + let mut set = ActiveSet::new(5000); + for index in indices { + set.set(index, true); + } + assert_eq!(set.iter().collect::>(), indices); + for index in indices { + assert_eq!(set.first(), Some(index)); + set.set(index, false); + } + assert!(set.first().is_none()); + let mut owners = vec![None; 5000]; + owners[63] = Some(0); + owners[4096] = Some(1); + let mut accumulator = Accumulator::::new(5000, 2); + accumulator.set(4096, F::from(7u64), &owners); + accumulator.set(63, F::from(3u64), &owners); + assert_eq!(accumulator.pivots.first(), Some(0)); + accumulator.set(63, F::from(0u64), &owners); + assert_eq!(accumulator.pivots.first(), Some(1)); + accumulator.set(4096, F::from(0u64), &owners); + accumulator.set(4999, F::from(9u64), &owners); + accumulator.clear(); + assert!(accumulator.columns.first().is_none()); + assert!(accumulator + .values + .iter() + .all(|value| *value == F::from(0u64))); + } + + #[test] + fn row_and_column_ordering_reduce_hub_fill_in() { + let matrices = std::array::from_fn(|bit| { + let rows = (0..32) + .map(|row| { + let index = row + bit * 32; + let mut values = vec![(0, F::from(1u64)), (index + 1, F::from(2u64))]; + if index % 3 == 0 { + values.push((100, F::from((index + 3) as u64))); + } + values + }) + .collect(); + Core::from_rows(128, rows) + }); + let layer = MbpLayer { matrices }; + let (old_left, old_right) = layer.reference_rank_factor(true); + let (left, right) = layer.rank_factor(true); + assert_eq!(right.num_rows(), old_right.num_rows()); + assert!(left.nnz() + right.nnz() < old_left.nnz() + old_right.nnz()); + assert_eq!( + dense(&left).multiply(&dense(&right)), + DenseMatrix::stack(&layer, true) + ); + } + + #[test] + #[ignore = "manual comparison with the previous merge-based sparse algorithm"] + fn benchmark_sparse_rank_fill_in() { + use std::time::Instant; + for case in ["hub", "hub-tail", "coupled"] { + let matrices = std::array::from_fn(|bit| { + let rows = (0..256) + .map(|row| { + let index = row + bit * 256; + if case == "coupled" { + let mut state = index as u64 + 1; + let mut entries = std::collections::BTreeMap::new(); + for _ in 0..32 { + state = state.wrapping_mul(6364136223846793005).wrapping_add(1); + entries.insert( + (state >> 32) as usize % 256, + F::from((state % 31 + 1) as u64), + ); + } + return entries.into_iter().collect(); + } + let mut entries = vec![ + (0, F::from((index % 13 + 1) as u64)), + (index + 1, F::from(1u64)), + ]; + if case == "hub-tail" { + for offset in 0..16 { + entries.push(( + 2048 + ((index * 17 + offset * 31) % 512), + F::from(((index + offset) % 31 + 1) as u64), + )); + } + } + entries.sort_unstable_by_key(|&(column, _)| column); + entries + }) + .collect(); + Core::from_rows(4096, rows) + }); + let layer = MbpLayer { matrices }; + for previous in [true, false] { + #[cfg(feature = "profiling-allocator")] + let (base, count) = (crate::ALLOCATOR.reset_max(), crate::ALLOCATOR.count()); + let start = Instant::now(); + let (left, right) = if previous { + layer.reference_rank_factor(true) + } else { + layer.rank_factor(true) + }; + let elapsed = start.elapsed(); + #[cfg(feature = "profiling-allocator")] + let (peak, allocations) = ( + crate::ALLOCATOR.max().saturating_sub(base), + crate::ALLOCATOR.count() - count, + ); + eprintln!( + "{case}, {}: {elapsed:?}, rank {}, nnz ({}, {})", + if previous { "previous" } else { "new" }, + right.num_rows(), + left.nnz(), + right.nnz() + ); + #[cfg(feature = "profiling-allocator")] + eprintln!("peak additional heap {peak} bytes, {allocations} allocations"); + let product = left.multiply(&right); + for row in 0..512 { + assert_eq!( + product.iter_row(row).collect::>(), + layer.sparse_stack_row(row, true) + ); + } + } + } + } + + #[test] + fn sparse_factorizations_match_dense_rank_for_all_small_ternary_inputs() { + // Wide padding keeps elimination sparse. Scrambled nonzero columns + // exercise basis IDs introduced in a different order from pivot columns. + for mut code in 0..729 { + let mut matrix = DenseMatrix::zeros(2, 11); + for row in 0..2 { + for column in [9, 2, 5] { + matrix[(row, column)] = F::from((code % 3) as i64 - 1); + code /= 3; + } + } + let layer = MbpLayer { + matrices: std::array::from_fn(|_| Core::force_sparse(matrix.clone())), + }; + for vertically in [false, true] { + assert_factorization(&layer, vertically); + } + } + } + + #[test] + fn nonunit_pivots_mixed_storage_fill_in_and_empty_shapes() { + for (height, width) in [(0, 0), (0, 5), (5, 0), (3, 17), (32, 32)] { + let mut seed = 29u64; + let matrices = std::array::from_fn(|_| { + let mut matrix = DenseMatrix::zeros(height, width); + for value in &mut matrix.values { + seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1); + *value = F::from(((seed >> 32) % 17) as i64 - 8); + } + matrix + }); + let sparse = MbpLayer { + matrices: matrices + .each_ref() + .map(|matrix| Core::force_sparse(matrix.clone())), + }; + let [first, second] = matrices; + let mixed = MbpLayer { + matrices: [Core::force_sparse(first), Core::Dense(second)], + }; + for vertically in [false, true] { + assert_factorization(&sparse, vertically); + assert_factorization(&mixed, vertically); + } + } + } + + #[test] + #[ignore = "manual sparse versus dense rank-factorization comparison"] + fn benchmark_sparse_rank_factorization() { + use std::time::Instant; + let matrices = std::array::from_fn(|bit| { + let rows = (0..64) + .map(|row| vec![(0, F::from(1u64)), (1 + row + bit * 64, F::from(2u64))]) + .collect(); + Core::from_rows(4096, rows) + }); + let layer = MbpLayer { matrices }; + assert!(layer.matrices.iter().all(Core::is_sparse)); + let start = Instant::now(); + let (left, right) = layer.rank_factor(true); + let sparse_time = start.elapsed(); + let start = Instant::now(); + let original = DenseMatrix::stack(&layer, true); + let (reference_left, reference_right) = original.clone().rank_factor(); + let dense_time = start.elapsed(); + assert_eq!(right.num_rows(), reference_right.num_rows); + assert_eq!(dense(&left.multiply(&right)), original); + assert_eq!(reference_left.multiply(&reference_right), original); + eprintln!( + "sparse rank {sparse_time:?}, dense including stack {dense_time:?}; rank {}, factor \ + nnz ({}, {}); avoided input {} bytes", + right.num_rows(), + left.nnz(), + right.nnz(), + 128 * 4096 * size_of::() + ); + } +} diff --git a/tooling/cli/src/cmd/extract/rescale.rs b/tooling/cli/src/cmd/extract/rescale.rs new file mode 100644 index 000000000..a3392b381 --- /dev/null +++ b/tooling/cli/src/cmd/extract/rescale.rs @@ -0,0 +1,397 @@ +use { + super::MatrixBranchingProgram, + ark_ff::{batch_inversion, Field}, + tracing::{info, instrument}, +}; + +// Bound the search and workspace even when almost every coefficient differs. +const CANDIDATES_PER_SIDE: usize = 4; +const CANDIDATES: usize = 2 * CANDIDATES_PER_SIDE; +const MAX_SWEEPS: usize = 4; + +/// Misra–Gries candidates, grouping coefficients up to sign. Exact scoring +/// below prevents approximate frequency estimates from accepting a regression. +#[derive(Clone)] +struct Frequent { + values: [F; CANDIDATES_PER_SIDE], + votes: [usize; CANDIDATES_PER_SIDE], +} + +impl Frequent { + fn new() -> Self { + Self { + values: [F::ZERO; CANDIDATES_PER_SIDE], + votes: [0; CANDIDATES_PER_SIDE], + } + } + + fn offer(&mut self, value: F) { + let negative = -value; + if let Some(index) = (0..CANDIDATES_PER_SIDE).find(|&i| { + self.votes[i] != 0 && (self.values[i] == value || self.values[i] == negative) + }) { + self.votes[index] += 1; + } else if let Some(index) = self.votes.iter().position(|&votes| votes == 0) { + self.values[index] = value; + self.votes[index] = 1; + } else { + for votes in &mut self.votes { + *votes -= 1; + } + } + } + + fn candidates(&self) -> impl Iterator + '_ { + self.values + .iter() + .zip(&self.votes) + .map(|(&value, &votes)| if votes == 0 { F::ZERO } else { value }) + } +} + +impl MatrixBranchingProgram { + /// Improve the number of +/-1 coefficients without changing widths or + /// support. At a cut, A_b -> A_b D and B_b -> D^-1 B_b for both branches. + /// Only strict improvements across all four cores are accepted. This is + /// bounded coordinate descent, not a globally optimal change of basis. + /// Each sweep takes O(stored entries + widths) time with a fixed candidate + /// count, and O(max width) workspace besides rebuilding one interned core. + #[instrument(skip_all)] + pub(super) fn rescale(&mut self) { + let cuts = self.layers.len().saturating_sub(1); + for sweep in 0..MAX_SWEEPS { + let mut changed = 0; + let mut improvement = 0; + for index in 0..cuts { + let cut = if sweep % 2 == 0 { + index + } else { + cuts - 1 - index + }; + let (coordinates, gain) = self.rescale_cut(cut); + changed += coordinates; + improvement += gain; + } + info!( + "MBP rescaling sweep {}: {} coordinates changed; other count reduced by {}.", + sweep + 1, + changed, + improvement + ); + if changed == 0 { + break; + } + } + } + + fn rescale_cut(&mut self, cut: usize) -> (usize, usize) { + let width = self.layers[cut].matrices[0].num_cols(); + let mut incoming = vec![Frequent::new(); width]; + let mut outgoing = vec![Frequent::new(); width]; + let mut baseline = vec![0usize; width]; + let mut has_other = false; + for core in &self.layers[cut].matrices { + for ((_, column), value) in core.iter() { + incoming[column].offer(value); + let unit = value == F::ONE || value == -F::ONE; + baseline[column] += usize::from(unit); + has_other |= !unit; + } + } + for core in &self.layers[cut + 1].matrices { + for ((row, _), value) in core.iter() { + outgoing[row].offer(value); + let unit = value == F::ONE || value == -F::ONE; + baseline[row] += usize::from(unit); + has_other |= !unit; + } + } + if !has_other { + return (0, 0); + } + + // Incoming a becomes a*d, so propose d=1/a. Outgoing b becomes + // b/d, so propose d=b. Signs do not affect the objective. + let mut incoming_scales: Vec = incoming.iter().flat_map(Frequent::candidates).collect(); + batch_inversion(&mut incoming_scales); + let mut scales = Vec::with_capacity(width * CANDIDATES); + for (column, frequencies) in outgoing.iter().enumerate() { + scales.extend_from_slice( + &incoming_scales[column * CANDIDATES_PER_SIDE..(column + 1) * CANDIDATES_PER_SIDE], + ); + scales.extend(frequencies.candidates()); + } + drop(incoming); + drop(outgoing); + drop(incoming_scales); + let mut inverses = scales.clone(); + batch_inversion(&mut inverses); + let mut scores = vec![0usize; scales.len()]; + for core in &self.layers[cut].matrices { + for ((_, column), value) in core.iter() { + let start = column * CANDIDATES; + let negative = -value; + for (score, &target) in scores[start..start + CANDIDATES] + .iter_mut() + .zip(&inverses[start..start + CANDIDATES]) + { + *score += usize::from(value == target || negative == target); + } + } + } + for core in &self.layers[cut + 1].matrices { + for ((row, _), value) in core.iter() { + let start = row * CANDIDATES; + let negative = -value; + for (score, &target) in scores[start..start + CANDIDATES] + .iter_mut() + .zip(&scales[start..start + CANDIDATES]) + { + *score += usize::from(value == target || negative == target); + } + } + } + let mut diagonal = vec![F::ONE; width]; + let mut inverse_diagonal = vec![F::ONE; width]; + let mut changed = 0; + let mut improvement = 0; + for column in 0..width { + let mut best = baseline[column]; + for candidate in column * CANDIDATES..(column + 1) * CANDIDATES { + if scores[candidate] > best { + best = scores[candidate]; + diagonal[column] = scales[candidate]; + inverse_diagonal[column] = inverses[candidate]; + } + } + if best > baseline[column] { + changed += 1; + improvement += best - baseline[column]; + } + } + drop(scales); + drop(inverses); + drop(scores); + drop(baseline); + if changed != 0 { + for core in &mut self.layers[cut].matrices { + core.scale_axis(&diagonal, true); + } + for core in &mut self.layers[cut + 1].matrices { + core.scale_axis(&inverse_diagonal, false); + } + } + (changed, improvement) + } +} + +#[cfg(test)] +mod tests { + use { + super::*, + crate::cmd::extract::{matrix::DenseMatrix, Core, MbpLayer}, + ark_ff::AdditiveGroup, + provekit_backend_bn254::FieldElement as F, + }; + + fn evaluate(program: &MatrixBranchingProgram, assignment: usize) -> Vec { + if let Some(constant) = program.constant { + return constant.to_vec(); + } + let mut values = vec![F::ONE]; + for (level, layer) in program.layers.iter().enumerate().rev() { + let core = &layer.matrices[(assignment >> level) & 1]; + let mut next = vec![F::ZERO; core.num_rows()]; + for ((row, column), value) in core.iter() { + next[row] += value * values[column]; + } + values = next; + } + values + } + + fn other_count(program: &MatrixBranchingProgram) -> usize { + program + .layers + .iter() + .flat_map(|layer| &layer.matrices) + .map(|core| core.value_counts()[2]) + .sum() + } + + fn support(program: &MatrixBranchingProgram) -> Vec<(usize, usize, Vec<(usize, usize)>)> { + program + .layers + .iter() + .flat_map(|layer| &layer.matrices) + .map(|core| { + ( + core.num_rows(), + core.num_cols(), + core.iter().map(|(index, _)| index).collect(), + ) + }) + .collect() + } + + fn assert_equivalent(program: &mut MatrixBranchingProgram) { + let table: Vec<_> = (0..1 << program.layers.len()) + .map(|assignment| evaluate(program, assignment)) + .collect(); + let before_support = support(program); + let before_other = other_count(program); + program.rescale(); + assert_eq!(support(program), before_support); + assert!(other_count(program) <= before_other); + for (assignment, expected) in table.iter().enumerate() { + assert_eq!( + &evaluate(program, assignment), + expected, + "assignment {assignment}" + ); + } + } + + #[test] + fn removes_diagonal_gauges_in_all_storage_formats() { + let widths = [3, 3, 4, 1]; + let diagonals: Vec> = vec![ + vec![F::ONE; 3], + [2u64, 3, 5].into_iter().map(F::from).collect(), + [7u64, 11, 13, 17].into_iter().map(F::from).collect(), + vec![F::ONE], + ]; + for representation in 0..4 { + let mut program = MatrixBranchingProgram { + index_bits: 0, + constant: None, + layers: (0..3) + .map(|level| MbpLayer { + matrices: std::array::from_fn(|bit| { + let mut dense = DenseMatrix::zeros(widths[level], widths[level + 1]); + for row in 0..dense.num_rows { + for column in 0..dense.num_cols { + let base = match (row + column + bit) % 3 { + 0 => F::ZERO, + 1 => F::ONE, + _ => -F::ONE, + }; + dense[(row, column)] = base + * diagonals[level][row].inverse().unwrap() + * diagonals[level + 1][column]; + } + } + match if representation == 3 { + (level + bit) % 3 + } else { + representation + } { + 0 => Core::Dense(dense), + 1 => Core::force_sparse(dense), + _ => Core::force_direct(dense), + } + }), + }) + .collect(), + }; + assert!(other_count(&program) > 0); + assert_equivalent(&mut program); + assert_eq!(other_count(&program), 0, "representation {representation}"); + } + } + + #[test] + fn scores_both_adjacent_layers_before_accepting_a_scale() { + let mut program = MatrixBranchingProgram { + index_bits: 1, + constant: None, + layers: vec![ + MbpLayer::from_dense(std::array::from_fn(|_| DenseMatrix { + num_rows: 3, + num_cols: 1, + values: vec![F::ONE; 3], + })), + MbpLayer::from_dense(std::array::from_fn(|_| DenseMatrix { + num_rows: 1, + num_cols: 1, + values: vec![F::from(2u64)], + })), + ], + }; + // Normalizing the two outgoing coefficients would spoil six incoming + // units, so the scale must be rejected. + assert_eq!(program.rescale_cut(0), (0, 0)); + assert_equivalent(&mut program); + assert_eq!(other_count(&program), 2); + } + + #[test] + fn mixed_coefficients_preserve_every_assignment_and_never_regress() { + for seed in 0..24 { + let widths = [3, 4, 3, 2, 1]; + let mut program = MatrixBranchingProgram { + index_bits: 2, + constant: None, + layers: (0..4) + .map(|level| MbpLayer { + matrices: std::array::from_fn(|bit| { + let dense = DenseMatrix { + num_rows: widths[level], + num_cols: widths[level + 1], + values: (0..widths[level] * widths[level + 1]) + .map(|index| { + F::from( + ((index * 17 + seed * 11 + level * 7 + bit * 3) % 13) + as i64 + - 6, + ) + }) + .collect(), + }; + match (seed + level + bit) % 3 { + 0 => Core::Dense(dense), + 1 => Core::force_sparse(dense), + _ => Core::force_direct(dense), + } + }), + }) + .collect(), + }; + // Check each cut's claimed gain against actual coefficient counts. + for cut in 0..3 { + let before = other_count(&program); + let (_, gain) = program.rescale_cut(cut); + assert_eq!(other_count(&program) + gain, before); + } + assert_equivalent(&mut program); + program.minimize(); + assert_equivalent(&mut program); + } + } + + #[test] + fn empty_widths_and_constant_programs_are_unchanged() { + let constant = [F::from(2u64), -F::ONE, F::ZERO]; + let mut program = MatrixBranchingProgram { + index_bits: 0, + layers: vec![], + constant: Some(constant), + }; + program.rescale(); + assert_eq!(program.constant, Some(constant)); + for widths in [[3, 0, 1], [3, 2, 1]] { + let mut program = MatrixBranchingProgram { + index_bits: 1, + constant: None, + layers: (0..2) + .map(|level| { + MbpLayer::from_dense(std::array::from_fn(|_| { + DenseMatrix::zeros(widths[level], widths[level + 1]) + })) + }) + .collect(), + }; + assert_equivalent(&mut program); + } + } +} diff --git a/tooling/cli/src/cmd/extract/ternary.rs b/tooling/cli/src/cmd/extract/ternary.rs new file mode 100644 index 000000000..5ef934a21 --- /dev/null +++ b/tooling/cli/src/cmd/extract/ternary.rs @@ -0,0 +1,1231 @@ +//! Exact, greedy sparsification with ternary internal cores. The first core +//! carries arbitrary A/B/C coefficients and is exempt from the alphabet check, +//! but its nonzeros are included in every cost comparison. +use { + super::{Core, MatrixBranchingProgram, MbpLayer}, + ark_ff::Field, + tracing::{info, instrument}, +}; + +mod patterns; +mod shears; + +type SparseVector = Vec<(usize, F)>; + +fn ternary(value: F) -> bool { + value.is_zero() || value == F::ONE || value == -F::ONE +} + +fn layer_nnz(layer: &MbpLayer) -> usize { + layer.matrices.iter().map(Core::nnz).sum() +} + +fn layer_is_ternary(layer: &MbpLayer) -> bool { + layer + .matrices + .iter() + .all(|core| core.value_counts()[2] == 0) +} + +impl MatrixBranchingProgram { + /// Minimize (total nnz, sum of internal widths) by guarded rank proposals, + /// paired elementary basis changes, and shared-pattern extraction. Every + /// accepted rewrite is exact. Sharing may increase width while reducing + /// nnz. This is a local heuristic, not a minimum-rank or minimum-nnz + /// guarantee. + /// + /// Each direction scans every cut once. Sparse indexes prune impossible + /// pairs while preserving the all-pairs trial order and accepted moves. + /// A dense worst-case scan takes O(k W^3) field operations. Further sweeps + /// are needed until the strictly decreasing integer cost stabilizes. + /// Rank proposals retain the existing rank factorizer's workspace costs. + #[instrument(skip_all)] + pub(super) fn minimize_ternary(&mut self) { + assert!( + self.layers.iter().skip(1).all(layer_is_ternary), + "ternary minimization requires ternary internal input cores" + ); + let cuts = self.layers.len().saturating_sub(1); + if cuts == 0 { + return; + } + let mut versions = vec![0usize; self.layers.len()]; + // A failed pass remains a fixed point until either adjacent core + // changes. Rank proposals depend on direction, so cache separately. + let mut clean = vec![[None; 2]; cuts]; + for sweep in 1usize.. { + let mut changed = false; + for forward in [true, false] { + let direction_index = usize::from(!forward); + let mut skipped = 0; + for index in 0..cuts { + let cut = if forward { index } else { cuts - 1 - index }; + let version = (versions[cut], versions[cut + 1]); + if clean[cut][direction_index] == Some(version) { + skipped += 1; + continue; + } + let before = self.ternary_cut_cost(cut); + let mut cut_changed = self.try_ternary_rank(cut, forward); + cut_changed |= self.sparsify_ternary_cut(cut); + changed |= cut_changed; + if cut_changed { + versions[cut] += 1; + versions[cut + 1] += 1; + } else { + clean[cut][direction_index] = Some(version); + } + let after = self.ternary_cut_cost(cut); + info!( + "Ternary sweep {sweep}, cut {cut}, forward={forward}: nnz {} -> {}, width \ + {} -> {}.", + before.0, after.0, before.1, after.1 + ); + } + let direction = if forward { "forward" } else { "backward" }; + info!("Ternary sweep {sweep} {direction}: skipped {skipped} unchanged cuts."); + self.log_stats(&format!("Ternary sweep {sweep} {direction}")); + } + if !changed { + break; + } + } + } + + fn ternary_cut_cost(&self, cut: usize) -> (usize, usize) { + ( + layer_nnz(&self.layers[cut]) + layer_nnz(&self.layers[cut + 1]), + self.layers[cut].matrices[0].num_cols(), + ) + } + + fn try_ternary_rank(&mut self, cut: usize, forward: bool) -> bool { + let (left, right) = if forward { + let rows = self.layers[cut].matrices[0].num_rows(); + let (l, r) = self.layers[cut].rank_factor(true); + let rank = r.num_rows(); + let left = MbpLayer { + matrices: std::array::from_fn(|bit| l.block(bit * rows, rows, 0, rank)), + }; + if cut != 0 && !layer_is_ternary(&left) { + return false; + } + let right = MbpLayer { + matrices: self.layers[cut + 1] + .matrices + .each_ref() + .map(|core| r.multiply(core)), + }; + (left, right) + } else { + let cols = self.layers[cut + 1].matrices[0].num_cols(); + let (l, r) = self.layers[cut + 1].rank_factor(false); + let rank = l.num_cols(); + let right = MbpLayer { + matrices: std::array::from_fn(|bit| r.block(0, rank, bit * cols, cols)), + }; + if !layer_is_ternary(&right) { + return false; + } + let left = MbpLayer { + matrices: self.layers[cut] + .matrices + .each_ref() + .map(|core| core.multiply(&l)), + }; + (left, right) + }; + let cost = ( + layer_nnz(&left) + layer_nnz(&right), + left.matrices[0].num_cols(), + ); + // Check emitted cores, including the absorbed transfer matrix. + if cost >= self.ternary_cut_cost(cut) + || (cut != 0 && !layer_is_ternary(&left)) + || !layer_is_ternary(&right) + { + return false; + } + self.layers[cut] = left; + self.layers[cut + 1] = right; + true + } + + fn sparsify_ternary_cut(&mut self, cut: usize) -> bool { + let mut pair = Cut::new(&self.layers[cut], &self.layers[cut + 1]); + let mut changed = pair.remove_dead(); + // Keep deleted slots empty during the scan, so indices remain stable. + // A later sweep revisits moves exposed after their position was passed. + changed |= pair.shear_sweep(cut != 0); + // New shared states are eligible for shears on the next sweep. + changed |= pair.extract_shared_patterns(cut != 0) != 0; + if changed { + let (left, right) = pair.into_layers(); + self.layers[cut] = left; + self.layers[cut + 1] = right; + } + changed + } +} + +/// Columns of A = [M_i0; M_i1] and rows of B = [M_(i+1)0 M_(i+1)1]. +/// Both orientations are sparse; no dense pair of cores is materialized. +struct Cut { + incoming: Vec>, + outgoing: Vec>, + left_rows: usize, + right_cols: usize, +} + +impl Cut { + fn new(left: &MbpLayer, right: &MbpLayer) -> Self { + let width = left.matrices[0].num_cols(); + let left_rows = left.matrices[0].num_rows(); + let right_cols = right.matrices[0].num_cols(); + let mut incoming = vec![Vec::new(); width]; + let mut outgoing = vec![Vec::new(); width]; + for bit in 0..2 { + assert_eq!(left.matrices[bit].num_rows(), left_rows); + assert_eq!(left.matrices[bit].num_cols(), width); + assert_eq!(right.matrices[bit].num_rows(), width); + assert_eq!(right.matrices[bit].num_cols(), right_cols); + for ((row, col), value) in left.matrices[bit].iter() { + incoming[col].push((bit * left_rows + row, value)); + } + for ((row, col), value) in right.matrices[bit].iter() { + outgoing[row].push((bit * right_cols + col, value)); + } + } + Self { + incoming, + outgoing, + left_rows, + right_cols, + } + } + + fn remove_dead(&mut self) -> bool { + let mut changed = false; + for (incoming, outgoing) in self.incoming.iter_mut().zip(&mut self.outgoing) { + if incoming.is_empty() || outgoing.is_empty() { + // Even a state with no edges consumes width. + changed = true; + incoming.clear(); + outgoing.clear(); + } + } + changed + } + + /// A[:,q] += s*A[:,p], B[p,:] -= s*B[q,:], where s is +/-1. + /// The two cross terms cancel, so every two-bit core product is unchanged. + /// Score without allocating; only materialize the two sums on acceptance. + fn try_shear(&mut self, p: usize, q: usize, subtract: bool, ternary_left: bool) -> bool { + self.try_shear_with(p, q, subtract, ternary_left, |_, _| {}) + } + + /// Notify the candidate index of the old/new outgoing support only after + /// acceptance. Rejected trials do not allocate or update index state. + fn try_shear_with( + &mut self, + p: usize, + q: usize, + subtract: bool, + ternary_left: bool, + on_accept: impl FnOnce(&[(usize, F)], &[(usize, F)]), + ) -> bool { + if self.incoming[p].is_empty() || self.incoming[q].is_empty() { + return false; + } + let sides = [ + (&self.incoming[q], &self.incoming[p], subtract, ternary_left), + (&self.outgoing[p], &self.outgoing[q], !subtract, true), + ]; + let totals = sides.map(|(lhs, rhs, ..)| lhs.len() + rhs.len()); + let before = totals[0] + totals[1]; + // If the new incoming column is zero, delete outgoing[q]; if the new + // outgoing row is zero, delete incoming[p]. Include those savings in + // every bound, so profitable state deletions are never pruned. + let retained = [self.outgoing[q].len(), self.incoming[p].len()]; + let cost = |nnz: [usize; 2]| { + nnz[0] + + nnz[1] + + if nnz[0] == 0 { 0 } else { retained[0] } + + if nnz[1] == 0 { 0 } else { retained[1] } + }; + let mut nnz = sides.map(|(lhs, rhs, ..)| combined_nnz_lower_bound(lhs, rhs)); + if cost(nnz) >= before { + return false; + } + + // Check the shorter side first. BDD outgoing rows are usually much + // shorter than incoming columns, and can reject an expensive trial. + let first = usize::from(totals[1] < totals[0]); + for side in [first, 1 - first] { + if nnz[side] == totals[side] { + continue; // Disjoint index ranges: the bound is exact. + } + let other = 1 - side; + let base = nnz[other] + if nnz[other] == 0 { 0 } else { retained[other] }; + // cost(nnz) < before implies base < before. A zero sum deletes + // its paired vector; a nonzero sum must also pay retained[side]. + let limit = (before - base - 1).saturating_sub(retained[side]); + let (lhs, rhs, subtract, require_ternary) = sides[side]; + let Some(count) = combined_nnz(lhs, rhs, subtract, require_ternary, limit) else { + return false; + }; + nnz[side] = count; + } + debug_assert!(cost(nnz) < before); + + let incoming = if nnz[0] == 0 { + Vec::new() + } else { + combine(&self.incoming[q], &self.incoming[p], subtract, ternary_left) + .expect("scored incoming sum") + }; + let outgoing = if nnz[1] == 0 { + Vec::new() + } else { + combine(&self.outgoing[p], &self.outgoing[q], !subtract, true) + .expect("scored outgoing sum") + }; + on_accept(&self.outgoing[p], &outgoing); + self.incoming[q] = incoming; + self.outgoing[p] = outgoing; + if self.outgoing[p].is_empty() { + self.incoming[p].clear(); + } + if self.incoming[q].is_empty() { + self.outgoing[q].clear(); + } + true + } + + fn into_layers(self) -> (MbpLayer, MbpLayer) { + let mut left_rows = vec![Vec::new(); 2 * self.left_rows]; + let mut right_rows: [Vec>; 2] = [Vec::new(), Vec::new()]; + let mut width = 0; + for (incoming, outgoing) in self.incoming.into_iter().zip(self.outgoing) { + if incoming.is_empty() || outgoing.is_empty() { + continue; + } + for (row, value) in incoming { + left_rows[row].push((width, value)); + } + let mut rows: [SparseVector; 2] = [Vec::new(), Vec::new()]; + for (column, value) in outgoing { + // With zero right width all outgoing vectors are empty. + rows[column / self.right_cols].push((column % self.right_cols, value)); + } + for (dest, row) in right_rows.iter_mut().zip(rows) { + dest.push(row); + } + width += 1; + } + let high_rows = left_rows.split_off(self.left_rows); + ( + MbpLayer { + matrices: [left_rows, high_rows].map(|rows| Core::from_rows(width, rows)), + }, + MbpLayer { + matrices: right_rows.map(|rows| Core::from_rows(self.right_cols, rows)), + }, + ) + } +} + +/// O(1) lower bound on the support size of a sum or difference. Exact when +/// either input is empty or their sorted index ranges do not overlap. +fn combined_nnz_lower_bound(lhs: &[(usize, F)], rhs: &[(usize, F)]) -> usize { + if lhs.is_empty() + || rhs.is_empty() + || lhs[lhs.len() - 1].0 < rhs[0].0 + || rhs[rhs.len() - 1].0 < lhs[0].0 + { + lhs.len() + rhs.len() + } else { + lhs.len().abs_diff(rhs.len()) + } +} + +/// Score a sparse sum without storing it. Inputs have sorted unique indices, +/// nonzero entries, and are already ternary when require_ternary is true. +/// Only intersecting entries need field arithmetic or alphabet checks. +/// Reject as soon as the finalized prefix alone exceeds the nonzero budget. +fn combined_nnz( + lhs: &[(usize, F)], + rhs: &[(usize, F)], + subtract: bool, + require_ternary: bool, + limit: usize, +) -> Option { + if lhs.len() / 8 > rhs.len() { + return sparse_update_nnz(lhs, rhs, subtract, require_ternary, limit); + } + if rhs.len() / 8 > lhs.len() { + // Swapping operands of a difference negates the result, preserving + // both support size and membership in {0, +/-1}. + return sparse_update_nnz(rhs, lhs, subtract, require_ternary, limit); + } + combined_nnz_linear(lhs, rhs, subtract, require_ternary, limit) +} + +fn combined_nnz_linear( + lhs: &[(usize, F)], + rhs: &[(usize, F)], + subtract: bool, + require_ternary: bool, + limit: usize, +) -> Option { + let (mut i, mut j, mut nnz) = (0, 0, 0); + while i < lhs.len() && j < rhs.len() { + match lhs[i].0.cmp(&rhs[j].0) { + std::cmp::Ordering::Less => { + i += 1; + nnz += 1; + } + std::cmp::Ordering::Greater => { + j += 1; + nnz += 1; + } + std::cmp::Ordering::Equal => { + let value = if subtract { + lhs[i].1 - rhs[j].1 + } else { + lhs[i].1 + rhs[j].1 + }; + if require_ternary && !ternary(value) { + return None; + } + nnz += usize::from(!value.is_zero()); + i += 1; + j += 1; + } + } + if nnz > limit { + return None; + } + } + nnz += lhs.len() - i + rhs.len() - j; + (nnz <= limit).then_some(nnz) +} + +/// Score a short update by jumping over untouched ranges of the longer vector. +fn sparse_update_nnz( + lhs: &[(usize, F)], + rhs: &[(usize, F)], + subtract: bool, + require_ternary: bool, + limit: usize, +) -> Option { + let (mut i, mut nnz) = (0, 0); + for &(column, value) in rhs { + let untouched = lhs[i..].partition_point(|&(index, _)| index < column); + i += untouched; + nnz += untouched; + if nnz > limit { + return None; + } + if i < lhs.len() && lhs[i].0 == column { + let value = if subtract { + lhs[i].1 - value + } else { + lhs[i].1 + value + }; + if require_ternary && !ternary(value) { + return None; + } + nnz += usize::from(!value.is_zero()); + i += 1; + } else { + nnz += 1; + } + if nnz > limit { + return None; + } + } + nnz += lhs.len() - i; + (nnz <= limit).then_some(nnz) +} + +/// Merge sorted sparse vectors. As with combined_nnz, inputs must already be +/// ternary if require_ternary is set; only overlaps can violate the alphabet. +fn combine( + lhs: &[(usize, F)], + rhs: &[(usize, F)], + subtract: bool, + require_ternary: bool, +) -> Option> { + if lhs.len() / 8 > rhs.len() { + let mut result = Vec::with_capacity(lhs.len() + rhs.len()); + let mut i = 0; + for &(column, value) in rhs { + let end = i + lhs[i..].partition_point(|&(index, _)| index < column); + result.extend_from_slice(&lhs[i..end]); + i = end; + let value = if i < lhs.len() && lhs[i].0 == column { + let value = if subtract { + lhs[i].1 - value + } else { + lhs[i].1 + value + }; + if require_ternary && !ternary(value) { + return None; + } + i += 1; + value + } else if subtract { + -value + } else { + value + }; + if !value.is_zero() { + result.push((column, value)); + } + } + result.extend_from_slice(&lhs[i..]); + return Some(result); + } + combine_linear(lhs, rhs, subtract, require_ternary) +} + +fn combine_linear( + lhs: &[(usize, F)], + rhs: &[(usize, F)], + subtract: bool, + require_ternary: bool, +) -> Option> { + let mut result = Vec::with_capacity(lhs.len() + rhs.len()); + let (mut i, mut j) = (0, 0); + while i < lhs.len() || j < rhs.len() { + let column = match (lhs.get(i), rhs.get(j)) { + (Some(a), Some(b)) => a.0.min(b.0), + (Some(a), None) => a.0, + (None, Some(b)) => b.0, + (None, None) => unreachable!(), + }; + let mut value = F::ZERO; + if i < lhs.len() && lhs[i].0 == column { + value += lhs[i].1; + i += 1; + } + if j < rhs.len() && rhs[j].0 == column { + value += if subtract { -rhs[j].1 } else { rhs[j].1 }; + j += 1; + } + if require_ternary && !ternary(value) { + return None; + } + if !value.is_zero() { + result.push((column, value)); + } + } + Some(result) +} + +#[cfg(test)] +mod tests { + use { + super::{ + super::{extract, extract_cqbdd, matrix::DenseMatrix, MinimizeStrategy}, + *, + }, + ark_ff::AdditiveGroup, + provekit_backend_bn254::FieldElement as F, + provekit_common::R1CS, + }; + + // Previous eager implementation, retained only as a test/benchmark oracle. + impl Cut { + fn reference_try_shear( + &mut self, + p: usize, + q: usize, + subtract: bool, + ternary_left: bool, + ) -> bool { + if self.incoming[p].is_empty() || self.incoming[q].is_empty() { + return false; + } + let Some(incoming) = + combine(&self.incoming[q], &self.incoming[p], subtract, ternary_left) + else { + return false; + }; + let Some(outgoing) = combine(&self.outgoing[p], &self.outgoing[q], !subtract, true) + else { + return false; + }; + let before = self.incoming[p].len() + + self.incoming[q].len() + + self.outgoing[p].len() + + self.outgoing[q].len(); + let after = incoming.len() + + outgoing.len() + + if outgoing.is_empty() { + 0 + } else { + self.incoming[p].len() + } + + if incoming.is_empty() { + 0 + } else { + self.outgoing[q].len() + }; + if after >= before { + return false; + } + self.incoming[q] = incoming; + self.outgoing[p] = outgoing; + if self.outgoing[p].is_empty() { + self.incoming[p].clear(); + } + if self.incoming[q].is_empty() { + self.outgoing[q].clear(); + } + true + } + } + + fn matrix(rows: usize, cols: usize, values: &[i64]) -> DenseMatrix { + assert_eq!(values.len(), rows * cols); + DenseMatrix { + num_rows: rows, + num_cols: cols, + values: values.iter().copied().map(F::from).collect(), + } + } + + fn dense_layer(rows: usize, cols: usize, values: [&[i64]; 2]) -> MbpLayer { + MbpLayer { + matrices: values.map(|values| Core::Dense(matrix(rows, cols, values))), + } + } + + // Independent evaluation of every assignment, without using the rewrites. + fn table(program: &MatrixBranchingProgram) -> Vec> { + (0..1usize << program.layers.len()) + .map(|assignment| { + if let Some(values) = program.constant { + return values.to_vec(); + } + let mut values = vec![F::ONE]; + for (level, layer) in program.layers.iter().enumerate().rev() { + let core = &layer.matrices[(assignment >> level) & 1]; + assert_eq!(core.num_cols(), values.len()); + let mut next = vec![F::ZERO; core.num_rows()]; + for ((row, column), value) in core.iter() { + next[row] += value * values[column]; + } + values = next; + } + values + }) + .collect() + } + + fn cost(program: &MatrixBranchingProgram) -> (usize, usize) { + ( + program.layers.iter().map(layer_nnz).sum(), + program + .layers + .iter() + .take(program.layers.len().saturating_sub(1)) + .map(|layer| layer.matrices[0].num_cols()) + .sum(), + ) + } + + fn assert_optimized(program: &mut MatrixBranchingProgram) { + let expected = table(program); + let before = cost(program); + let variables = program.layers.len(); + let constant = program.constant; + program.minimize_ternary(); + assert_eq!(table(program), expected); + assert_eq!(program.layers.len(), variables); + assert_eq!(program.constant, constant); + assert!(program.layers.iter().skip(1).all(layer_is_ternary)); + assert!(cost(program) <= before); + } + + #[test] + fn exhaustive_equivalence_for_generated_cores_and_storage_formats() { + for seed in 0..48u64 { + let mut state = seed + 1; + let mut next = || { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + state >> 32 + }; + let widths = [3, 1 + seed as usize % 5, 4, 3, 1]; + let layers = widths + .windows(2) + .enumerate() + .map(|(level, dims)| MbpLayer { + matrices: std::array::from_fn(|bit| { + let mut dense = DenseMatrix::zeros(dims[0], dims[1]); + for value in &mut dense.values { + *value = if level == 0 { + F::from((next() % 7) as i64 - 3) + } else { + match next() % 4 { + 0 => F::ONE, + 1 => -F::ONE, + _ => F::ZERO, + } + }; + } + match (seed as usize + level + bit) % 3 { + 0 => Core::Dense(dense), + 1 => Core::force_sparse(dense), + _ => Core::force_direct(dense), + } + }), + }) + .collect(); + let mut program = MatrixBranchingProgram { + index_bits: 2, + layers, + constant: None, + }; + assert_optimized(&mut program); + // A second invocation must already be at a fixed point. + let optimized = cost(&program); + assert_optimized(&mut program); + assert_eq!(cost(&program), optimized); + } + } + + #[test] + fn cached_sweeps_match_uncached_minimization() { + for seed in 0..24u64 { + let build = || { + let mut random = seed + 1; + let widths = [3, 40, 35, 8, 3, 1]; + let layers = widths + .windows(2) + .enumerate() + .map(|(level, dims)| MbpLayer { + matrices: std::array::from_fn(|bit| { + let mut dense = DenseMatrix::zeros(dims[0], dims[1]); + for value in &mut dense.values { + random = random + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + let code = (random >> 32) % 10; + *value = if level == 0 { + F::from(code as i64 - 5) + } else { + match code { + 0 => F::ONE, + 1 => -F::ONE, + _ => F::ZERO, + } + }; + } + match (seed as usize + level + bit) % 3 { + 0 => Core::Dense(dense), + 1 => Core::force_sparse(dense), + _ => Core::force_direct(dense), + } + }), + }) + .collect(); + MatrixBranchingProgram { + index_bits: 3, + layers, + constant: None, + } + }; + let mut cached = build(); + let mut reference = build(); + cached.minimize_ternary(); + loop { + let mut changed = false; + let cuts = reference.layers.len() - 1; + for forward in [true, false] { + for index in 0..cuts { + let cut = if forward { index } else { cuts - 1 - index }; + changed |= reference.try_ternary_rank(cut, forward); + changed |= reference.sparsify_ternary_cut(cut); + } + } + if !changed { + break; + } + } + assert_eq!(cost(&cached), cost(&reference)); + for (actual, expected) in cached.layers.iter().zip(&reference.layers) { + for bit in 0..2 { + let a = &actual.matrices[bit]; + let b = &expected.matrices[bit]; + assert_eq!((a.num_rows(), a.num_cols()), (b.num_rows(), b.num_cols())); + assert_eq!(a.iter().collect::>(), b.iter().collect::>()); + } + } + } + } + + #[test] + fn rejects_rank_proposal_when_absorbed_transfer_is_not_ternary() { + let mut program = MatrixBranchingProgram { + index_bits: 1, + layers: vec![ + dense_layer(3, 2, [&[1, 1, 0, 0, 0, 0], &[0; 6]]), + dense_layer(2, 1, [&[1, 1], &[1, 1]]), + ], + constant: None, + }; + // Forward factorization combines the two outgoing rows into [2, 2]. + let expected = table(&program); + let before = cost(&program); + assert!(!program.try_ternary_rank(0, true)); + assert_eq!(cost(&program), before); + assert_eq!(table(&program), expected); + + // Backward factorization can instead absorb the 2 into the weighted + // boundary. That is permitted, and reduces the number of nonzeros. + assert!(program.try_ternary_rank(0, false)); + assert_eq!(table(&program), expected); + assert!(cost(&program) < before); + assert!(layer_is_ternary(&program.layers[1])); + assert!(!layer_is_ternary(&program.layers[0])); + } + + #[test] + fn rejects_ternary_rank_reduction_that_increases_total_nonzeros() { + let mut program = MatrixBranchingProgram { + index_bits: 2, + layers: vec![ + dense_layer(3, 3, [&[1, 0, 1, 0, 1, 1, 0, 0, 0], &[0; 9]]), + dense_layer(3, 2, [&[0, 0, 0, 0, 1, 1], &[0, 0, 0, 0, 1, 1]]), + dense_layer(2, 1, [&[1, 1], &[0, 1]]), + ], + constant: None, + }; + // The third incoming column is the sum of the first two. Eliminating + // it saves two incoming entries but duplicates four outgoing entries. + assert_eq!(program.ternary_cut_cost(0), (8, 3)); + let expected = table(&program); + assert!(!program.try_ternary_rank(0, true)); + assert_eq!(program.ternary_cut_cost(0), (8, 3)); + assert_eq!(table(&program), expected); + } + + fn cut_product(cut: &Cut) -> DenseMatrix { + let mut result = DenseMatrix::zeros(2 * cut.left_rows, 2 * cut.right_cols); + for (incoming, outgoing) in cut.incoming.iter().zip(&cut.outgoing) { + for &(row, a) in incoming { + for &(column, b) in outgoing { + result[(row, column)] += a * b; + } + } + } + result + } + + fn pair_cost(cut: &Cut) -> usize { + cut.incoming.iter().chain(&cut.outgoing).map(Vec::len).sum() + } + + fn copy_cut(cut: &Cut) -> Cut { + Cut { + incoming: cut.incoming.clone(), + outgoing: cut.outgoing.clone(), + left_rows: cut.left_rows, + right_cols: cut.right_cols, + } + } + + fn sparse_vectors(alphabet: &[i64], width: usize) -> Vec> { + (0..alphabet.len().pow(width as u32)) + .map(|mut code| { + (0..width) + .filter_map(|column| { + let value = alphabet[code % alphabet.len()]; + code /= alphabet.len(); + (value != 0).then_some((column, F::from(value))) + }) + .collect() + }) + .collect() + } + + #[test] + fn bounded_scoring_matches_materialized_sums_exhaustively() { + for (alphabet, require_ternary) in + [(&[-1, 0, 1][..], true), (&[-2, -1, 0, 1, 2][..], false)] + { + let vectors = sparse_vectors(alphabet, 3); + for lhs in &vectors { + for rhs in &vectors { + for subtract in [false, true] { + let expected = combine(lhs, rhs, subtract, require_ternary); + for limit in 0..=6 { + assert_eq!( + combined_nnz(lhs, rhs, subtract, require_ternary, limit), + expected.as_ref().map(Vec::len).filter(|&nnz| nnz <= limit) + ); + } + if let Some(expected) = expected { + assert!(combined_nnz_lower_bound(lhs, rhs) <= expected.len()); + } + } + } + } + } + } + + #[test] + fn unbalanced_sparse_updates_match_linear_merges() { + for require_ternary in [false, true] { + for length in [0, 1, 8, 9, 32, 129, 1024] { + for seed in 0..20u64 { + let mut random = seed; + let mut coefficient = || { + random = random + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + let magnitude = if require_ternary { + 1 + } else { + 1 + (random >> 32) % 4 + }; + F::from(if random & 1 == 0 { + magnitude as i64 + } else { + -(magnitude as i64) + }) + }; + let long: Vec<_> = (0..length).map(|i| (2 * i, coefficient())).collect(); + // Empty updates, early/middle/late intersections, new + // coordinates, and insertion beyond the longer support. + let short: Vec<_> = [0, 1, length, 2 * length, 2 * length + 1] + .into_iter() + .take((seed % 6) as usize) + .collect::>() + .into_iter() + .map(|i| (i, coefficient())) + .collect(); + for (lhs, rhs) in [(&long, &short), (&short, &long)] { + for subtract in [false, true] { + let expected = combine_linear(lhs, rhs, subtract, require_ternary); + assert_eq!(combine(lhs, rhs, subtract, require_ternary), expected); + for limit in + [0, length / 2, length.saturating_sub(1), length, length + 5] + { + assert_eq!( + combined_nnz(lhs, rhs, subtract, require_ternary, limit), + expected.as_ref().map(Vec::len).filter(|&nnz| nnz <= limit) + ); + } + } + } + } + } + } + } + + #[test] + #[ignore = "manual unbalanced-update benchmark; run with --test-threads=1 --nocapture"] + fn benchmark_unbalanced_sparse_updates() { + use std::{hint::black_box, time::Instant}; + + let long: Vec<_> = (0..65_536).map(|i| (i, F::ONE)).collect(); + let short = [(0, F::ONE), (32_768, F::ONE), (65_535, F::ONE)]; + assert_eq!( + combine(&long, &short, true, true), + combine_linear(&long, &short, true, true) + ); + for (name, score) in [ + ( + "linear", + combined_nnz_linear:: + as fn(&[(usize, F)], &[(usize, F)], bool, bool, usize) -> Option, + ), + ("sparse", combined_nnz::), + ] { + let start = Instant::now(); + for _ in 0..200 { + assert_eq!( + score(black_box(&long), black_box(&short), true, true, long.len()), + Some(long.len() - short.len()) + ); + } + eprintln!("{name} scoring: {:?} per 200 updates", start.elapsed()); + } + for (name, merge) in [ + ( + "linear", + combine_linear:: + as fn(&[(usize, F)], &[(usize, F)], bool, bool) -> Option>, + ), + ("sparse", combine::), + ] { + let start = Instant::now(); + for _ in 0..200 { + black_box(merge(black_box(&long), black_box(&short), true, true).unwrap()); + } + eprintln!("{name} merge: {:?} per 200 updates", start.elapsed()); + } + } + + #[test] + fn lazy_shear_matches_eager_acceptance_and_state_exhaustively() { + let outgoing = sparse_vectors(&[-1, 0, 1], 2); + for ternary_left in [false, true] { + let incoming = if ternary_left { + sparse_vectors(&[-1, 0, 1], 2) + } else { + sparse_vectors(&[-2, -1, 0, 1, 2], 2) + }; + for ap in &incoming { + for aq in &incoming { + for bp in &outgoing { + for bq in &outgoing { + for subtract in [false, true] { + let mut actual = Cut { + incoming: vec![ap.clone(), aq.clone()], + outgoing: vec![bp.clone(), bq.clone()], + left_rows: 1, + right_cols: 1, + }; + let mut expected = copy_cut(&actual); + assert_eq!( + actual.try_shear(0, 1, subtract, ternary_left), + expected.reference_try_shear(0, 1, subtract, ternary_left) + ); + assert_eq!(actual.incoming, expected.incoming); + assert_eq!(actual.outgoing, expected.outgoing); + } + } + } + } + } + } + } + + #[test] + #[ignore = "manual rejected-shear benchmark; run with --test-threads=1 --nocapture"] + fn benchmark_rejected_shears() { + use std::{hint::black_box, time::Instant}; + + fn vector(start: usize, step: usize, count: usize, value: i64) -> SparseVector { + (0..count) + .map(|i| (start + step * i, F::from(value))) + .collect() + } + + let cases = [ + ( + "disjoint ranges", + [vector(0, 1, 64, 1), vector(64, 1, 64, 1)], + [vector(0, 1, 64, 1), vector(64, 1, 64, 1)], + ), + ( + "interleaved disjoint supports", + [vector(0, 2, 64, 1), vector(1, 2, 64, 1)], + [vector(0, 2, 64, 1), vector(1, 2, 64, 1)], + ), + ( + "early ternary conflict", + [vector(0, 1, 64, 1), vector(0, 1, 64, 1)], + [vector(0, 2, 64, 1), vector(1, 2, 64, 1)], + ), + ( + "short outgoing conflict", + [vector(0, 2, 512, 1), vector(1, 2, 512, 1)], + [vector(0, 1, 2, 1), vector(0, 1, 2, -1)], + ), + ( + "length bound", + [vector(0, 2, 512, 1), vector(1, 2, 8, 1)], + [vector(0, 2, 8, 1), vector(1, 2, 512, 1)], + ), + ("cost tie", [vector(0, 1, 1, 1), vector(1, 1, 1, 1)], [ + vector(0, 1, 2, 1), + vector(0, 1, 1, 1), + ]), + ( + "one-entry conflict", + [vector(0, 1, 1, 1), vector(0, 1, 1, 1)], + [vector(0, 1, 1, 1), vector(1, 1, 1, 1)], + ), + ]; + for (name, incoming, outgoing) in cases { + let mut cut = Cut { + incoming: incoming.into(), + outgoing: outgoing.into(), + left_rows: 512, + right_cols: 512, + }; + for (strategy, trial) in [ + ( + "eager", + Cut::reference_try_shear as fn(&mut Cut, usize, usize, bool, bool) -> bool, + ), + ("scored", Cut::try_shear), + ] { + for _ in 0..100 { + assert!(!trial(black_box(&mut cut), 0, 1, false, true)); + } + #[cfg(feature = "profiling-allocator")] + let count = crate::ALLOCATOR.count(); + let start = Instant::now(); + for _ in 0..100_000 { + assert!(!trial(black_box(&mut cut), 0, 1, false, true)); + } + let elapsed = start.elapsed(); + #[cfg(feature = "profiling-allocator")] + let allocations = crate::ALLOCATOR.count() - count; + eprintln!("{name}, {strategy}: {elapsed:?} per 100,000 rejected trials"); + #[cfg(feature = "profiling-allocator")] + eprintln!(" {allocations} allocations"); + } + } + } + + #[test] + fn paired_shear_reduces_nonzeros_without_changing_two_bit_products() { + for subtract in [false, true] { + let sign = if subtract { -F::ONE } else { F::ONE }; + let mut cut = Cut { + incoming: vec![vec![(0, F::ONE)], vec![(1, F::ONE)]], + outgoing: vec![vec![(0, sign), (1, sign), (2, F::ONE)], vec![ + (0, F::ONE), + (1, F::ONE), + ]], + left_rows: 1, + right_cols: 2, + }; + let expected = cut_product(&cut); + assert_eq!(pair_cost(&cut), 7); + assert!(cut.try_shear(0, 1, subtract, true)); + assert_eq!(pair_cost(&cut), 6); + assert_eq!(cut_product(&cut), expected); + let (left, right) = cut.into_layers(); + assert!(layer_is_ternary(&left)); + assert!(layer_is_ternary(&right)); + assert_eq!(cut_product(&Cut::new(&left, &right)), expected); + } + } + + #[test] + fn shear_rejects_nonternary_emissions_and_moves_without_savings() { + let mut cut = Cut { + incoming: vec![vec![(0, F::ONE)], vec![(0, F::ONE)]], + outgoing: vec![vec![(0, F::ONE), (1, F::ONE)]; 2], + left_rows: 1, + right_cols: 1, + }; + let expected = cut_product(&cut); + assert!(!cut.try_shear(0, 1, false, true)); // Incoming 2. + assert_eq!(cut_product(&cut), expected); + assert_eq!(pair_cost(&cut), 6); + assert!(cut.try_shear(0, 1, false, false)); // Allowed at the boundary. + assert_eq!(cut_product(&cut), expected); + assert!(cut.incoming[0].is_empty()); + let (left, right) = cut.into_layers(); + assert_eq!(left.matrices[0].num_cols(), 1); + assert!(!layer_is_ternary(&left)); + assert!(layer_is_ternary(&right)); + + let mut cut = Cut { + incoming: vec![vec![(0, F::ONE)], vec![(1, F::ONE)]], + outgoing: vec![vec![(0, F::ONE)], vec![(0, -F::ONE)]], + left_rows: 1, + right_cols: 1, + }; + assert!(!cut.try_shear(0, 1, false, true)); // Outgoing 2. + cut.outgoing[1] = vec![(1, F::ONE)]; + assert!(!cut.try_shear(0, 1, false, true)); // Adds nonzeros. + } + + #[test] + fn removes_dead_states_and_handles_zero_widths_and_constants() { + let mut program = MatrixBranchingProgram { + index_bits: 1, + layers: vec![ + dense_layer(3, 2, [&[1, 1, 2, 2, 3, 3]; 2]), + dense_layer(2, 1, [&[1, -1]; 2]), + ], + constant: None, + }; + assert_optimized(&mut program); + assert_eq!(cost(&program), (0, 0)); + assert_optimized(&mut program); + + let mut cut = Cut { + incoming: vec![vec![], vec![(0, F::ONE)], vec![]], + outgoing: vec![vec![(0, F::ONE)], vec![], vec![]], + left_rows: 1, + right_cols: 1, + }; + assert!(cut.remove_dead()); + let (left, right) = cut.into_layers(); + assert_eq!(left.matrices[0].num_cols(), 0); + assert_eq!(right.matrices[0].num_rows(), 0); + + let mut constant = MatrixBranchingProgram { + index_bits: 0, + layers: vec![], + constant: Some([F::from(7u64); 3]), + }; + assert_optimized(&mut constant); + let mut single = MatrixBranchingProgram { + index_bits: 0, + layers: vec![dense_layer(3, 1, [&[2, 3, 4], &[-5, 0, 7]])], + constant: None, + }; + assert_optimized(&mut single); + } + + #[test] + fn ternary_extraction_preserves_r1cs_coefficients_and_padded_domain() { + let mut r1cs = R1CS::::new(); + r1cs.add_witnesses(7); + for row in 0..5usize { + let a: Vec<_> = (0..7) + .map(|col| (F::from(((row + 1) * (col + 1)) as u64), col)) + .collect(); + r1cs.add_constraint(&a, &[(F::from(-3i64), row)], &[( + F::from(7u64), + (row + 1) % 7, + )]); + } + let original = extract_cqbdd(&r1cs).to_mbp(); + let optimized = extract(&r1cs, MinimizeStrategy::Ternary); + assert!(cost(&optimized) <= cost(&original)); + assert!(optimized.layers.iter().skip(1).all(layer_is_ternary)); + for (assignment, values) in table(&optimized).into_iter().enumerate() { + let mut row = 0; + let mut col = 0; + for bit in 0..3 { + row |= ((assignment >> (2 * bit)) & 1) << bit; + col |= ((assignment >> (2 * bit + 1)) & 1) << bit; + } + let expected = [&r1cs.a, &r1cs.b, &r1cs.c].map(|matrix| { + if row < matrix.num_rows && col < matrix.num_cols { + matrix + .get(row, col) + .map_or(F::ZERO, |value| r1cs.interner.get(value).unwrap()) + } else { + F::ZERO + } + }); + assert_eq!(values, expected, "row {row}, column {col}"); + } + let empty = extract(&R1CS::::new(), MinimizeStrategy::Ternary); + assert_eq!(empty.constant, Some([F::ZERO; 3])); + let mut constant = R1CS::::new(); + constant.add_witnesses(1); + constant.add_constraint(&[(F::from(7u64), 0)], &[], &[]); + assert_eq!( + extract(&constant, MinimizeStrategy::Ternary).constant, + Some([F::from(7u64), F::ZERO, F::ZERO]) + ); + } +} diff --git a/tooling/cli/src/cmd/extract/ternary/patterns.rs b/tooling/cli/src/cmd/extract/ternary/patterns.rs new file mode 100644 index 000000000..d2463528d --- /dev/null +++ b/tooling/cli/src/cmd/extract/ternary/patterns.rs @@ -0,0 +1,615 @@ +//! Share signed subexpressions across a cut: +//! L R = [L, L t] [R - t h; h]. The transpose shares incoming columns. +//! Candidate enumeration is bounded; every accepted rewrite is scored exactly. +use { + super::{ternary, Cut, SparseVector}, + ark_ff::Field, + std::{ + cmp::Reverse, + collections::{hash_map::Entry, HashMap}, + }, + tracing::{info, instrument}, +}; + +// At most C(8,2) + C(8,3) = 84 candidates per live state. Prefer coordinates +// used by many states; break ties by coordinate for deterministic discovery. +const MAX_PATTERN_ENTRIES: usize = 8; + +/// Sorted coordinates with first coefficient +1; remaining signs are relative +/// to the first. Pairs leave their last coordinate/sign zero. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)] +struct Pattern { + columns: [usize; 3], + negative: u8, + len: u8, +} + +impl Pattern { + fn new(entries: &[(usize, bool)]) -> Self { + let mut pattern = Self { + columns: [0; 3], + negative: 0, + len: entries.len() as u8, + }; + for (i, &(column, negative)) in entries.iter().enumerate() { + pattern.columns[i] = column; + pattern.negative |= u8::from(negative ^ entries[0].1) << i; + } + pattern + } + + fn columns(&self) -> &[usize] { + &self.columns[..usize::from(self.len)] + } + + fn value(&self, index: usize) -> F { + if self.negative & (1 << index) == 0 { + F::ONE + } else { + -F::ONE + } + } + + fn matches(&self, vector: &[(usize, F)], negative: bool) -> bool { + self.columns().iter().enumerate().all(|(i, column)| { + let Ok(index) = vector.binary_search_by_key(column, |&(column, _)| column) else { + return false; + }; + let expected = self.value::(i); + vector[index].1 == if negative { -expected } else { expected } + }) + } + + fn vector(&self) -> SparseVector { + self.columns() + .iter() + .enumerate() + .map(|(i, &column)| (column, self.value(i))) + .collect() + } +} + +#[derive(Clone, Copy)] +struct Occurrence { + state: usize, + negative: bool, +} + +/// Avoid a heap allocation for every pattern that occurs only once. +enum Occurrences { + One(Occurrence), + Many(Vec), +} + +fn record(table: &mut HashMap, entries: &[(usize, bool)], state: usize) { + let occurrence = Occurrence { + state, + negative: entries[0].1, + }; + match table.entry(Pattern::new(entries)) { + Entry::Vacant(entry) => { + entry.insert(Occurrences::One(occurrence)); + } + Entry::Occupied(mut entry) => { + let value = entry.get_mut(); + match value { + Occurrences::One(previous) => { + *value = Occurrences::Many(vec![*previous, occurrence]); + } + Occurrences::Many(occurrences) => occurrences.push(occurrence), + } + } + } +} + +fn candidates( + source: &[SparseVector], + other: &[SparseVector], + dimension: usize, +) -> Vec<(Pattern, Vec)> { + let mut frequencies = vec![0usize; dimension]; + for (row, opposite) in source.iter().zip(other) { + if opposite.is_empty() { + continue; + } + for &(column, value) in row { + if value == F::ONE || value == -F::ONE { + frequencies[column] += 1; + } + } + } + let mut table = HashMap::new(); + let mut entries = Vec::new(); + for (state, (row, opposite)) in source.iter().zip(other).enumerate() { + if opposite.is_empty() { + continue; + } + entries.clear(); + entries.extend(row.iter().filter_map(|&(column, value)| { + (frequencies[column] > 1 && (value == F::ONE || value == -F::ONE)) + .then_some((column, value != F::ONE)) + })); + if entries.len() > MAX_PATTERN_ENTRIES { + entries.select_nth_unstable_by_key(MAX_PATTERN_ENTRIES, |&(column, _)| { + (Reverse(frequencies[column]), column) + }); + entries.truncate(MAX_PATTERN_ENTRIES); + } + entries.sort_unstable_by_key(|&(column, _)| column); + for i in 0..entries.len() { + for j in i + 1..entries.len() { + record(&mut table, &[entries[i], entries[j]], state); + for k in j + 1..entries.len() { + record(&mut table, &[entries[i], entries[j], entries[k]], state); + } + } + } + } + let mut candidates: Vec<_> = table + .into_iter() + .filter_map(|(pattern, occurrences)| match occurrences { + Occurrences::One(_) => None, + Occurrences::Many(occurrences) => Some((pattern, occurrences)), + }) + .collect(); + // Prefer larger potential reuse. This is only a priority, not the actual + // gain: the opposite-side sum and dead-state savings are scored below. + candidates.sort_unstable_by_key(|(pattern, occurrences)| { + ( + Reverse((occurrences.len() - 1) * usize::from(pattern.len)), + *pattern, + ) + }); + candidates +} + +/// Reusable accumulator for sums of signed sparse vectors. Only touched slots +/// are cleared. Check the final sum: a partial +/-2 may subsequently cancel. +struct Sum { + values: Vec, + marked: Vec, + touched: Vec, +} + +impl Sum { + fn new(dimension: usize) -> Self { + Self { + values: vec![F::ZERO; dimension], + marked: vec![false; dimension], + touched: Vec::new(), + } + } + + fn add(&mut self, vector: &[(usize, F)], negative: bool) { + for &(column, value) in vector { + if !self.marked[column] { + self.marked[column] = true; + self.touched.push(column); + } + self.values[column] += if negative { -value } else { value }; + } + } + + fn nnz(&self, require_ternary: bool) -> Option { + let mut nnz = 0; + for &column in &self.touched { + let value = self.values[column]; + if require_ternary && !ternary(value) { + return None; + } + nnz += usize::from(!value.is_zero()); + } + Some(nnz) + } + + fn clear(&mut self) { + for column in self.touched.drain(..) { + self.values[column] = F::ZERO; + self.marked[column] = false; + } + } + + fn take(&mut self, nnz: usize) -> SparseVector { + let mut result = Vec::with_capacity(nnz); + self.touched.sort_unstable(); + for column in self.touched.drain(..) { + let value = std::mem::replace(&mut self.values[column], F::ZERO); + self.marked[column] = false; + if !value.is_zero() { + result.push((column, value)); + } + } + result + } +} + +impl Cut { + /// One discovery pass per orientation. New states and newly exposed + /// patterns are reconsidered by the enclosing optimizer's next sweep. + /// Each accepted extraction strictly reduces nnz, even if width grows. + #[instrument(skip_all)] + pub(super) fn extract_shared_patterns(&mut self, ternary_left: bool) -> usize { + let outgoing = extract( + &mut self.outgoing, + &mut self.incoming, + 2 * self.right_cols, + 2 * self.left_rows, + ternary_left, + "outgoing", + ); + let incoming = extract( + &mut self.incoming, + &mut self.outgoing, + 2 * self.left_rows, + 2 * self.right_cols, + true, + "incoming", + ); + outgoing + incoming + } +} + +/// Factor patterns out of source vectors and combine their opposite vectors. +/// The transpose uses the identical algebra. Source residuals only lose +/// entries, so only the new opposite vector can violate the coefficient +/// alphabet. +fn extract( + source: &mut Vec>, + other: &mut Vec>, + source_dimension: usize, + other_dimension: usize, + require_ternary: bool, + direction: &str, +) -> usize { + let candidates = candidates(source, other, source_dimension); + let count = candidates.len(); + if count == 0 { + return 0; + } + let mut sum = Sum::new(other_dimension); + let (mut accepted, mut saved) = (0, 0); + for (pattern, mut occurrences) in candidates { + // Earlier extractions may have removed shared entries or entire states. + // State IDs stay stable until Cut::into_layers compacts them. + occurrences.retain(|occurrence| { + !other[occurrence.state].is_empty() + && pattern.matches(&source[occurrence.state], occurrence.negative) + }); + if occurrences.len() < 2 { + continue; + } + let size = usize::from(pattern.len); + let mut removed = size * occurrences.len(); + for occurrence in &occurrences { + sum.add(&other[occurrence.state], occurrence.negative); + if source[occurrence.state].len() == size { + removed += other[occurrence.state].len(); + } + } + let Some(nnz) = sum.nnz(require_ternary) else { + sum.clear(); + continue; + }; + // A zero sum needs no helper state at all. + let added = if nnz == 0 { 0 } else { size + nnz }; + if added >= removed { + sum.clear(); + continue; + } + for occurrence in occurrences { + source[occurrence.state].retain(|&(column, _)| !pattern.columns().contains(&column)); + if source[occurrence.state].is_empty() { + other[occurrence.state].clear(); + } + } + if nnz == 0 { + sum.clear(); + } else { + source.push(pattern.vector()); + other.push(sum.take(nnz)); + } + accepted += 1; + saved += removed - added; + } + info!( + "Shared {direction} patterns: {count} candidates, {accepted} accepted, {saved} nonzeros \ + removed." + ); + saved +} + +#[cfg(test)] +mod tests { + use { + super::*, + ark_ff::{AdditiveGroup, Zero}, + provekit_backend_bn254::FieldElement as F, + }; + + fn vector(entries: &[(usize, i64)]) -> SparseVector { + entries + .iter() + .map(|&(column, value)| (column, F::from(value))) + .collect() + } + + fn example() -> Cut { + Cut { + incoming: (0..3).map(|i| vec![(i, F::ONE)]).collect(), + outgoing: (0..3) + .map(|i| vector(&[(0, 1), (1, 1), (i + 2, 1)])) + .collect(), + left_rows: 2, + right_cols: 3, + } + } + + fn transpose(cut: Cut) -> Cut { + Cut { + incoming: cut.outgoing, + outgoing: cut.incoming, + left_rows: cut.right_cols, + right_cols: cut.left_rows, + } + } + + fn nnz(cut: &Cut) -> usize { + cut.incoming.iter().chain(&cut.outgoing).map(Vec::len).sum() + } + + fn product(cut: &Cut) -> Vec> { + let mut result = vec![vec![F::ZERO; 2 * cut.right_cols]; 2 * cut.left_rows]; + for (incoming, outgoing) in cut.incoming.iter().zip(&cut.outgoing) { + for &(row, a) in incoming { + for &(column, b) in outgoing { + result[row][column] += a * b; + } + } + } + result + } + + fn check(cut: &mut Cut, ternary_left: bool) -> usize { + let before = nnz(cut); + let expected = product(cut); + let saved = cut.extract_shared_patterns(ternary_left); + assert_eq!(product(cut), expected); + assert_eq!(nnz(cut) + saved, before); + assert_eq!(cut.incoming.len(), cut.outgoing.len()); + for side in [&cut.incoming, &cut.outgoing] { + for row in side { + assert!(row.windows(2).all(|pair| pair[0].0 < pair[1].0)); + assert!(row.iter().all(|&(_, value)| !value.is_zero())); + } + } + assert!(cut + .outgoing + .iter() + .flatten() + .all(|&(_, value)| ternary(value))); + if ternary_left { + assert!(cut + .incoming + .iter() + .flatten() + .all(|&(_, value)| ternary(value))); + } + saved + } + + #[test] + fn shared_pair_escapes_single_shear_plateau_in_both_orientations() { + for reverse in [false, true] { + let mut cut = if reverse { + transpose(example()) + } else { + example() + }; + for p in 0..3 { + for q in 0..3 { + if p != q { + for subtract in [false, true] { + assert!(!cut.try_shear(p, q, subtract, true)); + } + } + } + } + assert_eq!(nnz(&cut), 12); + assert_eq!(check(&mut cut, true), 1); + assert_eq!(cut.incoming.len(), 4); + assert_eq!(nnz(&cut), 11); + let expected = product(&cut); + let (left, right) = cut.into_layers(); + assert_eq!(left.matrices[0].num_cols(), 4); + assert_eq!(product(&Cut::new(&left, &right)), expected); + } + } + + #[test] + fn signed_triples_are_shared_when_pairs_would_only_tie() { + let mut cut = Cut { + incoming: vec![vector(&[(0, 1)]), vector(&[(1, 1)])], + outgoing: vec![ + vector(&[(0, 1), (1, -1), (2, 1), (3, 1)]), + vector(&[(0, -1), (1, 1), (2, -1), (4, 1)]), + ], + left_rows: 1, + right_cols: 3, + }; + assert_eq!(nnz(&cut), 10); + assert_eq!(check(&mut cut, true), 1); + assert_eq!(cut.outgoing[2], vector(&[(0, 1), (1, -1), (2, 1)])); + assert_eq!(cut.incoming[2], vector(&[(0, 1), (1, -1)])); + } + + #[test] + fn scores_opposite_side_and_respects_boundary_exemption() { + let mut tie = example(); + tie.incoming.pop(); + tie.outgoing.pop(); + assert_eq!(check(&mut tie, true), 0); + + let mut conflict = example(); + conflict.incoming = vec![vector(&[(0, 1)]); 3]; + assert_eq!(check(&mut conflict, true), 0); + assert_eq!(check(&mut conflict, false), 3); + assert_eq!(conflict.incoming[3], vector(&[(0, 3)])); + + let mut transposed = example(); + transposed.incoming = vec![vector(&[(0, 1)]); 3]; + let mut transposed = transpose(transposed); + // Only the incoming (weighted) boundary can contain a general value; + // a shared incoming pattern may not introduce 3 on the outgoing side. + assert_eq!(check(&mut transposed, false), 0); + + let mut weighted = example(); + weighted.incoming = vec![vector(&[(0, 2)]), vector(&[(0, 3)]), vector(&[(0, -7)])]; + assert_eq!(check(&mut weighted, false), 3); + assert_eq!(weighted.incoming[3], vector(&[(0, -2)])); + } + + #[test] + fn checks_final_sum_after_cancellation_and_omits_zero_helpers() { + let mut cut = example(); + cut.incoming = vec![vector(&[(0, 1)]), vector(&[(0, 1)]), vector(&[(0, -1)])]; + assert_eq!(check(&mut cut, true), 3); // Partial sum 2, final sum 1. + assert_eq!(cut.incoming[3], vector(&[(0, 1)])); + + let mut cut = example(); + cut.incoming = vec![vector(&[(0, 1)]), vector(&[(0, -1)])]; + cut.outgoing.pop(); + assert_eq!(check(&mut cut, true), 4); + assert_eq!(cut.incoming.len(), 2); // No zero helper appended. + } + + #[test] + fn accounts_for_dead_states_and_rechecks_overlapping_candidates() { + let mut cut = example(); + for row in &mut cut.outgoing { + row.pop(); + } + assert_eq!(check(&mut cut, true), 4); + assert_eq!(nnz(&cut), 5); + let expected = product(&cut); + let (left, right) = cut.into_layers(); + assert_eq!(left.matrices[0].num_cols(), 1); + assert_eq!(product(&Cut::new(&left, &right)), expected); + + let mut cut = example(); + cut.outgoing = vec![vector(&[(0, 1), (1, 1), (2, 1), (3, 1)]); 3]; + // Removing one triple invalidates the other overlapping triples/pairs. + assert_eq!(check(&mut cut, true), 3); + assert_eq!(check(&mut cut, true), 0); + } + + #[test] + fn discovery_is_bounded_and_prefers_frequent_coordinates() { + let source = vec![(0..20).map(|i| (i, F::ONE)).collect::>(); 3]; + let other = (0..3).map(|i| vec![(i, F::ONE)]).collect::>(); + let found = candidates(&source, &other, 20); + assert_eq!(found.len(), 84); + assert!(found.iter().all(|(_, occurrences)| occurrences.len() == 3)); + + let mut source: Vec<_> = (0..3) + .map(|state| { + (0..20) + .map(|i| (state * 20 + i, F::ONE)) + .collect::>() + }) + .collect(); + for row in &mut source { + row.extend([(60, F::ONE), (61, -F::ONE), (62, F::ONE)]); + } + let found = candidates(&source, &other, 63); + assert_eq!(found.len(), 4); // Three pairs plus the shared triple. + assert!(found + .iter() + .any(|(p, _)| p.len == 3 && p.columns == [60, 61, 62])); + } + + #[test] + fn generated_cuts_preserve_every_two_bit_product_and_reach_a_fixed_point() { + let mut state = 1u64; + for case in 0..96 { + let mut random = || { + state = state.wrapping_mul(6364136223846793005).wrapping_add(1); + state >> 32 + }; + let width = case % 7; + let ternary_left = case % 2 == 0; + let mut vectors = |dimension, weighted| { + (0..width) + .map(|_| { + (0..dimension) + .filter_map(|column| { + let value = if weighted { + F::from((random() % 9) as i64 - 4) + } else { + F::from((random() % 3) as i64 - 1) + }; + (!value.is_zero()).then_some((column, value)) + }) + .collect() + }) + .collect() + }; + let mut cut = Cut { + incoming: vectors(6, !ternary_left), + outgoing: vectors(8, false), + left_rows: 3, + right_cols: 4, + }; + cut.remove_dead(); + let expected = product(&cut); + let bound = nnz(&cut); + let mut stopped = false; + for _ in 0..=bound { + if check(&mut cut, ternary_left) == 0 { + stopped = true; + break; + } + } + assert!(stopped); + let (left, right) = cut.into_layers(); + assert_eq!(product(&Cut::new(&left, &right)), expected); + } + } + + #[test] + #[ignore = "manual shared-pattern throughput; run with --test-threads=1 --nocapture"] + fn benchmark_shared_patterns() { + use std::time::Instant; + + let groups = 10_000; + let width = 3 * groups; + let mut cut = Cut { + incoming: (0..width).map(|i| vec![(i, F::ONE)]).collect(), + outgoing: (0..width) + .map(|i| vector(&[(2 * (i / 3), 1), (2 * (i / 3) + 1, 1), (2 * groups + i, 1)])) + .collect(), + left_rows: width / 2, + right_cols: (2 * groups + width) / 2, + }; + let before = nnz(&cut); + #[cfg(feature = "profiling-allocator")] + let (base, count) = (crate::ALLOCATOR.reset_max(), crate::ALLOCATOR.count()); + let start = Instant::now(); + let saved = cut.extract_shared_patterns(true); + let elapsed = start.elapsed(); + #[cfg(feature = "profiling-allocator")] + let (peak, allocations) = ( + crate::ALLOCATOR.max().saturating_sub(base), + crate::ALLOCATOR.count() - count, + ); + assert_eq!(saved, groups); + assert_eq!(before - nnz(&cut), saved); + assert_eq!(cut.incoming.len(), width + groups); + eprintln!( + "shared patterns: {elapsed:?}, width {width} -> {}, nnz {before} -> {}", + cut.incoming.len(), + nnz(&cut) + ); + #[cfg(feature = "profiling-allocator")] + eprintln!("peak additional heap {peak} bytes, {allocations} allocations"); + } +} diff --git a/tooling/cli/src/cmd/extract/ternary/shears.rs b/tooling/cli/src/cmd/extract/ternary/shears.rs new file mode 100644 index 000000000..468b347ae --- /dev/null +++ b/tooling/cli/src/cmd/extract/ternary/shears.rs @@ -0,0 +1,425 @@ +//! Exact candidate pruning for the ordered shear scan. +//! +//! A profitable shear must shrink either A[:,q] or B[p,:]. For ternary +//! vectors, shrinking dst +/- src requires overlap with more than half of +//! src's support. Every such overlap hits any ceil(|src|/2) source entries; +//! choose rare entries as anchors to avoid high-degree coordinates. +//! +//! At the weighted boundary A need not be ternary, but shrinking still +//! requires at least one exact cancellation. Index (coordinate, +/-value) +//! there instead. Both filters are necessary conditions, never heuristics. + +use { + super::{Cut, SparseVector}, + ark_ff::Field, + std::{ + collections::{BTreeSet, HashMap}, + hash::Hash, + time::Instant, + }, + tracing::info, +}; + +/// Sorted postings preserve the original p, q, sign trial order. Keys are +/// cached per state so an accepted rewrite can update just the affected lists. +struct Index { + postings: HashMap>, + keys: Vec>, +} + +impl Index { + fn new(keys: Vec>) -> Self { + let mut postings = HashMap::>::new(); + for (state, keys) in keys.iter().enumerate() { + for &key in keys { + postings.entry(key).or_default().insert(state); + } + } + Self { postings, keys } + } + + fn replace(&mut self, state: usize, keys: Vec) { + if self.keys[state] == keys { + return; + } + for key in &self.keys[state] { + let states = self.postings.get_mut(key).expect("indexed key"); + states.remove(&state); + if states.is_empty() { + self.postings.remove(key); + } + } + for &key in &keys { + self.postings.entry(key).or_default().insert(state); + } + self.keys[state] = keys; + } + + fn extend(&self, keys: &[K], from: usize, pending: &mut BTreeSet) { + for key in keys { + if let Some(states) = self.postings.get(key) { + pending.extend(states.range(from..).copied()); + } + } + } + + fn anchors(&self, state: usize, ternary: bool) -> Vec { + let mut keys = self.keys[state].clone(); + if ternary { + let count = keys.len().div_ceil(2); + if count < keys.len() { + keys.select_nth_unstable_by_key(count, |key| (self.postings[key].len(), *key)); + keys.truncate(count); + } + } + keys + } +} + +fn outgoing_anchors(vector: &SparseVector, degrees: &[usize]) -> Vec { + let mut keys: Vec<_> = vector.iter().map(|&(column, _)| column).collect(); + let count = keys.len().div_ceil(2); + if count < keys.len() { + keys.select_nth_unstable_by_key(count, |&column| (degrees[column], column)); + keys.truncate(count); + } + keys.sort_unstable(); + keys +} + +/// Coordinates already queried cannot expose a new future q: its anchors +/// stay fixed throughout this p scan. Only newly added support needs a query. +/// In particular, deleting entries from a large hub causes no posting walks. +fn extend_added( + index: &Index, + before: &[(usize, F)], + after: &[(usize, F)], + from: usize, + pending: &mut BTreeSet, +) { + let mut old = 0; + for &(column, _) in after { + while old < before.len() && before[old].0 < column { + old += 1; + } + if old == before.len() || before[old].0 != column { + index.extend(&[column], from, pending); + } + } +} + +impl Cut { + pub(super) fn shear_sweep(&mut self, ternary_left: bool) -> bool { + // Tiny cuts are cheaper to scan directly than to index. + if self.incoming.len() < 32 { + return self.shear_sweep_all_pairs(ternary_left); + } + self.shear_sweep_indexed(ternary_left) + } + + fn shear_sweep_all_pairs(&mut self, ternary_left: bool) -> bool { + let mut changed = false; + for p in 0..self.incoming.len() { + for q in 0..self.incoming.len() { + if p != q { + for subtract in [false, true] { + changed |= self.try_shear(p, q, subtract, ternary_left); + } + } + } + } + changed + } + + fn shear_sweep_indexed(&mut self, ternary_left: bool) -> bool { + if ternary_left { + self.shear_sweep_with_keys(true, |column, _| column) + } else { + self.shear_sweep_with_keys(false, |column, value: F| (column, value.min(-value))) + } + } + + fn shear_sweep_with_keys( + &mut self, + ternary_left: bool, + key: impl Fn(usize, F) -> K, + ) -> bool { + let start = Instant::now(); + let keys = |vector: &SparseVector| { + vector + .iter() + .map(|&(column, value)| key(column, value)) + .collect() + }; + let width = self.incoming.len(); + let mut incoming = Index::new(self.incoming.iter().map(keys).collect()); + // Fixed degrees suffice: rarity affects efficiency, not coverage. + let mut degrees = vec![0usize; 2 * self.right_cols]; + for vector in &self.outgoing { + for &(column, _) in vector { + degrees[column] += 1; + } + } + let mut outgoing = Index::new( + self.outgoing + .iter() + .map(|v| outgoing_anchors(v, °rees)) + .collect(), + ); + let mut pending = BTreeSet::new(); + let (mut trials, mut accepted) = (0usize, 0usize); + for p in 0..width { + if self.incoming[p].is_empty() { + continue; + } + pending.clear(); + incoming.extend(&incoming.anchors(p, ternary_left), 0, &mut pending); + // Reverse query: outgoing[q] is the source of B[p] -= s B[q]. + for &(column, _) in &self.outgoing[p] { + outgoing.extend(&[column], 0, &mut pending); + } + let mut changed_p = false; + while let Some(q) = pending.pop_first() { + if p == q || self.incoming[q].is_empty() { + continue; + } + let mut changed_q = false; + for subtract in [false, true] { + trials += 1; + if self.try_shear_with(p, q, subtract, ternary_left, |before, after| { + extend_added(&outgoing, before, after, q + 1, &mut pending); + }) { + accepted += 1; + changed_q = true; + } + } + if !changed_q { + continue; + } + changed_p = true; + incoming.replace(q, keys(&self.incoming[q])); + if self.incoming[q].is_empty() { + outgoing.replace(q, Vec::new()); + } + if self.incoming[p].is_empty() { + incoming.replace(p, Vec::new()); + break; + } + // A[:,p] and future A[:,q] are unchanged. The acceptance + // callback already enqueued any new B[p,:] candidates. + } + if changed_p { + outgoing.replace(p, outgoing_anchors(&self.outgoing[p], °rees)); + } + } + info!( + "Shear scan: {trials} of {} signed trials, {accepted} accepted, {:?}.", + width + .saturating_mul(width.saturating_sub(1)) + .saturating_mul(2), + start.elapsed() + ); + accepted != 0 + } +} + +#[cfg(test)] +mod tests { + use {super::*, provekit_backend_bn254::FieldElement as F}; + + fn copy_cut(cut: &Cut) -> Cut { + Cut { + incoming: cut.incoming.clone(), + outgoing: cut.outgoing.clone(), + left_rows: cut.left_rows, + right_cols: cut.right_cols, + } + } + + fn compare(mut cut: Cut, ternary_left: bool) { + cut.remove_dead(); + let mut expected = copy_cut(&cut); + for _ in 0..3 { + assert_eq!( + cut.shear_sweep_indexed(ternary_left), + expected.shear_sweep_all_pairs(ternary_left) + ); + assert_eq!(cut.incoming, expected.incoming, "different incoming cores"); + assert_eq!(cut.outgoing, expected.outgoing, "different outgoing cores"); + } + } + + fn vectors(alphabet: &[i64]) -> Vec> { + alphabet + .iter() + .flat_map(|&x| { + alphabet.iter().map(move |&y| { + [x, y] + .into_iter() + .enumerate() + .filter_map(|(column, value)| { + (value != 0).then_some((column, F::from(value))) + }) + .collect() + }) + }) + .collect() + } + + #[test] + fn indexed_scan_matches_all_pairs_exhaustively() { + let outgoing = vectors(&[-1, 0, 1]); + for ternary_left in [false, true] { + let incoming = vectors(if ternary_left { + &[-1, 0, 1] + } else { + &[-2, -1, 0, 1, 2] + }); + for ap in &incoming { + for aq in &incoming { + for bp in &outgoing { + for bq in &outgoing { + compare( + Cut { + incoming: vec![ap.clone(), aq.clone()], + outgoing: vec![bp.clone(), bq.clone()], + left_rows: 1, + right_cols: 1, + }, + ternary_left, + ); + } + } + } + } + } + } + + #[test] + fn indexed_scan_matches_mutating_scans_with_hubs_and_general_boundary() { + let mut seed = 783837_u64; + let mut random = || { + seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1); + seed >> 32 + }; + for _ in 0..2000 { + let width = 2 + random() as usize % 18; + let dimension = 2 * (1 + random() as usize % 8); + let ternary_left = random() % 2 == 0; + let mut vectors = |general| { + (0..width) + .map(|_| { + (0..dimension) + .filter_map(|column| { + let sample = random(); + let value = if general { + (sample % 9) as i64 - 4 + } else { + (sample % 3) as i64 - 1 + }; + (value != 0 && (column == 0 || sample % 5 < 2)) + .then_some((column, F::from(value))) + }) + .collect() + }) + .collect() + }; + compare( + Cut { + incoming: vectors(!ternary_left), + outgoing: vectors(false), + left_rows: dimension / 2, + right_cols: dimension / 2, + }, + ternary_left, + ); + } + } + + #[test] + fn index_updates_remove_old_keys_and_keep_postings_ordered() { + let mut index = Index::new(vec![vec![1, 4], vec![1, 2], vec![2, 3]]); + index.replace(1, vec![3, 4]); + index.replace(0, Vec::new()); + let mut pending = BTreeSet::new(); + index.extend(&[1, 2, 3, 4], 1, &mut pending); + assert_eq!(pending.into_iter().collect::>(), [1, 2]); + assert!(!index.postings.contains_key(&1)); + assert_eq!(index.postings[&4].iter().copied().collect::>(), [1]); + } + + #[test] + fn outgoing_refresh_queries_only_added_coordinates_and_future_states() { + let index = Index::new(vec![vec![3], vec![1], vec![2], vec![3]]); + let before = vec![(1, F::ONE), (2, F::ONE)]; + let after = vec![(2, -F::ONE), (3, F::ONE)]; + let mut pending = BTreeSet::new(); + extend_added(&index, &before, &after, 1, &mut pending); + assert_eq!(pending.into_iter().collect::>(), [3]); + } + + #[test] + #[ignore = "manual real-cut comparison; set PROVEKIT_BENCH_VERIFIER and optionally \ + PROVEKIT_BENCH_CUT"] + fn benchmark_target_shear_cut() { + use { + crate::cmd::extract::extract_cqbdd, provekit_backend_bn254::Verifier, + provekit_common::file::read, + }; + provekit_backend_bn254::register(); + let path = std::env::var("PROVEKIT_BENCH_VERIFIER").expect("path to a prepared PKV"); + let cut: usize = std::env::var("PROVEKIT_BENCH_CUT") + .unwrap_or_else(|_| "7".into()) + .parse() + .unwrap(); + let verifier: Verifier = read(std::path::Path::new(&path)).unwrap(); + let mut program = extract_cqbdd(&verifier.r1cs).to_mbp(); + program.try_ternary_rank(cut, true); + let mut original = Cut::new(&program.layers[cut], &program.layers[cut + 1]); + original.remove_dead(); + let mut indexed = copy_cut(&original); + let start = Instant::now(); + let changed = original.shear_sweep_all_pairs(cut != 0); + let baseline = start.elapsed(); + let start = Instant::now(); + assert_eq!(indexed.shear_sweep_indexed(cut != 0), changed); + eprintln!( + "Real cut {cut}, width {}: all-pairs {baseline:?}, indexed {:?}", + original.incoming.len(), + start.elapsed() + ); + assert_eq!(indexed.incoming, original.incoming); + assert_eq!(indexed.outgoing, original.outgoing); + } + + #[test] + #[ignore = "manual sparse scan benchmark; run with --test-threads=1 --nocapture"] + fn benchmark_sparse_shear_scan() { + // A large weighted boundary with distinct coefficients, as in the + // first sha256_35 cut. Outgoing rows have disjoint supports. + let width = 20_000; + let cut = Cut { + incoming: (0..width) + .map(|i| vec![(0, F::from((i + 1) as u64))]) + .collect(), + outgoing: (0..width) + .map(|i| vec![(2 * i, F::ONE), (2 * i + 1, F::ONE)]) + .collect(), + left_rows: 1, + right_cols: width, + }; + for indexed in [false, true] { + let mut actual = copy_cut(&cut); + let start = Instant::now(); + let changed = if indexed { + actual.shear_sweep_indexed(false) + } else { + actual.shear_sweep_all_pairs(false) + }; + eprintln!("indexed={indexed}, width={width}: {:?}", start.elapsed()); + assert!(!changed); + assert_eq!(actual.incoming, cut.incoming); + assert_eq!(actual.outgoing, cut.outgoing); + } + } +} diff --git a/tooling/cli/src/cmd/mod.rs b/tooling/cli/src/cmd/mod.rs index ccd2b6373..4390061d4 100644 --- a/tooling/cli/src/cmd/mod.rs +++ b/tooling/cli/src/cmd/mod.rs @@ -1,5 +1,6 @@ mod analyze_pkp; mod circuit_stats; +mod extract; mod generate_gnark_inputs; pub mod prepare; mod prove; diff --git a/tooling/cli/src/cmd/prepare.rs b/tooling/cli/src/cmd/prepare.rs index 639e8018f..e4014eee6 100644 --- a/tooling/cli/src/cmd/prepare.rs +++ b/tooling/cli/src/cmd/prepare.rs @@ -1,5 +1,5 @@ use { - super::{util::resolve_key_path, Command}, + super::{extract::MinimizeStrategy, util::resolve_key_path, Command}, anyhow::{anyhow, bail, Context as _, Result}, argh::FromArgs, mavros_artifacts::R1CS as MavrosR1CS, @@ -60,6 +60,14 @@ pub struct Args { #[argh(option, long = "compiler", default = "Compiler::Noir")] compiler: Compiler, + /// MBP optimization: "rank" (default) or "ternary" (noir only) + #[argh(option, long = "mbp-strategy", default = "MinimizeStrategy::Rank")] + mbp_strategy: MinimizeStrategy, + + /// skip Gaussian elimination optimization of the R1CS (noir only) + #[argh(switch, long = "skip-ge")] + skip_ge: bool, + /// name of the package to compile (noir only; default: enclosing package) #[argh(option)] package: Option, @@ -212,8 +220,16 @@ impl Args { )?; for (package, artifact) in binary_packages.iter().zip(artifacts) { - let scheme = NoirCompiler::from_program(artifact, hash_config, witness_mode) - .context("while building Noir proof scheme")?; + let scheme = NoirCompiler::from_program_with_ge( + artifact, + hash_config, + witness_mode, + !self.skip_ge, + ) + .context("while building Noir proof scheme")?; + + super::extract::extract(scheme.r1cs(), self.mbp_strategy); + let spark_setup = self.maybe_build_spark(&scheme, hash_config)?; let pkp_path = self .pkp_path @@ -480,3 +496,26 @@ pub fn build_spark_r1cs_mavros( Ok(SparkMatrix::new(row, col, val, 2 * row_cnt, 2 * col_cnt)) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_mbp_strategy_and_preserves_rank_default() { + let default = Args::from_args(&["prepare"], &[]).unwrap(); + assert_eq!(default.mbp_strategy, MinimizeStrategy::Rank); + assert!(!default.skip_ge); + let unoptimized = Args::from_args(&["prepare"], &["--skip-ge"]).unwrap(); + assert!(unoptimized.skip_ge); + for (name, expected) in [ + ("rank", MinimizeStrategy::Rank), + ("ternary", MinimizeStrategy::Ternary), + ] { + let args = Args::from_args(&["prepare"], &["--mbp-strategy", name]).unwrap(); + assert_eq!(args.mbp_strategy, expected); + } + let error = Args::from_args(&["prepare"], &["--mbp-strategy", "unknown"]).unwrap_err(); + assert!(error.output.contains("Unknown MBP strategy")); + } +} diff --git a/tooling/cli/src/main.rs b/tooling/cli/src/main.rs index d273bace7..9a025d141 100644 --- a/tooling/cli/src/main.rs +++ b/tooling/cli/src/main.rs @@ -13,7 +13,9 @@ use { anyhow::Result, span_stats::SpanStats, tracing::subscriber, - tracing_subscriber::{self, filter::LevelFilter, layer::SubscriberExt as _, Layer, Registry}, + tracing_subscriber::{ + self, filter::LevelFilter, layer::SubscriberExt as _, EnvFilter, Layer, Registry, + }, }; #[cfg(feature = "profiling-allocator")] @@ -28,7 +30,11 @@ fn main() -> Result<()> { let level = LevelFilter::TRACE; #[cfg(not(debug_assertions))] let level = LevelFilter::INFO; - let subscriber = Registry::default().with(SpanStats.with_filter(level)); + // RUST_LOG can override the default level and filter individual crates. + let filter = EnvFilter::builder() + .with_default_directive(level.into()) + .from_env()?; + let subscriber = Registry::default().with(SpanStats.with_filter(filter)); #[cfg(feature = "tracy")] let subscriber = {