New optimization: Move non-mutable array of Copy type to .rodata

[tesuji]

Continue the discussion in #73780 (comment).

nbdd0121 said that:

let bar: [u32; 10] = [0,10,20,30,40,50,60,70,80,90];

Semantically just says: put this on the stack. Rust is doing exactly what you ask it to do;
it cannot place foo inside rodata because it is a value.

---

Consider this code: https://rust.godbolt.org/z/6TBstc

pub fn foo(x: usize) -> i32 {
    let base: [i32; 64] = [
        67, 754, 860, 559, 368, 870, 548, 972,
        141, 731, 351, 664, 32, 4, 996, 741,
        203, 292, 237, 480, 151, 940, 777, 540,
        143, 587, 747, 65, 152, 517, 882, 880,
        712, 595, 370, 901, 237, 53, 789, 785,
        912, 650, 896, 367, 316, 392, 62, 473,
        675, 691, 281, 192, 445, 970, 225, 425,
        628, 324, 322, 206, 912, 867, 462, 92
    ];
    base[x % 64]
}

The Rust compiler puts base on the stack in both opt-level 2 and 3:

The same code in C (I know that C array is mostly pointer): https://godbolt.org/z/WYbKs97E6

#include <stdlib.h>
#include <stdint.h>

uint32_t square(size_t x) {
    uint32_t bases[64] = {
        67, 754, 860, 559, 368, 870, 548, 972,
        141, 731, 351, 664, 32, 4, 996, 741,
        203, 292, 237, 480, 151, 940, 777, 540,
        143, 587, 747, 65, 152, 517, 882, 880,
        712, 595, 370, 901, 237, 53, 789, 785,
        912, 650, 896, 367, 316, 392, 62, 473,
        675, 691, 281, 192, 445, 970, 225, 425,
        628, 324, 322, 206, 912, 867, 462, 92
    };
    return bases[x % 64];
}

clang (with -O1) optimizes out the bases array to .rodata and use that static array directly.
gcc also fails to optimized. In -Oz gcc optimizes it to a memcpy.

---

Could the same optimization be done in Rust?

[nbdd0121]

The optimisation is certainly possible in many cases but it is not trivial. E.g. A non-mutable array can still be moved. In C you cannot move an array.

[Mark-Simulacrum]

Interestingly, it looks like the array is put in rodata in Rust too, if the array type is change to i32.

[tesuji]

Yeah, if the array type is changed to i32/u32, the array is put in .rodata. But it still pulls that array and pushes
it into stack:

.LCPI0_0:
        .long   1
        .long   3
        .long   4
        .long   5
example::foo:
        sub     rsp, 24
        movaps  xmm0, xmmword ptr [rip + .LCPI0_0]
        movaps  xmmword ptr [rsp], xmm0
        mov     dword ptr [rsp + 16], 2
        movabs  rcx, -3689348814741910323
        mov     rax, rdi
        mul     rcx
        shr     rdx, 2
        lea     rax, [rdx + 4*rdx]
        sub     rdi, rax
        mov     eax, dword ptr [rsp + 4*rdi]
        add     rsp, 24
        ret

[nbdd0121]

@lzutao That's not putting the array on rodata. You can see it contains just 4 elements; it is indeed just LLVM groups 4 u32 stores into one u128 store and that needs the constant on memory. It is very clear from LLVM IR:

Here's the LLVM IR for the debug version:

define i32 @_ZN10playground3foo17ha3dfdbe3e3d033d9E(i64 %x) unnamed_addr #0 !dbg !6 {
start:
  %x.dbg.spill = alloca i64, align 8
  %base = alloca [5 x i32], align 4
  store i64 %x, i64* %x.dbg.spill, align 8
  call void @llvm.dbg.declare(metadata i64* %x.dbg.spill, metadata !14, metadata !DIExpression()), !dbg !20
  call void @llvm.dbg.declare(metadata [5 x i32]* %base, metadata !15, metadata !DIExpression()), !dbg !21
  %0 = bitcast [5 x i32]* %base to i32*, !dbg !22
  store i32 1, i32* %0, align 4, !dbg !22
  %1 = getelementptr inbounds [5 x i32], [5 x i32]* %base, i32 0, i32 1, !dbg !22
  store i32 3, i32* %1, align 4, !dbg !22
  %2 = getelementptr inbounds [5 x i32], [5 x i32]* %base, i32 0, i32 2, !dbg !22
  store i32 4, i32* %2, align 4, !dbg !22
  %3 = getelementptr inbounds [5 x i32], [5 x i32]* %base, i32 0, i32 3, !dbg !22
  store i32 5, i32* %3, align 4, !dbg !22
  %4 = getelementptr inbounds [5 x i32], [5 x i32]* %base, i32 0, i32 4, !dbg !22
  store i32 2, i32* %4, align 4, !dbg !22
  %_3 = urem i64 %x, 5, !dbg !23
  %_6 = icmp ult i64 %_3, 5, !dbg !24
  %5 = call i1 @llvm.expect.i1(i1 %_6, i1 true), !dbg !24
  br i1 %5, label %bb1, label %panic, !dbg !24

bb1:                                              ; preds = %start
  %6 = getelementptr inbounds [5 x i32], [5 x i32]* %base, i64 0, i64 %_3, !dbg !24
  %7 = load i32, i32* %6, align 4, !dbg !24
  ret i32 %7, !dbg !25

panic:                                            ; preds = %start
; call core::panicking::panic_bounds_check
  call void @_ZN4core9panicking18panic_bounds_check17h4b3d0dcda831e378E(i64 %_3, i64 5, %"core::panic::Location"* noalias readonly align 8 dereferenceable(24) bitcast (<{ i8*, [16 x i8] }>* @alloc3 to %"core::panic::Location"*)), !dbg !24
  unreachable, !dbg !24
}

for the release version

define i32 @_ZN10playground3foo17hf1f5714ed87b6704E(i64 %x) unnamed_addr #0 {
start:
  %base = alloca [5 x i32], align 16
  %0 = bitcast [5 x i32]* %base to i8*
  call void @llvm.lifetime.start.p0i8(i64 20, i8* nonnull %0)
  %1 = bitcast [5 x i32]* %base to <4 x i32>*
  store <4 x i32> <i32 1, i32 3, i32 4, i32 5>, <4 x i32>* %1, align 16
  %2 = getelementptr inbounds [5 x i32], [5 x i32]* %base, i64 0, i64 4
  store i32 2, i32* %2, align 16
  %_3 = urem i64 %x, 5
  %3 = getelementptr inbounds [5 x i32], [5 x i32]* %base, i64 0, i64 %_3
  %4 = load i32, i32* %3, align 4
  call void @llvm.lifetime.end.p0i8(i64 20, i8* nonnull %0)
  ret i32 %4
}

So clearly the IR we generate initialises array elements one by one. On the other hand, Clang generates this in O0:

@__const.foo.bases = private unnamed_addr constant [5 x i32] [i32 1, i32 3, i32 4, i32 5, i32 2], align 16

define dso_local i32 @foo(i64 %0) #0 !dbg !7 {
  %2 = alloca i64, align 8
  %3 = alloca [5 x i32], align 16
  store i64 %0, i64* %2, align 8
  call void @llvm.dbg.declare(metadata i64* %2, metadata !19, metadata !DIExpression()), !dbg !20
  call void @llvm.dbg.declare(metadata [5 x i32]* %3, metadata !21, metadata !DIExpression()), !dbg !25
  %4 = bitcast [5 x i32]* %3 to i8*, !dbg !25
  call void @llvm.memcpy.p0i8.p0i8.i64(i8* align 16 %4, i8* align 16 bitcast ([5 x i32]* @__const.foo.bases to i8*), i64 20, i1 false), !dbg !25
  %5 = load i64, i64* %2, align 8, !dbg !26
  %6 = urem i64 %5, 5, !dbg !27
  %7 = getelementptr inbounds [5 x i32], [5 x i32]* %3, i64 0, i64 %6, !dbg !28
  %8 = load i32, i32* %7, align 4, !dbg !28
  ret i32 %8, !dbg !29
}

So it seems that clang has take care of the const promotion of arrays. LLVM is smart enough to observe that memcpy is not necessary and optimise that away, but it is not yet smart enough to do a const promotion.

[nbdd0121]

Idea: we can just add an optimisation pass to turn any expr: [T;N] into *&expr when T:Copy

Example:

pub fn foo(x: usize) -> u32 {
    let base: [u32; 5] = [1,3,4,5,2];
    base[x % 5]
}

turns into:

pub fn foo(x: usize) -> u32 {
    let base: [u32; 5] = *&[1,3,4,5,2];
    base[x % 5]
}

This will enable the optimisation!

[nbdd0121]

Probably this is an optimisation that we can make for all large const values. I can totally see large structs being benefits by this as memcpy should be faster than initialising fields one by one. I am not sure whether this is something Rust or LLVM should do. It might be easier to implement on Rust side though because we have more information.

One caveat would be const arrays like [42; 1000], which is cheaper to initialise than copying from rodata.

[bdonlan]

Counterpoint: Do rust semantics guarantee that "stack-allocated" arrays have different pointer locations? That is, is this code guaranteed to not panic?

fn foo() {
  let x: [u32; 5] = [1,2,3,4,5];
  let y: [u32; 5] = [1,2,3,4,5];

  assert_ne!(&x as *const [u32; 5], &y as *const [u32; 5])
}

If so, the proposed optimization pass must take this into account and ensure that, if the pointer location is observed, then the array is actually reified into a copy on the stack.

[nbdd0121]

@bdonlan My proposed optimisation will turn that into

fn foo() {
  let x: [u32; 5] = *&[1,2,3,4,5];
  let y: [u32; 5] = *&[1,2,3,4,5];

  assert_ne!(&x as *const [u32; 5], &y as *const [u32; 5])
}

and that shouldn't change the semantics.

[tesuji]

There are missed-optimizations like #75978. One has to borrow the static array to make optimization happens.
But there are plans to demote static slice,which may cause optimization losses.
Not sure what to do, but I will cc @RalfJung about MIR.

@rustbot modify labels: A-mir-opt A-LLVM