Function operating on arrays performs a lot worse for arrays of length 2

[Canleskis]

Hello, after some discussion on Stack Overflow, I was suggested to file an issue here.

The following function performs unexpectedly poorly with N = 2.

fn fold<const N: usize>(v: Vec<Array<N>>) -> Vec<Array<N>> {
    let result = v.iter().map(|a1| {
        v.iter().fold(Array::default(), |acc, a2| {
            let d = *a2 - *a1;

            acc + d
        })
    });

    result.collect()
}

Array is a simple wrapper implementing Add and Sub around an array of f32s. The results are the same without that wrapper.

Here is a graph generated with criterion of the benchmark of that function:

Someone on the Stack Overflow discussion suggested that this was because:
"The shl, shr and or operations on rdx and rsi suggest that for N = 2 the two floats are stored in one 64 bit general purpose register, whereas in the other cases the value of a1 is persisted in N separate xmm registers".
Another user suggested that this wrong optimization was caused by opt-level=3, so on the LLVM side.

Here is the assembly code where a user marked important areas.

Here is a repository of the full code for the benchmark.

This was observed both in stable and nightly.

[Rageking8]

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

[nikic]

This probably happens because the Rust ABI passes [f32; 2] as i64 to functions, resulting in lots of packing an unpacking that messes up optimization.

[nikic]

So yeah, we're doing something like this:

  %.val.i.i.i.i.i.i.i = load i64, ptr %self.sroa.0.010.i.i.i.i.i.i.i, align 4, !dbg !242, !noalias !243
  %rhs.sroa.0.0.extract.trunc.i.i.i.i.i.i.i.i.i.i.i.i = trunc i64 %.val.i.i.i.i.i.i.i to i32
  %10 = bitcast i32 %rhs.sroa.0.0.extract.trunc.i.i.i.i.i.i.i.i.i.i.i.i to float
  %rhs.sroa.2.0.extract.shift.i.i.i.i.i.i.i.i.i.i.i.i = lshr i64 %.val.i.i.i.i.i.i.i, 32
  %rhs.sroa.2.0.extract.trunc.i.i.i.i.i.i.i.i.i.i.i.i = trunc i64 %rhs.sroa.2.0.extract.shift.i.i.i.i.i.i.i.i.i.i.i.i to i32
  %11 = bitcast i32 %rhs.sroa.2.0.extract.trunc.i.i.i.i.i.i.i.i.i.i.i.i to float

Where we extract two floats from an i64 using shifts, truncates and bitcasts.

LLVM should probably try to split those loads, but it would certainly help if rustc used a less dumb ABI in the first place.

[alex]

The code generated for 2 on 1.83 looks better: https://rust.godbolt.org/z/vcxrhq348 but I haven't actually measured.