[ARM64EC] Fix thunks for C++ methods returning structs. (#95876)

For C++ methods, the first argument is "this", and the second is the
sret argument, which needs to be returned indirectly. Add handling for
this case.

2 files changed, 64 insertions, 8 deletions

diff --git a/llvm/lib/Target/AArch64/AArch64Arm64ECCallLowering.cpp b/llvm/lib/Target/AArch64/AArch64Arm64ECCallLowering.cpp
index 0ec15d3..218201f 100644
--- a/llvm/lib/Target/AArch64/AArch64Arm64ECCallLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64Arm64ECCallLowering.cpp
@@ -209,27 +209,36 @@ void AArch64Arm64ECCallLowering::getThunkRetType(
 #endif
   if (T->isVoidTy()) {
     if (FT->getNumParams()) {
-      auto SRetAttr = AttrList.getParamAttr(0, Attribute::StructRet);
-      auto InRegAttr = AttrList.getParamAttr(0, Attribute::InReg);
-      if (SRetAttr.isValid() && InRegAttr.isValid()) {
+      Attribute SRetAttr0 = AttrList.getParamAttr(0, Attribute::StructRet);
+      Attribute InRegAttr0 = AttrList.getParamAttr(0, Attribute::InReg);
+      Attribute SRetAttr1, InRegAttr1;
+      if (FT->getNumParams() > 1) {
+        // Also check the second parameter (for class methods, the first
+        // parameter is "this", and the second parameter is the sret pointer.)
+        // It doesn't matter which one is sret.
+        SRetAttr1 = AttrList.getParamAttr(1, Attribute::StructRet);
+        InRegAttr1 = AttrList.getParamAttr(1, Attribute::InReg);
+      }
+      if ((SRetAttr0.isValid() && InRegAttr0.isValid()) ||
+          (SRetAttr1.isValid() && InRegAttr1.isValid())) {
         // sret+inreg indicates a call that returns a C++ class value. This is
         // actually equivalent to just passing and returning a void* pointer
-        // as the first argument. Translate it that way, instead of trying
-        // to model "inreg" in the thunk's calling convention, to simplify
-        // the rest of the code.
+        // as the first or second argument. Translate it that way, instead of
+        // trying to model "inreg" in the thunk's calling convention; this
+        // simplfies the rest of the code, and matches MSVC mangling.
         Out << "i8";
         Arm64RetTy = I64Ty;
         X64RetTy = I64Ty;
         return;
       }
-      if (SRetAttr.isValid()) {
+      if (SRetAttr0.isValid()) {
         // FIXME: Sanity-check the sret type; if it's an integer or pointer,
         // we'll get screwy mangling/codegen.
         // FIXME: For large struct types, mangle as an integer argument and
         // integer return, so we can reuse more thunks, instead of "m" syntax.
         // (MSVC mangles this case as an integer return with no argument, but
         // that's a miscompile.)
-        Type *SRetType = SRetAttr.getValueAsType();
+        Type *SRetType = SRetAttr0.getValueAsType();
         Align SRetAlign = AttrList.getParamAlignment(0).valueOrOne();
         Type *Arm64Ty, *X64Ty;
         canonicalizeThunkType(SRetType, SRetAlign, /*Ret*/ true, ArgSizeBytes,
diff --git a/llvm/test/CodeGen/AArch64/arm64ec-entry-thunks.ll b/llvm/test/CodeGen/AArch64/arm64ec-entry-thunks.ll
index e9556b9..0cf678f 100644
--- a/llvm/test/CodeGen/AArch64/arm64ec-entry-thunks.ll
+++ b/llvm/test/CodeGen/AArch64/arm64ec-entry-thunks.ll
@@ -444,6 +444,50 @@ define %T2 @simple_struct(%T1 %0, %T2 %1, %T3, %T4) nounwind {
   ret %T2 %1
 }
 
+define void @cxx_method(ptr noundef nonnull align 8 dereferenceable(8) %0, ptr dead_on_unwind inreg noalias writable sret(i64) align 8 %1) {
+; CHECK-LABEL:    .def    $ientry_thunk$cdecl$i8$i8i8;
+; CHECK: .section        .wowthk$aa,"xr",discard,$ientry_thunk$cdecl$i8$i8i8
+; CHECK:          // %bb.0:
+; CHECK-NEXT:     stp     q6, q7, [sp, #-176]!            // 32-byte Folded Spill
+; CHECK-NEXT:     .seh_save_any_reg_px    q6, 176
+; CHECK-NEXT:     stp     q8, q9, [sp, #32]               // 32-byte Folded Spill
+; CHECK-NEXT:     .seh_save_any_reg_p     q8, 32
+; CHECK-NEXT:     stp     q10, q11, [sp, #64]             // 32-byte Folded Spill
+; CHECK-NEXT:     .seh_save_any_reg_p     q10, 64
+; CHECK-NEXT:     stp     q12, q13, [sp, #96]             // 32-byte Folded Spill
+; CHECK-NEXT:     .seh_save_any_reg_p     q12, 96
+; CHECK-NEXT:     stp     q14, q15, [sp, #128]            // 32-byte Folded Spill
+; CHECK-NEXT:     .seh_save_any_reg_p     q14, 128
+; CHECK-NEXT:     stp     x29, x30, [sp, #160]            // 16-byte Folded Spill
+; CHECK-NEXT:     .seh_save_fplr  160
+; CHECK-NEXT:     add     x29, sp, #160
+; CHECK-NEXT:     .seh_add_fp     160
+; CHECK-NEXT:     .seh_endprologue
+; CHECK-NEXT:     blr     x9
+; CHECK-NEXT:     adrp    x8, __os_arm64x_dispatch_ret
+; CHECK-NEXT:     ldr     x1, [x8, :lo12:__os_arm64x_dispatch_ret]
+; CHECK-NEXT:     mov     x8, x0
+; CHECK-NEXT:     .seh_startepilogue
+; CHECK-NEXT:     ldp     x29, x30, [sp, #160]            // 16-byte Folded Reload
+; CHECK-NEXT:     .seh_save_fplr  160
+; CHECK-NEXT:     ldp     q14, q15, [sp, #128]            // 32-byte Folded Reload
+; CHECK-NEXT:     .seh_save_any_reg_p     q14, 128
+; CHECK-NEXT:     ldp     q12, q13, [sp, #96]             // 32-byte Folded Reload
+; CHECK-NEXT:     .seh_save_any_reg_p     q12, 96
+; CHECK-NEXT:     ldp     q10, q11, [sp, #64]             // 32-byte Folded Reload
+; CHECK-NEXT:     .seh_save_any_reg_p     q10, 64
+; CHECK-NEXT:     ldp     q8, q9, [sp, #32]               // 32-byte Folded Reload
+; CHECK-NEXT:     .seh_save_any_reg_p     q8, 32
+; CHECK-NEXT:     ldp     q6, q7, [sp], #176              // 32-byte Folded Reload
+; CHECK-NEXT:     .seh_save_any_reg_px    q6, 176
+; CHECK-NEXT:     .seh_endepilogue
+; CHECK-NEXT:     br      x1
+; CHECK-NEXT:     .seh_endfunclet
+; CHECK-NEXT:     .seh_endproc
+  ret void
+}
+
+
 ; Verify the hybrid bitmap
 ; CHECK-LABEL:    .section        .hybmp$x,"yi"
 ; CHECK-NEXT:     .symidx "#no_op"
@@ -476,3 +520,6 @@ define %T2 @simple_struct(%T1 %0, %T2 %1, %T3, %T4) nounwind {
 ; CHECK-NEXT:     .symidx "#simple_struct"
 ; CHECK-NEXT:     .symidx $ientry_thunk$cdecl$m8$i8m8m16m24
 ; CHECK-NEXT:     .word   1
+; CHECK-NEXT:     .symidx "#cxx_method"
+; CHECK-NEXT:     .symidx $ientry_thunk$cdecl$i8$i8i8
+; CHECK-NEXT:     .word   1