[SDAGBuilder] Handle multi-part arguments in argument copy elision (PR63430)

[nikic]

Summary

When eliding an argument copy, we need to update the chain to ensure the argument reads are performed before later writes. However, the code doing this only handled this for the first part of the argument. If the argument had multiple parts, the chains of the later parts were dropped. Make sure we preserve all chains.

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

[arsenm]

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

10928

For some reason this is split into two hooks for the register type and the number of registers. Wouldn't you need to account for the split and promoted / demoted case?

[nikic]

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

10928

Not sure which case you have in mind here, but this is doing the same as the code below (line 10954).

[pengfei]

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

10928

IIUC, we cannot assume all ValueVTs are the same VT. So it doesn't look correct to sum them together.
How about iterate Ins[i] and check for isSplitEnd?

[nikic]

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

10928

This doesn't assume that the VTs are the same, it just sums of the parts for all the (potentially different) VTs. We don't care about the types, just how many there are.

[pengfei]

LGTM.

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

10928

Thinking twice, it doesn't matter if ValueVTs is the same or not. We can see them as a whole for conservative reason.

[nikic]

Landed as https://github.com/llvm/llvm-project/commit/81ec494c363d4934e692e8b35e0b3fbbc3de1c2b, forgot the tag.

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

 
 /// Try to elide argument copies from memory into a local alloca. Succeeds if
 /// ArgVal is a load from a suitable fixed stack object.
 static void tryToElideArgumentCopy(
     FunctionLoweringInfo &FuncInfo, SmallVectorImpl<SDValue> &Chains,
     DenseMap<int, int> &ArgCopyElisionFrameIndexMap,
     SmallPtrSetImpl<const Instruction *> &ElidedArgCopyInstrs,
     ArgCopyElisionMapTy &ArgCopyElisionCandidates, const Argument &Arg,
-    SDValue ArgVal, bool &ArgHasUses) {
+    ArrayRef<SDValue> ArgVals, bool &ArgHasUses) {
   // Check if this is a load from a fixed stack object.
-  auto *LNode = dyn_cast<LoadSDNode>(ArgVal);
+  auto *LNode = dyn_cast<LoadSDNode>(ArgVals[0]);
   if (!LNode)
     return;
   auto *FINode = dyn_cast<FrameIndexSDNode>(LNode->getBasePtr().getNode());
   if (!FINode)
     return;
 
   // Check that the fixed stack object is the right size and alignment.
   // Look at the alignment that the user wrote on the alloca instead of looking
LLVM_DEBUG({
     dbgs() << "Eliding argument copy from " << Arg << " to " << *AI << '\n'
            << "  Replacing frame index " << OldIndex << " with " << FixedIndex
            << '\n';
   });
   MFI.RemoveStackObject(OldIndex);
   MFI.setIsImmutableObjectIndex(FixedIndex, false);
   AllocaIndex = FixedIndex;
   ArgCopyElisionFrameIndexMap.insert({OldIndex, FixedIndex});
+  for (SDValue ArgVal : ArgVals)
     Chains.push_back(ArgVal.getValue(1));
 
   // Avoid emitting code for the store implementing the copy.
   const StoreInst *SI = ArgCopyIter->second.second;
   ElidedArgCopyInstrs.insert(SI);
 
   // Check for uses of the argument again so that we can avoid exporting ArgVal
   // if it is't used by anything other than the store.
   for (const Value *U : Arg.users()) {
for (const Argument &Arg : F.args()) {
     if (NumValues == 0)
       continue;
 
     bool ArgHasUses = !Arg.use_empty();
 
     // Elide the copying store if the target loaded this argument from a
     // suitable fixed stack object.
     if (Ins[i].Flags.isCopyElisionCandidate()) {
+      unsigned NumParts = 0;
+      for (EVT VT : ValueVTs)
+        NumParts += TLI->getNumRegistersForCallingConv(*CurDAG->getContext(),
+                                                       F.getCallingConv(), VT);
+
       tryToElideArgumentCopy(*FuncInfo, Chains, ArgCopyElisionFrameIndexMap,
                              ElidedArgCopyInstrs, ArgCopyElisionCandidates, Arg,
-                             InVals[i], ArgHasUses);
+                             ArrayRef(&InVals[i], NumParts), ArgHasUses);
     }
 
     // If this argument is unused then remember its value. It is used to generate
     // debugging information.
     bool isSwiftErrorArg =
         TLI->supportSwiftError() &&
         Arg.hasAttribute(Attribute::SwiftError);
     if (!ArgHasUses && !isSwiftErrorArg) {

llvm/test/CodeGen/X86/pr63430.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --no_x86_scrub_sp --version 2
 ; RUN: llc -mtriple=x86_64-unknown-linux < %s | FileCheck %s
 
-; TODO: This is a miscompile.
+; Make sure the argument is read before the stack slot is over-written.
 define i1 @test(ptr %a0, ptr %a1, ptr %a2, ptr %a3, ptr %a4, ptr %a5, i128 %x) {
 ; CHECK-LABEL: test:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    movq 8(%rsp), %rax
 ; CHECK-NEXT:    xorps %xmm0, %xmm0
-; CHECK-NEXT:    movaps %xmm0, 8(%rsp)
 ; CHECK-NEXT:    andq 16(%rsp), %rax
+; CHECK-NEXT:    movaps %xmm0, 8(%rsp)
 ; CHECK-NEXT:    cmpq $-1, %rax
 ; CHECK-NEXT:    sete %al
 ; CHECK-NEXT:    retq
   %alloca = alloca i128
   store i128 %x, ptr %alloca
   store i128 0, ptr %alloca
   %cmp = icmp eq i128 %x, -1
   ret i1 %cmp
 }