binius_math/

inner_product.rs

// Copyright 2024-2025 Irreducible Inc.
// Copyright 2026 The Binius Developers

//! Inner products `sum_i a_i * b_i` over sequences of field elements.

use std::{iter, ops::Deref};

use binius_field::{ExtensionField, Field, FieldOps, PackedField, WideMul};
use binius_utils::rayon::{
	prelude::*,
	task_size::{IndexedParallelIteratorExt, WorkPerItem},
};

use crate::FieldBuffer;

/// Computes the inner product `sum_i a_i * b_i`.
///
/// Arithmetic is element-wise.
/// A packed element therefore yields one independent inner product per lane, not one scalar.
///
/// # Panics
///
/// Panics if the two sequences have different lengths.
#[inline]
pub fn inner_product<F: FieldOps>(
	a: impl IntoIterator<Item = F>,
	b: impl IntoIterator<Item = F>,
) -> F {
	itertools::zip_eq(a, b).map(|(a_i, b_i)| b_i * a_i).sum()
}

/// Computes the inner product `sum_i a_i * b_i`.
///
/// The left side lies in a subfield of the right side's field.
///
/// # Panics
///
/// Panics if the two sequences have different lengths.
#[inline]
pub fn inner_product_subfield<F, FSub>(
	a: impl IntoIterator<Item = FSub>,
	b: impl IntoIterator<Item = F>,
) -> F
where
	F: Field + ExtensionField<FSub>,
	FSub: Field,
{
	itertools::zip_eq(a, b).map(|(a_i, b_i)| b_i * a_i).sum()
}

/// Sums the coefficient-by-coefficient products of two multilinears.
///
/// ```text
/// result = sum_i a_i * b_i
/// ```
///
/// This is not the product polynomial; it is a single field element.
/// Pairing a polynomial with an equality indicator expansion evaluates it at that point.
///
/// ## Preconditions
///
/// * the two buffers must hold the same number of coefficients
#[inline]
pub fn inner_product_buffers<F, P, DataA, DataB>(
	a: &FieldBuffer<P, DataA>,
	b: &FieldBuffer<P, DataB>,
) -> F
where
	F: Field,
	P: PackedField<Scalar = F>,
	DataA: Deref<Target = [P]>,
	DataB: Deref<Target = [P]>,
{
	// Both buffers hold the same number of packed words, so the packed form applies directly.
	inner_product_packed(a.log_len(), a.iter_packed().copied(), b.iter_packed().copied())
}

/// Sums the coefficient-by-coefficient products of two multilinears, across threads.
///
/// The value matches the single-threaded pairing.
/// The coefficient range is split into tasks, each summing its own partial products.
///
/// ## Preconditions
///
/// * the two buffers must hold the same number of coefficients
#[inline]
pub fn inner_product_par<F, P, DataA, DataB>(
	a: &FieldBuffer<P, DataA>,
	b: &FieldBuffer<P, DataB>,
) -> F
where
	F: Field,
	P: PackedField<Scalar = F>,
	DataA: Deref<Target = [P]>,
	DataB: Deref<Target = [P]>,
{
	// Below one packed word the buffer still occupies a whole word, so only the live lanes count.
	let n = a.len();

	// Accumulate the products in unreduced (wide) form and reduce a single time at the end.
	// For packed `GF(2^128)` fields this amortizes the reduction cost across all products.
	// For every other field the widening multiply is the trivial one, so this is a plain
	// product-then-sum.
	let wide_sum = a
		.as_ref()
		.par_iter()
		.zip_eq(b.as_ref().par_iter())
		.with_min_task(WorkPerItem::FieldMuls)
		.map(|(&a_i, &b_i)| P::wide_mul(a_i, b_i))
		.sum::<<P as WideMul>::Output>();
	P::reduce(wide_sum).into_iter().take(n).sum()
}

/// Computes the inner product of two scalar sequences generated by iterators of packed elements.
///
/// ## Preconditions
///
/// * `a` and `b` have length `1 << log_n.saturating_sub(P::LOG_WIDTH)`
#[inline]
pub fn inner_product_packed<F, P>(
	log_n: usize,
	a: impl ExactSizeIterator<Item = P>,
	b: impl ExactSizeIterator<Item = P>,
) -> F
where
	F: Field,
	P: PackedField<Scalar = F>,
{
	assert_eq!(a.len(), 1 << log_n.saturating_sub(P::LOG_WIDTH)); // pre-condition
	assert_eq!(b.len(), 1 << log_n.saturating_sub(P::LOG_WIDTH)); // pre-condition

	// Accumulate the products in unreduced (wide) form and reduce a single time at the end.
	// For packed `GF(2^128)` fields this amortizes the reduction cost across all products.
	// For every other field the widening multiply is the trivial one, so this is a plain
	// product-then-sum.
	let wide_sum = iter::zip(a, b)
		.map(|(a_i, b_i)| P::wide_mul(a_i, b_i))
		.sum::<<P as WideMul>::Output>();
	P::reduce(wide_sum).into_iter().take(1 << log_n).sum()
}

#[cfg(test)]
mod tests {
	use binius_field::{Divisible, Ghash128b, PackedField, PackedGhash4x128b, Random};
	use proptest::prelude::*;
	use rand::{SeedableRng, rngs::StdRng};

	use super::*;
	use crate::test_utils::random_field_buffer;

	type P = PackedGhash4x128b;
	type F = Ghash128b;

	// The packing width is four scalars, so this range straddles it in both directions.
	const MAX_VARS: usize = 8;

	proptest! {
		#[test]
		fn the_two_buffer_inner_products_agree_with_the_scalar_reference(
			n_vars in 0..=MAX_VARS,
			seed: u64,
		) {
			let mut rng = StdRng::seed_from_u64(seed);
			let a = random_field_buffer::<P>(&mut rng, n_vars);
			let b = random_field_buffer::<P>(&mut rng, n_vars);

			// Both forms defer the field reduction, so a scalar sum is the reference.
			let reference = (0..a.len()).map(|i| a.get(i) * b.get(i)).sum::<F>();

			prop_assert_eq!(inner_product_buffers(&a, &b), reference);
			prop_assert_eq!(inner_product_par(&a, &b), reference);
		}
	}

	#[test]
	fn test_inner_product_over_packed_is_lane_wise() {
		let mut rng = StdRng::seed_from_u64(7);

		// Invariant: a packed element is a vector of independent lanes.
		// Multiplication and addition never cross a lane boundary.
		// So lane `j` of the result depends only on lane `j` of the two inputs.
		//
		// Fixture state: 8 packed elements per side, 4 lanes each.
		// Every lane therefore carries its own sequence of 8 scalars.
		//
		//     a: [ a_0[0] | a_0[1] | a_0[2] | a_0[3] ], [ a_1[0] | ... ], ...
		//     b: [ b_0[0] | b_0[1] | b_0[2] | b_0[3] ], [ b_1[0] | ... ], ...
		//     -> lane j = sum_i a_i[j] * b_i[j]
		let a = (0..8).map(|_| P::random(&mut rng)).collect::<Vec<P>>();
		let b = (0..8).map(|_| P::random(&mut rng)).collect::<Vec<P>>();

		let packed = inner_product(a.iter().copied(), b.iter().copied());

		// Rebuild each lane's own scalar inner product and compare it against that lane.
		for lane in 0..P::WIDTH {
			let expected = iter::zip(&a, &b)
				.map(|(a_i, b_i)| a_i.get(lane) * b_i.get(lane))
				.sum::<F>();
			assert_eq!(packed.get(lane), expected, "mismatch in lane {lane}");
		}
	}

	#[test]
	fn test_inner_product_matches_subfield_at_the_full_field() {
		let mut rng = StdRng::seed_from_u64(11);

		// Invariant: every field is a degree-1 extension of itself.
		// The same-field form and the subfield form therefore agree on that overlap.
		//
		// Fixture state: 16 random elements per side, drawn from the same field.
		let a = (0..16).map(|_| F::random(&mut rng)).collect::<Vec<F>>();
		let b = (0..16).map(|_| F::random(&mut rng)).collect::<Vec<F>>();

		// The turbofish forces the subfield to be the full field, which is the overlap.
		// Pin the two forms equal there so their separate bodies cannot drift apart.
		assert_eq!(
			inner_product(a.iter().copied(), b.iter().copied()),
			inner_product_subfield::<F, F>(a.iter().copied(), b.iter().copied()),
		);
	}

	#[test]
	fn test_inner_product_packed_matches_naive() {
		let mut rng = StdRng::seed_from_u64(42);

		// The packed inner product defers the `GF(2^128)` reduction (widening multiply).
		// Check the reduced result against a naive scalar-by-scalar reference.
		for log_n in [4, 8, 12] {
			let n = 1 << log_n;
			let a = (0..n / P::WIDTH)
				.map(|_| P::random(&mut rng))
				.collect::<Vec<P>>();
			let b = (0..n / P::WIDTH)
				.map(|_| P::random(&mut rng))
				.collect::<Vec<P>>();

			let naive = iter::zip(&a, &b)
				.flat_map(|(a_i, b_i)| iter::zip(a_i.iter(), b_i.iter()))
				.map(|(a_j, b_j)| a_j * b_j)
				.sum::<F>();

			let packed: F = inner_product_packed(log_n, a.iter().copied(), b.iter().copied());
			assert_eq!(packed, naive, "mismatch at log_n={log_n}");
		}
	}

	#[test]
	fn test_inner_product_packed_below_the_packing_width() {
		let mut rng = StdRng::seed_from_u64(0);

		// Below the packing width the trailing scalars of the one word must not contribute.
		let a = P::random(&mut rng);
		let b = P::random(&mut rng);

		let packed: F = inner_product_packed(0, iter::once(a), iter::once(b));
		assert_eq!(packed, a.iter().next().unwrap() * b.iter().next().unwrap());
	}
}