Branchless Rust: Making a Filter 4x Faster by Removing an If

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 
    }

Worth noting that as written the "trick" results in memory usage proportional to the size of the input rather than the output. If the filter rejects most of the input the difference could be quite noticeable.

Great explanation of why a branchless approach results in such a speed up. I've never really had to deal with performance optimization at this level. Generally it's probably best not to get too involved letting the CPU black box do its thing.

I do wonder, would the performance characteristics of branchless vs branching be consistent across different CPUs/architectures? If you had a CPU that wasn't trying to be fancy with branch prediction, would the regular algo be faster?

I've been doing leetcode in Janet in a (sometimes) tacit (variabless), branchless way:

    (def find-shared-gcd
      (comp
       (fn [e] (max ;(map (fn [d] (* d ;(map |(- 1 (min 1 (mod $ d))) e)))
                         (range 1 (+ 1 (min ;e))))))
       |((juxt* max min) ;$)))

   
    (defn max-diff `where elements increase` [& numbs]
      (reduce max
              -1 (filter |(< 0 $) # strip 0s and add -1 in case (= true (apply > numbs))
                             (map - numbs (accumulate2 min numbs)))))

Thanks for sharing, optimisations like these are what keeps the fun in programming. I have been optimising my JSONLogic evaluator in rust and used arena allocator and preallocation tricks that gave me good jump in tuning. Let me see if branchless programming techniques can get any further in my case

Would PGO figure this out?

This article is 100% AI written. The data was interesting, the commentary overly verbose and hard to gain useful insights from.