[llvm] [win/arm64] Enable tail call with inreg arguments when possible (PR #134671)

[Hans Wennborg via llvm-commits]

https://github.com/zmodem created https://github.com/llvm/llvm-project/pull/134671

Tail calls were disabled from callers with inreg parameters in 5dc8aeb with a fixme to check if the callee also takes an inreg parameter.

The issue is that inreg parameters (which are passed in x0 or x1 for free and member functions respectively) are supposed to be returned (in x0) at the end of the function. In case of a tail call, that means the callee needs to return the same value as the caller would.

We can check for that case, and it's not as niche as it sounds, as that's how Clang will lower one function with an sret return value calling another, such as:

```
struct T { int x; };
struct S {
    T foo();
    T bar();
};
T S::foo() { return bar(); } // foo's sret argument will get passed directly to bar
```

Fixes #133098

>From 62a35c522ddbf9e28c2b70ffcb7c84bbbc2b79b8 Mon Sep 17 00:00:00 2001
From: Hans Wennborg <hans at chromium.org>
Date: Wed, 2 Apr 2025 16:04:50 +0200
Subject: [PATCH] [win/arm64] Enable tail call with inreg arguments when
 possible

---
 .../Target/AArch64/AArch64ISelLowering.cpp    | 20 ++++++++++------
 .../CodeGen/AArch64/arm64-windows-tailcall.ll | 24 +++++++++++++++++++
 2 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 0534d2d546325..90ae1a7e57cdb 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -8703,13 +8703,19 @@ bool AArch64TargetLowering::isEligibleForTailCallOptimization(
       return false;
 
     // On Windows, "inreg" attributes signify non-aggregate indirect returns.
-    // In this case, it is necessary to save/restore X0 in the callee. Tail
-    // call opt interferes with this. So we disable tail call opt when the
-    // caller has an argument with "inreg" attribute.
-
-    // FIXME: Check whether the callee also has an "inreg" argument.
-    if (i->hasInRegAttr())
-      return false;
+    // In this case, it is necessary to save X0/X1 in the callee and return it
+    // in X0. Tail call opt may interfere with this, so we disable tail call
+    // opt when the caller has an "inreg" attribute -- except if the callee
+    // also has that attribute on the same argument, and the same value is
+    // passed.
+    if (i->hasInRegAttr()) {
+      unsigned ArgNum = i - CallerF.arg_begin();
+      if (!CLI.CB || CLI.CB->arg_size() <= ArgNum ||
+          !CLI.CB->getParamAttr(ArgNum, Attribute::InReg).isValid() ||
+          CLI.CB->getArgOperand(ArgNum) != i) {
+        return false;
+      }
+    }
   }
 
   if (canGuaranteeTCO(CalleeCC, getTargetMachine().Options.GuaranteedTailCallOpt))
diff --git a/llvm/test/CodeGen/AArch64/arm64-windows-tailcall.ll b/llvm/test/CodeGen/AArch64/arm64-windows-tailcall.ll
index 55799d0dcb2d2..64491abd1bad7 100644
--- a/llvm/test/CodeGen/AArch64/arm64-windows-tailcall.ll
+++ b/llvm/test/CodeGen/AArch64/arm64-windows-tailcall.ll
@@ -16,3 +16,27 @@ entry:
 }
 
 declare dso_local void @"?foo"(ptr dereferenceable(4))
+
+
+declare void @inreg_callee(ptr, ptr inreg sret(%class.C))
+
+define void @inreg_caller_1(ptr %a, ptr inreg sret(%class.C) %b) {
+; A different value is passed to the inreg parameter, so tail call is not possible.
+; CHECK-LABEL: inreg_caller_1
+; CHECK: mov x19, x1
+; CHECK: bl inreg_callee
+; CHECK: mov x0, x19
+
+  tail call void @inreg_callee(ptr %b, ptr inreg sret(%class.C) %a)
+  ret void
+}
+
+define void @inreg_caller_2(ptr %a, ptr inreg sret(%class.C) %b) {
+; The inreg attribute and value lines up between caller and callee, so it can
+; be tail called.
+; CHECK-LABEL: inreg_caller_2
+; CHECK: b inreg_callee
+
+  tail call void @inreg_callee(ptr %a, ptr inreg sret(%class.C) %b)
+  ret void
+}