From 0a334f9c87ece5ddf6690086837846be535874e5 Mon Sep 17 00:00:00 2001 From: "Sergey \"Shnatsel\" Davidoff" Date: Sat, 29 Aug 2026 18:16:58 +0100 Subject: [PATCH] Use a more optimal formulation of swizzle_dyn for AVX2 --- crates/core_simd/src/swizzle_dyn.rs | 23 +++++++++++++++-------- 1 file changed, 15 insertions(+), 8 deletions(-) diff --git a/crates/core_simd/src/swizzle_dyn.rs b/crates/core_simd/src/swizzle_dyn.rs index da91e01801c..a4985a219fd 100644 --- a/crates/core_simd/src/swizzle_dyn.rs +++ b/crates/core_simd/src/swizzle_dyn.rs @@ -265,19 +265,26 @@ unsafe fn avx2_pshufb(bytes: Simd, idxs: Simd) -> Simd { use x86::_mm256_shuffle_epi8 as avx2_half_pshufb; // SAFETY: Caller promised AVX2 unsafe { - let lolo = avx2_cross_shuffle::<0x00>(bytes.into(), bytes.into()); - let hihi = avx2_cross_shuffle::<0x11>(bytes.into(), bytes.into()); + let bytes = bytes.into(); + let indices = idxs.into(); + let swapped = avx2_cross_shuffle::<0x01>(bytes, bytes); // Adding 0x60 preserves the low nibble and bit 4 for valid // indices 0..=31. Larger indices get their high bit set, so // VPSHUFB supplies the required out-of-bounds zeroing. - let control = x86::_mm256_adds_epu8(idxs.into(), x86::_mm256_set1_epi8(0x60)); + let control = x86::_mm256_adds_epu8(indices, x86::_mm256_set1_epi8(0x60)); - // Move index bit 4 into each byte's sign bit for VPBLENDVB. - let select_high = x86::_mm256_slli_epi16::<3>(control); - let from_low = avx2_half_pshufb(lolo, control); - let from_high = avx2_half_pshufb(hihi, control); - x86::_mm256_blendv_epi8(from_low, from_high, select_high).into() + let local = avx2_half_pshufb(bytes, control); + let remote = avx2_half_pshufb(swapped, control); + + // In the low lane, adding 0x10 moves the valid index's bit 4 + // into the sign bit. The high lane has the opposite + // local/remote mapping, so adding 0x90 flips the selection. + // Out-of-range indices already zeroed both shuffle results, + // making the blend selection irrelevant for them. + let select_bias = x86::_mm256_set_m128i(x86::_mm_set1_epi8(-112), x86::_mm_set1_epi8(16)); + let select_remote = x86::_mm256_add_epi8(control, select_bias); + x86::_mm256_blendv_epi8(local, remote, select_remote).into() } }