[SelectionDAG] Fix argument copy elision with irregular types

[LemonBoy]

Summary

D29668 enabled to avoid a useless copy of the argument value into an alloca if the caller places it in memory (as it often happens on x86) by directly forwarding the pointer to it. This optimization is illegal if the type contains padding bytes: if a truncating store into the alloca is replaced the upper bits are filled with garbage and produce code misbehaving at runtime.

[LemonBoy]

Ping?

[rnk]

I don't see why this optimization should be declared illegal. Why isn't it the caller's responsibility to store these narrow integers canonically? C doesn't have narrow integers, but GCC doesn't do this. If a variable is passed in memory, it just uses the memory that was passed into it.

llvm/test/CodeGen/X86/arg-copy-elide.ll

80

We don't want this, do we? Seems like it would be a regression.

[LemonBoy]

In D102153#2764464, @rnk wrote:

I don't see why this optimization should be declared illegal. Why isn't it the caller's responsibility to store these narrow integers canonically? C doesn't have narrow integers, but GCC doesn't do this. If a variable is passed in memory, it just uses the memory that was passed into it.

AFAICS it's the callee's responsibility to chop the extra bits off.
In this (compile at -O0) example you can see that @expect is correctly truncating the i3 value passed in-register, but fails to do so for the third argument that's passed in-memory.

[rnk]

Thinking about this more, I think this is a good change. Currently we use the zext and sext attributes to document if the high bits of an argument are initialized, and if so, to what. If those attributes aren't present, we shouldn't do copy elision. I have a minor concern that needs to be addressed, though.

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

9895

Don't we still need this check? I don't see any other code that checks that the allocated type and argument type match. Otherwise we could try to copy elide an i32 stored into an i64 alloca, for example. Please add a test for that, sorry for not already having one.

llvm/test/CodeGen/X86/arg-copy-elide.ll

80

This code turns out to not be representative of what clang generates, so I don't think we need to worry about it. See this bool example here:
https://gcc.godbolt.org/z/efPavanfb

void escape(bool *);
void bool_arg(bool b) { escape(&b); }

We fail to copy elide because the argument is i1, but the alloca is i8. I think this will be true generally for non-byte-sized integers, so there is really no harm to this change to copy elision.

[LemonBoy]

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

9895

In tryToElideArgumentCopy there's a check between the object size pointed by the old and new frame indices, in case of a partially initialized i64 being forwarded as a i32 the check will fail.
No strong feeling about this, I can put the check back if you prefer.

[LemonBoy]

Add a test w/ partially-initialized alloca.

[rnk]

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

9895

I think we should keep the store size vs. alloca size check here. If I were debugging this code later, I would be confused as to how an alloca like this became a copy elision candidate that gets rejected later in tryToElideArgumentCopy.

[LemonBoy]

Reinstate the alloca initialization check.

[rnk]

lgtm, thanks.

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

if (!Info)
       continue;
     const AllocaInst *AI = cast<AllocaInst>(Dst);
 
     // Skip allocas that have been initialized or clobbered.
     if (*Info != StaticAllocaInfo::Unknown)
       continue;
 
     // Check if the stored value is an argument, and that this store fully
-    // initializes the alloca. Don't elide copies from the same argument twice.
+    // initializes the alloca.
+    // If the argument type has padding bits we can't directly forward a pointer
+    // as the upper bits may contain garbage.
+    // Don't elide copies from the same argument twice.
     const Value *Val = SI->getValueOperand()->stripPointerCasts();
     const auto *Arg = dyn_cast<Argument>(Val);
     if (!Arg || Arg->hasPassPointeeByValueCopyAttr() ||
         Arg->getType()->isEmptyTy() ||
         DL.getTypeStoreSize(Arg->getType()) !=
             DL.getTypeAllocSize(AI->getAllocatedType()) ||
+        !DL.typeSizeEqualsStoreSize(Arg->getType()) ||
         ArgCopyElisionCandidates.count(Arg)) {
       *Info = StaticAllocaInfo::Clobbered;
       continue;
     }
 
     LLVM_DEBUG(dbgs() << "Found argument copy elision candidate: " << *AI
                       << '\n');
 

llvm/test/CodeGen/X86/arg-copy-elide.ll

entry:
   call void @addrof_i64(i64* %x.addr)
   ret i64 %x
 }
 
 define i1 @i1_arg(i1 %x) {
 ; CHECK-LABEL: i1_arg:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    pushl %ebx
+; CHECK-NEXT:    pushl %eax
 ; CHECK-NEXT:    movb {{[0-9]+}}(%esp), %bl
+; CHECK-NEXT:    movl %ebx, %eax
+; CHECK-NEXT:    andb $1, %al
+; CHECK-NEXT:    movb %al, {{[0-9]+}}(%esp)
 ; CHECK-NEXT:    leal {{[0-9]+}}(%esp), %eax
 ; CHECK-NEXT:    pushl %eax
 ; CHECK-NEXT:    calll _addrof_i1
 ; CHECK-NEXT:    addl $4, %esp
 ; CHECK-NEXT:    movl %ebx, %eax
+; CHECK-NEXT:    addl $4, %esp
 ; CHECK-NEXT:    popl %ebx
 ; CHECK-NEXT:    retl
   %x.addr = alloca i1
   store i1 %x, i1* %x.addr
   call void @addrof_i1(i1* %x.addr)
   ret i1 %x
 }
 

 ; CHECK-NEXT:    retl
   %v.p = alloca i32
   store i32 %v, i32* %v.p
   call void @addrof_i32(i32* %v.p)
   store i32 %v, i32* %sret
   ret void
 }
 
-; FIXME: The argument elision is not legal in presence of irregular types.
+define void @avoid_partially_initialized_alloca(i32 %x) {
+; CHECK-LABEL: avoid_partially_initialized_alloca:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    pushl %ebp
+; CHECK-NEXT:    movl %esp, %ebp
+; CHECK-NEXT:    andl $-8, %esp
+; CHECK-NEXT:    subl $8, %esp
+; CHECK-NEXT:    movl 8(%ebp), %eax
+; CHECK-NEXT:    movl %eax, (%esp)
+; CHECK-NEXT:    movl %esp, %eax
+; CHECK-NEXT:    pushl %eax
+; CHECK-NEXT:    calll _addrof_i32
+; CHECK-NEXT:    addl $4, %esp
+; CHECK-NEXT:    movl %ebp, %esp
+; CHECK-NEXT:    popl %ebp
+; CHECK-NEXT:    retl
+  %a = alloca i64
+  %p = bitcast i64* %a to i32*
+  store i32 %x, i32* %p
+  call void @addrof_i32(i32* %p)
+  ret void
+}
+
+; Ensure no copy elision happens as the two i3 values fed into icmp may have
+; garbage in the upper bits, a truncation is needed.
 
 define i1 @use_i3(i3 %a1, i3 %a2) {
 ; CHECK-LABEL: use_i3:
 ; CHECK:       # %bb.0:
+; CHECK-NEXT:    pushl %eax
 ; CHECK-NEXT:    movb {{[0-9]+}}(%esp), %al
 ; CHECK-NEXT:    andb $7, %al
-; CHECK-NEXT:    cmpb {{[0-9]+}}(%esp), %al
+; CHECK-NEXT:    movb {{[0-9]+}}(%esp), %cl
+; CHECK-NEXT:    andb $7, %cl
+; CHECK-NEXT:    movb %cl, {{[0-9]+}}(%esp)
+; CHECK-NEXT:    cmpb %cl, %al
 ; CHECK-NEXT:    sete %al
+; CHECK-NEXT:    popl %ecx
 ; CHECK-NEXT:    retl
-  %tmp = alloca i3, align 4
-  store i3 %a2, i3* %tmp, align 4
-  %l = load i3, i3* %tmp
-  %res = icmp eq i3 %a1, %l
+  %tmp = alloca i3
+  store i3 %a2, i3* %tmp
+  %val = load i3, i3* %tmp
+  %res = icmp eq i3 %a1, %val
   ret i1 %res
 }