perf(verifier): batch reduced setup rows that share one functional

crates/akita-verifier/src/setup_contribution/scan/reduced.rs

@@ -196,6 +196,13 @@ where
     let setup_flat = setup_view.as_slice();
     let job_rings = super::super::segments::SETUP_SCAN_JOB_RINGS;
     let num_jobs = required.div_ceil(job_rings);
+    // A single base-ring functional is independent of the physical ring index.
+    // Reuse the shared-functional kernel already used by lifted scans. Mixed
+    // dimensions or distinct functionals retain the general per-class path.
+    let shared_functional = match functional_classes {
+        [class] if class.ratio == 1 => Some(class.functional),
+        _ => None,
+    };
     cfg_try_fold_reduce!(
         0..num_jobs,
         E::zero,
@@ -206,6 +213,16 @@ where
                 .ok_or(AkitaError::InvalidProof)?
                 .min(required);
             let setup = setup_flat.get(lo..hi).ok_or(AkitaError::InvalidProof)?;
+            let mut shared_rows = Vec::new();
+            let mut shared_weights = Vec::new();
+            if shared_functional.is_some() {
+                shared_rows.try_reserve_exact(setup.len()).map_err(|_| {
+                    AkitaError::InvalidSetup("reduced scan row allocation failed".into())
+                })?;
+                shared_weights.try_reserve_exact(setup.len()).map_err(|_| {
+                    AkitaError::InvalidSetup("reduced scan weight allocation failed".into())
+                })?;
+            }
             let mut segment_cursors = partitions
                 .iter()
                 .map(|partition| {
@@ -241,15 +258,25 @@ where
                 for (class, scalar) in functional_classes.iter().zip(&mut class_scalars) {
                     let scalar = std::mem::replace(scalar, A::zero()).finish();
                     if !scalar.is_zero() {
-                        let functional = class.chunk::<BASE_D>(base_idx)?;
-                        term.add_product(
-                            MulBaseUnreduced::dot_base(functional, ring.coefficients()),
-                            scalar,
-                        );
+                        if shared_functional.is_some() {
+                            shared_rows.push(ring.coefficients().as_slice());
+                            shared_weights.push(scalar);
+                        } else {
+                            let functional = class.chunk::<BASE_D>(base_idx)?;
+                            term.add_product(
+                                MulBaseUnreduced::dot_base(functional, ring.coefficients()),
+                                scalar,
+                            );
+                        }
                     }
                 }
             }
-            Ok(acc + term.finish())
+            let contribution = if let Some(functional) = shared_functional {
+                E::weighted_dot_base_rows(&shared_rows, &shared_weights, functional)
+            } else {
+                term.finish()
+            };
+            Ok(acc + contribution)
         },
         |lhs, rhs| Ok(lhs + rhs)
     )

crates/akita-verifier/src/setup_contribution/tests.rs

@@ -20,6 +20,7 @@ mod direct_evaluation_regressions;
 mod fused_scan;
 mod packed_direct;
 mod prepare;
+mod reduced_shared;
 mod span_evaluators;
 
 type F = Prime128OffsetA7F7;

crates/akita-verifier/src/setup_contribution/tests/fused_scan.rs

@@ -89,7 +89,7 @@ fn naive_physical_b_weights<E: Field>(
 }
 
 #[allow(clippy::too_many_arguments)]
-fn reduced_direct_literal_oracle<E>(
+pub(super) fn reduced_direct_literal_oracle<E>(
     plan: &SetupContributionPlan<E>,
     setup: &AkitaExpandedSetup<F>,
     inputs: &TestSetupInputs,

crates/akita-verifier/src/setup_contribution/tests/reduced_shared.rs

@@ -0,0 +1,140 @@
+use super::fused_scan::reduced_direct_literal_oracle;
+use super::*;
+use jolt_field::Ring;
+
+fn check_shared_reduced_scan<E>(blocks: usize, scalar: impl Fn(u64) -> E, zero_rows: bool)
+where
+    E: ExtField<F> + MulBaseUnreduced<F>,
+{
+    let role_dims = CommitmentRingDims {
+        inner: 64,
+        outer: 64,
+        opening: 64,
+    };
+    let (inputs, groups, layout, base_plan, _, base_relation_point, fold_gadget) =
+        structured_weight_fixture_with_outgoing(blocks, &[blocks], role_dims, 64);
+    let geometry = base_plan.relation_address_geometry();
+    let relation_point = base_relation_point
+        .iter()
+        .enumerate()
+        .map(|(index, _)| scalar(601 + index as u64))
+        .collect::<Vec<_>>();
+    let coefficient_point = (0..geometry.relation_coefficient_variable_count())
+        .map(|index| scalar(701 + index as u64))
+        .collect::<Vec<_>>();
+    let row_weights = inputs
+        .eq_tau1
+        .iter()
+        .copied()
+        .map(|weight| {
+            if zero_rows {
+                E::zero()
+            } else {
+                E::lift_base(weight)
+            }
+        })
+        .collect::<Vec<_>>();
+    let alpha = scalar(809);
+    let plan = SetupContributionPlan::<E>::prepare::<F>(
+        &inputs.level_params,
+        &inputs.opening_batch,
+        1,
+        row_weights.clone().into(),
+        &layout,
+        &groups,
+        PreparedRelationAddress::new(&relation_point).unwrap(),
+        Some(&fold_gadget),
+        geometry,
+    )
+    .unwrap();
+    let scan = DirectScan::new(
+        plan,
+        PreparedCoefficientFunctional::reduced_evaluation(alpha, &coefficient_point, geometry)
+            .unwrap(),
+    )
+    .unwrap();
+    let plan = scan.plan();
+    // Check that the fixture exercises the single functional, ratio-one case.
+    let DirectScanMode::Reduced { groups: direct, .. } = &scan.mode else {
+        panic!("expected reduced fixture");
+    };
+    assert_eq!(plan.projection_geometry().base_ring_dim(), 64);
+    let first = &direct[0].roles[0].functional;
+    assert_eq!(first.len(), 64);
+    for group in direct {
+        for role in &group.roles {
+            assert_eq!(role.functional.as_ref(), first.as_ref());
+        }
+    }
+    if blocks == 131 {
+        // Multiple 2048-ring jobs and a partial last job, without padding away
+        // the unequal ownership or physical-prefix boundaries.
+        assert!(plan.required() > super::super::segments::SETUP_SCAN_JOB_RINGS);
+        assert_ne!(
+            plan.required() % super::super::segments::SETUP_SCAN_JOB_RINGS,
+            0
+        );
+    }
+    let setup_coefficients = plan.required() * 64;
+    let setup = AkitaExpandedSetup::from_trusted_seed_derived_parts_unchecked(
+        AkitaSetupDescriptor {
+            max_num_vars: 0,
+            max_num_batched_polys: 0,
+            num_field_elements: setup_coefficients,
+            setup_seed: [0u8; 32].into(),
+        },
+        FlatMatrix::from_flat_data(
+            (0..setup_coefficients)
+                .map(|index| test_scalar(901 + index as u128))
+                .collect(),
+        ),
+    );
+    let expected = reduced_direct_literal_oracle(
+        plan,
+        &setup,
+        &inputs,
+        &groups,
+        &layout,
+        &fold_gadget,
+        &row_weights,
+        &coefficient_point,
+        alpha,
+    );
+    assert_eq!(scan.evaluate_direct::<F>(&setup).unwrap(), expected);
+    if zero_rows {
+        assert_eq!(expected, E::zero());
+    }
+
+    // A truncated public setup is still rejected at the existing boundary.
+    let truncated = AkitaExpandedSetup::from_trusted_seed_derived_parts_unchecked(
+        AkitaSetupDescriptor {
+            max_num_vars: 0,
+            max_num_batched_polys: 0,
+            num_field_elements: setup_coefficients - 64,
+            setup_seed: [0u8; 32].into(),
+        },
+        FlatMatrix::from_flat_data(vec![F::zero(); setup_coefficients - 64]),
+    );
+    assert!(scan.evaluate_direct::<F>(&truncated).is_err());
+}
+
+#[test]
+fn shared_reduced_scan_matches_literal_oracle_across_jobs() {
+    check_shared_reduced_scan(131, F::from_u64, false);
+}
+
+#[test]
+fn shared_reduced_scan_handles_zero_weights_and_boolean_points() {
+    check_shared_reduced_scan(5, |value| F::from_u64(value % 2), true);
+    check_shared_reduced_scan(5, |value| F::from_u64(value % 2), false);
+}
+
+#[test]
+fn shared_reduced_scan_matches_literal_oracle_over_extension_field() {
+    type X = jolt_field::Ext2<F>;
+    check_shared_reduced_scan(
+        5,
+        |value| X::from_base_slice(&[F::from_u64(value), F::from_u64(value + 10_000)]),
+        false,
+    );
+}