[X86] Don't elide argument copies for scalarized vectors (PR63475)

[nikic]

Summary

When eliding argument copies, the memory layout between a plain store of the type and the layout of the argument lowering on the stack must match. For multi-part argument lowerings, this is not necessarily the case.

The code already tried to prevent this optimization for "scalarized and extended" vectors, but the check for "extends" was incomplete. While a scalarized vector of i32s stores i32 values on the stack, these are stored in 8 byte stack slots (on x86_64), so effectively have padding.

Rather than trying to add more special cases to handle this (which is not straightforward), I'm going in the other direction and exclude scalarized vectors from this optimization entirely. This seems like a rare case that is not worth the hassle -- the complete lack of test coverage is not reassuring either.

Fixes https://github.com/llvm/llvm-project/issues/63475.

[nikic]

Ping

[pengfei]

Why do we change the distance from 4 to 8 in callee rather than change 8 to 4 for caller?
Although we don't have a standard calling converion for illegal types in IR, and what makes it more corner is its power-of-2 vector <8 x i32> is illegal too (https://godbolt.org/z/ds9WGPTeP), I think we still should make it close to practical converion, i.e., https://godbolt.org/z/YoMbzxzr8

[nikic]

In D154078#4476524, @pengfei wrote:

Why do we change the distance from 4 to 8 in callee rather than change 8 to 4 for caller?
Although we don't have a standard calling converion for illegal types in IR, and what makes it more corner is its power-of-2 vector <8 x i32> is illegal too (https://godbolt.org/z/ds9WGPTeP), I think we still should make it close to practical converion, i.e., https://godbolt.org/z/YoMbzxzr8

My understanding of how this works is that we legalize the vector into (legally typed) parts, and then those parts get assigned to registers until we run out of argument registers. At that point the remainder is passed on the stack. How that looks like is determined by the calling convention -- on x86_64 it means that the 32-bit registers get pushed into 64-bit stack slots.

I don't think we can make the 32-bit registers spill into 32-bit stack slots for the specific case that they come from an illegal vector -- and it wouldn't make a lot of sense, if you consider that only parts of the vector might get passed via stack, and the rest via registers, so you wouldn't get a complete vector on the stack anyway.

The other thing we can do is to chose a better argument passing legalization for the vector, and e.g. pass it via combination of widening and splitting as two <4 x i32> vectors instead. This would of course be the more efficient calling convention.

The problem here is that SDAG argument lowering currently doesn't support this. It supports widening a vector and it supports splitting a vector, but it doesn't support doing both. I took a look, and supporting this doesn't look particularly easy.

In any case, I still think that the fix here makes sense, in that if a vector gets scalarized, we should disable the argument copy elision optimization. The question whether the vector gets scalarized or not is ultimately a separate one...

[nikic]

Ping :)

[pengfei]

LGTM. I don't think the calling convention of illegal type matters much and we can always change it in the future.

llvm/test/CodeGen/X86/pr63475.ll

99–100

Is it out of bound?

[nikic]

llvm/test/CodeGen/X86/pr63475.ll

99–100

It is... fixed in https://github.com/llvm/llvm-project/commit/f78a06ef111d619d160a92d1fd7553337a1921a1.

llvm/lib/Target/X86/X86ISelLowering.cpp

if (Flags.isByVal()) {
     // can be improved with deeper analysis.
     int FI = MFI.CreateFixedObject(Bytes, VA.getLocMemOffset(), isImmutable,
                                    /*isAliased=*/true);
     return DAG.getFrameIndex(FI, PtrVT);
   }
 
   EVT ArgVT = Ins[i].ArgVT;
 
-  // If this is a vector that has been split into multiple parts, and the
-  // scalar size of the parts don't match the vector element size, then we can't
-  // elide the copy. The parts will have padding between them instead of being
-  // packed like a vector.
-  bool ScalarizedAndExtendedVector =
-      ArgVT.isVector() && !VA.getLocVT().isVector() &&
-      VA.getLocVT().getSizeInBits() != ArgVT.getScalarSizeInBits();
+  // If this is a vector that has been split into multiple parts, don't elide
+  // the copy. The layout on the stack may not match the packed in-memory
+  // layout.
+  bool ScalarizedVector = ArgVT.isVector() && !VA.getLocVT().isVector();
 
   // This is an argument in memory. We might be able to perform copy elision.
   // If the argument is passed directly in memory without any extension, then we
   // can perform copy elision. Large vector types, for example, may be passed
   // indirectly by pointer.
   if (Flags.isCopyElisionCandidate() &&
       VA.getLocInfo() != CCValAssign::Indirect && !ExtendedInMem &&
-      !ScalarizedAndExtendedVector) {
+      !ScalarizedVector) {
     SDValue PartAddr;
     if (Ins[i].PartOffset == 0) {
       // If this is a one-part value or the first part of a multi-part value,
       // create a stack object for the entire argument value type and return a
       // load from our portion of it. This assumes that if the first part of an
       // argument is in memory, the rest will also be in memory.
       int FI = MFI.CreateFixedObject(ArgVT.getStoreSize(), VA.getLocMemOffset(),
                                      /*IsImmutable=*/false);

llvm/test/CodeGen/X86/pr63475.ll

 ; CHECK-NEXT:    callq callee@PLT
 ; CHECK-NEXT:    addq $64, %rsp
 ; CHECK-NEXT:    popq %rax
 ; CHECK-NEXT:    retq
   call void @callee(ptr null, ptr null, ptr null, ptr null, ptr null, ptr null, <7 x i32> <i32 42, i32 42, i32 42, i32 42, i32 42, i32 42, i32 42>)
   ret void
 }
 
-; FIXME: This is a miscompile.
+; Make sure the stack offsets are correct. The distance between them should
+; be 8, not 4.
 define void @callee(ptr %p0, ptr %p1, ptr %p2, ptr %p3, ptr %p4, ptr %p5, <7 x i32> %arg) nounwind {
 ; CHECK-LABEL: callee:
 ; CHECK:       # %bb.0: # %start
 ; CHECK-NEXT:    pushq %rbp
 ; CHECK-NEXT:    pushq %r15
 ; CHECK-NEXT:    pushq %r14
 ; CHECK-NEXT:    pushq %r13
 ; CHECK-NEXT:    pushq %r12
 ; CHECK-NEXT:    pushq %rbx
-; CHECK-NEXT:    pushq %rax
-; CHECK-NEXT:    movl 112(%rsp), %ebx
-; CHECK-NEXT:    movl 104(%rsp), %ebp
-; CHECK-NEXT:    movl 96(%rsp), %r14d
-; CHECK-NEXT:    movl 76(%rsp), %r15d
-; CHECK-NEXT:    movl 72(%rsp), %r12d
-; CHECK-NEXT:    movl 64(%rsp), %edi
-; CHECK-NEXT:    movl 68(%rsp), %r13d
-; CHECK-NEXT:    callq use@PLT
-; CHECK-NEXT:    movl %r13d, %edi
-; CHECK-NEXT:    callq use@PLT
-; CHECK-NEXT:    movl %r12d, %edi
+; CHECK-NEXT:    subq $40, %rsp
+; CHECK-NEXT:    movl 120(%rsp), %ebx
+; CHECK-NEXT:    movd %ebx, %xmm0
+; CHECK-NEXT:    movl 112(%rsp), %ebp
+; CHECK-NEXT:    movd %ebp, %xmm1
+; CHECK-NEXT:    punpckldq {{.*#+}} xmm1 = xmm1[0],xmm0[0],xmm1[1],xmm0[1]
+; CHECK-NEXT:    movl 104(%rsp), %r15d
+; CHECK-NEXT:    movd %r15d, %xmm0
+; CHECK-NEXT:    movl 96(%rsp), %edi
+; CHECK-NEXT:    movd %edi, %xmm2
+; CHECK-NEXT:    punpckldq {{.*#+}} xmm2 = xmm2[0],xmm0[0],xmm2[1],xmm0[1]
+; CHECK-NEXT:    punpcklqdq {{.*#+}} xmm2 = xmm2[0],xmm1[0]
+; CHECK-NEXT:    movl 136(%rsp), %r14d
+; CHECK-NEXT:    movd %r14d, %xmm0
+; CHECK-NEXT:    movl 128(%rsp), %r12d
+; CHECK-NEXT:    movd %r12d, %xmm1
+; CHECK-NEXT:    punpckldq {{.*#+}} xmm1 = xmm1[0],xmm0[0],xmm1[1],xmm0[1]
+; CHECK-NEXT:    movl 144(%rsp), %r13d
+; CHECK-NEXT:    movl %r13d, 36(%rsp)
+; CHECK-NEXT:    movq %xmm1, 28(%rsp)
+; CHECK-NEXT:    movdqu %xmm2, 12(%rsp)
 ; CHECK-NEXT:    callq use@PLT
 ; CHECK-NEXT:    movl %r15d, %edi
 ; CHECK-NEXT:    callq use@PLT
-; CHECK-NEXT:    movl %r14d, %edi
-; CHECK-NEXT:    callq use@PLT
 ; CHECK-NEXT:    movl %ebp, %edi
 ; CHECK-NEXT:    callq use@PLT
 ; CHECK-NEXT:    movl %ebx, %edi
 ; CHECK-NEXT:    callq use@PLT
-; CHECK-NEXT:    addq $8, %rsp
+; CHECK-NEXT:    movl %r12d, %edi
+; CHECK-NEXT:    callq use@PLT
+; CHECK-NEXT:    movl %r14d, %edi
+; CHECK-NEXT:    callq use@PLT
+; CHECK-NEXT:    movl %r13d, %edi
+; CHECK-NEXT:    callq use@PLT
+; CHECK-NEXT:    addq $40, %rsp
 ; CHECK-NEXT:    popq %rbx
 ; CHECK-NEXT:    popq %r12
 ; CHECK-NEXT:    popq %r13
 ; CHECK-NEXT:    popq %r14
 ; CHECK-NEXT:    popq %r15
 ; CHECK-NEXT:    popq %rbp
 ; CHECK-NEXT:    retq
 start:
   %alloca = alloca [7 x i32], align 4
   store <7 x i32> %arg, ptr %alloca, align 4
   %extract0 = extractelement <7 x i32> %arg, i64 0
   call void @use(i32 %extract0)
   %extract1 = extractelement <7 x i32> %arg, i64 1
   call void @use(i32 %extract1)
   %extract2 = extractelement <7 x i32> %arg, i64 2
   call void @use(i32 %extract2)
   %extract3 = extractelement <7 x i32> %arg, i64 3
   call void @use(i32 %extract3)
   %extract4 = extractelement <7 x i32> %arg, i64 4
   call void @use(i32 %extract4)
   %extract5 = extractelement <7 x i32> %arg, i64 5
   call void @use(i32 %extract5)
   %extract6 = extractelement <7 x i32> %arg, i64 6
   call void @use(i32 %extract6)
   ret void
 }
 
 declare void @use(i32)