Two multi-group root fixes + downstream preset/catalog enablement (from aerie)

crates/akita-config/src/proof_optimized.rs

@@ -23,9 +23,9 @@ use std::sync::{LazyLock, Mutex};
 /// smallest reachable `2`) keeps the shrink strong enough that every preset — dense
 /// and small-field included — supports the full `nv` range, and matches the value
 /// the unpinned planner already favored at the root.
-pub(crate) const PROOF_OPTIMIZED_LOG_BASIS_MIN: u32 = 3;
+pub const PROOF_OPTIMIZED_LOG_BASIS_MIN: u32 = 3;
 /// Maximum proof-optimized log-basis.
-pub(crate) const PROOF_OPTIMIZED_LOG_BASIS_MAX: u32 = 6;
+pub const PROOF_OPTIMIZED_LOG_BASIS_MAX: u32 = 6;
 /// Explicit sparse-binary chunk size used by standard one-hot presets.
 ///
 /// Smaller/nonstandard chunking is represented by a separately named preset
@@ -48,7 +48,7 @@ const DEFAULT_GROUP_BATCH_MAX_PRECOMMITTED_GROUPS: usize = 2;
 /// [`akita_challenges::SparseChallengeConfig::production_for_ring_dim`].
 /// A preset's `D` is fixed across all schedule levels, so both the planner DP
 /// and the generated-table expansion call the per-`Cfg` hook with `d == Cfg::D`.
-pub(crate) fn proof_optimized_ring_challenge_config(
+pub fn proof_optimized_ring_challenge_config(
     d: usize,
 ) -> Result<akita_challenges::SparseChallengeConfig, AkitaError> {
     let cfg =
@@ -60,7 +60,7 @@ pub(crate) fn proof_optimized_ring_challenge_config(
     Ok(cfg)
 }
 
-pub(crate) fn proof_optimized_schedule_key(
+pub fn proof_optimized_schedule_key(
     layout: &OpeningClaimsLayout,
 ) -> Result<AkitaScheduleLookupKey, AkitaError> {
     layout.check()?;
@@ -91,7 +91,7 @@ type SetupMatrixCapacityCache =
 static SETUP_MATRIX_CAPACITY_CACHE: SetupMatrixCapacityCache =
     LazyLock::new(|| Mutex::new(HashMap::new()));
 
-pub(crate) fn proof_optimized_setup_matrix_capacity<Cfg: CommitmentConfig>(
+pub fn proof_optimized_setup_matrix_capacity<Cfg: CommitmentConfig>(
     max_num_vars: usize,
     max_num_batched_polys: usize,
 ) -> Result<SetupMatrixCapacity, AkitaError> {
@@ -402,6 +402,13 @@ fn matrix_coefficient_len(
 /// shared ring-challenge policy, the shared setup-matrix sizer, and the
 /// `[PROOF_OPTIMIZED_LOG_BASIS_MIN, MAX]` basis range, so those are not
 /// parameters.
+///
+/// Exported so downstream workspaces can define their own application
+/// `CommitmentConfig` presets with identical semantics; the expansion
+/// references `akita_types`, `akita_challenges`, `akita_schedules`, and
+/// `akita_field` by name, so invokers must have those crates as direct
+/// dependencies.
+#[macro_export]
 macro_rules! impl_proof_optimized_preset {
     (@selection_policy default) => {
         fn selection_policy() -> akita_schedules::SelectionPolicyId {
@@ -435,41 +442,41 @@ macro_rules! impl_proof_optimized_preset {
         const RING_DIMENSION_CANDIDATES: &'static [akita_types::CommitmentRingDims] = $candidates;
     };
     ($cfg:ident, $field:ty, $ext_field:ty, $family:expr, $d:expr, $field_bits:expr, $log_commit_bound:expr, fold_norms = $fold_norms:expr) => {
-        impl_proof_optimized_preset!(@core $cfg, $field, $ext_field, $family, $d, $field_bits, $log_commit_bound, $fold_norms, none, default);
+        $crate::impl_proof_optimized_preset!(@core $cfg, $field, $ext_field, $family, $d, $field_bits, $log_commit_bound, $fold_norms, none, default);
     };
     ($cfg:ident, $field:ty, $ext_field:ty, $family:expr, $d:expr, $field_bits:expr, $log_commit_bound:expr, fold_norms = $fold_norms:expr, schedules = ($feat:literal, $family_name:literal, $table:ident)) => {
-        impl_proof_optimized_preset!(@core $cfg, $field, $ext_field, $family, $d, $field_bits, $log_commit_bound, $fold_norms, table, $feat, $family_name, $table, default);
+        $crate::impl_proof_optimized_preset!(@core $cfg, $field, $ext_field, $family, $d, $field_bits, $log_commit_bound, $fold_norms, table, $feat, $family_name, $table, default);
     };
     ($cfg:ident, $field:ty, $ext_field:ty, $family:expr, $d:expr, $field_bits:expr, $log_commit_bound:expr, fold_norms = $fold_norms:expr, schedules = ($feat:literal, $family_name:literal, $table:ident), selection_policy = $selection_policy:expr) => {
-        impl_proof_optimized_preset!(@core $cfg, $field, $ext_field, $family, $d, $field_bits, $log_commit_bound, $fold_norms, table, $feat, $family_name, $table, selection_policy = $selection_policy);
+        $crate::impl_proof_optimized_preset!(@core $cfg, $field, $ext_field, $family, $d, $field_bits, $log_commit_bound, $fold_norms, table, $feat, $family_name, $table, selection_policy = $selection_policy);
     };
     ($cfg:ident, $field:ty, $ext_field:ty, $family:expr, $d:expr, $field_bits:expr, $log_commit_bound:expr, fold_norms = $fold_norms:expr, schedules = ($feat:literal, $family_name:literal, $table:ident), ring_dimension_candidates = $candidates:expr) => {
-        impl_proof_optimized_preset!(@core $cfg, $field, $ext_field, $family, $d, $field_bits, $log_commit_bound, $fold_norms, table, $feat, $family_name, $table, ring_dimension_candidates = $candidates);
+        $crate::impl_proof_optimized_preset!(@core $cfg, $field, $ext_field, $family, $d, $field_bits, $log_commit_bound, $fold_norms, table, $feat, $family_name, $table, ring_dimension_candidates = $candidates);
     };
     ($cfg:ident, $field:ty, $ext_field:ty, $family:expr, $d:expr, $field_bits:expr, $log_commit_bound:expr, fold_norms = $fold_norms:expr, schedules = ($feat:literal, $family_name:literal, $table:ident), selection_policy = $selection_policy:expr, ring_dimension_candidates = $candidates:expr) => {
-        impl_proof_optimized_preset!(@core $cfg, $field, $ext_field, $family, $d, $field_bits, $log_commit_bound, $fold_norms, table, $feat, $family_name, $table, selection_policy = $selection_policy, ring_dimension_candidates = $candidates);
+        $crate::impl_proof_optimized_preset!(@core $cfg, $field, $ext_field, $family, $d, $field_bits, $log_commit_bound, $fold_norms, table, $feat, $family_name, $table, selection_policy = $selection_policy, ring_dimension_candidates = $candidates);
     };
     (@options default) => {
-        impl_proof_optimized_preset!(@selection_policy default);
+        $crate::impl_proof_optimized_preset!(@selection_policy default);
     };
     (@options selection_policy = $selection_policy:expr) => {
-        impl_proof_optimized_preset!(@selection_policy $selection_policy);
+        $crate::impl_proof_optimized_preset!(@selection_policy $selection_policy);
     };
     (@options ring_dimension_candidates = $candidates:expr) => {
-        impl_proof_optimized_preset!(@ring_dimension_candidates $candidates);
-        impl_proof_optimized_preset!(@selection_policy default);
+        $crate::impl_proof_optimized_preset!(@ring_dimension_candidates $candidates);
+        $crate::impl_proof_optimized_preset!(@selection_policy default);
     };
     (@options selection_policy = $selection_policy:expr, ring_dimension_candidates = $candidates:expr) => {
-        impl_proof_optimized_preset!(@ring_dimension_candidates $candidates);
-        impl_proof_optimized_preset!(@selection_policy $selection_policy);
+        $crate::impl_proof_optimized_preset!(@ring_dimension_candidates $candidates);
+        $crate::impl_proof_optimized_preset!(@selection_policy $selection_policy);
     };
     (@core $cfg:ident, $field:ty, $ext_field:ty, $family:expr, $d:expr, $field_bits:expr, $log_commit_bound:expr, $fold_norms:expr, none, $($options:tt)*) => {
         impl $crate::CommitmentConfig for $cfg {
             type Field = $field;
             type ExtField = $ext_field;
             const D: usize = $d;
 
-            impl_proof_optimized_preset!(@options $($options)*);
+            $crate::impl_proof_optimized_preset!(@options $($options)*);
 
             fn decomposition() -> akita_types::DecompositionParams {
                 akita_types::DecompositionParams {
@@ -537,7 +544,7 @@ macro_rules! impl_proof_optimized_preset {
                 )?)
             }
 
-            impl_proof_optimized_preset!(@schedule_catalog none);
+            $crate::impl_proof_optimized_preset!(@schedule_catalog none);
         }
     };
     (@core $cfg:ident, $field:ty, $ext_field:ty, $family:expr, $d:expr, $field_bits:expr, $log_commit_bound:expr, $fold_norms:expr, table, $feat:literal, $family_name:literal, $table:ident, $($options:tt)*) => {
@@ -546,7 +553,7 @@ macro_rules! impl_proof_optimized_preset {
             type ExtField = $ext_field;
             const D: usize = $d;
 
-            impl_proof_optimized_preset!(@options $($options)*);
+            $crate::impl_proof_optimized_preset!(@options $($options)*);
 
             fn decomposition() -> akita_types::DecompositionParams {
                 akita_types::DecompositionParams {
@@ -614,7 +621,7 @@ macro_rules! impl_proof_optimized_preset {
                 )?)
             }
 
-            impl_proof_optimized_preset!(@schedule_catalog ($feat, $family_name, $table));
+            $crate::impl_proof_optimized_preset!(@schedule_catalog ($feat, $family_name, $table));
         }
     };
 }

crates/akita-prover/src/protocol/sumcheck/digit_range/direct_range_leaf/live_prefix.rs

@@ -165,7 +165,7 @@ impl<E: FieldCore + FromPrimitiveInt + HasUnreducedOps> LowBasisRangeCheckProver
         &self,
         compact_range_image: &[V],
     ) -> EqFactoredUniPoly<E> {
-        debug_assert!(self.rounds_completed < self.col_bits);
+        debug_assert!(self.rounds_completed < self.num_vars);
         debug_assert_eq!(
             compact_range_image.len(),
             self.live_x_cols * (1usize << (self.num_vars - self.col_bits))
@@ -249,7 +249,7 @@ impl<E: FieldCore + FromPrimitiveInt + HasUnreducedOps> LowBasisRangeCheckProver
         &self,
         range_image: &[E],
     ) -> EqFactoredUniPoly<E> {
-        debug_assert!(self.rounds_completed < self.col_bits);
+        debug_assert!(self.rounds_completed < self.num_vars);
         let y_len = range_image.len() / self.live_x_cols;
         let (e_first, e_second) = self.split_eq.remaining_eq_tables();
         let num_first = e_first.len();

crates/akita-prover/src/protocol/sumcheck/relation_range_image/evaluation_trace.rs

@@ -159,6 +159,22 @@ impl<E: FieldCore> PreparedProverEvaluationTrace<E> {
         Ok(self.get(0, 0, 1))
     }
 
+    /// A trace of the given geometry whose weight function is identically zero.
+    ///
+    /// Used by virtual-only stage-2 instances that carry no committed
+    /// evaluation-trace term: every lane has empty support, so `get` returns
+    /// zero everywhere, coefficient/lane folds are no-ops over the empty
+    /// source set, and [`Self::final_value`] resolves to zero once folding
+    /// completes.
+    pub(crate) fn zero(live_lane_count: usize, coeff_count: usize) -> Self {
+        Self {
+            lane_terms: vec![Vec::new(); live_lane_count],
+            sources: Vec::new(),
+            live_lane_count,
+            coeff_count,
+        }
+    }
+
     #[cfg(test)]
     pub(crate) fn from_dense(dense: Vec<E>, live_lane_count: usize, coeff_count: usize) -> Self {
         assert_eq!(dense.len(), live_lane_count * coeff_count);

crates/akita-prover/src/protocol/sumcheck/relation_range_image/lifecycle.rs

@@ -1,6 +1,56 @@
 use super::*;
 
 impl<E: FieldCore + FromPrimitiveInt + HasUnreducedOps> RelationRangeImageProver<E> {
+    /// Create a stage-2 instance containing only the virtual range-image term.
+    ///
+    /// This is the standalone companion to
+    /// [`DigitRangeProver`](crate::protocol::sumcheck::DigitRangeProver):
+    /// stage 1 proves that the compact balanced-digit table is pointwise in
+    /// range, while this sumcheck links its carried range-image claim
+    /// `S(r) = range_image_evaluation` to an opening of the same digit table
+    /// through `S = w(w + 1)`. No relation or evaluation-trace term is
+    /// included.
+    pub fn new_virtual_only(
+        w_evals_compact: impl Into<std::sync::Arc<[i8]>>,
+        stage1_point: &[E],
+        range_image_evaluation: E,
+        b: usize,
+        live_lane_count: usize,
+        lane_bits: usize,
+        coefficient_bits: usize,
+    ) -> Result<Self, AkitaError> {
+        let lane_capacity = 1usize
+            .checked_shl(
+                u32::try_from(lane_bits).map_err(|_| {
+                    AkitaError::InvalidInput("stage-2 lane width overflow".to_string())
+                })?,
+            )
+            .ok_or_else(|| AkitaError::InvalidInput("stage-2 lane width overflow".to_string()))?;
+        let coeff_count = 1usize
+            .checked_shl(u32::try_from(coefficient_bits).map_err(|_| {
+                AkitaError::InvalidInput("stage-2 coefficient width overflow".to_string())
+            })?)
+            .ok_or_else(|| {
+                AkitaError::InvalidInput("stage-2 coefficient width overflow".to_string())
+            })?;
+        Self::new(
+            E::one(),
+            w_evals_compact,
+            stage1_point,
+            range_image_evaluation,
+            b,
+            vec![E::zero(); coeff_count],
+            vec![E::zero(); lane_capacity],
+            live_lane_count,
+            lane_bits,
+            coefficient_bits,
+            E::zero(),
+            PreparedProverEvaluationTrace::zero(live_lane_count, coeff_count),
+            E::zero(),
+            None,
+        )
+    }
+
     /// Create a fused stage-2 virtual-claim + relation sumcheck prover.
     #[allow(clippy::too_many_arguments)]
     #[tracing::instrument(skip_all, name = "RelationRangeImageProver::new")]

crates/akita-schedules/src/generated/walk.rs

@@ -6,7 +6,7 @@
 //! expand every typed fold once and recompute witness transitions and
 //! proof-byte totals.
 
-use akita_challenges::SparseChallengeConfig;
+use akita_challenges::{SparseChallengeConfig, TensorChallengeShape};
 use akita_field::{AkitaError, Prime128OffsetA7F7};
 use akita_types::{
     extension_opening_reduction_level_bytes, level_proof_bytes, terminal_response_bytes,
@@ -94,7 +94,17 @@ pub(crate) fn walk_generated_schedule_entry(
     root_params.witness_chunk =
         partition_to_chunk(entry.root.witness_partition, distributed_levels)?;
     let root_output_len = if is_multi_group {
-        root_params.output_witness_len::<Prime128OffsetA7F7>(&key.opening_layout()?)?
+        // The residual decomposition depth depends on the catalog family's
+        // field width; a concrete field type is not available here, so pass
+        // the policy's width (hardcoding a 128-bit field silently inflated
+        // the recomputed multi-group root output for 32/64-bit families and
+        // rejected their valid generated rows).
+        root_params.output_witness_len_for_field_bits(
+            &key.opening_layout()?,
+            u32::try_from(field_bits).map_err(|_| {
+                AkitaError::InvalidSetup("field width does not fit u32".to_string())
+            })?,
+        )?
     } else {
         planned_next_witness_len(
             field_bits,

crates/akita-sumcheck/Cargo.toml

@@ -20,11 +20,12 @@ akita-algebra = { version = "0.1.0", path = "../akita-algebra" }
 akita-field = { version = "0.1.0", path = "../akita-field" }
 akita-serialization = { version = "0.1.0", path = "../akita-serialization" }
 akita-transcript = { version = "0.1.0", path = "../akita-transcript" }
+rayon = { version = "1.10", optional = true }
 tracing = "0.1"
 
 [features]
 default = []
-parallel = ["akita-algebra/parallel", "akita-field/parallel"]
+parallel = ["dep:rayon", "akita-algebra/parallel", "akita-field/parallel"]
 
 [lints]
 workspace = true

crates/akita-sumcheck/src/batched_sumcheck.rs

@@ -14,6 +14,8 @@ use akita_field::{CanonicalField, FieldCore, FromPrimitiveInt, HalvingField};
 use akita_serialization::AkitaSerialize;
 use akita_transcript::labels;
 use akita_transcript::Transcript;
+#[cfg(feature = "parallel")]
+use rayon::prelude::*;
 
 fn mul_pow_2<E: FieldCore>(x: E, k: usize) -> E {
     let mut result = x;
@@ -74,14 +76,14 @@ pub struct BatchedSumcheckRoundResult<E: FieldCore> {
 /// Returns an error if the field inverse of 2 does not exist.
 #[tracing::instrument(skip_all, name = "prove_batched_sumcheck")]
 pub fn prove_batched_sumcheck<F, T, E, S>(
-    mut instances: Vec<&mut dyn SumcheckInstanceProver<E>>,
+    mut instances: Vec<&mut (dyn SumcheckInstanceProver<E> + Send)>,
     transcript: &mut T,
     mut sample_challenge: S,
 ) -> Result<(SumcheckProof<E>, Vec<E>), AkitaError>
 where
     F: FieldCore + CanonicalField,
     T: Transcript<F>,
-    E: FieldCore + FromPrimitiveInt + HalvingField + AkitaSerialize,
+    E: FieldCore + FromPrimitiveInt + HalvingField + AkitaSerialize + Send + Sync,
     S: FnMut(&mut T) -> E,
 {
     if instances.is_empty() {
@@ -129,10 +131,8 @@ where
     let mut challenges = Vec::with_capacity(max_num_rounds);
 
     for round in 0..max_num_rounds {
-        let univariate_polys: Vec<UniPoly<E>> = instances
-            .iter_mut()
-            .zip(individual_claims.iter())
-            .map(|(inst, previous_claim)| {
+        let compute_univariate =
+            |(inst, previous_claim): (&mut &mut (dyn SumcheckInstanceProver<E> + Send), &E)| {
                 let n = inst.num_rounds();
                 let offset = max_num_rounds - n;
                 let active = round >= offset && round < offset + n;
@@ -141,7 +141,21 @@ where
                 } else {
                     UniPoly::from_coeffs(vec![previous_claim.half()])
                 }
-            })
+            };
+        // With many instances (the fused selector batch carries dozens), the
+        // per-instance round computations dominate late rounds whose domains
+        // are too small for intra-instance parallelism; fan the instances out.
+        #[cfg(feature = "parallel")]
+        let univariate_polys: Vec<UniPoly<E>> = instances
+            .par_iter_mut()
+            .zip(individual_claims.par_iter())
+            .map(compute_univariate)
+            .collect();
+        #[cfg(not(feature = "parallel"))]
+        let univariate_polys: Vec<UniPoly<E>> = instances
+            .iter_mut()
+            .zip(individual_claims.iter())
+            .map(compute_univariate)
             .collect();
 
         let batched_poly = linear_combination(&univariate_polys, &batching_coeffs);
@@ -172,14 +186,18 @@ where
         }
 
         // Ingest challenge into each active instance.
-        for inst in instances.iter_mut() {
+        let ingest = |inst: &mut &mut (dyn SumcheckInstanceProver<E> + Send)| {
             let n = inst.num_rounds();
             let offset = max_num_rounds - n;
             let active = round >= offset && round < offset + n;
             if active {
                 inst.ingest_challenge(round - offset, r_j);
             }
-        }
+        };
+        #[cfg(feature = "parallel")]
+        instances.par_iter_mut().for_each(ingest);
+        #[cfg(not(feature = "parallel"))]
+        instances.iter_mut().for_each(ingest);
 
         round_polys.push(compressed);
     }

crates/akita-types/src/layout/params.rs

@@ -926,6 +926,24 @@ impl CommittedGroupParams {
     pub fn output_witness_len<F: CanonicalField>(
         &self,
         opening_batch: &OpeningClaimsLayout,
+    ) -> Result<usize, AkitaError> {
+        let modulus = crate::detect_field_modulus::<F>();
+        let field_bits = 128 - (modulus.saturating_sub(1)).leading_zeros();
+        self.output_witness_len_for_field_bits(opening_batch, field_bits)
+    }
+
+    /// [`Self::output_witness_len`] with the field width supplied as a value.
+    ///
+    /// Catalog expansion validates generated rows without naming a concrete
+    /// field type; the residual decomposition depth depends only on the
+    /// family's field width, so passing the policy's `field_bits` keeps the
+    /// recomputed layout identical to the runtime layout for every field
+    /// width (a hardcoded 128-bit field silently inflates the residual depth
+    /// for 32/64-bit families).
+    pub fn output_witness_len_for_field_bits(
+        &self,
+        opening_batch: &OpeningClaimsLayout,
+        field_bits: u32,
     ) -> Result<usize, AkitaError> {
         opening_batch.check()?;
         self.witness_chunk.validate()?;
@@ -934,7 +952,7 @@ impl CommittedGroupParams {
             self,
             opening_batch,
             self.witness_chunk.num_chunks,
-            crate::r_decomp_levels::<F>(self.log_basis_open),
+            crate::sis::compute_num_digits_field_width(field_bits, self.log_basis_open),
         )?;
         Ok(witness_layout.live_coeff_len())
     }