[COFF, ARM64] Fix ABI implementation of struct returns

[mgrang]

Summary

Refer the ABI doc at: https://docs.microsoft.com/en-us/cpp/build/arm64-windows-abi-conventions?view=vs-2019#return-values

Related clang patch: D60349

[mgrang]

Refer https://bugs.llvm.org/show_bug.cgi?id=41135 for a detailed discussion on the issue. Please review this and the related clang patch and let me know if this implementation is something we can proceed with.

Note: In case, for instance methods (when "this" is passed in x0), we would need to do a mov from x1 to get the sret parameter. I have not yet implemented this. Will do so once the reviewers +1 the implementation.

[mgrang]

The new patch now also handles struct returns for instance methods.

[kristof.beyls]

Hi Mandeep,

I had a quick look through and wrote up some thoughts that triggered immediately.
I'll try to have a further look when I find more time.

Thanks!

lib/Target/AArch64/AArch64ISelLowering.cpp

3212

As said elsewhere, I don't think the note is that helpful after this gets committed.

3219

I wonder if the a more exact word could be used than "non-POD"?
The updated documentation at https://docs.microsoft.com/en-us/cpp/build/arm64-windows-abi-conventions?view=vs-2019#return-values doesn't describe the rules in terms of POD anymore. I wonder if a more exact term could be used here than "non-POD"? I realize that may be hard... Maybe just drop the word "non-POD", as it probably wouldn't remove anything from the clarity of the explanation in the comment?

3223

I needed to browse around the surrounding code quite a bit to understand what "I == 1" really means.
I guess it means "this is the second argument to the function". If my guess is correct, maybe it'd be useful to spell it out in the comment?

3226

I wonder if it'd be a good idea to add a pointer in the comment to the relevant MSVC ABI documentation indicating that the struct pointer must be returned in X0?
I think that would be https://docs.microsoft.com/en-us/cpp/build/arm64-windows-abi-conventions?view=vs-2019#return-values?

3988

This note is helpful for review, indicating it's based on X86 implementation. However, I don't think it adds too much value if it'd be committed as is to the code base.

lib/Target/AArch64/AArch64MachineFunctionInfo.h

94–98

It seems this is also present in X86MachineFunctionInfo.
I wonder if it wouldn't be better to move this to the MachineFunctionInfo base class so we don't need to add an identical interface in 2 derived classes?

[mgrang]

lib/Target/AArch64/AArch64MachineFunctionInfo.h

94–98

MachineFunctionInfo base only has a factory function for creating objects. All fields/methods are in derived classes. So not sure if we should move SRetReturnReg to the base.

[kristof.beyls]

Hi Mandeep,

I spent a bit more time understanding the code in more depth now.
I have a few more remarks that I hope will help to make the code simpler and more understandable.
Overall, the design seems acceptable to me - especially as it is in line with how things are done in the X86 backend. It's good that we don't unnecessarily use more different ways to implement things.
One big remark though: I think this needs regression tests to be added that should cover the lines of code you're introducing in this patch.

I expect to be offline for the next 7 days, so if anyone else would want to review/approve an updated version of this patch: I'd also be happy with that if the above comments have been taken into account.

lib/Target/AArch64/AArch64CallingConvention.td

19–23

Is CCIfNotSubtarget used anywhere? If not, this can be removed?

59–60

Based on my understanding of this code, I think the following comment would make it easier to understand, by explaining what is aimed for without making the comment only about windows systems:

// In AAPCS, an SRet is passed in X8, not X0 like a normal pointer parameter.
// However, on windows, in some circumstances, the SRet is passed in in X0 or X1 instead.
// The presence of the inreg attribute indicates that SRet is passed in the alternative register (X0 or X1), not X8:
//  - X0 for non-instance methods.
//  - X1 for instance methods.

62–67

Assuming the proposed extended comment/documentation is correct, I think it would be easier to follow the logic if "CCIfInReg" rather than "CCIfNotInReg" was used to express the logic (one negation less). Do you think that would be possible?

62–67

It seems to me that the case where the SRet parameter is passed in X0 isn't explicitly handled here.
Should it be?
If not, I think the comment above should be extended to explain how that case is handled elsewhere.

lib/Target/AArch64/AArch64ISelLowering.cpp

3211–3212

I get the impression that this doesn't need to be protected by an "IsWin64" test.
The sret inreg combination will only be produced by a front-end when targeting windows. So maybe best to just test for "sret" and "inreg" attributes to alter behaviour, not also on the "IsWin64" predicate?

Or maybe this block of code is there to make sure that on Windows the Sret pointer does get returned (in X0). And there is no other way to get the information that the sret pointer must be returned, apart from there being an "sret inreg" function parameter present?
If that's the case, I think that extending the comment "Struct returns on Windows" to e.g. "On windows, Sret pointers must be returned, so record the pointer in a virtual register at the start of the function so it can be returned in the epilogue", makes the intent of this code block a bit easier to understand. Having looked into X86ISelLowering for equivalent functionality, I saw the following words were used in a comment there, that I thought was especially clear: "Save the argument into a virtual register so that we can access it from the return points.".

3217–3225

Assuming my understanding on the intent of this code is correct (i.e. "Save the argument into a virtual register so that we can access it from the return points."), there is no need to make a distinction between whether the sret pointer is passed in in X0 or X1 - it needs to be returned via X0 in both cases.
So, why not just have a very simple loop here that looks for the presence of an argument with both "sret" and "inreg" attributes and save the value of that argument in the getSRetReturnReg?
I think the loop could look almost identical to how similar functionality is implemented in X86ISelLowering with quite a bit less code - making this quite a bit easier to understand and maintain?

3987–3990

Similar comment to the above: maybe it's not necessary to check for IsWin64 here - just checking on whether getSRetReturnReg is set could be enough?
It would make the code less complex.
The comment could be extended a bit to make this a bit more clear - I thought the equivalent comment on X86IselLowering was pretty clear. Maybe something like:

// Windows AArch64 ABIs require that for returning structs by value we copy
// the sret argument into X0 for the return.
// We saved the argument into a virtual register in the entry block,
// so now we copy the value out and into X0.

lib/Target/AArch64/AArch64MachineFunctionInfo.h

94–98

I see - thanks for explaining!

[mgrang]

Thanks @kristof.beyls for the review. We are still working on the combining the clang patches which should fix PR41135 and PR41136. I will then add unit tests to this patch.

[rnk]

In D60348#1466139, @kristof.beyls wrote:

I expect to be offline for the next 7 days, so if anyone else would want to review/approve an updated version of this patch: I'd also be happy with that if the above comments have been taken into account.

Fair enough, I just got back to work, so I'll look into it. :)

From what I can tell, the last update implements all the suggested simplifications. I have a suggestion, but maybe it cuts against the direction you were suggesting.

Overall, I'm happy with this.

lib/Target/AArch64/AArch64CallingConvention.td

39–51

I feel like the whole discussion of free function vs. C++ instance method (X0/X1) could be avoided by falling through to the normal i64 argument handling below. IIUC, these are roughly the conditions we want:

if (IsSRet && (IsInReg || !IsWin64))
  // assign to X8

I'm not sure how to express that in our calling conv tablegen, though.

lib/Target/AArch64/AArch64MachineFunctionInfo.h

94–98

If we did want to avoid this duplication, I would say just toss it directly on MachineFunction. But, for now, I think it's best to do things like we normally do: duplicate some logic between targets but try to keep them aligned.

[efriedma]

Do you need to modify AArch64TargetLowering::isEligibleForTailCallOptimization to prevent a tail call in cases where the tail call would return the wrong value?

[mgrang]

In D60348#1474684, @efriedma wrote:

Do you need to modify AArch64TargetLowering::isEligibleForTailCallOptimization to prevent a tail call in cases where the tail call would return the wrong value?

Thanks Eli. I have prevented tail call for arm64-windows when sret arguments are present.

[efriedma]

lib/Target/AArch64/AArch64ISelLowering.cpp

3432

Can we restrict this to only arguments with both sret and inreg?

Maybe leave a FIXME to check whether the callee also has an sret+inreg argument. That optimization isn't critical, but nice to have at some point.

[mgrang]

lib/Target/AArch64/AArch64ISelLowering.cpp

3432

You cannot have inreg without sret. So I think just checking inreg should be enough. Will change this to check for inreg. Thanks.

[efriedma]

Is it possible to enable fast-isel and/or global-isel for ARM64 Windows? Do we need additional code for them? (I'm okay with making those changes separate patches, but it should be easy to at least add testcases now.)

lib/Target/AArch64/AArch64CallingConvention.td

57

Instead of the separate CCIfValNo check, can you just write CCAssignToRegWithShadow<[X0, X1], [W0, W1]>, like we do for normal i64 arguments?

lib/Target/AArch64/AArch64ISelLowering.cpp

3218

We should call LowerFormalArguments exactly once... maybe make this an assertion?

3432

Should we also check for inreg in other places the patch checks for sret?

[mstorsjo]

In D60348#1476017, @efriedma wrote:

Is it possible to enable fast-isel and/or global-isel for ARM64 Windows? Do we need additional code for them? (I'm okay with making those changes separate patches, but it should be easy to at least add testcases now.)

GlobalISEL is the default for unoptimized builds, and iirc FastISEL should work as well (it was pretty well working already when GlobalISEL was enabled by default for aarch64).

[rnk]

lib/IR/Function.cpp

149 ↗

(On Diff #196152)

Why have this check? inreg is typically applied to regular integer i32 arguments and vectors.

[rnk]

lib/Target/AArch64/AArch64CallingConvention.td

39–51

Nice, I think this is simpler.

lib/Target/AArch64/AArch64ISelLowering.cpp

3428

Hang on, why are we checking for inreg at all? Won't we still erroneously perform TCO on this example that we currently TCO?

struct Foo {
  Foo();
  Foo(const Foo &o);
  Foo(int, int);
  ~Foo();
  void *p;
};

void bar(Foo *, int, int);

Foo getfoo(int x, int y) {
  Foo f;
  bar(&f, x, y);
  return f;
}

In this example, NRVO fires, and &f is the sret pointer passed in. bar isn't guaranteed to return &f, but we tail call it, and inreg isn't involved at all.

I think the precise check that we are looking for is that, if the caller has an arg marked sret, then that argument must be passed to the callee and be marked sret. It doesn't matter if the callee receives the sret pointer in a different argument (X0, X1, X8), what matters is that it's returned in X0, and that matches the caller function's requirements. It might also be worth enabling TCO on constructors, which use the returned convention in IR, so something like this:

Foo getfoo(int x, int y) {
  return Foo(x, y);
}

Bailing on the caller using sret is conservatively correct of course.

[mgrang]

lib/Target/AArch64/AArch64ISelLowering.cpp

3428

Thanks Reid. I see in other targets like RISCV we disable TCO if caller or callee use StructRet. Should we do the same here? Also should we restrict that to only Win64?

[rnk]

lib/Target/AArch64/AArch64ISelLowering.cpp

3428

I suppose yes, restricting it to Win64 makes sense, since that's where the preservation requirement comes from.

[rnk]

I went over other reviewers comments, and I think this was the only open one:

In D60348#1476017, @efriedma wrote:

Is it possible to enable fast-isel and/or global-isel for ARM64 Windows? Do we need additional code for them? (I'm okay with making those changes separate patches, but it should be easy to at least add testcases now.)

Based on the comments in the test, I take it this will be addressed in a subsequent patch. Sounds good.

lgtm

[efriedma]

lib/Target/AArch64/AArch64ISelLowering.cpp

3428

Hang on, why are we checking for inreg at all?

"inreg" specifically marks the calls where we have to restore x0. For normal C calls, we follow the usual AArch64 ABI rules, and don't need to preserve the sret address.

[mgrang]

lib/Target/AArch64/AArch64ISelLowering.cpp

3428

I agree with Eli. We need to avoid TCO when we have to restore X0. The "sret" attribute is also used for aggregates larger than 16 bytes where the address is passed in X8 but no register needs to be restored.

Like here:

struct S3 { int a, b, c, d, e; };
S3 f3() { return S3{}; }

define dso_local void @"?f3@@YA?AUS3@@XZ"(%struct.S3* noalias nocapture sret %agg.result)

The "inreg" attribute specifies that X0 needs to be restored. So I guess we should indeed check for "inreg" (not "sret").

[rnk]

lib/Target/AArch64/AArch64ISelLowering.cpp

3428

Sorry, my misunderstanding. I thought the sret pointer was returned in all cases: aggregate, non-aggregate, instance method or otherwise, but I see it's only for non-aggregates.

[mgrang]

@rnk @efriedma Is this patch good to commit?

[efriedma]

lib/Target/AArch64/AArch64ISelLowering.cpp

3418

Unnecessary change.

test/CodeGen/AArch64/arm64-windows-returns.ll

72 ↗

(On Diff #196544)

Please add a testcase for a non-instance method, where the first argument is inreg.

[efriedma]

LGTM. (Please merge this at the same time as the clang patch.)

lib/Target/AArch64/AArch64CallingConvention.td

 //===----------------------------------------------------------------------===//
 
 /// CCIfAlign - Match of the original alignment of the arg
 class CCIfAlign<string Align, CCAction A> :
   CCIf<!strconcat("ArgFlags.getOrigAlign() == ", Align), A>;
 /// CCIfBigEndian - Match only if we're in big endian mode.
 class CCIfBigEndian<CCAction A> :
   CCIf<"State.getMachineFunction().getDataLayout().isBigEndian()", A>;
+// CCIfValNo - Match if this is argument number N.
+class CCIfValNo<string N, CCAction A> : CCIf<!strconcat("ValNo == ", N), A>;
 
 //===----------------------------------------------------------------------===//
 // ARM AAPCS64 Calling Convention
 //===----------------------------------------------------------------------===//
 
 let Entry = 1 in
 def CC_AArch64_AAPCS : CallingConv<[
   CCIfType<[iPTR], CCBitConvertToType<i64>>,
   CCIfType<[v2f32], CCBitConvertToType<v2i32>>,
   CCIfType<[v2f64, v4f32], CCBitConvertToType<v2i64>>,
 
   // Big endian vectors must be passed as if they were 1-element vectors so that
   // their lanes are in a consistent order.
   CCIfBigEndian<CCIfType<[v2i32, v2f32, v4i16, v4f16, v8i8],
                          CCBitConvertToType<f64>>>,
   CCIfBigEndian<CCIfType<[v2i64, v2f64, v4i32, v4f32, v8i16, v8f16, v16i8],
                          CCBitConvertToType<f128>>>,
 
-  // An SRet is passed in X8, not X0 like a normal pointer parameter.
+  // In AAPCS, an SRet is passed in X8, not X0 like a normal pointer parameter.
+  // However, on windows, in some circumstances, the SRet is passed in X0 or X1
+  // instead.  The presence of the inreg attribute indicates that SRet is
+  // passed in the alternative register (X0 or X1), not X8:
+  // - X0 for non-instance methods.
+  // - X1 for instance methods.
+
+  // The "sret" attribute identifies indirect returns.
+  // The "inreg" attribute identifies non-aggregate types.
+  // The position of the "sret" attribute identifies instance/non-instance
+  // methods.
+  // "sret" on argument 0 means non-instance methods.
+  // "sret" on argument 1 means instance methods.
+  CCIfValNo<"1",
+    CCIfInReg<CCIfType<[i64],
+      CCIfSRet<CCIfType<[i64], CCAssignToRegWithShadow<[X1], [W1]>>>>>>,
+
+  CCIfInReg<CCIfType<[i64],
+    CCIfSRet<CCIfType<[i64], CCAssignToRegWithShadow<[X0], [W0]>>>>>,
+
   CCIfSRet<CCIfType<[i64], CCAssignToRegWithShadow<[X8], [W8]>>>,
 
   // Put ByVal arguments directly on the stack. Minimum size and alignment of a
   // slot is 64-bit.
   CCIfByVal<CCPassByVal<8, 8>>,
 
   // The 'nest' parameter, if any, is passed in X18.
   // Darwin uses X18 as the platform register and hence 'nest' isn't currently
   // supported there.
   CCIfNest<CCAssignToReg<[X18]>>,
 
   // Pass SwiftSelf in a callee saved register.
   CCIfSwiftSelf<CCIfType<[i64], CCAssignToRegWithShadow<[X20], [W20]>>>,
 
   // A SwiftError is passed in X21.
   CCIfSwiftError<CCIfType<[i64], CCAssignToRegWithShadow<[X21], [W21]>>>,
 

lib/Target/AArch64/AArch64ISelLowering.cpp

if (MFI.hasMustTailInVarArgFunc()) {
       // Compute the set of forwarded registers. The rest are scratch.
       SmallVectorImpl<ForwardedRegister> &Forwards =
                                        FuncInfo->getForwardedMustTailRegParms();
       CCInfo.analyzeMustTailForwardedRegisters(Forwards, RegParmTypes,
                                                CC_AArch64_AAPCS);
     }
   }
 
+  // On Windows, SRet pointers must be returned, so record the pointer in a
+  // virtual register at the start of the function so it can be returned in the
+  // epilogue.
+  for (unsigned I = 0, E = Ins.size(); I != E; ++I) {
+    if (Ins[I].Flags.isSRet()) {
+      unsigned Reg = FuncInfo->getSRetReturnReg();
+      if (!Reg) {
+        MVT PtrTy = getPointerTy(DAG.getDataLayout());
+        Reg = MF.getRegInfo().createVirtualRegister(getRegClassFor(PtrTy));
+        FuncInfo->setSRetReturnReg(Reg);
+      }
+
+      SDValue Copy = DAG.getCopyToReg(DAG.getEntryNode(), DL, Reg, InVals[I]);
+      Chain = DAG.getNode(ISD::TokenFactor, DL, MVT::Other, Copy, Chain);
+      break;
+    }
+  }
+
   unsigned StackArgSize = CCInfo.getNextStackOffset();
   bool TailCallOpt = MF.getTarget().Options.GuaranteedTailCallOpt;
   if (DoesCalleeRestoreStack(CallConv, TailCallOpt)) {
     // This is a non-standard ABI so by fiat I say we're allowed to make full
     // use of the stack area to be popped, which must be aligned to 16 bytes in
     // any case:
     StackArgSize = alignTo(StackArgSize, 16);
 
bool AArch64TargetLowering::isEligibleForTailCallOptimization(
     const SmallVectorImpl<ISD::InputArg> &Ins, SelectionDAG &DAG) const {
   if (!mayTailCallThisCC(CalleeCC))
     return false;
 
   MachineFunction &MF = DAG.getMachineFunction();
   const Function &CallerF = MF.getFunction();
   CallingConv::ID CallerCC = CallerF.getCallingConv();
   bool CCMatch = CallerCC == CalleeCC;
 
   // Byval parameters hand the function a pointer directly into the stack area
   // we want to reuse during a tail call. Working around this *is* possible (see
   // X86) but less efficient and uglier in LowerCall.
   for (Function::const_arg_iterator i = CallerF.arg_begin(),
                                     e = CallerF.arg_end();
        i != e; ++i)
     if (i->hasByValAttr())
       return false;
 
   if (getTargetMachine().Options.GuaranteedTailCallOpt)
     return canGuaranteeTCO(CalleeCC) && CCMatch;
 
   // Externally-defined functions with weak linkage should not be
   // tail-called on AArch64 when the OS does not support dynamic
   // pre-emption of symbols, as the AAELF spec requires normal calls
   // to undefined weak functions to be replaced with a NOP or jump to the
   // next instruction. The behaviour of branch instructions in this
   // situation (as used for tail calls) is implementation-defined, so we
   // cannot rely on the linker replacing the tail call with a return.
   if (GlobalAddressSDNode *G = dyn_cast<GlobalAddressSDNode>(Callee)) {
     const GlobalValue *GV = G->getGlobal();
     const Triple &TT = getTargetMachine().getTargetTriple();

 }
 
 SDValue
 AArch64TargetLowering::LowerReturn(SDValue Chain, CallingConv::ID CallConv,
                                    bool isVarArg,
                                    const SmallVectorImpl<ISD::OutputArg> &Outs,
                                    const SmallVectorImpl<SDValue> &OutVals,
                                    const SDLoc &DL, SelectionDAG &DAG) const {
+  auto &MF = DAG.getMachineFunction();
+  auto *FuncInfo = MF.getInfo<AArch64FunctionInfo>();
+
   CCAssignFn *RetCC = CallConv == CallingConv::WebKit_JS
                           ? RetCC_AArch64_WebKit_JS
                           : RetCC_AArch64_AAPCS;
   SmallVector<CCValAssign, 16> RVLocs;
   CCState CCInfo(CallConv, isVarArg, DAG.getMachineFunction(), RVLocs,
                  *DAG.getContext());
   CCInfo.AnalyzeReturn(Outs, RetCC);
 
case CCValAssign::BCvt:
       Arg = DAG.getNode(ISD::BITCAST, DL, VA.getLocVT(), Arg);
       break;
     }
 
     Chain = DAG.getCopyToReg(Chain, DL, VA.getLocReg(), Arg, Flag);
     Flag = Chain.getValue(1);
     RetOps.push_back(DAG.getRegister(VA.getLocReg(), VA.getLocVT()));
   }
+
+  // Windows AArch64 ABIs require that for returning structs by value we copy
+  // the sret argument into X0 for the return.
+  // We saved the argument into a virtual register in the entry block,
+  // so now we copy the value out and into X0.
+  if (unsigned SRetReg = FuncInfo->getSRetReturnReg()) {
+    SDValue Val = DAG.getCopyFromReg(RetOps[0], DL, SRetReg,
+                                     getPointerTy(MF.getDataLayout()));
+
+    unsigned RetValReg = AArch64::X0;
+    Chain = DAG.getCopyToReg(Chain, DL, RetValReg, Val, Flag);
+    Flag = Chain.getValue(1);
+
+    RetOps.push_back(
+      DAG.getRegister(RetValReg, getPointerTy(DAG.getDataLayout())));
+  }
+
   const AArch64RegisterInfo *TRI = Subtarget->getRegisterInfo();
   const MCPhysReg *I =
       TRI->getCalleeSavedRegsViaCopy(&DAG.getMachineFunction());
   if (I) {
     for (; *I; ++I) {
       if (AArch64::GPR64RegClass.contains(*I))
         RetOps.push_back(DAG.getRegister(*I, MVT::i64));
       else if (AArch64::FPR64RegClass.contains(*I))

lib/Target/AArch64/AArch64MachineFunctionInfo.h

class AArch64FunctionInfo final : public MachineFunctionInfo {
   /// True when the stack gets realigned dynamically because the size of stack
   /// frame is unknown at compile time. e.g., in case of VLAs.
   bool StackRealigned = false;
 
   /// True when the callee-save stack area has unused gaps that may be used for
   /// other stack allocations.
   bool CalleeSaveStackHasFreeSpace = false;
 
+  /// SRetReturnReg - sret lowering includes returning the value of the
+  /// returned struct in a register. This field holds the virtual register into
+  /// which the sret argument is passed.
+  unsigned SRetReturnReg = 0;
+
   /// Has a value when it is known whether or not the function uses a
   /// redzone, and no value otherwise.
   /// Initialized during frame lowering, unless the function has the noredzone
   /// attribute, in which case it is set to false at construction.
   Optional<bool> HasRedZone;
 
   /// ForwardedMustTailRegParms - A list of virtual and physical registers
   /// that must be forwarded to every musttail call.
public:
   void setVarArgsGPRSize(unsigned Size) { VarArgsGPRSize = Size; }
 
   int getVarArgsFPRIndex() const { return VarArgsFPRIndex; }
   void setVarArgsFPRIndex(int Index) { VarArgsFPRIndex = Index; }
 
   unsigned getVarArgsFPRSize() const { return VarArgsFPRSize; }
   void setVarArgsFPRSize(unsigned Size) { VarArgsFPRSize = Size; }
 
+  unsigned getSRetReturnReg() const { return SRetReturnReg; }
+  void setSRetReturnReg(unsigned Reg) { SRetReturnReg = Reg; }
+
   unsigned getJumpTableEntrySize(int Idx) const {
     auto It = JumpTableEntryInfo.find(Idx);
     if (It != JumpTableEntryInfo.end())
       return It->second.first;
     return 4;
   }
   MCSymbol *getJumpTableEntryPCRelSymbol(int Idx) const {
     return JumpTableEntryInfo.find(Idx)->second.second;