Arc::new duplicates stack memory

[marc0246]

Running the following through Godbolt, using rustc 1.71.0-nightly (cba1407 2023-05-10) with -C opt-level=3:

use std::sync::Arc;

pub fn test1(x: u64) -> Arc<[u64]> {
    let array = [x; 1000];
    Arc::new(array)
}

I expected rustc to create an 8 KB stack allocation and a memcpy from it to the heap allocation, instead it creates 2 identical stack allocations and 2 memcpys to go along with them:

example::test1:
        push    rbx
        sub     rsp, 4096
        mov     qword ptr [rsp], 0
        sub     rsp, 4096
        mov     qword ptr [rsp], 0
        sub     rsp, 4096
        mov     qword ptr [rsp], 0
        sub     rsp, 3712
        movq    xmm0, rdi
        pshufd  xmm0, xmm0, 68
        mov     eax, 18
.LBB0_1:
        movdqu  xmmword ptr [rsp + 8*rax - 144], xmm0
        movdqu  xmmword ptr [rsp + 8*rax - 128], xmm0
        movdqu  xmmword ptr [rsp + 8*rax - 112], xmm0
        movdqu  xmmword ptr [rsp + 8*rax - 96], xmm0
        movdqu  xmmword ptr [rsp + 8*rax - 80], xmm0
        movdqu  xmmword ptr [rsp + 8*rax - 64], xmm0
        movdqu  xmmword ptr [rsp + 8*rax - 48], xmm0
        movdqu  xmmword ptr [rsp + 8*rax - 32], xmm0
        movdqu  xmmword ptr [rsp + 8*rax - 16], xmm0
        movdqu  xmmword ptr [rsp + 8*rax], xmm0
        add     rax, 20
        cmp     rax, 1018
        jne     .LBB0_1
        lea     rdi, [rsp + 8000]
        mov     rsi, rsp
        mov     edx, 8000
        call    qword ptr [rip + memcpy@GOTPCREL]
        mov     edi, 8016
        mov     esi, 8
        call    qword ptr [rip + __rust_alloc@GOTPCREL]
        test    rax, rax
        je      .LBB0_3
        mov     rbx, rax
        mov     qword ptr [rax], 1
        mov     qword ptr [rax + 8], 1
        mov     rdi, rax
        add     rdi, 16
        lea     rsi, [rsp + 8000]
        mov     edx, 8000
        call    qword ptr [rip + memcpy@GOTPCREL]
        mov     edx, 1000
        mov     rax, rbx
        add     rsp, 16000
        pop     rbx
        ret
.LBB0_3:
        mov     edi, 8
        mov     esi, 8016
        call    qword ptr [rip + alloc::alloc::handle_alloc_error@GOTPCREL]
        ud2

Trying to narrow the problem down further, when eliminating the stack allocations entirely:

#![feature(get_mut_unchecked, new_uninit)]

use std::sync::Arc;

pub fn test2(x: u64) -> Arc<[u64]> {
    let mut arc = Arc::new_uninit_slice(1000);
    for elem in unsafe { Arc::get_mut_unchecked(&mut arc) } {
        elem.write(x);
    }
    unsafe { arc.assume_init() }
}

The code duplication still persists, and is limited to the arcinner_layout_for_value_layout call:

example::test2:
        push    r15
        push    r14
        push    r12
        push    rbx
        push    rax
        mov     rbx, rdi
        mov     r12, qword ptr [rip + alloc::sync::arcinner_layout_for_value_layout@GOTPCREL]
        mov     edi, 8
        mov     esi, 8000
        call    r12
        mov     r14, rax
        mov     r15, rdx
        mov     edi, 8
        mov     esi, 8000
        call    r12
        test    rdx, rdx
        je      .LBB0_2
        mov     rdi, rdx
        mov     rsi, rax
        call    qword ptr [rip + __rust_alloc@GOTPCREL]
.LBB0_2:
        test    rax, rax
        je      .LBB0_6
        mov     qword ptr [rax], 1
        mov     qword ptr [rax + 8], 1
        movq    xmm0, rbx
        pshufd  xmm0, xmm0, 68
        mov     ecx, 20
.LBB0_4:
        movdqu  xmmword ptr [rax + 8*rcx - 144], xmm0
        movdqu  xmmword ptr [rax + 8*rcx - 128], xmm0
        movdqu  xmmword ptr [rax + 8*rcx - 112], xmm0
        movdqu  xmmword ptr [rax + 8*rcx - 96], xmm0
        movdqu  xmmword ptr [rax + 8*rcx - 80], xmm0
        movdqu  xmmword ptr [rax + 8*rcx - 64], xmm0
        movdqu  xmmword ptr [rax + 8*rcx - 48], xmm0
        movdqu  xmmword ptr [rax + 8*rcx - 32], xmm0
        movdqu  xmmword ptr [rax + 8*rcx - 16], xmm0
        movdqu  xmmword ptr [rax + 8*rcx], xmm0
        add     rcx, 20
        cmp     rcx, 1020
        jne     .LBB0_4
        mov     edx, 1000
        add     rsp, 8
        pop     rbx
        pop     r12
        pop     r14
        pop     r15
        ret
.LBB0_6:
        mov     rdi, r14
        mov     rsi, r15
        call    qword ptr [rip + alloc::alloc::handle_alloc_error@GOTPCREL]
        ud2

[marc0246]

@rustbot label +A-codegen +I-slow

[nikic]

Godbolt: https://rust.godbolt.org/z/GKWY8nT78

[marc0246]

It seems this is not unique to Arc::new, happens using Rc::new as well.

It may be useful to note that without the let binding, rustc emits a single alloca: https://godbolt.org/z/9j7cevvjs

[caojoshua]

This applies to a lot of things that takes an array as input. For example, Box and Vec https://rust.godbolt.org/z/rzoPneT3q

I have a potential fix in https://reviews.llvm.org/D151638. TLDR, marking a tail call makes Alias Analysis smarter, a memcpy can be optimized, and a dead alloca can be removed. I'll elaborate below.

1. rustc copies by-ref arguments, explaining why there are two allocas. rustc generated code basically memcpy's from alloca1 to alloca2, and then another memcpy from alloca2 to the Arc/Vec/Box.
2. TailCallElim fails to mark rust_alloc as a tail call. The tail mark means that the callee will not access the callers local memory. The call cannot be marked because the function iterates over GEPs, and the loop latch condition uses the GEP in a icmp, escaping the alloca. The patch I submitted makes sure icmps are not escaping.
3. MemCpyOptimizer tries to optimize memcpys, but can only do so if the source alloca is not written to since the first memcpy. BasicAA can determine a callee cannot read/write an alloca if the call is a tail call. rust_alloc does not have the tail mark from the previous step, and BasicAA deduces that rust_alloc may read/write caller allocas.

Interesting is that if you take the optimized code that @nikic posted and run it through -passes=memcpyopt,instcombine, the extra alloca is removed https://rust.godbolt.org/z/WcK48WG9c. In the -O3 pipeline, LoopVectorize changes the loop from iterating over GEPS to indicies, removing the icmp against a GEP. TailCallElim is called after, marking rust_alloc as tail, but MemCpyOptimizer is not run again.

[caojoshua]

The duplicate alloca could also have been removed if rust_alloc was set to only access inaccessible memory. It is set for libc malloc here. If its correct, maybe rustc can set it for rust_alloc (Big emphasis on if its correct. I don't know about the internals).

[caojoshua]

I also want to note that this is the rustc code that emits code for the arrays. It intentionally iterates from the start to end pointer. It would probably not be hard to change the code to iterate over array indices instead, which would help LLVM optimizations, but I'm not sure if that is a desirable change.

[the8472]

The call cannot be marked because the function iterates over GEPs, and the loop latch condition uses the GEP in a icmp, escaping the alloca. The patch I submitted makes sure icmps are not escaping.

Huh, I was under the impression that iterating over GEPs (ptr = ptr.add(1)) and breaking on pointer equality is the preferred way to do these kinds of loops. We're doing that in several places.

[caojoshua]

@the8472 I am not 100% which is better, maybe there are drawbacks to iterating over indices that I am not aware of.

Could you share some examples of where else there might be GEP iterating? I have a change in this LLVM PR which should make LLVM smarter about comparisons against GEP inductions, but I'm questioning whether we should merge it. Examples could help next steps.

[the8472]

It's used for slices and vecs.

[the8472]

This lead me to experiment with various loop shapes for slice::Iter::fold in #106343 and an index-based do-while loop + a zero-len precheck did yield the largest wins out of the approaches I tried.