Add swizzle_dyn_precise, mirroring swizzle_dyn from std::simd - #276
Conversation
|
I've also submitted similar changes upstream to
One PR is already merged, the other is awaiting review. |
|
The implementation has evolved somewhat since the PR description was written, here's a blog post that goes into detail and is more representative of the current state: https://shnatsel.github.io/improving-std-simd-swizzle-dyn/ |
…on to see what happens" This reverts commit de4a2c7.
…ecise and we can assert it has to always match the scalar fallback
… it's precise and we can assert it has to always match the scalar fallback" to avoid adding extra dependencies This reverts commit 4e8a56a.
…eric split/combine since there's no difference on benchmarks on real Zen2 hardware, which was the best-case scenario for this implementation according to llvm-mca
… llvm-mca results: Haswell, Broadwell and Skylake effectively unchanged for the 256-bit case; Zen1, Tiger Lake show 23% improvement, Zen3 7% to 20% improvement. The 512-bit case that decomposes into this is about the same, but with a 20% improvement on Skylake and a clear 20% improvement on Zen3.
0a551a7 to
8aee113
Compare
# Conflicts: # fearless_simd_tests/tests/harness/lm_generated.rs
…e from vpcmpgtb + vpblendvb to vpaddb + vpor
|
anxsan12 suggested a further optimization that turned out to be worthwhile, so I applied it. It is also merged upstream: rust-lang/portable-simd#543 |
Doubles the throughput on both Ice Lake and Zen4. Improves latency on Ice Lake from 23 to 13 cycles, on Zen4 from 17 to 15 cycles. Improves encoding performance of a toy base64 implementation on Zen4 by -20% (time) +25% (throughput), no change on decoding benchmark
|
@DJMcNab could you take a quick look? All the complex parts (x86 implementations, 2x vector width trick) have already been reviewed upstream in
WASM and NEON implementations are trivial. There are unit tests and random testing to ensure correctness, so I'm quite confident in this code. |
LaurenzV
left a comment
There was a problem hiding this comment.
Only took a brief look.
There was a problem hiding this comment.
It's a bit unfortunate that it makes the file that much longer. :/ Have you checked whether for loops create similar assembly? If so, maybe that would be preferable (same for Fallback).
There was a problem hiding this comment.
The original motivation for fully unrolled forms is #29 (comment)
In a nutshell, LLVM is good at autovectorizing things on x86. But other architectures aren't nearly as sophisticated, and many generally useful optimizations still only exist in the x86 backend. The default Aarch64 cost model in LLVM also has very little loop unrolling, despite ample ILP on recent Aarch64 chips.
I'll double-check what the fallback for this particular op does on SSE2 and on something obscure like POWER.
There was a problem hiding this comment.
I didn't look at the fallback very closely until you asked, but I did manage to find a better fallback formulation after a while. It even improves large swizzles on SSE4.2, which is a somewhat relevant target. Thanks!
… into one cmov per element instead of branches on SSE2 and SSE4.2 and autovectorizes on RISC-V
Change `load_interleaved_128` that returns a 512-bit vector to `load_four_interleaved` that returns four 128-bit vectors. This maps to the real-world usage better (in Vello, all uses of this function save one go through manual split/combine first), eliminates pointless split/combine in x86 implementations, and most importantly allows us to add wrappers for vld3q later which would perform the same operation for RGB rather than RGBA. Here's how this change affects vello: linebender/vello#1796 Also supports interleaved load/store for all widths, addressing a part of linebender#297 The implementations are untouched in this PR save for removing redundant split/combine. AVX2 leaves performance on the table because it uses the sse2 unpack tree for most types, but I don't want to write a bespoke implementation for what is just a swizzle. Instead let's merge linebender#276, then I'll add a relaxed version for x86 that doesn't zero out-of-bounds indices (update: that's linebender#304), and then we'll just lower interleaving into that on x86 since it would call literally all the same intrinsics.
A variant of linebender#276 that produces implementation-defined values instead of 0 on out-of-bounds indices. Only native-size vectors get the relaxed behavior; swizzle_dyn_precise is reused for 2x the native vector width because zeroing allows cheaply joining the halves. Improves performance on x86 and fallback compared to swizzle_dyn_precise. Reduces register pressure, which is important on x86. WASM also improves for 512-bit vectors where it routes through fallback. Part of linebender#29 ### Performance details Most AVX-512 machines get -1 cycle of latency (icelake-server 4->3 cycles, zen4 6->5 cycles), except for Sapphire Rapids where this doubles throughput and halves the latency. AVX2 ~~is unchanged~~ slightly improves throughput at the cost of worse latency (18->21 cycles on Haswell and Skylake, 20->21 on Zen3) and uses one less register. SSE4.2 results: Intel [Westmere](https://en.wikipedia.org/wiki/Westmere_(microarchitecture)) gets double the throughput and latency drops from 3 to 1 cycle. AMD Jaguar (btver2) drops latency from 4 to 2 cycles, throughput improves by 50%. SSE4.2-only Intel [Tremont](https://en.wikipedia.org/wiki/Tremont_(microarchitecture)) launched as recently as 2021 improves latency from 7 to 5 cycles, throughput is unchanged. Fallback (also used by SSE4.2 and WASM for 512-bit vectors) improves 50% in both latency and throughput.
We only had widen/narrow between u8 and u16 before. This PR removes that and builds a generic widen/narrow from the ground up, using the API from linebender#127. It covers all integer and float types. Implementations for floats, and for integers on AVX-512, NEON and WASM are a breeze. Very straightforward, love them. SSE4.2 and AVX2 only have saturating conversions but not truncating ones for integers, so truncating ones have to be emulated on top of saturating ones. There are also no conversions to/from 64-bit integers, so those have to be emulated too. I included saturating conversions in addition to truncating ones in the API, since all hardware has them natively, even older x86. This did end up adding a bit of complexity because 64-bit saturating casts have to be emulated on AVX2 and earlier. I think it's still worth keeping because they don't add that much. All the other complexity (truncating conversion emulation on AVX2 and earlier) is unavoidable. I've axed SSE2-specific kernels since autovectorization works okay for them, and I'm not willing to complicate this even further. Since this is still quite a lot of complexity due to all the x86 emulation, I've added extensive tests on concrete values as well as a random test under `#[ignore]` to be run in release mode used to validate it, same as for other non-trivial ops like linebender#276 Performance-wise rust-lang/rust#159464 bites us here (regression in Rust 1.96, 1.95 and earlier work fine) but it's not so bad that scalar wins, so nothing we can do about it right now (other than inline assembly but i'm not willing to go there). This is [expected to be fixed](rust-lang/rust#159464 (comment)) in the upgrade to LLVM 23 which is coming Soon™. In the meantime this PR is tuned for 1.95 on llvm-mca. ## Generic API considerations There's an API decision on how to expose these ops to generic code, specifically the int/float distinction. It's a tension between flexibility and brevity. In this PR I implemented both regular and saturating for floats and they just do the usual rounding, same as `as` casts. This abstracts over vector type if all you want is widen/narrow. However, if you want to operate on the widened vector afterwards, you need an additional generic bound, e.g. `V::Widened: SimdInt<S>`: <details><summary>Snippet with an additional bound</summary> <p> ```rust use fearless_simd::{SimdFrom, prelude::*, u8x16}; #[inline(always)] fn add_widened_halves<S, V>(value: V) -> V::Widened where S: Simd, V: SimdInt<S> + SimdWiden<S>, V::Widened: SimdInt<S>, { let (low, high) = value.widen(); low + high } #[inline(always)] fn fixed_width<S: Simd>(simd: S) -> [u16; 8] { let input = u8x16::simd_from( simd, [1, 2, 3, 4, 5, 6, 7, 8, 10, 20, 30, 40, 50, 60, 70, 80], ); add_widened_halves(input).into() } #[inline(always)] fn native_width<S: Simd>(simd: S) -> u16 { let input = S::u8s::from_fn(simd, |i| i as u8 + 1); add_widened_halves(input)[0] } ``` </p> </details> This is not a problem on one function, but this extra bound is infectious, so all callers have to require it too until a concrete type is reached. So you have `SimdWiden<S> + V::Widened: SimdInt<S>` propagated all the way up the call stack. Alternatively, one could argue floats really don't have the saturating variant and maybe they should only implement `narrow` but not `saturating_narrow`. So we could make two traits, `SimdNarrowInt<S>` and `SimdNarrowFloat<S>`. So instead of sometimes carrying the `SimdNarrow<S> + V::Narrowed: SimdInt<S>` bound your entire call stack you only carry `SimdNarrowInt<S>`, one having `saturating_*` and the other not, at the expense of not being able to abstract over ints and floats in widen/narrow ops anymore. Abstracting over ints and floats isn't that useful right now since `SimdBase` doesn't have anything on it, but this is something we can and probably should change now that masks aren't `SimdBase`, see linebender#301. Or we could have both `SimdWiden` and `SimdWidenInt`/`SimdWidenFloat`, with the current behavior and `SimdIntWiden`/`SimdFloatWiden` being shorthands to avoid carrying extra bounds - at the expense of more than one trait providing `narrow` and `widen` for each op, and that being potentially confusing. Part of linebender#297 Supersedes linebender#127
I wondered how
std::simdlowersswizzle_dyn- not the limited within_blocks version we added in #266 but the full arbitrary version - and the answer turned out to be "poorly". So I got nerd-sniped, badly. It's 4:30 AM as I write this.The in-tree std::simd is trivially suboptimal, not even taking the most obvious optimization avenues like using 512-bit table lookups on NEON. The NEON gap is fixed upstream but not synced into std yet.
They also lower into a slow, branchy scalar fallback for all cases that don't have a native shuffle op available. There is a better way: if all you have is 128-bit vectors, you can get a 256-bit vector shuffle with four 128-bit shuffles and a couple bitwise ORs.
This PR fixes all of that:
Adds a further optimized 512-bit shuffle on AVX2 which is better than the generic decompositionI've written a toy SIMD base64 decoder using 512-bit shuffles to benchmark this. On Zen4 the specialized AVX2 implementation improves the benchmark by 5%, but it'll never use it because it has AVX-512. On Zen2, where the specialized 512-bit AVX2 is supposed to help the most, there was no difference on benchmarks on real hardware at all, so I went ahead and deleted the specialized 512-bit AVX2 implementation.
Hand-written unit tests are included. Tests on random data that check the optimized SIMD impls against a trivial scalar impl were used to verify the implementation but aren't included in the final PR to avoid extra deps and slow tests (they're preserved in history for posterity but reverted).
I've implemented the *_precise variant first because it is genuinely important for some algorithms, mirrors the std::simd implementation exactly, and the "out of bounds is zero" property actually allows for cool optimizations for non-native vector widths.
This probably deserves a blog post.