Nice post!
You can do even a bit better if you're willing to use intrinsics. In particular this kind of operation is well-suited for compress-type operations, available as a first-class operation in at least AVX512, SVE and RVV; you can also emulate them reasonably quickly on NEON and AVX2.
Here's an example, building on the OP's work:
pub fn filter_compress(input: &[f64], threshold: f64) -> Vec<f64> {
use std::arch::x86_64::*;
let mut out = vec![0.0; input.len()];
let mut n = 0usize;
let (head, tail) = input.as_chunks::<8>();
for chunk in head {
unsafe {
let p = _mm512_loadu_pd(chunk.as_ptr());
let m = _mm512_cmpnle_pd_mask(p, _mm512_set1_pd(threshold));
let compress = _mm512_maskz_compress_pd(m, p);
_mm512_storeu_pd(out.as_mut_ptr().wrapping_add(n), compress);
n += m.count_ones() as usize;
}
}
for &x in tail {
out[n] = x;
n += (x > threshold) as usize;
}
out.truncate(n);
out
}
Nice! I saw the code and thought, I bet there’s a way to do some SIMD here… never touched intrinsics in Rust before so I really appreciate you writing it up!
See the following pdf for example on how to do this with SSSE3 (pages 104-133) or even SSE2 (pages 151-173)
https://deplinenoise.files.wordpress.com/2015/03/gdc2015_afr...
Great link, nice find!
This is interesting. So at a certain scale, CPU optimization becomes irrelevant because you're just waiting for new data to come in?
I heard there's a way to arrange code such that the compiler can autobectorize easier. I wonder if there's a way to do that here?
Would probably have to pass `-C target-cpu=native` to cargo so that llvm is allowed to use AVX512.
Good question. I personally doubt that the compress instruction is easy to coax compilers into generating, as there are many edge cases to consider.
For example, you'll notice here that we perform a full vector store of 8 elements unconditionally, even if only a few of the elements are active. This is safe, though, because the output buffer is as large as the input buffer, and we're chunking by 8, so we'll never trash memory past the end; but this is a tricky analysis. Performance-wise, we rely on the CPU's store buffer to make these overlapping stores cheap.
Instead, you might think that you could just store the elements which are actually active, using a masked store. In fact there is also an intrinsic for this purpose (_mm512_mask_compressstoreu_pd), but it is extremely slow on some CPUs, namely Zen 4, so it's dangerous to use unless you know exactly what CPU you're using. (In my testing, there also seems to be some weird hazard on Zen 5 where multiple memory-destination compress instructions to nearby, even non-overlapping, addresses are serialized. But I haven't looked closer at this.)
I think WUFFS "iterate loops" might help. WUFFS requires that processing a chunk of N items has code to process one at a time, which means it'll work for any N. However you can optionally provide specialisations for doing K at a time and the compiler is responsible for carving the input up as appropriate so e.g. N = K + K + 1 + 1 + 1 your K-at-a-time code runs twice, the extras are handled 1-at-a-time.
So this divides up the problem, the compiler can vectorize your 8-at-a-time code without needing to handle edge cases where N isn't a multiple of 8, and if a later pass notices we actually never end up using those edge cases they're dead code, if it doesn't they're just a rarely-taken branch once.
I believe the bounds check, in particular, is devastating for autovectorization. There are ways around it, but it requires additional code in safe rust.
Edit: actually looks like autovectorization is in play here [1]. Doesn't look like the bounds check gets in the way at all.
[1] https://godbolt.org/z/af4qGba5o
There's no autovectorization there; scalar f64-s just are always stored in xmm registers. And the bounds check is still there.
Compress patterns aren't recognized by any open-source compiler autovectorizer as far as I'm aware of. (I think intel's proprietary C/C++ compiler can?)
Thank you for sharing this. How would you emulate this kind of operation on avx2?
Once you have a mask of the positions you want to compress, you can generate a shuffle index vector from that mask to place the desired elements in the low part of the vector. You can expand the mask into nibble-sized indices using pext/pdep and some magic constants, then expand those nibble-sized indices into a vector of indices to use as the shuffle indices.
Yes, that's one approach. Another reasonable approach is to get out a mask from the comparison using `vmovmskpd` and use that to look up a shuffle constant, since there are only 16 possibilities. This also works well on NEON, although I wonder there whether it'd make more sense to find the shuffle dynamically rather than loading it.
Keep in mind it's UB to be:
> Executing code compiled with target features that the current thread of execution does not support
I.e. calling AVX512 on Neon architecture.
You need to wrap it in target attributes to even dream of it being safe.
This particular UB is not one of the subtle cases. You will almost certainly get illegal instruction signals if you mess this up.
Is rust UB different from C UB? C UB must be avoided at all costs even if you think you know the actual behavior.
It is not.
[dead]