binius_field/arch/
strategies.rs

1// Copyright 2024-2025 Irreducible Inc.
2// Copyright 2026 The Binius Developers
3
4use std::{
5	array,
6	iter::Sum,
7	marker::PhantomData,
8	ops::{Add, AddAssign, Sub, SubAssign},
9};
10
11use bytemuck::TransparentWrapper;
12
13use crate::{
14	BinaryField,
15	arithmetic_traits::{InvertOrZero, Square, WideMul},
16	divisible::Divisible,
17	packed_fields::primitive::PackedPrimitiveType,
18	underlier::Underlier,
19};
20
21/// Strategy that splits the underlier into `SubU`-sized lanes, applies the sub-packing
22/// `PackedPrimitiveType<SubU, F>`'s op to each lane, and recombines — a generic fallback for
23/// packings that lack a specialized full-width [`Square`], [`InvertOrZero`], or [`WideMul`]. The
24/// sub-underlier `SubU` is a `PhantomData` parameter so the packing type `T` stays last for the
25/// macro's `Divide<SubU, $name, N>` form.
26///
27/// `N` is the lane count: callers always pass `N = <U as Divisible<SubU>>::N` (or the literal it
28/// works out to). `Square`/`InvertOrZero` stream through [`Divisible`] and ignore `N`, but it is
29/// still required so every `Divide` instantiation names its lane count explicitly. `WideMul` must
30/// defer reduction, so it materializes one unreduced product per lane in an `N`-element
31/// [`LaneWideProduct`] — and an associated const can't be an array length without
32/// `generic_const_exprs`, which is why `N` is a const generic rather than read from `Divisible`.
33#[repr(transparent)]
34#[derive(TransparentWrapper)]
35#[transparent(T)]
36pub struct Divide<SubU, T, const N: usize>(T, PhantomData<SubU>);
37
38impl<U, SubU, F, const N: usize> Square for Divide<SubU, PackedPrimitiveType<U, F>, N>
39where
40	U: Underlier + Divisible<SubU>,
41	SubU: Underlier,
42	F: BinaryField,
43	PackedPrimitiveType<SubU, F>: Square,
44{
45	#[inline]
46	fn square(self) -> Self {
47		let val = Self::peel(self);
48		let squared = Divisible::<SubU>::value_iter(val.to_underlier()).map(|lane| {
49			PackedPrimitiveType::<SubU, F>::from_underlier(lane)
50				.square()
51				.to_underlier()
52		});
53		Self::wrap(PackedPrimitiveType::from_underlier(Divisible::<SubU>::from_iter(squared)))
54	}
55}
56
57impl<U, SubU, F, const N: usize> InvertOrZero for Divide<SubU, PackedPrimitiveType<U, F>, N>
58where
59	U: Underlier + Divisible<SubU>,
60	SubU: Underlier,
61	F: BinaryField,
62	PackedPrimitiveType<SubU, F>: InvertOrZero,
63{
64	#[inline]
65	fn invert_or_zero(self) -> Self {
66		let val = Self::peel(self);
67		let inverted = Divisible::<SubU>::value_iter(val.to_underlier()).map(|lane| {
68			PackedPrimitiveType::<SubU, F>::from_underlier(lane)
69				.invert_or_zero()
70				.to_underlier()
71		});
72		Self::wrap(PackedPrimitiveType::from_underlier(Divisible::<SubU>::from_iter(inverted)))
73	}
74}
75
76/// One independent deferred wide product per `SubU` lane of a [`Divide`] widening multiply. Lanes
77/// accumulate (`Add`/`Sub`/`Sum`) and reduce independently, mirroring the packing structure, so a
78/// sum of products is reduced only once per lane. `N` is the lane count.
79#[derive(Clone, Copy, Debug)]
80pub struct LaneWideProduct<O, const N: usize>(pub [O; N]);
81
82impl<O: Copy + Default, const N: usize> Default for LaneWideProduct<O, N> {
83	#[inline]
84	fn default() -> Self {
85		Self([O::default(); N])
86	}
87}
88
89impl<O: Copy + Add<Output = O>, const N: usize> Add for LaneWideProduct<O, N> {
90	type Output = Self;
91
92	#[inline]
93	fn add(self, rhs: Self) -> Self {
94		Self(array::from_fn(|i| self.0[i] + rhs.0[i]))
95	}
96}
97
98impl<O: Copy + Add<Output = O>, const N: usize> AddAssign for LaneWideProduct<O, N> {
99	#[inline]
100	fn add_assign(&mut self, rhs: Self) {
101		*self = *self + rhs;
102	}
103}
104
105impl<O: Copy + Sub<Output = O>, const N: usize> Sub for LaneWideProduct<O, N> {
106	type Output = Self;
107
108	#[inline]
109	fn sub(self, rhs: Self) -> Self {
110		Self(array::from_fn(|i| self.0[i] - rhs.0[i]))
111	}
112}
113
114impl<O: Copy + Sub<Output = O>, const N: usize> SubAssign for LaneWideProduct<O, N> {
115	#[inline]
116	fn sub_assign(&mut self, rhs: Self) {
117		*self = *self - rhs;
118	}
119}
120
121impl<O: Copy + Default + Add<Output = O>, const N: usize> Sum for LaneWideProduct<O, N> {
122	#[inline]
123	fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
124		iter.fold(Self::default(), |acc, x| acc + x)
125	}
126}
127
128impl<U, SubU, F, const N: usize> WideMul for Divide<SubU, PackedPrimitiveType<U, F>, N>
129where
130	U: Underlier + Divisible<SubU>,
131	SubU: Underlier,
132	F: BinaryField,
133	PackedPrimitiveType<SubU, F>: WideMul,
134	<PackedPrimitiveType<SubU, F> as WideMul>::Output: Copy + Default,
135{
136	type Output = LaneWideProduct<<PackedPrimitiveType<SubU, F> as WideMul>::Output, N>;
137
138	#[inline]
139	fn wide_mul(a: Self, b: Self) -> Self::Output {
140		debug_assert_eq!(N, <U as Divisible<SubU>>::N, "N must equal Divisible<SubU>::N");
141
142		let a = Self::peel(a).to_underlier();
143		let b = Self::peel(b).to_underlier();
144
145		let mut lanes = [<PackedPrimitiveType<SubU, F> as WideMul>::Output::default(); N];
146		for (slot, (lhs, rhs)) in lanes
147			.iter_mut()
148			.zip(Divisible::<SubU>::value_iter(a).zip(Divisible::<SubU>::value_iter(b)))
149		{
150			*slot = <PackedPrimitiveType<SubU, F> as WideMul>::wide_mul(
151				PackedPrimitiveType::from_underlier(lhs),
152				PackedPrimitiveType::from_underlier(rhs),
153			);
154		}
155		LaneWideProduct(lanes)
156	}
157
158	#[inline]
159	fn reduce(wide: Self::Output) -> Self {
160		let lanes = wide.0.into_iter().map(|product| {
161			<PackedPrimitiveType<SubU, F> as WideMul>::reduce(product).to_underlier()
162		});
163		Self::wrap(PackedPrimitiveType::from_underlier(Divisible::<SubU>::from_iter(lanes)))
164	}
165}
binius_field/arch/strategies.rs

binius_field/arch/
strategies.rs
