binius_field/arch/portable/
packed_arithmetic.rs

1// Copyright 2024-2025 Irreducible Inc.
2
3use crate::{
4	PackedField, TowerField,
5	binary_field::BinaryField,
6	linear_transformation::{FieldLinearTransformation, Transformation},
7	packed::PackedBinaryField,
8	underlier::{UnderlierWithBitOps, WithUnderlier},
9};
10
11pub trait UnderlierWithBitConstants: UnderlierWithBitOps
12where
13	Self: 'static,
14{
15	const INTERLEAVE_EVEN_MASK: &'static [Self];
16	const INTERLEAVE_ODD_MASK: &'static [Self];
17
18	/// Interleave with the given bit size
19	fn interleave(self, other: Self, log_block_len: usize) -> (Self, Self) {
20		// There are 2^7 = 128 bits in a u128
21		assert!(log_block_len < Self::INTERLEAVE_EVEN_MASK.len());
22
23		let block_len = 1 << log_block_len;
24
25		// See Hacker's Delight, Section 7-3.
26		// https://dl.acm.org/doi/10.5555/2462741
27		let t = ((self >> block_len) ^ other) & Self::INTERLEAVE_EVEN_MASK[log_block_len];
28		let c = self ^ t << block_len;
29		let d = other ^ t;
30
31		(c, d)
32	}
33
34	/// Transpose with the given bit size
35	fn transpose(mut self, mut other: Self, log_block_len: usize) -> (Self, Self) {
36		// There are 2^7 = 128 bits in a u128
37		assert!(log_block_len < Self::INTERLEAVE_EVEN_MASK.len());
38
39		for log_block_len in (log_block_len..Self::LOG_BITS).rev() {
40			(self, other) = self.interleave(other, log_block_len);
41		}
42
43		(self, other)
44	}
45}
46
47/// Generate the mask with ones in the odd packed element positions and zeros in even
48macro_rules! interleave_mask_even {
49	($underlier:ty, $tower_level:literal) => {{
50		let scalar_bits = 1 << $tower_level;
51
52		let mut mask: $underlier = (1 << scalar_bits) - 1;
53		let log_width = <$underlier>::LOG_BITS - $tower_level;
54		let mut i = 1;
55		while i < log_width {
56			mask |= mask << (scalar_bits << i);
57			i += 1;
58		}
59
60		mask
61	}};
62}
63
64pub(crate) use interleave_mask_even;
65
66/// Generate the mask with ones in the even packed element positions and zeros in odd
67macro_rules! interleave_mask_odd {
68	($underlier:ty, $tower_level:literal) => {
69		interleave_mask_even!($underlier, $tower_level) << (1 << $tower_level)
70	};
71}
72
73pub(crate) use interleave_mask_odd;
74
75/// View the inputs as vectors of packed binary tower elements and transpose as 2x2 square matrices.
76/// Given vectors <a_0, a_1, a_2, a_3, ...> and <b_0, b_1, b_2, b_3, ...>, returns a tuple with
77/// <a0, b0, a2, b2, ...> and <a1, b1, a3, b3>.
78fn interleave<U: UnderlierWithBitConstants, F: TowerField>(a: U, b: U) -> (U, U) {
79	let mask = U::INTERLEAVE_EVEN_MASK[F::TOWER_LEVEL];
80
81	let block_len = F::N_BITS;
82	let t = ((a >> block_len) ^ b) & mask;
83	let c = a ^ (t << block_len);
84	let d = b ^ t;
85
86	(c, d)
87}
88
89/// View the input as a vector of packed binary tower elements and add the adjacent ones.
90/// Given a vector <a_0, a_1, a_2, a_3, ...>, returns <a0 + a1, a0 + a1, a2 + a3, a2 + a3, ...>.
91fn xor_adjacent<U: UnderlierWithBitConstants, F: TowerField>(a: U) -> U {
92	let mask = U::INTERLEAVE_EVEN_MASK[F::TOWER_LEVEL];
93
94	let block_len = F::N_BITS;
95	let t = ((a >> block_len) ^ a) & mask;
96
97	t ^ (t << block_len)
98}
99
100/// Packed transformation implementation.
101/// Stores bases in a form of:
102/// [
103///     [<base vec 1> ... <base vec 1>],
104///     ...
105///     [<base vec N> ... <base vec N>]
106/// ]
107/// Transformation complexity is `N*log(N)` where `N` is `OP::Scalar::DEGREE`.
108pub struct PackedTransformation<OP> {
109	bases: Vec<OP>,
110}
111
112impl<OP> PackedTransformation<OP>
113where
114	OP: PackedBinaryField,
115{
116	pub fn new<Data: AsRef<[OP::Scalar]> + Sync>(
117		transformation: &FieldLinearTransformation<OP::Scalar, Data>,
118	) -> Self {
119		Self {
120			bases: transformation
121				.bases()
122				.iter()
123				.map(|base| OP::broadcast(*base))
124				.collect(),
125		}
126	}
127}
128
129/// Broadcast lowest field for each element, e.g. `[<0001><0000>] -> [<1111><0000>]`
130fn broadcast_lowest_bit<U: UnderlierWithBitOps>(mut data: U, log_packed_bits: usize) -> U {
131	for i in 0..log_packed_bits {
132		data |= data << (1 << i)
133	}
134
135	data
136}
137
138impl<U, IP, OP, IF, OF> Transformation<IP, OP> for PackedTransformation<OP>
139where
140	IP: PackedField<Scalar = IF> + WithUnderlier<Underlier = U>,
141	OP: PackedField<Scalar = OF> + WithUnderlier<Underlier = U>,
142	IF: BinaryField,
143	OF: BinaryField,
144	U: UnderlierWithBitOps,
145{
146	fn transform(&self, input: &IP) -> OP {
147		let mut result = OP::zero();
148		let ones = OP::one().to_underlier();
149		let mut input = input.to_underlier();
150
151		for base in &self.bases {
152			let base_component = input & ones;
153			// contains ones at positions which correspond to non-zero components
154			let mask = broadcast_lowest_bit(base_component, OF::LOG_DEGREE);
155			result += OP::from_underlier(mask & base.to_underlier());
156			input = input >> 1;
157		}
158
159		result
160	}
161}
162
163#[cfg(test)]
164mod tests {
165
166	use super::*;
167	use crate::BinaryField1b;
168
169	const NUM_TESTS: u64 = 100;
170
171	fn check_interleave<F: TowerField>(a: u128, b: u128, c: u128, d: u128) {
172		assert_eq!(interleave::<u128, F>(a, b), (c, d));
173		assert_eq!(interleave::<u128, F>(c, d), (a, b));
174	}
175
176	#[test]
177	fn test_interleave() {
178		check_interleave::<BinaryField1b>(
179			0x0000000000000000FFFFFFFFFFFFFFFF,
180			0xFFFFFFFFFFFFFFFF0000000000000000,
181			0xAAAAAAAAAAAAAAAA5555555555555555,
182			0xAAAAAAAAAAAAAAAA5555555555555555,
183		);
184	}
185}