Strange ASM pessimizations as a result of algorithmic optimization

[HadrienG2]

So, I was trying to sum a slice of f32s quickly on stable Rust.

But like pretty much all floating-point reductions, the naive algorithm (floats.iter().sum::<f32>()) does not autovectorize because its "natural" summation order introduces a serial dependency between successive sums. Which makes SIMD parallelization illegal in the eye of a compiler that guarantees bitwise floating point reproducibility like rustc does. Fair enough.

I saw this as a good motivation to move to explicit SIMD programming, but did not want to lose hardware portability (or, more precisely, wanted to keep it easy), so I tried to see how close I could get to stdsimd on stable Rust with only autovectorization and a pinch of hardware-specific vector size tuning.

---

Some hours of trial and error later, I got into a reasonably satisfactory state. In particular, the core of the algorithm...

    // Perform concurrent SIMD accumulation
    let mut accumulators = [SimdF32::default(); CONCURRENCY];
    for chunk in chunks {
        for (accumulator, &vector) in accumulators.iter_mut().zip(chunk) {
            *accumulator += vector;
        }
    }

    // Merge the SIMD accumulators into one
    let mut stride = CONCURRENCY / 2;
    while stride > 0 {
        for i in 0..stride {
            accumulators[i] += accumulators[i + stride];
        }
        stride /= 2;
    }
    let mut accumulator = accumulators[0];

...translated pretty much into the assembly that I would have written by hand, which made me very happy...

.LBB0_17:
	addps	(%rdi), %xmm1
	addps	16(%rdi), %xmm3
	addps	32(%rdi), %xmm2
	addps	48(%rdi), %xmm4
	addq	$64, %rdi
	addq	$4, %rcx
	jne	.LBB0_17
	jmp	.LBB0_18

...

.LBB0_18:
	addps	%xmm4, %xmm3
	addps	%xmm2, %xmm1
	addps	%xmm3, %xmm1

---

Nitpicky as I am, however, I was still a little bit unhappy about the part afterwards, which introduced a chain of serial dependencies that could become a bit long if I were to use a lot of accumulators...

    for &vector in remainder {
        accumulator += vector;
    }

.LBB0_31:
	addps	(%rcx), %xmm1
	addq	$16, %rcx
	incq	%rax
	jne	.LBB0_31

...because I knew that in this particular case, there should be an easy way to avoid that, which is to interleave the SIMD accumulator merging with the summation of remaining data.

    // Reinterprete input as aligned SIMD vectors + some extra floats.
    let (peel, mut vectors, tail) = unsafe { input.align_to::<SimdF32>() };
    
    // Perform concurrent SIMD accumulation, starting at maximal concurrency and
    // decreasing once the number of input vectors gets too small
    let mut accumulators = [SimdF32::default(); MAX_CONCURRENCY];
    let mut concurrency = MAX_CONCURRENCY;
    while concurrency > 0 {
        // Set up accumulators and chunked SIMD data according to current concurrency
        let accumulators = &mut accumulators[..concurrency];
        let chunks = vectors.chunks_exact(concurrency);
        let remainder = chunks.remainder();

        // Perform concurrent SIMD accumulation
        for chunk in chunks {
            for (accumulator, &vector) in accumulators.iter_mut().zip(chunk) {
                *accumulator += vector;
            }
        }
        
        // Halve SIMD concurrency to accumulate remaining data
        concurrency /= 2;
        for i in 0..concurrency {
            accumulators[i] += accumulators[i+concurrency];
        }
        vectors = remainder;
    }
    let accumulator = accumulators[0];

However, much to my surprise, performing this algorithmic optimization leads rustc to heavily pessimize the inner loop code by spilling all but one accumulator on every iteration:

.LBB0_16:
	addps	(%rax), %xmm1
	movaps	16(%rsp), %xmm2
	movaps	32(%rsp), %xmm3
	movaps	48(%rsp), %xmm4
	addps	16(%rax), %xmm2
	movaps	%xmm2, 16(%rsp)
	addps	32(%rax), %xmm3
	addps	48(%rax), %xmm4
	movaps	%xmm3, 32(%rsp)
	movaps	%xmm4, 48(%rsp)
	addq	$64, %rax
	addq	$4, %rdi
	jne	.LBB0_16

Why would that happen? The only explanation I have is that rustc is somehow unable to prove that the accumulators slice does not alias with the vectors/remainder slices, and thus spills to memory just in case accumulator changes would affect the input of the next computations.

But this sounds like a bug to me: given that I have an &mut to the accumulators, my understanding is that rustc should be able to infer that no other code can see the accumulators, and thus they can remain resident in registers for the entire duration of the accumulation loop.

Can someone with more knowledge of how rustc and LLVM do their optimization magic cross-check this and tell if my understanding is correct or if the register spills are truly necessary to preserve the semantics of my code?

---

Also, this is on stable release 1.57.0. On beta and nightly, the generated code becomes even stranger:

.LBB0_16:
	movq	(%rdx), %rcx
	movq	8(%rdx), %rbx
	movd	%ebx, %xmm5
	shrq	$32, %rbx
	movd	%ecx, %xmm6
	shrq	$32, %rcx
	movd	%ecx, %xmm7
	punpckldq	%xmm6, %xmm7
	movd	%ebx, %xmm6
	punpckldq	%xmm5, %xmm6
	punpcklqdq	%xmm7, %xmm6
	addps	%xmm6, %xmm2
	addps	16(%rdx), %xmm1
	addps	32(%rdx), %xmm4
	addps	48(%rdx), %xmm3
	addq	$64, %rdx
	addq	$4, %rax
	jne	.LBB0_16

Here, rustc generates the code I would expect for the last three accumulators, but then it goes crazy with the first accumulator and generates the least efficient SSE load I have ever seen.

So it seems the aliasing issue got resolved, but was replaced by another issue beyond my comprehension... Here again, compiler expert help would be useful.

[hkratz]

@rustbot label +T-compiler +I-slow +A-codegen +A-llvm

[AngelicosPhosphoros]

Your splitting by CONCURENCY threads is suboptimal.
You should split your vectors by chunks with size len() / CONCURENCY + 1 instead of chunks with size CONCURENCY. Your current version interacts badly with caches.

[HadrienG2]

@AngelicosPhosphoros : Nice catch, this might be the reason why I did not manage so far to get this code above ~50% of the CPU's peak load throughput (with inputs in L1 cache). I'll try this suggestion out, among other ideas.

[HadrienG2]

@AngelicosPhosphoros On a Zen 2 CPU at least, increasing load stride makes no difference to throughput. What did make a difference, however, was remembering to add codegen-units = 1 to my bench cargo profile so that rustc stops generating bad code in unpredictable ways. Given extra tuning, this gave me a peak throughput at 85% of theoretical hardware maximum in SSE mode and 73% of hardware maximum in AVX mode, which is good enough for me.

[AngelicosPhosphoros]

Yeah, I really love codegen-units=1 and lto="fat" for production builds.