fix avx512 (panicked on small instances) (#193)

* fix avx512 (panicked on small instances)

* add `test_aggregation`

---------

Co-authored-by: Tom Wambsgans <TomWambsgans@users.noreply.github.com>

Author: TomWambsgans

crates/backend/poly/src/utils.rs

@@ -61,6 +61,10 @@ pub const fn packing_width<EF: Field>() -> usize {
     PFPacking::<EF>::WIDTH
 }
 
+pub const fn must_unpack_multilinears<EF: Field>(n_vars: usize) -> bool {
+    n_vars <= 1 + packing_log_width::<EF>()
+}
+
 pub fn batch_fold_multilinears<
     EF: PrimeCharacteristicRing + Copy + Send + Sync,
     IF: Copy + Sub<Output = IF> + Send + Sync,

crates/backend/sumcheck/src/prove.rs

@@ -115,11 +115,8 @@ where
 
     let mut challenges = Vec::new();
     for _ in 0..n_rounds {
-        // If Packing is enabled, and there are too little variables, we unpack everything:
-        if multilinears.by_ref().is_packed() && n_vars <= 1 + packing_log_width::<EF>() {
-            // unpack
+        if multilinears.by_ref().is_packed() && must_unpack_multilinears::<EF>(n_vars) {
             multilinears = multilinears.by_ref().unpack().as_owned_or_clone().into();
-            // SplitEq handles unpacking transparently via get_unpacked
         }
 
         let ps = compute_and_send_polynomial(

crates/backend/sumcheck/src/sc_computation.rs

@@ -386,17 +386,25 @@ where
             |sc, pf, ed| sc.eval_packed_extension(&pf, ed),
             packing_unpack_sum,
         ),
-        MleGroupRef::BasePacked(multilinears) => sumcheck_compute_core(
-            multilinears,
-            degree,
-            |i| split_eq.map(|seq| seq.get_packed(i)),
-            computation,
-            extra_data,
-            missing_mul_factor,
-            packed_fold_size,
-            |sc, pf, ed| sc.eval_packed_base(&pf, ed),
-            packing_unpack_sum,
-        ),
+        MleGroupRef::BasePacked(multilinears) => {
+            if let Some(seq) = split_eq {
+                assert!(
+                    !seq.is_remainder_mode(),
+                    "BasePacked sumcheck received SplitEq in remainder mode"
+                );
+            }
+            sumcheck_compute_core(
+                multilinears,
+                degree,
+                |i| split_eq.map(|seq| seq.get_packed(i)),
+                computation,
+                extra_data,
+                missing_mul_factor,
+                packed_fold_size,
+                |sc, pf, ed| sc.eval_packed_base(&pf, ed),
+                packing_unpack_sum,
+            )
+        }
         MleGroupRef::Base(multilinears) => sumcheck_compute_core(
             multilinears,
             degree,
@@ -507,6 +515,12 @@ where
             )
         }
         MleGroupRef::BasePacked(multilinears) => {
+            if let Some(seq) = split_eq {
+                assert!(
+                    !seq.is_remainder_mode(),
+                    "BasePacked fold-and-compute received SplitEq in remainder mode"
+                );
+            }
             let prev_folded_size = multilinears[0].len() / 2;
             let prev_folding_factor_packed = EFPacking::<EF>::from(prev_folding_factor);
             sumcheck_fold_and_compute_core(

crates/backend/sumcheck/src/split_eq.rs

@@ -14,7 +14,7 @@ impl<EF: ExtensionField<PF<EF>>> SplitEq<EF> {
     pub fn new(eq_point: &[EF]) -> Self {
         let n = eq_point.len();
 
-        if n <= packing_log_width::<EF>() * 2 {
+        if must_unpack_multilinears::<EF>(n + 1) {
             return Self {
                 eq_lo: vec![EF::ONE],
                 eq_hi_packed: Vec::new(),
@@ -77,7 +77,7 @@ impl<EF: ExtensionField<PF<EF>>> SplitEq<EF> {
 
     #[inline(always)]
     pub fn get_packed(&self, i: usize) -> EFPacking<EF> {
-        debug_assert!(!self.is_remainder_mode(), "get_packed called in remainder mode");
+        assert!(!self.is_remainder_mode(), "get_packed called in remainder mode");
         let packed_hi = self.eq_hi_packed.len();
         if self.eq_lo.len() > 1 {
             EFPacking::<EF>::from(self.eq_lo[i >> self.log_packed_hi]) * self.eq_hi_packed[i & (packed_hi - 1)]

crates/sub_protocols/src/air_sumcheck.rs

@@ -79,7 +79,7 @@ where
     }
 
     fn compute_bare_round_poly(&mut self) -> DensePolynomial<EF> {
-        if self.multilinears.is_packed() && self.multilinears.n_vars() <= 1 + packing_log_width::<EF>() {
+        if self.multilinears.is_packed() && must_unpack_multilinears::<EF>(self.multilinears.n_vars()) {
             let old = std::mem::replace(
                 &mut self.multilinears,
                 MleGroup::Owned(MleGroupOwned::Extension(vec![])),

tests/test_lean_multisig.rs

@@ -1,5 +1,6 @@
-use lean_multisig::{AggregatedXMSS, setup_prover, xmss_aggregate, xmss_verify_aggregation};
+use lean_multisig::{AggregatedXMSS, AggregationTopology, setup_prover, xmss_aggregate, xmss_verify_aggregation};
 use rand::{RngExt, SeedableRng, rngs::StdRng};
+use rec_aggregation::benchmark::run_aggregation_benchmark;
 use xmss::{
     signers_cache::{BENCHMARK_SLOT, get_benchmark_signatures, message_for_benchmark},
     xmss_key_gen, xmss_sign, xmss_verify,
@@ -18,6 +19,18 @@ fn test_xmss_signature() {
     xmss_verify(&pub_key, &msg, &signature, slot).unwrap();
 }
 
+#[test]
+fn test_aggregation() {
+    for n_signatures in [1, 2, 4, 8, 16, 32, 64, 128] {
+        let topology = AggregationTopology {
+            raw_xmss: n_signatures,
+            children: vec![],
+            log_inv_rate: 1,
+        };
+        run_aggregation_benchmark(&topology, 0, false);
+    }
+}
+
 #[test]
 fn test_recursive_aggregation() {
     setup_prover();