Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
62 changes: 62 additions & 0 deletions src/simd.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
///
Expand Down
32 changes: 32 additions & 0 deletions src/simd_avx2.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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] };
Comment on lines +1745 to +1746

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Reject unsupported exchange granularities

Because this public const-generic API does not constrain G to the documented values 1, 2, 4, and 8, calls such as exchange::<3> silently duplicate and discard lanes rather than performing an exchange, while exchange::<16> indexes beyond the 16-lane arrays and panics. The tests exercise only valid powers of two, so add a const/runtime assertion for the supported set to this implementation and the five replicated backend implementations.

Useful? React with 👍 / 👎.

}
(Self::from_array(nl), Self::from_array(nh))
}
}

impl U32x16 {
Expand Down
32 changes: 32 additions & 0 deletions src/simd_avx512.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
32 changes: 32 additions & 0 deletions src/simd_neon.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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")]
Expand Down
32 changes: 32 additions & 0 deletions src/simd_nightly/u_word_types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
32 changes: 32 additions & 0 deletions src/simd_scalar.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
32 changes: 32 additions & 0 deletions src/simd_wasm.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
Loading