[llvm] [arm64ec] Fix missing sret return in Arm64EC entry thunks for large struct returns (PR #185452)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-platform-windows

Author: Daniel Paoliello (dpaoliello)

<details>
<summary>Changes</summary>

When an Arm64EC function returns a struct by value that is too large for x64's `RAX` (>8 bytes), the entry thunk synthesizes a hidden sret pointer parameter for the x64 side. However, this
parameter was never marked with the sret attribute, so ISel did not copy its value into `x8` (the Arm64EC mapping of `RAX`) on return. This caused the x64 caller to see a garbage pointer in `RAX` instead of the return buffer address.

The change adds the sret attribute to the thunk's synthesized pointer parameter, so that `LowerFormalArguments` saves it and `LowerReturn` restores it to `x8` before the tail call to `__os_arm64x_dispatch_ret`.

Fixes #<!-- -->185390

---
Full diff: https://github.com/llvm/llvm-project/pull/185452.diff


2 Files Affected:

- (modified) llvm/lib/Target/AArch64/AArch64Arm64ECCallLowering.cpp (+5) 
- (modified) llvm/test/CodeGen/AArch64/arm64ec-entry-thunks.ll (+2) 


``````````diff
diff --git a/llvm/lib/Target/AArch64/AArch64Arm64ECCallLowering.cpp b/llvm/lib/Target/AArch64/AArch64Arm64ECCallLowering.cpp
index c27a693ceecc1..866e2f9c4218c 100644
--- a/llvm/lib/Target/AArch64/AArch64Arm64ECCallLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64Arm64ECCallLowering.cpp
@@ -593,6 +593,11 @@ Function *AArch64Arm64ECCallLowering::buildEntryThunk(Function *F) {
 
   Value *RetVal = Call;
   if (TransformDirectToSRet) {
+    // The x64 side returns this value indirectly via a hidden pointer (sret).
+    // Mark the thunk's pointer arg with sret so that ISel saves it and copies
+    // it into x8 (RAX) on return, matching the x64 calling convention.
+    Thunk->addParamAttr(
+        1, Attribute::getWithStructRetType(M->getContext(), RetTy));
     IRB.CreateStore(RetVal, Thunk->getArg(1));
   } else if (X64RetType != RetTy) {
     Value *CastAlloca = IRB.CreateAlloca(X64RetType);
diff --git a/llvm/test/CodeGen/AArch64/arm64ec-entry-thunks.ll b/llvm/test/CodeGen/AArch64/arm64ec-entry-thunks.ll
index 2c1b735ffe28c..09fe884940c6a 100644
--- a/llvm/test/CodeGen/AArch64/arm64ec-entry-thunks.ll
+++ b/llvm/test/CodeGen/AArch64/arm64ec-entry-thunks.ll
@@ -359,6 +359,7 @@ define [3 x i64] @large_array([3 x i64] %0, [2 x double], [2 x [2 x i64]]) nounw
 ; CHECK-NEXT:     adrp    x8, __os_arm64x_dispatch_ret
 ; CHECK-NEXT:     str     x2, [x19, #16]
 ; CHECK-NEXT:     ldr     x0, [x8, :lo12:__os_arm64x_dispatch_ret]
+; CHECK-NEXT:     mov     x8, x19
 ; CHECK-NEXT:     .seh_startepilogue
 ; CHECK-NEXT:     ldp     x29, x30, [sp, #168]            // 16-byte Folded Reload
 ; CHECK-NEXT:     .seh_save_fplr  168
@@ -568,6 +569,7 @@ define <8 x i16> @large_vector(<8 x i16> %0) {
 ; CHECK-NEXT:     adrp	x8, __os_arm64x_dispatch_ret
 ; CHECK-NEXT:     str	q0, [x19]
 ; CHECK-NEXT:     ldr	x0, [x8, :lo12:__os_arm64x_dispatch_ret]
+; CHECK-NEXT:     mov	x8, x19
 ; CHECK-NEXT:     .seh_startepilogue
 ; CHECK-NEXT:     ldp	x29, x30, [sp, #168]            // 16-byte Folded Reload
 ; CHECK-NEXT:     .seh_save_fplr	168

``````````

</details>


https://github.com/llvm/llvm-project/pull/185452