Optimize swizzle_dyn for double the native vector width

[Shnatsel]

It is quite trivial to decompose a shuffle of double the native vector width into four native-width shuffles. This adds a generic helper that does this, and wires it up to ssse3, AVX2 and WASM.

The strict, portable behavior of always setting out-of-bounds indices to 0 rather than following the hardware behavior actually helps us here: we don't need to separately track which indices were out of bounds for each half of the lookup table, and can combine the results from both halves with a simple bitwise or because each element will be 0 in at least one half of the output.

This is a port/upstreaming of my work on fearless_simd: linebender/fearless_simd#276

In fearless_simd I've measured a 4x performance improvement on a toy base64 implementation from this change. I have not benchmarked SSE4.2 in fearless_simd directly but llvm-mca says it's a big improvement.

I have not disassembled or benchmarked this implementation to confirm the performance improvement because CONTRIBUTING.md does not explain how to do that on this crate as opposed to its re-export from std.

I've run tests on ssse3 and AVX2, which pass.

I have not run tests on WASM, because the tests fail to run in wasmtime with the wasm32-wasip1 target. I hope CI covers that.

I have not applied this to LoongArch in this PR, since I have no means of testing it and I doubt CI has either.

[programmerjake]

I have not applied this to LoongArch in this PR, since I have no means of testing it and I doubt CI has either.

there is actually a CI job for loongarch64 with default target features -- it's run using qemu.

maybe it would be better to refactor the code to something like:

const fn native_swizzle<const N: usize>(f: fn(src: Simd<u8, N>, idxs: Simd<u8, N>) -> Simd<u8, N>) -> (usize, unsafe fn(!)) {
    (N, core::mem::transmute(f))
}

const NATIVE_SWIZZLES: &[&[(usize, unsafe fn(!))]] = &[
    #[cfg(all(any(target_arch = "x86", target_arch = "x86_64"), target_feature = "ssse3"))]
    &[native_swizzle::<16>(|src, idxs| unsafe {
        transize(x86::_mm_shuffle_epi8, src, zeroing_idxs(idxs))
    })],
    #[cfg(all(
        any(target_arch = "aarch64", target_arch = "arm64ec"),
        target_feature = "neon",
        target_endian = "little"
    ))]
    &[
        native_swizzle::<8>(aarch64_swizzle),
        native_swizzle::<16>(aarch64_swizzle),
        native_swizzle::<24>(aarch64_swizzle),
        native_swizzle::<32>(aarch64_swizzle),
        native_swizzle::<48>(aarch64_swizzle),
        native_swizzle::<64>(aarch64_swizzle),
    ],
    ...
];

const fn find_native_swizzle(n: usize) -> Option<fn(!)> {
    let mut i = 0;
    while i < NATIVE_SWIZZLES.len() {
        let mut j = 0;
        while j < NATIVE_SWIZZLES[i].len() {
            let (cur_n, f) = NATIVE_SWIZZLES[i][j];
            if cur_n == n {
                return Some(f);
            }
            j += 1;
        }
        i += 1;
    }
    None
}

impl<const N: usize> Simd<u8, N> {
    pub fn swizzle_dyn(self, idxs: Self) -> Self {
        unsafe {
            if let Some(f) = const { find_native_swizzle(N) } {
                let f: fn(Self, Self) -> Self = core::mem::transmute(f);
                return f(self, idxs);
            }
            macro_rules! check_split {
                ($n:literal) => {
                    if const { N == $n && find_native_swizzle($n / 2).is_some() } {
                        return transize(swizzle_dyn_split::<$n, { $n / 2 }>, self, idxs);
                    }
                };
            }
            check_split!(16);
            check_split!(32);
            check_split!(64);
            // fallback scalar code:
            ...
        }
    }
}

[Shnatsel]

I've applied the same pattern to LSX and LASX. CI passed.

I admit that proposed refactoring goes over my head. I feel it would make the code less approachable to non-expert contributors such as myself. I'd prefer to keep it out of scope of this PR.

[RalfJung]

If either AVX2 or AVX512 is available we'd use one of the "32" arms above, so this seems redundant?

View changes since the review

[Shnatsel]

Without this rustc complains about an unreachable pattern:

warning: unreachable pattern
  --> crates/core_simd/src/swizzle_dyn.rs:78:17
   |
62 |                 32 => transize(avx2_pshufb, self, idxs),
   |                 -- matches all the relevant values
...
78 |                 32 => transize(swizzle_dyn_split::<32, 16>, self, idxs),
   |                 ^^ no value can reach this
   |
   = note: `#[warn(unreachable_patterns)]` (part of `#[warn(unused)]`) on by default

[RalfJung]

IMO that warning is worth suppressing, it arises form the heavy cfg logic which the "unreachable pattern" lint knows nothing about.

[Shnatsel]

I've added #[allow(unreachable_patterns)], removed all #[cfg(not(...)) and refactored the code to select the best implementation using match precedence.

[programmerjake]

other than that extra comma, looks good to me

View changes since this review