Xjyuk

When running a sum loop over an array in Rust, I noticed a huge performance difference when CAPACITY >= 240. CAPACITY = 239 is about 80 times faster.

Is there special compilation optimization Rust is doing for "short" arrays?

Compiled with rustc -C opt-level=3.

use std::time::Instant;

const CAPACITY: usize = 240;
const IN_LOOPS: usize = 500000;

fn main() 
 let mut arr = [0; CAPACITY];
 for i in 0..CAPACITY 
 arr[i] = i;
 
 let mut sum = 0;
 let now = Instant::now();
 for _ in 0..IN_LOOPS 
 let mut s = 0;
 for i in 0..arr.len() 
 s += arr[i];
 
 sum += s;
 
 println!("sum: time::?", sum, now.elapsed());

You found a magic threshold above which LLVM stops performing certain optimizations. The threshold is 8 bytes * 240 = 1920 bytes (your array is an array of usizes, therefore the length is multiplied with 8 bytes, assuming x86-64 CPU). In this benchmark, one specific optimization -- only performed for length 239 -- is responsible for the huge speed difference. But let's start slowly:

(All code in this answer is compiled with -C opt-level=3)

pub fn foo() -> usize 
 let arr = [0; 240];
 let mut s = 0;
 for i in 0..arr.len() 
 s += arr[i];
 
 s

This simple code will produce roughly the assembly one would expect: a loop adding up elements. However, if you change 240 to 239, the emitted assembly differs quite a lot. See it on Godbolt Compiler Explorer. I only paste a small part of the assembly here:

movdqa xmm1, xmmword ptr [rsp + 32]
movdqa xmm0, xmmword ptr [rsp + 48]
paddq xmm1, xmmword ptr [rsp]
paddq xmm0, xmmword ptr [rsp + 16]
paddq xmm1, xmmword ptr [rsp + 64]
; more stuff omitted here ...
paddq xmm0, xmmword ptr [rsp + 1840]
paddq xmm1, xmmword ptr [rsp + 1856]
paddq xmm0, xmmword ptr [rsp + 1872]
paddq xmm0, xmm1
pshufd xmm1, xmm0, 78
paddq xmm1, xmm0

This is what's called loop unrolling: LLVM pastes the loop body a bunch of time to avoid having to execute all those "loop management instructions", i.e. the increment of loop variable, check if loop ended and the jump.

(In case you're wondering: the paddq and similar instructions are SIMD instructions which allow summing up multiple values in parallel. More even, two SIMD register (xmm0 and xmm1) are used in parallel so that instruction-level parallelism of the CPU can basically executed to of these instructions at the same time. After all, they are independent of one another. In the end, both registers are combined.)

But loop unrolling is not responsible for a performance difference of factor 80! At least not loop unrolling alone. Let's take a look at the actual benchmarking code, which puts the one loop inside another one:

const CAPACITY: usize = 239;
const IN_LOOPS: usize = 500000;

pub fn foo() -> usize 
 let mut arr = [0; CAPACITY];
 for i in 0..CAPACITY 
 arr[i] = i;
 

 let mut sum = 0;
 for _ in 0..IN_LOOPS 
 let mut s = 0;
 for i in 0..arr.len() 
 s += arr[i];
 
 sum += s;
 

 sum

(On Godbolt Compiler Explorer)

The assembly for CAPACITY = 240 looks normal: two nested loops. (At the start of the function there is quite some code just for initializing, which we will ignore.) For 239, however, it looks very different! We see that the initializing loop and the inner loop got unrolled: so far so expected.

The important difference is that for 239 LLVM was able to figure out that the result of the inner loop does not depend on the outer loop! As a consequence, LLVM emits code that basically first executes only the inner loop (calculating the sum) and then simulates the outer loop by adding up sum a bunch of times!

First we see almost the same assembly as above (the assembly representing the inner loop). Afterwards we see this (I commented to explain the assembly, the comments with * are especially important):

 ; at the start of the function, `rbx` was set to 0

 movq rax, xmm1 ; result of SIMD summing up stored in `rax`
 add rax, 711 ; add up missing terms from loop unrolling
 mov ecx, 500000 ; * init loop variable outer loop
.LBB0_1:
 add rbx, rax ; * rbx += rax
 add rcx, -1 ; * decrement loop variable
 jne .LBB0_1 ; * if loop variable != 0 jump to LBB0_1
 mov rax, rbx ; move rbx (the sum) back to rax
 ; two unimportant instructions omitted
 ret ; the return value is stored in `rax`

As you can see here, the result of the inner loop is taken, added up as often as the outer loop would have ran and then returned. LLVM can only perform this optimization because it understood that the inner loop is independent of the outer one.

This means the runtime changes from CAPACITY * IN_LOOPS to CAPACITY + IN_LOOPS. And this is responsible for the huge performance difference.

An additional note: can you do anything about this? Not really LLVM has to have such magic thresholds as without them LLVM-optimizations could take forever to complete on certain code. But we can also agree that this code was highly artificial. In practice, I doubt that such a huge difference would occur. The difference due to full loop unrolling is usually not even factor 2 in these cases. So no need to worry about real use cases.

As a last note about idiomatic Rust code: arr.iter().sum() is a better way to sum up all elements of an array. And changing this in the second example does not lead to a any notable differences in emitted assembly. You should use short and idiomatic versions unless you measured that it hurts performance.