logoalt Hacker News

anematodetoday at 4:03 AM0 repliesview on HN

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 
    }
For me it's about 25% less time than the branchless version with 1,000,000 elements, and 60% less with 10,000 elements where memory bandwidth effects are less relevant.