GlobalISel: Restructure argument lowering loop in handleAssignments

[arsenm]

Summary

This was structured in a way that implied every split argument is in
memory, or in registers. It is possible for a pass a original argument
partially in registers, and partially in memory. Transpose the logic
here to only consider a single piece at a time. Every individual
CCValAssign should be treated independently, and any merge to original
value needs to be handled later.

This is in preparation for merging some preprocessing hacks in the
AMDGPU calling convention lowering into the generic code. This was
intended to be NFC, but it does partially address a FIXME in the
memloc handling.

As a result, this does slightly change AArch64 handling of some
promoted arguments passed on the stack. The store will be emitted as
the smaller, piece type rather than a wider store of an anyext
value. I think this exposes a failure to merge stores later, as the
change in swifterror replaces a single 64-bit stp with 2 4-byte str.

I'm also not sure what the correct behavior for memlocs where the
promoted size is larger than the original value. I've opted to clamp
the memory access size to not exceed the value register to avoid the
explicit trunc/extend/vector widen/vector extract instruction. This
happens for AMDGPU for i8 arguments that end up stack passed, which
are promoted to i16 (I think this is a preexisting DAG bug though, and
they should not really be promoted when in memory).

[aemerson]

I don't think we should be changing the extending behavior for memlocs.

Just for the varargs case in Darwin, there's lots of code out there which incorrectly try to interpret a sub 64bit incoming varargs parameter as a 64 bit value. Although it should be technically correct to emit a smaller store, what happens in practice is that this code breaks for very hard to detect reasons (i.e. you no longer get a free zeroing of the upper bits of the stack slot). In arm64 we could force this to always explicitly zero-extend to 64 bits but that incurs a penalty at the call site.

It's unfortunate but copying the DAG behavior here is likely to cause less pain.

[arsenm]

Use LocVT, which this code should have been using in the first place which avoids shrinking the store.

I also found the 2 forms of assignValueToAddress confusing, and the full form interface can't handle the case where a single stack slot covers multiple registers (although it seems unlikely this would ever be needed)

[aemerson]

llvm/test/CodeGen/AArch64/GlobalISel/call-lowering-i128-on-stack.ll

11 ↗

(On Diff #277264)

This test is actually falling back to SDAG. If you add -global-isel-abort=1 it crashes.

[arsenm]

Remove no longer fixed aarch64 case

llvm/test/CodeGen/AArch64/GlobalISel/call-lowering-i128-on-stack.ll

11 ↗

(On Diff #277264)

I think enabling the fallback is a bad default for llc

[aemerson]

llvm/test/CodeGen/AArch64/GlobalISel/call-lowering-i128-on-stack.ll

11 ↗

(On Diff #277264)

Probably true.

[arsenm]

ping

[arsenm]

ping

[aemerson]

Sorry, didn’t realise you’d updated the patch.

[arsenm]

b98f902f1877c3d679f77645a267edc89ffcd5d6

llvm/include/llvm/CodeGen/GlobalISel/CallLowering.h

virtual void assignValueToAddress(Register ValVReg, Register Addr,
                                       uint64_t Size, MachinePointerInfo &MPO,
                                       CCValAssign &VA) = 0;
 
     /// An overload which takes an ArgInfo if additional information about
     /// the arg is needed.
     virtual void assignValueToAddress(const ArgInfo &Arg, Register Addr,
                                       uint64_t Size, MachinePointerInfo &MPO,
                                       CCValAssign &VA) {
+      assert(Arg.Regs.size() == 1);
       assignValueToAddress(Arg.Regs[0], Addr, Size, MPO, VA);
     }
 
     /// Handle custom values, which may be passed into one or more of \p VAs.
     /// \return The number of \p VAs that have been assigned after the first
     ///         one, and which should therefore be skipped from further
     ///         processing.
     virtual unsigned assignCustomValue(const ArgInfo &Arg,

llvm/lib/CodeGen/GlobalISel/CallLowering.cpp

for (unsigned i = 0, e = Args.size(), j = 0; i != e; ++i, ++j) {
 
     // FIXME: Pack registers if we have more than one.
     Register ArgReg = Args[i].Regs[0];
 
     EVT OrigVT = EVT::getEVT(Args[i].Ty);
     EVT VAVT = VA.getValVT();
     const LLT OrigTy = getLLTForType(*Args[i].Ty, DL);
 
-    if (VA.isRegLoc()) {
-      if (Handler.isIncomingArgumentHandler() && VAVT != OrigVT) {
-        if (VAVT.getSizeInBits() < OrigVT.getSizeInBits()) {
-          // Expected to be multiple regs for a single incoming arg.
-          unsigned NumArgRegs = Args[i].Regs.size();
-          if (NumArgRegs < 2)
-            return false;
+    // There should be Regs.size() ArgLocs per argument.
 
-          assert((j + (NumArgRegs - 1)) < ArgLocs.size() &&
-                 "Too many regs for number of args");
-          for (unsigned Part = 0; Part < NumArgRegs; ++Part) {
-            // There should be Regs.size() ArgLocs per argument.
-            VA = ArgLocs[j + Part];
-            Handler.assignValueToReg(Args[i].Regs[Part], VA.getLocReg(), VA);
-          }
-          j += NumArgRegs - 1;
-          // Merge the split registers into the expected larger result vreg
-          // of the original call.
-          MIRBuilder.buildMerge(Args[i].OrigRegs[0], Args[i].Regs);
-          continue;
-        }
-        const LLT VATy(VAVT.getSimpleVT());
-        Register NewReg =
-            MIRBuilder.getMRI()->createGenericVirtualRegister(VATy);
-        Handler.assignValueToReg(NewReg, VA.getLocReg(), VA);
-        // If it's a vector type, we either need to truncate the elements
-        // or do an unmerge to get the lower block of elements.
-        if (VATy.isVector() &&
-            VATy.getNumElements() > OrigVT.getVectorNumElements()) {
-          // Just handle the case where the VA type is 2 * original type.
-          if (VATy.getNumElements() != OrigVT.getVectorNumElements() * 2) {
-            LLVM_DEBUG(dbgs()
-                       << "Incoming promoted vector arg has too many elts");
-            return false;
-          }
-          auto Unmerge = MIRBuilder.buildUnmerge({OrigTy, OrigTy}, {NewReg});
-          MIRBuilder.buildCopy(ArgReg, Unmerge.getReg(0));
-        } else {
-          MIRBuilder.buildTrunc(ArgReg, {NewReg}).getReg(0);
-        }
-      } else if (!Handler.isIncomingArgumentHandler()) {
-        assert((j + (Args[i].Regs.size() - 1)) < ArgLocs.size() &&
-               "Too many regs for number of args");
-        // This is an outgoing argument that might have been split.
-        for (unsigned Part = 0; Part < Args[i].Regs.size(); ++Part) {
-          // There should be Regs.size() ArgLocs per argument.
-          VA = ArgLocs[j + Part];
-          Handler.assignValueToReg(Args[i].Regs[Part], VA.getLocReg(), VA);
-        }
-        j += Args[i].Regs.size() - 1;
-      } else {
-        Handler.assignValueToReg(ArgReg, VA.getLocReg(), VA);
-      }
-    } else if (VA.isMemLoc()) {
-      // Don't currently support loading/storing a type that needs to be split
-      // to the stack. Should be easy, just not implemented yet.
-      if (Args[i].Regs.size() > 1) {
-        LLVM_DEBUG(
-            dbgs()
-            << "Load/store a split arg to/from the stack not implemented yet");
-        return false;
+      if (VA.isMemLoc()) {
+        // Don't currently support loading/storing a type that needs to be split
+        // to the stack. Should be easy, just not implemented yet.
+        if (NumArgRegs > 1) {
+          LLVM_DEBUG(
+            dbgs()
+            << "Load/store a split arg to/from the stack not implemented yet\n");
+          return false;
+
+        // FIXME: Use correct address space for pointer size
+        EVT LocVT = VA.getValVT();
+        unsigned MemSize = LocVT == MVT::iPTR ? DL.getPointerSize()
+                                              : LocVT.getStoreSize();
+        unsigned Offset = VA.getLocMemOffset();
+        MachinePointerInfo MPO;
+        Register StackAddr = Handler.getStackAddress(MemSize, Offset, MPO);
+        Handler.assignValueToAddress(Args[i], StackAddr,
+                                     MemSize, MPO, VA);
+        continue;
+      }
+
+      assert(VA.isRegLoc() && "custom loc should have been handled already");
+
+      if (OrigVT.getSizeInBits() >= VAVT.getSizeInBits() ||
+          !Handler.isIncomingArgumentHandler()) {
+        // This is an argument that might have been split. There should be
+        // Regs.size() ArgLocs per argument.
+
+        // Insert the argument copies. If VAVT < OrigVT, we'll insert the merge
+        // to the original register after handling all of the parts.
+        Handler.assignValueToReg(Args[i].Regs[Part], VA.getLocReg(), VA);
+
+      // This ArgLoc covers multiple pieces, so we need to split it.
+
+    // Now that all pieces have been handled, re-pack any arguments into any
+    // wider, original registers.
+    if (Handler.isIncomingArgumentHandler()) {
+      if (VAVT.getSizeInBits() < OrigVT.getSizeInBits()) {
+        assert(NumArgRegs >= 2);
+
+        // Merge the split registers into the expected larger result vreg
+        // of the original call.
+        MIRBuilder.buildMerge(Args[i].OrigRegs[0], Args[i].Regs);
       }
-      MVT VT = MVT::getVT(Args[i].Ty);
-      unsigned Size = VT == MVT::iPTR ? DL.getPointerSize()
-                                      : alignTo(VT.getSizeInBits(), 8) / 8;
-      unsigned Offset = VA.getLocMemOffset();
-      MachinePointerInfo MPO;
-      Register StackAddr = Handler.getStackAddress(Size, Offset, MPO);
-      Handler.assignValueToAddress(Args[i], StackAddr, Size, MPO, VA);
-    } else {
-      // FIXME: Support byvals and other weirdness
-      return false;
     }
+
+    j += NumArgRegs - 1;
   }
+
   return true;
 }
 
 bool CallLowering::analyzeArgInfo(CCState &CCState,
                                   SmallVectorImpl<ArgInfo> &Args,
                                   CCAssignFn &AssignFnFixed,
                                   CCAssignFn &AssignFnVarArg) const {
   for (unsigned i = 0, e = Args.size(); i < e; ++i) {

llvm/lib/Target/AArch64/GISel/AArch64CallLowering.cpp

void assignValueToAddress(const CallLowering::ArgInfo &Arg, Register Addr,
                             uint64_t Size, MachinePointerInfo &MPO,
                             CCValAssign &VA) override {
     unsigned MaxSize = Size * 8;
     // For varargs, we always want to extend them to 8 bytes, in which case
     // we disable setting a max.
     if (!Arg.IsFixed)
       MaxSize = 0;
 
+    assert(Arg.Regs.size() == 1);
+
     Register ValVReg = VA.getLocInfo() != CCValAssign::LocInfo::FPExt
                            ? extendRegister(Arg.Regs[0], VA, MaxSize)
                            : Arg.Regs[0];
 
     // If we extended we might need to adjust the MMO's Size.
     const LLT RegTy = MRI.getType(ValVReg);
     if (RegTy.getSizeInBytes() > Size)
       Size = RegTy.getSizeInBytes();

llvm/lib/Target/AMDGPU/AMDGPUCallLowering.cpp

case CCValAssign::LocInfo::AExt: {
       break;
     }
     default:
       MIRBuilder.buildCopy(ValVReg, PhysReg);
       break;
     }
   }
 
-  void assignValueToAddress(Register ValVReg, Register Addr, uint64_t Size,
+  void assignValueToAddress(Register ValVReg, Register Addr, uint64_t MemSize,
                             MachinePointerInfo &MPO, CCValAssign &VA) override {
     MachineFunction &MF = MIRBuilder.getMF();
 
+    // The reported memory location may be wider than the value.
+    const LLT RegTy = MRI.getType(ValVReg);
+    MemSize = std::min(static_cast<uint64_t>(RegTy.getSizeInBytes()), MemSize);
+
     // FIXME: Get alignment
     auto MMO = MF.getMachineMemOperand(
-        MPO, MachineMemOperand::MOLoad | MachineMemOperand::MOInvariant, Size,
+        MPO, MachineMemOperand::MOLoad | MachineMemOperand::MOInvariant, MemSize,
         inferAlignFromPtrInfo(MF, MPO));
     MIRBuilder.buildLoad(ValVReg, Addr, *MMO);
   }
 
   /// How the physical register gets marked varies between formal
   /// parameters (it's a basic-block live-in), and a call instruction
   /// (it's an implicit-def of the BL).
   virtual void markPhysRegUsed(unsigned PhysReg) = 0;