Merge pull request #270 from AdaWorldAPI/claude/u32x16-exchange

simd: exchange<const G> on U32x16 — the transpose surface, all six backends

Author: AdaWorldAPI

src/simd.rs

@@ -898,6 +898,68 @@ mod tests {
         );
     }
 
+    /// `exchange::<G>` composes a complete 16x16 transpose over four stages.
+    ///
+    /// This is the control the codegen oracle's driver applies to
+    /// `transpose_16x16_composed`, brought into the test suite so the library
+    /// method is checked and not merely its probe-local twin. A packed-but-
+    /// wrong shuffle network would pass a codegen histogram and fail here.
+    ///
+    /// Note the semantics differ from `interleave_*` / `concat_*` above: those
+    /// mirror x86's per-128-bit-lane `unpack`, whereas `exchange` pairs lane
+    /// `c` with lane `c ^ G` across the whole vector. They are different
+    /// permutations and neither substitutes for the other.
+    #[test]
+    fn u32x16_exchange_stages_compose_a_transpose() {
+        fn stage<const G: usize>(m: &mut [U32x16; 16]) {
+            for r in 0..16 {
+                if r & G == 0 {
+                    let (lo, hi) = m[r].exchange::<G>(m[r | G]);
+                    m[r] = lo;
+                    m[r | G] = hi;
+                }
+            }
+        }
+
+        // Row r, lane c = r*16 + c. A transpose must yield row r, lane c = c*16 + r.
+        let mut m: [U32x16; 16] =
+            core::array::from_fn(|r| U32x16::from_array(core::array::from_fn(|c| (r * 16 + c) as u32)));
+
+        stage::<1>(&mut m);
+        stage::<2>(&mut m);
+        stage::<4>(&mut m);
+        stage::<8>(&mut m);
+
+        for (r, row) in m.iter().enumerate() {
+            let want: [u32; 16] = core::array::from_fn(|c| (c * 16 + r) as u32);
+            assert_eq!(row.to_array(), want, "row {r} after four exchange stages");
+        }
+    }
+
+    /// Each granularity in isolation, pinned lane-by-lane, so a backend that
+    /// gets one stage wrong is not masked by the composition above.
+    #[test]
+    fn u32x16_exchange_is_lane_exact_per_granularity() {
+        let a: [u32; 16] = core::array::from_fn(|i| 10 + i as u32);
+        let b: [u32; 16] = core::array::from_fn(|i| 30 + i as u32);
+        let (va, vb) = (U32x16::from_array(a), U32x16::from_array(b));
+
+        // lo[c] = a[c] when c & G == 0 else b[c ^ G]
+        // hi[c] = b[c] when c & G != 0 else a[c ^ G]
+        for g in [1usize, 2, 4, 8] {
+            let (lo, hi) = match g {
+                1 => va.exchange::<1>(vb),
+                2 => va.exchange::<2>(vb),
+                4 => va.exchange::<4>(vb),
+                _ => va.exchange::<8>(vb),
+            };
+            let want_lo: [u32; 16] = core::array::from_fn(|c| if c & g == 0 { a[c] } else { b[c ^ g] });
+            let want_hi: [u32; 16] = core::array::from_fn(|c| if c & g != 0 { b[c] } else { a[c ^ g] });
+            assert_eq!(lo.to_array(), want_lo, "exchange::<{g}> lo");
+            assert_eq!(hi.to_array(), want_hi, "exchange::<{g}> hi");
+        }
+    }
+
     /// `U32x8`'s shuffle + rotate surface, checked against the REAL x86
     /// intrinsics it claims to reproduce.
     ///

src/simd_avx2.rs

@@ -1715,6 +1715,38 @@ impl U32x16 {
         }
         Self::from_array(o)
     }
+
+    /// One butterfly exchange at block granularity `G` elements — the general
+    /// form of the whole unpack / lane-exchange family, parameterized by
+    /// granularity instead of one method per width.
+    ///
+    /// `G = 1` is the 32-bit unpack, `G = 2` the 64-bit unpack, `G = 4` the
+    /// 128-bit lane exchange, and `G = 8` the 256-bit half exchange that a
+    /// 512-bit lane additionally needs and for which no AVX2 intrinsic exists.
+    /// `G` is a const parameter, so every shuffle pattern is compile-time
+    /// constant — the same property a hand-written intrinsic has, and the
+    /// precondition for LLVM to select a shuffle rather than an indexed copy.
+    ///
+    /// Four stages over this (`G` = 1, 2, 4, 8, pairing row `r` with `r | G`)
+    /// compose a complete 16x16 transpose. Measured as
+    /// `transpose_16x16_composed`: **79 packed / 0 scalar-lane-arith**, 19 of
+    /// them real shuffles. The same transpose written as one monolithic index
+    /// loop measures **0 packed** — 1088 bytes of stack and a 256-iteration
+    /// scalar copy. The spelling is the entire difference. See
+    /// `.claude/knowledge/blake3-on-ndarray-simd.md`.
+    ///
+    /// No intrinsic override is earned: the generic form does not fail.
+    #[inline(always)]
+    pub fn exchange<const G: usize>(self, other: Self) -> (Self, Self) {
+        let (l, h) = (self.to_array(), other.to_array());
+        let mut nl = [0u32; 16];
+        let mut nh = [0u32; 16];
+        for c in 0..16 {
+            nl[c] = if c & G == 0 { l[c] } else { h[c ^ G] };
+            nh[c] = if c & G != 0 { h[c] } else { l[c ^ G] };
+        }
+        (Self::from_array(nl), Self::from_array(nh))
+    }
 }
 
 impl U32x16 {

src/simd_avx512.rs

@@ -1521,6 +1521,38 @@ impl U32x16 {
         }
         Self::from_array(o)
     }
+
+    /// One butterfly exchange at block granularity `G` elements — the general
+    /// form of the whole unpack / lane-exchange family, parameterized by
+    /// granularity instead of one method per width.
+    ///
+    /// `G = 1` is the 32-bit unpack, `G = 2` the 64-bit unpack, `G = 4` the
+    /// 128-bit lane exchange, and `G = 8` the 256-bit half exchange that a
+    /// 512-bit lane additionally needs and for which no AVX2 intrinsic exists.
+    /// `G` is a const parameter, so every shuffle pattern is compile-time
+    /// constant — the same property a hand-written intrinsic has, and the
+    /// precondition for LLVM to select a shuffle rather than an indexed copy.
+    ///
+    /// Four stages over this (`G` = 1, 2, 4, 8, pairing row `r` with `r | G`)
+    /// compose a complete 16x16 transpose. Measured as
+    /// `transpose_16x16_composed`: **79 packed / 0 scalar-lane-arith**, 19 of
+    /// them real shuffles. The same transpose written as one monolithic index
+    /// loop measures **0 packed** — 1088 bytes of stack and a 256-iteration
+    /// scalar copy. The spelling is the entire difference. See
+    /// `.claude/knowledge/blake3-on-ndarray-simd.md`.
+    ///
+    /// No intrinsic override is earned: the generic form does not fail.
+    #[inline(always)]
+    pub fn exchange<const G: usize>(self, other: Self) -> (Self, Self) {
+        let (l, h) = (self.to_array(), other.to_array());
+        let mut nl = [0u32; 16];
+        let mut nh = [0u32; 16];
+        for c in 0..16 {
+            nl[c] = if c & G == 0 { l[c] } else { h[c ^ G] };
+            nh[c] = if c & G != 0 { h[c] } else { l[c ^ G] };
+        }
+        (Self::from_array(nl), Self::from_array(nh))
+    }
 }
 
 impl U32x16 {

src/simd_neon.rs

@@ -1726,6 +1726,38 @@ impl U32x16 {
         }
         Self::from_array(o)
     }
+
+    /// One butterfly exchange at block granularity `G` elements — the general
+    /// form of the whole unpack / lane-exchange family, parameterized by
+    /// granularity instead of one method per width.
+    ///
+    /// `G = 1` is the 32-bit unpack, `G = 2` the 64-bit unpack, `G = 4` the
+    /// 128-bit lane exchange, and `G = 8` the 256-bit half exchange that a
+    /// 512-bit lane additionally needs and for which no AVX2 intrinsic exists.
+    /// `G` is a const parameter, so every shuffle pattern is compile-time
+    /// constant — the same property a hand-written intrinsic has, and the
+    /// precondition for LLVM to select a shuffle rather than an indexed copy.
+    ///
+    /// Four stages over this (`G` = 1, 2, 4, 8, pairing row `r` with `r | G`)
+    /// compose a complete 16x16 transpose. Measured as
+    /// `transpose_16x16_composed`: **79 packed / 0 scalar-lane-arith**, 19 of
+    /// them real shuffles. The same transpose written as one monolithic index
+    /// loop measures **0 packed** — 1088 bytes of stack and a 256-iteration
+    /// scalar copy. The spelling is the entire difference. See
+    /// `.claude/knowledge/blake3-on-ndarray-simd.md`.
+    ///
+    /// No intrinsic override is earned: the generic form does not fail.
+    #[inline(always)]
+    pub fn exchange<const G: usize>(self, other: Self) -> (Self, Self) {
+        let (l, h) = (self.to_array(), other.to_array());
+        let mut nl = [0u32; 16];
+        let mut nh = [0u32; 16];
+        for c in 0..16 {
+            nl[c] = if c & G == 0 { l[c] } else { h[c ^ G] };
+            nh[c] = if c & G != 0 { h[c] } else { l[c ^ G] };
+        }
+        (Self::from_array(nl), Self::from_array(nh))
+    }
 }
 
 #[cfg(target_arch = "aarch64")]

src/simd_nightly/u_word_types.rs

@@ -459,6 +459,38 @@ impl U32x16 {
         }
         Self::from_array(o)
     }
+
+    /// One butterfly exchange at block granularity `G` elements — the general
+    /// form of the whole unpack / lane-exchange family, parameterized by
+    /// granularity instead of one method per width.
+    ///
+    /// `G = 1` is the 32-bit unpack, `G = 2` the 64-bit unpack, `G = 4` the
+    /// 128-bit lane exchange, and `G = 8` the 256-bit half exchange that a
+    /// 512-bit lane additionally needs and for which no AVX2 intrinsic exists.
+    /// `G` is a const parameter, so every shuffle pattern is compile-time
+    /// constant — the same property a hand-written intrinsic has, and the
+    /// precondition for LLVM to select a shuffle rather than an indexed copy.
+    ///
+    /// Four stages over this (`G` = 1, 2, 4, 8, pairing row `r` with `r | G`)
+    /// compose a complete 16x16 transpose. Measured as
+    /// `transpose_16x16_composed`: **79 packed / 0 scalar-lane-arith**, 19 of
+    /// them real shuffles. The same transpose written as one monolithic index
+    /// loop measures **0 packed** — 1088 bytes of stack and a 256-iteration
+    /// scalar copy. The spelling is the entire difference. See
+    /// `.claude/knowledge/blake3-on-ndarray-simd.md`.
+    ///
+    /// No intrinsic override is earned: the generic form does not fail.
+    #[inline(always)]
+    pub fn exchange<const G: usize>(self, other: Self) -> (Self, Self) {
+        let (l, h) = (self.to_array(), other.to_array());
+        let mut nl = [0u32; 16];
+        let mut nh = [0u32; 16];
+        for c in 0..16 {
+            nl[c] = if c & G == 0 { l[c] } else { h[c ^ G] };
+            nh[c] = if c & G != 0 { h[c] } else { l[c ^ G] };
+        }
+        (Self::from_array(nl), Self::from_array(nh))
+    }
 }
 
 impl U32x16 {

src/simd_scalar.rs

@@ -1269,6 +1269,38 @@ impl U32x16 {
         }
         Self::from_array(o)
     }
+
+    /// One butterfly exchange at block granularity `G` elements — the general
+    /// form of the whole unpack / lane-exchange family, parameterized by
+    /// granularity instead of one method per width.
+    ///
+    /// `G = 1` is the 32-bit unpack, `G = 2` the 64-bit unpack, `G = 4` the
+    /// 128-bit lane exchange, and `G = 8` the 256-bit half exchange that a
+    /// 512-bit lane additionally needs and for which no AVX2 intrinsic exists.
+    /// `G` is a const parameter, so every shuffle pattern is compile-time
+    /// constant — the same property a hand-written intrinsic has, and the
+    /// precondition for LLVM to select a shuffle rather than an indexed copy.
+    ///
+    /// Four stages over this (`G` = 1, 2, 4, 8, pairing row `r` with `r | G`)
+    /// compose a complete 16x16 transpose. Measured as
+    /// `transpose_16x16_composed`: **79 packed / 0 scalar-lane-arith**, 19 of
+    /// them real shuffles. The same transpose written as one monolithic index
+    /// loop measures **0 packed** — 1088 bytes of stack and a 256-iteration
+    /// scalar copy. The spelling is the entire difference. See
+    /// `.claude/knowledge/blake3-on-ndarray-simd.md`.
+    ///
+    /// No intrinsic override is earned: the generic form does not fail.
+    #[inline(always)]
+    pub fn exchange<const G: usize>(self, other: Self) -> (Self, Self) {
+        let (l, h) = (self.to_array(), other.to_array());
+        let mut nl = [0u32; 16];
+        let mut nh = [0u32; 16];
+        for c in 0..16 {
+            nl[c] = if c & G == 0 { l[c] } else { h[c ^ G] };
+            nh[c] = if c & G != 0 { h[c] } else { l[c ^ G] };
+        }
+        (Self::from_array(nl), Self::from_array(nh))
+    }
 }
 
 impl U32x16 {

src/simd_wasm.rs

@@ -1052,6 +1052,38 @@ pub mod wasm32_simd {
             }
             Self::from_array(o)
         }
+
+        /// One butterfly exchange at block granularity `G` elements — the general
+        /// form of the whole unpack / lane-exchange family, parameterized by
+        /// granularity instead of one method per width.
+        ///
+        /// `G = 1` is the 32-bit unpack, `G = 2` the 64-bit unpack, `G = 4` the
+        /// 128-bit lane exchange, and `G = 8` the 256-bit half exchange that a
+        /// 512-bit lane additionally needs and for which no AVX2 intrinsic exists.
+        /// `G` is a const parameter, so every shuffle pattern is compile-time
+        /// constant — the same property a hand-written intrinsic has, and the
+        /// precondition for LLVM to select a shuffle rather than an indexed copy.
+        ///
+        /// Four stages over this (`G` = 1, 2, 4, 8, pairing row `r` with `r | G`)
+        /// compose a complete 16x16 transpose. Measured as
+        /// `transpose_16x16_composed`: **79 packed / 0 scalar-lane-arith**, 19 of
+        /// them real shuffles. The same transpose written as one monolithic index
+        /// loop measures **0 packed** — 1088 bytes of stack and a 256-iteration
+        /// scalar copy. The spelling is the entire difference. See
+        /// `.claude/knowledge/blake3-on-ndarray-simd.md`.
+        ///
+        /// No intrinsic override is earned: the generic form does not fail.
+        #[inline(always)]
+        pub fn exchange<const G: usize>(self, other: Self) -> (Self, Self) {
+            let (l, h) = (self.to_array(), other.to_array());
+            let mut nl = [0u32; 16];
+            let mut nh = [0u32; 16];
+            for c in 0..16 {
+                nl[c] = if c & G == 0 { l[c] } else { h[c ^ G] };
+                nh[c] = if c & G != 0 { h[c] } else { l[c ^ G] };
+            }
+            (Self::from_array(nl), Self::from_array(nh))
+        }
     }
 
     impl U32x16 {