Skip to content

Faster AVX2 swizzle_dyn(_precise) #362

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Shnatsel merged 2 commits into from
Aug 31, 2026
Merged

Faster AVX2 swizzle_dyn(_precise) #362

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4ffafc3
Use a more optimal formulation of swizzle_dyn_precise for AVX2
Shnatsel Aug 29, 2026
7440c4a
Optimize AVX2 swizzle_dyn
Shnatsel Aug 29, 2026
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
28 changes: 14 additions & 14 deletions fearless_simd/src/generated/avx2.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6369,12 +6369,12 @@ impl Simd for Avx2 {
let bytes = Bytes::to_bytes(a).val.0;
let indices = indices.into();
let swapped = _mm256_permute2x128_si256::<0x01>(bytes, bytes);
let local = _mm256_shuffle_epi8(bytes, indices);
let remote = _mm256_shuffle_epi8(swapped, indices);
let select_remote = _mm256_slli_epi16::<3>(indices);
let flip_high_lane = _mm256_set_m128i(_mm_set1_epi8(i8::MIN), _mm_setzero_si128());
let select_remote = _mm256_xor_si256(select_remote, flip_high_lane);
let result = _mm256_blendv_epi8(local, remote, select_remote);
let lane_bias = _mm256_set_m128i(_mm_set1_epi8(-16), _mm_set1_epi8(112));
let local_control = _mm256_add_epi8(indices, lane_bias);
let remote_control = _mm256_xor_si256(local_control, _mm256_set1_epi8(i8::MIN));
let local = _mm256_shuffle_epi8(bytes, local_control);
let remote = _mm256_shuffle_epi8(swapped, remote_control);
let result = _mm256_or_si256(local, remote);
Bytes::from_bytes(u8x32 {
val: crate::support::Aligned256(result),
simd: token,
Expand All @@ -6389,14 +6389,14 @@ impl Simd for Avx2 {
#[inline(always)]
fn kernel(token: Avx2, a: u8x32<Avx2>, indices: u8x32<Avx2>) -> u8x32<Avx2> {
let bytes = Bytes::to_bytes(a);
let idxs = indices;
let lolo = _mm256_permute2x128_si256::<0x00>(bytes.val.0, bytes.val.0);
let hihi = _mm256_permute2x128_si256::<0x11>(bytes.val.0, bytes.val.0);
let control = _mm256_adds_epu8(idxs.into(), _mm256_set1_epi8(0x60));
let select_high = _mm256_slli_epi16::<3>(control);
let from_low = _mm256_shuffle_epi8(lolo, control);
let from_high = _mm256_shuffle_epi8(hihi, control);
let result = _mm256_blendv_epi8(from_low, from_high, select_high);
let indices = indices.into();
let swapped = _mm256_permute2x128_si256::<0x01>(bytes.val.0, bytes.val.0);
let control = _mm256_adds_epu8(indices, _mm256_set1_epi8(0x60));
let local = _mm256_shuffle_epi8(bytes.val.0, control);
let remote = _mm256_shuffle_epi8(swapped, control);
let select_bias = _mm256_set_m128i(_mm_set1_epi8(-112), _mm_set1_epi8(16));
let select_remote = _mm256_add_epi8(control, select_bias);
let result = _mm256_blendv_epi8(local, remote, select_remote);
let result_bytes = u8x32 {
val: crate::support::Aligned256(result),
simd: token,
Expand Down
55 changes: 33 additions & 22 deletions fearless_simd_gen/src/mk_x86.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3703,19 +3703,22 @@ impl X86 {
let bytes = Bytes::to_bytes(a).val.0;
let indices = indices.into();
let swapped = _mm256_permute2x128_si256::<0x01>(bytes, bytes);
let local = _mm256_shuffle_epi8(bytes, indices);
let remote = _mm256_shuffle_epi8(swapped, indices);

// Move index bit 4 into each byte's sign bit for VPBLENDVB.
// The high output lane has the opposite local/remote mapping,
// so invert its blend controls.
let select_remote = _mm256_slli_epi16::<3>(indices);
let flip_high_lane = _mm256_set_m128i(
_mm_set1_epi8(i8::MIN),
_mm_setzero_si128(),

// For an in-range index, set the sign bit in the shuffle
// control for the table half that does not contain the
// requested byte. The high output lane has the opposite
// local/remote mapping.
let lane_bias = _mm256_set_m128i(
_mm_set1_epi8(-16),
_mm_set1_epi8(112),
);
let select_remote = _mm256_xor_si256(select_remote, flip_high_lane);
let result = _mm256_blendv_epi8(local, remote, select_remote);
let local_control = _mm256_add_epi8(indices, lane_bias);
let remote_control =
_mm256_xor_si256(local_control, _mm256_set1_epi8(i8::MIN));

let local = _mm256_shuffle_epi8(bytes, local_control);
let remote = _mm256_shuffle_epi8(swapped, remote_control);
let result = _mm256_or_si256(local, remote);
},
(Self::Avx512, 128 | 256 | 512) => {
let permute = intrinsic_ident("permutexvar", "epi8", vec_ty.n_bits());
Expand Down Expand Up @@ -3762,20 +3765,28 @@ impl X86 {
}
(Self::Avx2, 256) => quote! {
let bytes = Bytes::to_bytes(a);
let idxs = indices;
let lolo = _mm256_permute2x128_si256::<0x00>(bytes.val.0, bytes.val.0);
let hihi = _mm256_permute2x128_si256::<0x11>(bytes.val.0, bytes.val.0);
let indices = indices.into();
let swapped = _mm256_permute2x128_si256::<0x01>(bytes.val.0, bytes.val.0);

// 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 = _mm256_adds_epu8(idxs.into(), _mm256_set1_epi8(0x60));

// Move index bit 4 into each byte's sign bit for VPBLENDVB.
let select_high = _mm256_slli_epi16::<3>(control);
let from_low = _mm256_shuffle_epi8(lolo, control);
let from_high = _mm256_shuffle_epi8(hihi, control);
let result = _mm256_blendv_epi8(from_low, from_high, select_high);
let control = _mm256_adds_epu8(indices, _mm256_set1_epi8(0x60));

let local = _mm256_shuffle_epi8(bytes.val.0, control);
let remote = _mm256_shuffle_epi8(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 = _mm256_set_m128i(
_mm_set1_epi8(-112),
_mm_set1_epi8(16),
);
let select_remote = _mm256_add_epi8(control, select_bias);
let result = _mm256_blendv_epi8(local, remote, select_remote);
let result_bytes = #bytes { val: #wrapper(result), simd: #token };
},
(Self::Avx512, 128 | 256 | 512) => {
Expand Down
18 changes: 18 additions & 0 deletions fearless_simd_tests/tests/harness/ops/swizzle_dyn.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,24 @@ fn swizzle_dyn_u8x32_crosses_blocks<S: Simd>(simd: S) {
assert_swizzle_dyn(bytes, indices, *result);
}

#[simd_test]
fn swizzle_dyn_u8x32_every_valid_index_in_both_lanes<S: Simd>(simd: S) {
// Every value has four bits set, so ORing two distinct table bytes cannot
// accidentally produce either input byte.
let bytes = [
15, 23, 27, 29, 30, 39, 43, 45, 46, 51, 53, 54, 57, 58, 60, 71, 75, 77, 78, 83, 85, 86, 89,
90, 92, 99, 101, 102, 105, 106, 108, 113,
];
let value = u8x32::simd_from(simd, bytes);

for index in 0_u8..32 {
let index_vec = u8x32::simd_from(simd, [index; 32]);
let result = value.swizzle_dyn(index_vec);

assert_eq!(*result, [bytes[usize::from(index)]; 32], "index {index}");
}
}

#[simd_test]
fn swizzle_dyn_u8x64_crosses_blocks<S: Simd>(simd: S) {
let bytes: [u8; 64] = core::array::from_fn(|i| u8::try_from(i + 1).unwrap());
Expand Down
Loading