binius_field/
transpose.rs1use binius_utils::checked_arithmetics::log2_strict_usize;
4
5use super::packed::PackedField;
6use crate::{ExtensionField, Field, PackedExtension};
7
8pub 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 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}