[AArch64][GlobalISel] Support returned argument with multiple registers

[nikic]

Summary

The call lowering code assumed that a returned argument could only consist of one register. Pass an ArrayRef<Register> instead of Register to make sure that all parts get assigned.

Not familiar with this code, so not sure whether this is the right approach -- I could imagine that the proper fix is to disable the returned handling for this case, rather than supporting it.

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

[0x59616e]

This comment was removed by 0x59616e.

[0x59616e]

This comment was removed by 0x59616e.

[0x59616e]

ok I was thinking in the wrong way. Sorry for the incorrect information.

[arsenm]

LGTM. This should be done for fixing returned, but I think something is off here (at least for AArch64). isTypeIsValidForThisReturn claims this is only valid for 64-bit types, and I'm guessing the callbacks intent was to express the original type, not the legalized split type

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

protected:
   bool determineAssignments(ValueAssigner &Assigner,
                             SmallVectorImpl<ArgInfo> &Args,
                             CCState &CCInfo) const;
 
   /// Invoke ValueAssigner::assignArg on each of the given \p Args and then use
   /// \p Handler to move them to the assigned locations.
   ///
   /// \return True if everything has succeeded, false otherwise.
-  bool determineAndHandleAssignments(ValueHandler &Handler,
-                                     ValueAssigner &Assigner,
-                                     SmallVectorImpl<ArgInfo> &Args,
-                                     MachineIRBuilder &MIRBuilder,
-                                     CallingConv::ID CallConv, bool IsVarArg,
-                                     Register ThisReturnReg = Register()) const;
+  bool
+  determineAndHandleAssignments(ValueHandler &Handler, ValueAssigner &Assigner,
+                                ArrayRef<Register> ThisReturnRegs = None) const;
 
   /// Use \p Handler to insert code to handle the argument/return values
   /// represented by \p Args. It's expected determineAssignments previously
   /// processed these arguments to populate \p CCState and \p ArgLocs.
   bool handleAssignments(ValueHandler &Handler, SmallVectorImpl<ArgInfo> &Args,
                          CCState &CCState,
                          SmallVectorImpl<CCValAssign> &ArgLocs,
                          MachineIRBuilder &MIRBuilder,
-                         Register ThisReturnReg = Register()) const;
+                         ArrayRef<Register> ThisReturnRegs = None) const;
 
   /// Check whether parameters to a call that are passed in callee saved
   /// registers are the same as from the calling function.  This needs to be
   /// checked for tail call eligibility.
   bool parametersInCSRMatch(const MachineRegisterInfo &MRI,
                             const uint32_t *CallerPreservedMask,
                             const SmallVectorImpl<CCValAssign> &ArgLocs,
                             const SmallVectorImpl<ArgInfo> &OutVals) const;

llvm/lib/CodeGen/GlobalISel/CallLowering.cpp

if (LCMTy.isVector() && CoveringSize != SrcSize)
     UnmergeSrc = B.buildPadVectorWithUndefElements(LCMTy, SrcReg).getReg(0);
 
   B.buildUnmerge(DstRegs, UnmergeSrc);
 }
 
 bool CallLowering::determineAndHandleAssignments(
     ValueHandler &Handler, ValueAssigner &Assigner,
     SmallVectorImpl<ArgInfo> &Args, MachineIRBuilder &MIRBuilder,
-    CallingConv::ID CallConv, bool IsVarArg, Register ThisReturnReg) const {
+    CallingConv::ID CallConv, bool IsVarArg,
+    ArrayRef<Register> ThisReturnRegs) const {
   MachineFunction &MF = MIRBuilder.getMF();
   const Function &F = MF.getFunction();
   SmallVector<CCValAssign, 16> ArgLocs;
 
   CCState CCInfo(CallConv, IsVarArg, MF, ArgLocs, F.getContext());
   if (!determineAssignments(Assigner, Args, CCInfo))
     return false;
 
   return handleAssignments(Handler, Args, CCInfo, ArgLocs, MIRBuilder,
-                           ThisReturnReg);
+                           ThisReturnRegs);
 }
 
 static unsigned extendOpFromFlags(llvm::ISD::ArgFlagsTy Flags) {
   if (Flags.isSExt())
     return TargetOpcode::G_SEXT;
   if (Flags.isZExt())
     return TargetOpcode::G_ZEXT;
   return TargetOpcode::G_ANYEXT;
bool CallLowering::determineAssignments(ValueAssigner &Assigner,
   return true;
 }
 
 bool CallLowering::handleAssignments(ValueHandler &Handler,
                                      SmallVectorImpl<ArgInfo> &Args,
                                      CCState &CCInfo,
                                      SmallVectorImpl<CCValAssign> &ArgLocs,
                                      MachineIRBuilder &MIRBuilder,
-                                     Register ThisReturnReg) const {
+                                     ArrayRef<Register> ThisReturnRegs) const {
   MachineFunction &MF = MIRBuilder.getMF();
   MachineRegisterInfo &MRI = MF.getRegInfo();
   const Function &F = MF.getFunction();
   const DataLayout &DL = F.getParent()->getDataLayout();
 
   const unsigned NumArgs = Args.size();
 
   // Stores thunks for outgoing register assignments. This is used so we delay
for (unsigned Part = 0; Part < NumParts; ++Part) {
                                      DstMPO, DstAlign, SrcMPO, SrcAlign,
                                      MemSize, VA);
         }
         continue;
       }
 
       assert(!VA.needsCustom() && "custom loc should have been handled already");
 
-      if (i == 0 && ThisReturnReg.isValid() &&
+      if (i == 0 && !ThisReturnRegs.empty() &&
           Handler.isIncomingArgumentHandler() &&
           isTypeIsValidForThisReturn(ValVT)) {
-        Handler.assignValueToReg(Args[i].Regs[i], ThisReturnReg, VA);
+        Handler.assignValueToReg(ArgReg, ThisReturnRegs[Part], VA);
         continue;
       }
 
       if (Handler.isIncomingArgumentHandler())
         Handler.assignValueToReg(ArgReg, VA.getLocReg(), VA);
       else {
         DelayedOutgoingRegAssignments.emplace_back([=, &Handler]() {
           Handler.assignValueToReg(ArgReg, VA.getLocReg(), VA);

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

bool UsingReturnedArg =
         !OutArgs.empty() && OutArgs[0].Flags[0].isReturned();
 
     AArch64OutgoingValueAssigner Assigner(RetAssignFn, RetAssignFn, Subtarget,
                                           /*IsReturn*/ false);
     ReturnedArgCallReturnHandler ReturnedArgHandler(MIRBuilder, MRI, MIB);
     if (!determineAndHandleAssignments(
             UsingReturnedArg ? ReturnedArgHandler : Handler, Assigner, InArgs,
             MIRBuilder, Info.CallConv, Info.IsVarArg,
-            UsingReturnedArg ? OutArgs[0].Regs[0] : Register()))
+            UsingReturnedArg ? makeArrayRef(OutArgs[0].Regs) : None))
       return false;
   }
 
   if (Info.SwiftErrorVReg) {
     MIB.addDef(AArch64::X21, RegState::Implicit);
     MIRBuilder.buildCopy(Info.SwiftErrorVReg, Register(AArch64::X21));
   }
 

llvm/test/CodeGen/AArch64/pr53315-returned-i128.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -O0 -mtriple=aarch64-unknown-linux-gnu -verify-machineinstrs < %s | FileCheck %s
+
+define void @test() nounwind {
+; CHECK-LABEL: test:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    sub sp, sp, #32
+; CHECK-NEXT:    str x30, [sp, #16] // 8-byte Folded Spill
+; CHECK-NEXT:    mov x1, xzr
+; CHECK-NEXT:    str x1, [sp, #8] // 8-byte Folded Spill
+; CHECK-NEXT:    mov x0, x1
+; CHECK-NEXT:    bl returns_arg
+; CHECK-NEXT:    ldr x1, [sp, #8] // 8-byte Folded Reload
+; CHECK-NEXT:    mov x0, x1
+; CHECK-NEXT:    bl accepts_arg
+; CHECK-NEXT:    ldr x30, [sp, #16] // 8-byte Folded Reload
+; CHECK-NEXT:    add sp, sp, #32
+; CHECK-NEXT:    ret
+  %x = call i128 @returns_arg(i128 0)
+  call void @accepts_arg(i128 %x)
+  ret void
+}
+
+declare i128 @returns_arg(i128 returned)
+declare void @accepts_arg(i128)