p3_binary_field/bitslice/transpose.rs
1//! The square bit transpose: row `i` of the result is column `i` of the input.
2//!
3//! A `2n x 2n` matrix is four `n x n` quadrants, and transposing swaps the off-diagonal pair:
4//!
5//! ```text
6//! | A B | | A^T C^T |
7//! | C D | -> | B^T D^T |
8//! ```
9//!
10//! Doing that swap at every quadrant size costs a shift, a mask and two exclusive ors a pair.
11//! So `log2(B)` passes over the rows leave nothing left to transpose.
12//!
13//! This is Hacker's Delight figure 7-6, mirrored because column `0` is the lowest bit here.
14//!
15//! A `nB x nB` matrix is an `n x n` grid of `B x B` blocks, so the identity applies a level up.
16//! Phase 1 transposes each block in place, phase 2 swaps block `(i, j)` with block `(j, i)`.
17//!
18//! Phase 1 treats every word alike, so it sweeps a run of rows flat, at the vector width.
19//!
20//! Only square, only a power of two, one row per lane, and in place.
21
22use core::slice;
23
24use super::packing::{PackedGf2, block_mask};
25use super::underlier::{Underlier, Word};
26
27/// The backing words of a run of packings, as one flat slice.
28///
29/// A run of `k` packings is `k * n` contiguous words, so a pass sweeps it straight through.
30#[inline]
31const fn words_of<U: Underlier>(rows: &mut [PackedGf2<U>]) -> &mut [U::Word] {
32 // size exactly the backing words, so no padding between packings
33 // align at least a word's, so every packing starts word-aligned
34 const {
35 assert!(size_of::<PackedGf2<U>>() == U::WORDS * size_of::<U::Word>());
36 assert!(align_of::<PackedGf2<U>>() >= align_of::<U::Word>());
37 }
38
39 let len = rows.len() * U::WORDS;
40
41 // SAFETY: by the assertions the run is exactly `len` contiguous aligned words.
42 // They lie in one allocation, and every bit pattern of a word is valid.
43 unsafe { slice::from_raw_parts_mut(rows.as_mut_ptr().cast::<U::Word>(), len) }
44}
45
46/// Exchange the off-diagonal quadrants of every `2 * STEP` block, in every word of a row.
47///
48/// The quadrant side is a constant, so the compiler knows every trip count.
49/// That is what stops it vectorising the short passes blindly.
50///
51/// A quadrant as wide as a word belongs to the grid exchange, so that call folds away.
52// The chunk size `2 * STEP` is not a valid const generic argument, so the lint's fix is out.
53#[allow(clippy::chunks_exact_to_as_chunks)]
54#[inline(always)]
55fn stage<U: Underlier, const STEP: usize>(group: &mut [PackedGf2<U>]) {
56 if STEP >= <U::Word as Word>::BITS {
57 return;
58 }
59
60 // The column range of the lower-left quadrant of each block.
61 let mask = U::Word::from_u64(block_mask(STEP));
62
63 // Upper quadrants in the first half of each run, lower ones in the second.
64 for run in group.chunks_exact_mut(2 * STEP) {
65 let (top, bottom) = run.split_at_mut(STEP);
66
67 // The halves pair off word for word, so the sweep stays flat.
68 //
69 // upper right bits [STEP, 2*STEP) of a
70 // lower left bits [0, STEP) of b
71 for (a, b) in words_of(top).iter_mut().zip(words_of(bottom)) {
72 // Their difference, aligned at the bottom, then added back to each side.
73 let diff = ((*a >> STEP) ^ *b) & mask;
74
75 *a = *a ^ (diff << STEP);
76 *b = *b ^ diff;
77 }
78 }
79}
80
81impl<U: Underlier> PackedGf2<U> {
82 /// Replace the square bit matrix whose row `i` is `rows[i]` by its transpose.
83 ///
84 /// Lane `c` of row `r` becomes lane `r` of row `c`.
85 ///
86 /// # Panics
87 /// Panics unless the matrix has exactly as many rows as the packing has lanes.
88 pub fn transpose(rows: &mut [Self]) {
89 assert_eq!(
90 rows.len(),
91 Self::WIDTH,
92 "a bit transpose needs one row per lane"
93 );
94 let block = Self::WORD_BITS;
95
96 // Phase 1, a row-group at a time; a pass wider than the word is skipped inside it.
97 for group in rows.chunks_exact_mut(block) {
98 stage::<U, 32>(group);
99 stage::<U, 16>(group);
100 stage::<U, 8>(group);
101 stage::<U, 4>(group);
102 stage::<U, 2>(group);
103 stage::<U, 1>(group);
104 }
105
106 // Phase 2: block `(i, j)` is word `j` of rows `block*i ..`, so no bits move.
107 // A single-word row has no grid, and both loops are empty.
108 for i in 0..U::WORDS {
109 for j in (i + 1)..U::WORDS {
110 // Disjoint row ranges, because `i < j`.
111 let (above, below) = rows.split_at_mut(block * j);
112
113 for r in 0..block {
114 let upper = above[block * i + r].words()[j];
115 let lower = below[r].words()[i];
116 above[block * i + r].words_mut()[j] = lower;
117 below[r].words_mut()[i] = upper;
118 }
119 }
120 }
121 }
122}
123
124#[cfg(test)]
125mod tests {
126 use p3_field::PrimeCharacteristicRing;
127 use proptest::prelude::*;
128
129 use crate::{Gf2, PackedGf2x8, PackedGf2x64, PackedGf2x128};
130
131 /// One test module per matrix side, against the definition read lane by lane.
132 macro_rules! transpose_tests {
133 ($module:ident, $name:ident, $underlier:ty, $word:ty, $words:literal, $width:literal) => {
134 mod $module {
135 use super::*;
136 use crate::bitslice::underlier::Underlier;
137 use crate::$name;
138
139 fn from_words(words: [$word; $words]) -> $name {
140 $name::new(<$underlier>::from_words_fn(|i| words[i]))
141 }
142
143 /// Row `r` of the result is column `r` of the input, one lane at a time.
144 fn reference(rows: &[$name; $width]) -> [$name; $width] {
145 core::array::from_fn(|r| $name::from_fn(|c| rows[c].get(r)))
146 }
147
148 /// Entry `(r, c)` depends on its position, so any misplacement shows.
149 fn fixture() -> [$name; $width] {
150 core::array::from_fn(|r| {
151 $name::from_fn(|c| Gf2::from_bool((r * 7 + c * 13 + r * c) % 3 == 0))
152 })
153 }
154
155 #[test]
156 fn transpose_matches_the_definition() {
157 let mut rows = fixture();
158 let expected = reference(&rows);
159
160 $name::transpose(&mut rows);
161
162 assert_eq!(rows, expected);
163 }
164
165 #[test]
166 fn transposing_twice_is_the_identity() {
167 // An involution: two passes restore the input.
168 let original = fixture();
169 let mut rows = original;
170
171 $name::transpose(&mut rows);
172 $name::transpose(&mut rows);
173
174 assert_eq!(rows, original);
175 }
176
177 #[test]
178 fn the_extremes_are_fixed_or_swapped_as_expected() {
179 // Symmetric matrices are left alone.
180 let mut zeros = [$name::ZERO; $width];
181 $name::transpose(&mut zeros);
182 assert_eq!(zeros, [$name::ZERO; $width]);
183
184 let mut ones = [$name::ONE; $width];
185 $name::transpose(&mut ones);
186 assert_eq!(ones, [$name::ONE; $width]);
187
188 // Entry `(0, W-1)` moves to the mirrored corner, `(W-1, 0)`.
189 let mut corner = [$name::ZERO; $width];
190 corner[0].set($width - 1, Gf2::ONE);
191 $name::transpose(&mut corner);
192
193 let mut expected = [$name::ZERO; $width];
194 expected[$width - 1].set(0, Gf2::ONE);
195 assert_eq!(corner, expected);
196 }
197
198 #[test]
199 fn the_identity_matrix_is_its_own_transpose() {
200 // The diagonal is the fixed set, so this must not change at all.
201 let identity: [$name; $width] =
202 core::array::from_fn(|r| $name::from_fn(|c| Gf2::from_bool(r == c)));
203 let mut rows = identity;
204
205 $name::transpose(&mut rows);
206
207 assert_eq!(rows, identity);
208 }
209
210 proptest! {
211 #![proptest_config(ProptestConfig::with_cases(16))]
212
213 #[test]
214 fn transpose_matches_the_definition_on_random_matrices(
215 seed in any::<u64>(),
216 ) {
217 // A xorshift fills the matrix, so the printed seed reproduces it.
218 let mut state = seed | 1;
219 let mut next = move || {
220 state ^= state << 13;
221 state ^= state >> 7;
222 state ^= state << 17;
223 state
224 };
225 let mut rows: [$name; $width] = core::array::from_fn(|_| {
226 from_words(core::array::from_fn(|_| next() as $word))
227 });
228
229 let expected = reference(&rows);
230 $name::transpose(&mut rows);
231
232 prop_assert_eq!(rows, expected);
233 }
234 }
235 }
236 };
237 }
238
239 transpose_tests!(x8, PackedGf2x8, u8, u8, 1, 8);
240 transpose_tests!(x16, PackedGf2x16, u16, u16, 1, 16);
241 transpose_tests!(x32, PackedGf2x32, u32, u32, 1, 32);
242 transpose_tests!(x64, PackedGf2x64, u64, u64, 1, 64);
243 transpose_tests!(x128, PackedGf2x128, crate::M128, u64, 2, 128);
244 transpose_tests!(x256, PackedGf2x256, crate::M256, u64, 4, 256);
245 transpose_tests!(x512, PackedGf2x512, crate::M512, u64, 8, 512);
246
247 #[test]
248 fn a_hand_written_eight_by_eight_transposes_as_drawn() {
249 // Column 0 is drawn leftmost and is the lowest bit, so the literals read backwards.
250 //
251 // | input output
252 // | c0....c7 c0....c7
253 // | r0 11111111 10000000
254 // | r1 00000000 10000001
255 // | r2 00000000 10000001
256 // | r3 00010000 10010001
257 // | r4 00000000 10000001
258 // | r5 00000000 10000001
259 // | r6 00000000 10000001
260 // | r7 01111111 10000001
261 //
262 // Row 0 is every column, so it becomes column 0: every output row's bit 0.
263 // The lone entry at `(3, 3)` is on the diagonal and stays put.
264 let input = [
265 0b1111_1111u8,
266 0b0000_0000,
267 0b0000_0000,
268 0b0000_1000,
269 0b0000_0000,
270 0b0000_0000,
271 0b0000_0000,
272 0b1111_1110,
273 ];
274 let expected = [
275 0b0000_0001u8,
276 0b1000_0001,
277 0b1000_0001,
278 0b1000_1001,
279 0b1000_0001,
280 0b1000_0001,
281 0b1000_0001,
282 0b1000_0001,
283 ];
284
285 let mut rows: [PackedGf2x8; 8] = core::array::from_fn(|r| PackedGf2x8::new(input[r]));
286 PackedGf2x8::transpose(&mut rows);
287
288 let got: [u8; 8] = core::array::from_fn(|r| rows[r].words()[0]);
289 assert_eq!(got, expected);
290 }
291
292 #[test]
293 fn the_transpose_reads_lanes_the_way_the_packings_write_them() {
294 // Invariant: the transpose and the lane accessors cannot drift apart.
295 //
296 // before: row r, lane c
297 // after: row c, lane r, and nothing else
298 for r in [0usize, 1, 31, 62, 63] {
299 for c in [0usize, 1, 32, 63] {
300 let mut rows = [PackedGf2x64::ZERO; 64];
301 rows[r].set(c, Gf2::ONE);
302
303 PackedGf2x64::transpose(&mut rows);
304
305 let mut expected = [PackedGf2x64::ZERO; 64];
306 expected[c].set(r, Gf2::ONE);
307 assert_eq!(rows, expected, "entry ({r}, {c})");
308 }
309 }
310 }
311
312 #[test]
313 fn a_wide_transpose_moves_entries_across_word_boundaries() {
314 // Invariant: an entry crosses blocks, the step a per-word kernel would miss.
315 //
316 // entry (10, 100) block (0, 1), word 1 of row 10
317 // entry (100, 10) block (1, 0), word 0 of row 100
318 let mut rows = [PackedGf2x128::ZERO; 128];
319 rows[10].set(100, Gf2::ONE);
320
321 PackedGf2x128::transpose(&mut rows);
322
323 let mut expected = [PackedGf2x128::ZERO; 128];
324 expected[100].set(10, Gf2::ONE);
325 assert_eq!(rows, expected);
326 }
327}