[AArch64][GlobalISel] Support the 'returned' parameter attribute.

[aemerson]

Summary

On AArch64 (which seems to be the only target that supports it), this attribute allows codegen to avoid saving/restoring the value in x0 across a call.

Gives a 0.1% geomean -Os code size improvement on CTMark.

[arsenm]

llvm/lib/CodeGen/GlobalISel/CallLowering.cpp

231

I think it would be better to move this to a proper function in CallLowering, and take a calling convention argument (also not a fan of using EVT, but I guess we're stuck with it for a while)

[paquette]

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

271

Can we document this parameter?

llvm/lib/CodeGen/GlobalISel/CallLowering.cpp

174

Comment would be useful here

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

1012

Would it make sense to make this a function, since it's the same as the code in lowerTailCall?

(The tail call/standard call code really needs to be refactored, so maybe it's better to leave it. I'm not sure.)

llvm/test/CodeGen/AArch64/arm64-this-return.ll

2

Would it make sense to add a MIR test for this as well? Or would that just be unnecessary test duplication?

[aemerson]

Address comments.

[paquette]

LGTM

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

protected:
   ///        \p PackedTy, as returned by computeValueLLTs.
   void unpackRegs(ArrayRef<Register> DstRegs, Register SrcReg, Type *PackedTy,
                   MachineIRBuilder &MIRBuilder) const;
 
   /// Invoke Handler::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 handleAssignments(MachineIRBuilder &MIRBuilder,
-                         SmallVectorImpl<ArgInfo> &Args,
-                         ValueHandler &Handler) const;
-  bool handleAssignments(CCState &CCState,
-                         SmallVectorImpl<CCValAssign> &ArgLocs,
-                         MachineIRBuilder &MIRBuilder,
-                         SmallVectorImpl<ArgInfo> &Args,
-                         ValueHandler &Handler) const;
+  bool handleAssignments(
+      MachineIRBuilder &MIRBuilder, SmallVectorImpl<ArgInfo> &Args,
+      ValueHandler &Handler,
+      std::function<bool(EVT)> TypeIsValidForThisReturn =
+          [](EVT T) { return false; },
+      Register ThisReturnReg = Register()) const;
+  bool handleAssignments(
+      CCState &CCState, SmallVectorImpl<CCValAssign> &ArgLocs,
+      MachineIRBuilder &MIRBuilder, SmallVectorImpl<ArgInfo> &Args,
+      ValueHandler &Handler,
+      std::function<bool(EVT)> TypeIsValidForThisReturn =
+          [](EVT T) { return false; },
+      Register ThisReturnReg = Register()) const;
 
   /// Analyze passed or returned values from a call, supplied in \p ArgInfo,
   /// incorporating info about the passed values into \p CCState.
   ///
   /// Used to check if arguments are suitable for tail call lowering.
   bool analyzeArgInfo(CCState &CCState, SmallVectorImpl<ArgInfo> &Args,
                       CCAssignFn &AssignFnFixed,
                       CCAssignFn &AssignFnVarArg) const;

llvm/include/llvm/CodeGen/TargetCallingConv.h

public:
 
     bool isSecArgPass() const { return IsSecArgPass; }
     void setSecArgPass() { IsSecArgPass = 1; }
 
     bool isNest() const { return IsNest; }
     void setNest() { IsNest = 1; }
 
     bool isReturned() const { return IsReturned; }
-    void setReturned() { IsReturned = 1; }
+    void setReturned(bool V = true) { IsReturned = V; }
 
     bool isInConsecutiveRegs()  const { return IsInConsecutiveRegs; }
     void setInConsecutiveRegs(bool Flag = true) { IsInConsecutiveRegs = Flag; }
 
     bool isInConsecutiveRegsLast() const { return IsInConsecutiveRegsLast; }
     void setInConsecutiveRegsLast(bool Flag = true) {
       IsInConsecutiveRegsLast = Flag;
     }

llvm/lib/CodeGen/GlobalISel/CallLowering.cpp

if (Flags.isByVal() || Flags.isInAlloca() || Flags.isPreallocated()) {
     Align FrameAlign;
     if (auto ParamAlign = FuncInfo.getParamAlign(OpIdx - 2))
       FrameAlign = *ParamAlign;
     else
       FrameAlign = Align(getTLI()->getByValTypeAlignment(ElementTy, DL));
     Flags.setByValAlign(FrameAlign);
   }
   Flags.setOrigAlign(DL.getABITypeAlign(Arg.Ty));
+
+  if (Flags.isSwiftSelf())
+    Flags.setReturned(false);
 }
 
 template void
 CallLowering::setArgFlags<Function>(CallLowering::ArgInfo &Arg, unsigned OpIdx,
                                     const DataLayout &DL,
                                     const Function &FuncInfo) const;
 
 template void
void CallLowering::unpackRegs(ArrayRef<Register> DstRegs, Register SrcReg,
   SmallVector<uint64_t, 8> Offsets;
   computeValueLLTs(DL, *PackedTy, LLTs, &Offsets);
   assert(LLTs.size() == DstRegs.size() && "Regs / types mismatch");
 
   for (unsigned i = 0; i < DstRegs.size(); ++i)
     MIRBuilder.buildExtract(DstRegs[i], SrcReg, Offsets[i]);
 }
 
-bool CallLowering::handleAssignments(MachineIRBuilder &MIRBuilder,
-                                     SmallVectorImpl<ArgInfo> &Args,
-                                     ValueHandler &Handler) const {
+bool CallLowering::handleAssignments(
+    MachineIRBuilder &MIRBuilder, SmallVectorImpl<ArgInfo> &Args,
+    ValueHandler &Handler, std::function<bool(EVT)> TypeIsValidForThisReturn,
+    Register ThisReturnReg) const {
   MachineFunction &MF = MIRBuilder.getMF();
   const Function &F = MF.getFunction();
   SmallVector<CCValAssign, 16> ArgLocs;
   CCState CCInfo(F.getCallingConv(), F.isVarArg(), MF, ArgLocs, F.getContext());
-  return handleAssignments(CCInfo, ArgLocs, MIRBuilder, Args, Handler);
+  return handleAssignments(CCInfo, ArgLocs, MIRBuilder, Args, Handler,
+                           TypeIsValidForThisReturn, ThisReturnReg);
 }
 
-bool CallLowering::handleAssignments(CCState &CCInfo,
-                                     SmallVectorImpl<CCValAssign> &ArgLocs,
-                                     MachineIRBuilder &MIRBuilder,
-                                     SmallVectorImpl<ArgInfo> &Args,
-                                     ValueHandler &Handler) const {
+bool CallLowering::handleAssignments(
+    CCState &CCInfo, SmallVectorImpl<CCValAssign> &ArgLocs,
+    MachineIRBuilder &MIRBuilder, SmallVectorImpl<ArgInfo> &Args,
+    ValueHandler &Handler, std::function<bool(EVT)> TypeIsValidForThisReturn,
+    Register ThisReturnReg) const {
   MachineFunction &MF = MIRBuilder.getMF();
   const Function &F = MF.getFunction();
   const DataLayout &DL = F.getParent()->getDataLayout();
 
   unsigned NumArgs = Args.size();
   for (unsigned i = 0; i != NumArgs; ++i) {
     EVT CurVT = EVT::getEVT(Args[i].Ty);
     if (CurVT.isSimple() &&
if (Handler.isIncomingArgumentHandler()) {
         ISD::ArgFlagsTy Flags = OrigFlags;
         if (PartIdx == 0) {
           Flags.setSplit();
         } else {
           Flags.setOrigAlign(Align(1));
           if (PartIdx == NumParts - 1)
             Flags.setSplitEnd();
         }
+
+        // TODO: Also check if there is a valid extension that preserves the
+        // bits. However currently this call lowering doesn't support non-exact
+        // split parts, so that can't be tested.
+        if (OrigFlags.isReturned() &&
+            (NumParts * NewVT.getSizeInBits() != CurVT.getSizeInBits())) {
+          Flags.setReturned(false);
+        }
+
         Args[i].Regs.push_back(Unmerge.getReg(PartIdx));
         Args[i].Flags.push_back(Flags);
         if (Handler.assignArg(i, NewVT, NewVT, CCValAssign::Full,
                               Args[i], Args[i].Flags[PartIdx], CCInfo))
           return false;
       }
     }
   }
for (unsigned Part = 0; Part < NumArgRegs; ++Part) {
         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 (i == 0 && ThisReturnReg.isValid() &&
+          Handler.isIncomingArgumentHandler() &&
+          TypeIsValidForThisReturn(VAVT)) {
+        Handler.assignValueToReg(Args[i].Regs[i], ThisReturnReg, VA);
+        continue;
+      }
+
       // GlobalISel does not currently work for scalable vectors.
       if (OrigVT.getFixedSizeInBits() >= VAVT.getFixedSizeInBits() ||
           !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.

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

struct CallReturnHandler : public IncomingArgHandler {
 
   void markPhysRegUsed(MCRegister PhysReg) override {
     MIB.addDef(PhysReg, RegState::Implicit);
   }
 
   MachineInstrBuilder MIB;
 };
 
+/// A special return arg handler for "returned" attribute arg calls.
+struct ReturnedArgCallReturnHandler : public CallReturnHandler {
+  ReturnedArgCallReturnHandler(MachineIRBuilder &MIRBuilder,
+                               MachineRegisterInfo &MRI,
+                               MachineInstrBuilder MIB, CCAssignFn *AssignFn)
+      : CallReturnHandler(MIRBuilder, MRI, MIB, AssignFn) {}
+
+  void markPhysRegUsed(MCRegister PhysReg) override {}
+};
+
 struct OutgoingArgHandler : public CallLowering::OutgoingValueHandler {
   OutgoingArgHandler(MachineIRBuilder &MIRBuilder, MachineRegisterInfo &MRI,
                      MachineInstrBuilder MIB, CCAssignFn *AssignFn,
                      CCAssignFn *AssignFnVarArg, bool IsTailCall = false,
                      int FPDiff = 0)
       : OutgoingValueHandler(MIRBuilder, MRI, AssignFn), MIB(MIB),
         AssignFnVarArg(AssignFnVarArg), IsTailCall(IsTailCall), FPDiff(FPDiff),
         StackSize(0), SPReg(0) {}
bool AArch64CallLowering::lowerTailCall(
   const auto &Forwards = FuncInfo->getForwardedMustTailRegParms();
 
   // Do the actual argument marshalling.
   OutgoingArgHandler Handler(MIRBuilder, MRI, MIB, AssignFnFixed,
                              AssignFnVarArg, true, FPDiff);
   if (!handleAssignments(MIRBuilder, OutArgs, Handler))
     return false;
 
+  if (!OutArgs.empty() && OutArgs[0].Flags[0].isReturned()) {
+    // For 'this' returns, use the X0-preserving mask if applicable
+    Mask = TRI->getThisReturnPreservedMask(MF, Info.CallConv);
+    if (!Mask) {
+      OutArgs[0].Flags[0].setReturned(false);
+      Mask = TRI->getCallPreservedMask(MF, Info.CallConv);
+    }
+  } else {
+    Mask = TRI->getCallPreservedMask(MF, Info.CallConv);
+  }
+
   if (Info.IsVarArg && Info.IsMustTailCall) {
     // Now we know what's being passed to the function. Add uses to the call for
     // the forwarded registers that we *aren't* passing as parameters. This will
     // preserve the copies we build earlier.
     for (const auto &F : Forwards) {
       Register ForwardedReg = F.PReg;
       // If the register is already passed, or aliases a register which is
       // already being passed, then skip it.
bool AArch64CallLowering::lowerCall(MachineIRBuilder &MIRBuilder,
   // Create a temporarily-floating call instruction so we can add the implicit
   // uses of arg registers.
   unsigned Opc = getCallOpcode(MF, Info.Callee.isReg(), false);
 
   auto MIB = MIRBuilder.buildInstrNoInsert(Opc);
   MIB.add(Info.Callee);
 
   // Tell the call which registers are clobbered.
-  auto TRI = MF.getSubtarget<AArch64Subtarget>().getRegisterInfo();
-  const uint32_t *Mask = TRI->getCallPreservedMask(MF, Info.CallConv);
-  if (MF.getSubtarget<AArch64Subtarget>().hasCustomCallingConv())
-    TRI->UpdateCustomCallPreservedMask(MF, &Mask);
-  MIB.addRegMask(Mask);
-
-  if (TRI->isAnyArgRegReserved(MF))
-    TRI->emitReservedArgRegCallError(MF);
+  const uint32_t *Mask;
+  const auto *TRI = MF.getSubtarget<AArch64Subtarget>().getRegisterInfo();
 
   // Do the actual argument marshalling.
   OutgoingArgHandler Handler(MIRBuilder, MRI, MIB, AssignFnFixed,
                              AssignFnVarArg, false);
   if (!handleAssignments(MIRBuilder, OutArgs, Handler))
     return false;
 
+  bool UsingReturnedArg = !OutArgs.empty() && OutArgs[0].Flags[0].isReturned();
+  if (UsingReturnedArg) {
+    // For 'this' returns, use the X0-preserving mask if applicable
+    Mask = TRI->getThisReturnPreservedMask(MF, Info.CallConv);
+    if (!Mask) {
+      OutArgs[0].Flags[0].setReturned(false);
+      Mask = TRI->getCallPreservedMask(MF, Info.CallConv);
+    }
+  } else {
+    Mask = TRI->getCallPreservedMask(MF, Info.CallConv);
+  }
+
+  if (MF.getSubtarget<AArch64Subtarget>().hasCustomCallingConv())
+    TRI->UpdateCustomCallPreservedMask(MF, &Mask);
+  MIB.addRegMask(Mask);
+
+  if (TRI->isAnyArgRegReserved(MF))
+    TRI->emitReservedArgRegCallError(MF);
+
   // Now we can add the actual call instruction to the correct basic block.
   MIRBuilder.insertInstr(MIB);
 
   // If Callee is a reg, since it is used by a target specific
   // instruction, it must have a register class matching the
   // constraint of that instruction.
   if (Info.Callee.isReg())
     constrainOperandRegClass(MF, *TRI, MRI, *MF.getSubtarget().getInstrInfo(),
                              *MF.getSubtarget().getRegBankInfo(), *MIB,
                              MIB->getDesc(), Info.Callee, 0);
 
   // Finally we can copy the returned value back into its virtual-register. In
   // symmetry with the arguments, the physical register must be an
   // implicit-define of the call instruction.
   if (!Info.OrigRet.Ty->isVoidTy()) {
     CCAssignFn *RetAssignFn = TLI.CCAssignFnForReturn(Info.CallConv);
     CallReturnHandler Handler(MIRBuilder, MRI, MIB, RetAssignFn);
-    if (!handleAssignments(MIRBuilder, InArgs, Handler))
+    ReturnedArgCallReturnHandler ReturnedArgHandler(MIRBuilder, MRI, MIB,
+                                                    RetAssignFn);
+    if (!handleAssignments(
+            MIRBuilder, InArgs, UsingReturnedArg ? ReturnedArgHandler : Handler,
+            [](EVT VT) { return VT.getSizeInBits() == 64; },
+            UsingReturnedArg ? OutArgs[0].Regs[0] : Register()))
       return false;
   }
 
   if (Info.SwiftErrorVReg) {
     MIB.addDef(AArch64::X21, RegState::Implicit);
     MIRBuilder.buildCopy(Info.SwiftErrorVReg, Register(AArch64::X21));
   }
 

llvm/test/CodeGen/AArch64/GlobalISel/arm64-irtranslator-switch.ll

define void @jt_multiple_jump_tables(%1* %arg, i32 %arg1, i32* %arg2) {
   ; CHECK:   [[MUL:%[0-9]+]]:_(s64) = G_MUL [[PHI]], [[C111]]
   ; CHECK:   [[PTR_ADD:%[0-9]+]]:_(p0) = G_PTR_ADD [[GV]], [[MUL]](s64)
   ; CHECK:   [[C112:%[0-9]+]]:_(s64) = G_CONSTANT i64 8
   ; CHECK:   [[PTR_ADD1:%[0-9]+]]:_(p0) = G_PTR_ADD [[PTR_ADD]], [[C112]](s64)
   ; CHECK:   [[LOAD:%[0-9]+]]:_(p0) = G_LOAD [[PTR_ADD1]](p0) :: (load 8 from %ir.tmp59)
   ; CHECK:   ADJCALLSTACKDOWN 0, 0, implicit-def $sp, implicit $sp
   ; CHECK:   $x0 = COPY [[COPY]](p0)
   ; CHECK:   $x1 = COPY [[LOAD]](p0)
-  ; CHECK:   BL @wibble, csr_aarch64_aapcs, implicit-def $lr, implicit $sp, implicit $x0, implicit $x1, implicit-def $x0
-  ; CHECK:   [[COPY3:%[0-9]+]]:_(p0) = COPY $x0
+  ; CHECK:   BL @wibble, csr_aarch64_aapcs_thisreturn, implicit-def $lr, implicit $sp, implicit $x0, implicit $x1
+  ; CHECK:   [[COPY3:%[0-9]+]]:_(p0) = COPY [[COPY]](p0)
   ; CHECK:   ADJCALLSTACKUP 0, 0, implicit-def $sp, implicit $sp
   ; CHECK:   G_BR %bb.59
   ; CHECK: bb.57.bb62:
   ; CHECK:   successors: %bb.59(0x80000000)
   ; CHECK:   ADJCALLSTACKDOWN 0, 0, implicit-def $sp, implicit $sp
   ; CHECK:   $x0 = COPY [[COPY]](p0)
   ; CHECK:   $x1 = COPY [[COPY2]](p0)
-  ; CHECK:   BL @wibble, csr_aarch64_aapcs, implicit-def $lr, implicit $sp, implicit $x0, implicit $x1, implicit-def $x0
-  ; CHECK:   [[COPY4:%[0-9]+]]:_(p0) = COPY $x0
+  ; CHECK:   BL @wibble, csr_aarch64_aapcs_thisreturn, implicit-def $lr, implicit $sp, implicit $x0, implicit $x1
+  ; CHECK:   [[COPY4:%[0-9]+]]:_(p0) = COPY [[COPY]](p0)
   ; CHECK:   ADJCALLSTACKUP 0, 0, implicit-def $sp, implicit $sp
   ; CHECK:   G_BR %bb.59
   ; CHECK: bb.58.bb64:
   ; CHECK:   successors: %bb.59(0x80000000)
   ; CHECK:   [[COPY5:%[0-9]+]]:_(p0) = COPY [[FRAME_INDEX]](p0)
   ; CHECK:   ADJCALLSTACKDOWN 0, 0, implicit-def $sp, implicit $sp
   ; CHECK:   $w0 = COPY [[COPY1]](s32)
   ; CHECK:   $x1 = COPY [[COPY5]](p0)
   ; CHECK:   BL @baz, csr_aarch64_aapcs, implicit-def $lr, implicit $sp, implicit $w0, implicit $x1
   ; CHECK:   ADJCALLSTACKUP 0, 0, implicit-def $sp, implicit $sp
   ; CHECK:   ADJCALLSTACKDOWN 0, 0, implicit-def $sp, implicit $sp
   ; CHECK:   $x0 = COPY [[COPY]](p0)
   ; CHECK:   $x1 = COPY [[COPY5]](p0)
-  ; CHECK:   BL @wibble, csr_aarch64_aapcs, implicit-def $lr, implicit $sp, implicit $x0, implicit $x1, implicit-def $x0
-  ; CHECK:   [[COPY6:%[0-9]+]]:_(p0) = COPY $x0
+  ; CHECK:   BL @wibble, csr_aarch64_aapcs_thisreturn, implicit-def $lr, implicit $sp, implicit $x0, implicit $x1
+  ; CHECK:   [[COPY6:%[0-9]+]]:_(p0) = COPY [[COPY]](p0)
   ; CHECK:   ADJCALLSTACKUP 0, 0, implicit-def $sp, implicit $sp
   ; CHECK: bb.59.bb68:
   ; CHECK:   RET_ReallyLR
   ; CHECK: bb.60.bb69:
   ; CHECK:   successors: %bb.58(0x40000000), %bb.57(0x40000000)
   ; CHECK:   [[ICMP55:%[0-9]+]]:_(s1) = G_ICMP intpred(eq), [[COPY2]](p0), [[C55]]
   ; CHECK:   G_BRCOND [[ICMP55]](s1), %bb.58
   ; CHECK:   G_BR %bb.57

llvm/test/CodeGen/AArch64/arm64-this-return.ll

 ; RUN: llc < %s -mtriple=arm64-eabi | FileCheck %s
+; RUN: llc < %s -mtriple=arm64-eabi -global-isel -global-isel-abort=1 | FileCheck %s
 
 %struct.A = type { i8 }
 %struct.B = type { i32 }
 %struct.C = type { %struct.B }
 %struct.D = type { %struct.B }
 %struct.E = type { %struct.B, %struct.B }
 
 declare %struct.A* @A_ctor_base(%struct.A* returned)