binius_field/
transpose.rs

1// Copyright 2023-2025 Irreducible Inc.
2
3use binius_utils::checked_arithmetics::log2_strict_usize;
4
5use super::packed::PackedField;
6use crate::{ExtensionField, Field, PackedExtension};
7
8/// Transpose square blocks of elements within packed field elements in place.
9///
10/// The input elements are interpreted as a rectangular matrix with height `n = 2^n` in row-major
11/// order. This matrix is interpreted as a vector of square matrices of field elements, and each
12/// square matrix is transposed in-place.
13///
14/// # Arguments
15///
16/// * `log_n`: The base-2 logarithm of the dimension of the n x n square matrix. Must be less than
17///   or equal to the base-2 logarithm of the packing width.
18/// * `elems`: The packed field elements, length is a power-of-two multiple of `1 << log_n`.
19///
20/// # Preconditions
21///
22/// * `log_n` must be at most `P::LOG_WIDTH`.
23/// * `elems.len()` must be a power of two and at least `2^log_n`.
24pub fn square_transpose<P: PackedField>(log_n: usize, elems: &mut [P]) {
25	assert!(P::LOG_WIDTH >= log_n, "dimension n of square blocks must divide packing width");
26
27	let size = elems.len();
28	assert!(size.is_power_of_two(), "elems length must be a power of two, got {size}");
29	let log_size = log2_strict_usize(size);
30	assert!(
31		log_size >= log_n,
32		"elems must have length at least 2^log_n = {}, got {size}",
33		1 << log_n
34	);
35
36	let log_w = log_size - log_n;
37
38	// See Hacker's Delight, Section 7-3.
39	// https://dl.acm.org/doi/10.5555/2462741
40	for i in 0..log_n {
41		for j in 0..1 << (log_n - i - 1) {
42			for k in 0..1 << (log_w + i) {
43				let idx0 = (j << (log_w + i + 1)) | k;
44				let idx1 = idx0 | (1 << (log_w + i));
45
46				let v0 = elems[idx0];
47				let v1 = elems[idx1];
48				let (v0, v1) = v0.interleave(v1, i);
49				elems[idx0] = v0;
50				elems[idx1] = v1;
51			}
52		}
53	}
54}
55
56pub fn square_transforms_extension_field<F: Field, FE: ExtensionField<F> + PackedExtension<F>>(
57	values: &mut [FE],
58) {
59	square_transpose(FE::LOG_DEGREE, FE::cast_bases_mut(values))
60}
61
62#[cfg(test)]
63mod tests {
64	use super::*;
65	use crate::PackedBinaryField128x1b;
66
67	#[test]
68	fn test_square_transpose_128x1b() {
69		let mut elems = [
70			PackedBinaryField128x1b::from(0x00000000000000000000000000000000u128),
71			PackedBinaryField128x1b::from(0x00000000000000000000000000000000u128),
72			PackedBinaryField128x1b::from(0x00000000000000000000000000000000u128),
73			PackedBinaryField128x1b::from(0x00000000000000000000000000000000u128),
74			PackedBinaryField128x1b::from(0xffffffffffffffffffffffffffffffffu128),
75			PackedBinaryField128x1b::from(0xffffffffffffffffffffffffffffffffu128),
76			PackedBinaryField128x1b::from(0xffffffffffffffffffffffffffffffffu128),
77			PackedBinaryField128x1b::from(0xffffffffffffffffffffffffffffffffu128),
78		];
79		square_transpose(3, &mut elems);
80
81		let expected = [
82			PackedBinaryField128x1b::from(0xf0f0f0f0f0f0f0f0f0f0f0f0f0f0f0f0u128),
83			PackedBinaryField128x1b::from(0xf0f0f0f0f0f0f0f0f0f0f0f0f0f0f0f0u128),
84			PackedBinaryField128x1b::from(0xf0f0f0f0f0f0f0f0f0f0f0f0f0f0f0f0u128),
85			PackedBinaryField128x1b::from(0xf0f0f0f0f0f0f0f0f0f0f0f0f0f0f0f0u128),
86			PackedBinaryField128x1b::from(0xf0f0f0f0f0f0f0f0f0f0f0f0f0f0f0f0u128),
87			PackedBinaryField128x1b::from(0xf0f0f0f0f0f0f0f0f0f0f0f0f0f0f0f0u128),
88			PackedBinaryField128x1b::from(0xf0f0f0f0f0f0f0f0f0f0f0f0f0f0f0f0u128),
89			PackedBinaryField128x1b::from(0xf0f0f0f0f0f0f0f0f0f0f0f0f0f0f0f0u128),
90		];
91		assert_eq!(elems, expected);
92	}
93
94	#[test]
95	fn test_square_transpose_128x1b_multi_row() {
96		let mut elems = [
97			PackedBinaryField128x1b::from(0x00000000000000000000000000000000u128),
98			PackedBinaryField128x1b::from(0x00000000000000000000000000000000u128),
99			PackedBinaryField128x1b::from(0x00000000000000000000000000000000u128),
100			PackedBinaryField128x1b::from(0x00000000000000000000000000000000u128),
101			PackedBinaryField128x1b::from(0xffffffffffffffffffffffffffffffffu128),
102			PackedBinaryField128x1b::from(0xffffffffffffffffffffffffffffffffu128),
103			PackedBinaryField128x1b::from(0xffffffffffffffffffffffffffffffffu128),
104			PackedBinaryField128x1b::from(0xffffffffffffffffffffffffffffffffu128),
105		];
106		square_transpose(1, &mut elems);
107
108		let expected = [
109			PackedBinaryField128x1b::from(0xaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaau128),
110			PackedBinaryField128x1b::from(0xaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaau128),
111			PackedBinaryField128x1b::from(0xaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaau128),
112			PackedBinaryField128x1b::from(0xaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaau128),
113			PackedBinaryField128x1b::from(0xaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaau128),
114			PackedBinaryField128x1b::from(0xaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaau128),
115			PackedBinaryField128x1b::from(0xaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaau128),
116			PackedBinaryField128x1b::from(0xaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaau128),
117		];
118		assert_eq!(elems, expected);
119	}
120}