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 
    }

This problem is called stream compaction and there is a wealth of research on it. The best methods use prefix scan. They first efficiently compute the index in the output array of each element that satisfies the predicate and then they gather them in one linear operation.

Also, I can tell that you are a good writer. You didn't need the LLM to "polish" your text.

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?

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.

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.