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

[Hans Wennborg via llvm-commits]

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

>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 1/4] [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
+}

>From 8aff76cdc789a1a7939754fbf972003575e1eed5 Mon Sep 17 00:00:00 2001
From: Hans Wennborg <hans at chromium.org>
Date: Tue, 8 Apr 2025 09:25:55 +0200
Subject: [PATCH 2/4] add sret checks

---
 llvm/lib/Target/AArch64/AArch64ISelLowering.cpp | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 90ae1a7e57cdb..deb05b0dcb97c 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -8712,6 +8712,8 @@ bool AArch64TargetLowering::isEligibleForTailCallOptimization(
       unsigned ArgNum = i - CallerF.arg_begin();
       if (!CLI.CB || CLI.CB->arg_size() <= ArgNum ||
           !CLI.CB->getParamAttr(ArgNum, Attribute::InReg).isValid() ||
+          !i->hasStructRetAttr() ||
+          !CLI.CB->getParamAttr(ArgNum, Attribute::StructRet).isValid() ||
           CLI.CB->getArgOperand(ArgNum) != i) {
         return false;
       }

>From 1e66e3e5340bc511442ad79f15dd541957962166 Mon Sep 17 00:00:00 2001
From: Hans Wennborg <hans at chromium.org>
Date: Tue, 8 Apr 2025 09:30:01 +0200
Subject: [PATCH 3/4] typo

---
 llvm/test/CodeGen/AArch64/arm64-windows-tailcall.ll | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/test/CodeGen/AArch64/arm64-windows-tailcall.ll b/llvm/test/CodeGen/AArch64/arm64-windows-tailcall.ll
index 64491abd1bad7..cd0a77a280aec 100644
--- a/llvm/test/CodeGen/AArch64/arm64-windows-tailcall.ll
+++ b/llvm/test/CodeGen/AArch64/arm64-windows-tailcall.ll
@@ -32,7 +32,7 @@ define void @inreg_caller_1(ptr %a, ptr inreg sret(%class.C) %b) {
 }
 
 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
+; The inreg attribute and value line up between caller and callee, so it can
 ; be tail called.
 ; CHECK-LABEL: inreg_caller_2
 ; CHECK: b inreg_callee

>From a50dc332e7093af6e9eb9a6ffd2c751d6e52095d Mon Sep 17 00:00:00 2001
From: Hans Wennborg <hans at chromium.org>
Date: Tue, 8 Apr 2025 09:46:14 +0200
Subject: [PATCH 4/4] break up the if statement

---
 llvm/lib/Target/AArch64/AArch64ISelLowering.cpp | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index deb05b0dcb97c..4092680d3234f 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -8709,12 +8709,13 @@ bool AArch64TargetLowering::isEligibleForTailCallOptimization(
     // 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() ||
-          !i->hasStructRetAttr() ||
-          !CLI.CB->getParamAttr(ArgNum, Attribute::StructRet).isValid() ||
-          CLI.CB->getArgOperand(ArgNum) != i) {
+      unsigned ArgIdx = i - CallerF.arg_begin();
+      if (!CLI.CB || CLI.CB->arg_size() <= ArgIdx)
+        return false;
+      AttributeSet Attrs = CLI.CB->getParamAttributes(ArgIdx);
+      if (!Attrs.hasAttribute(Attribute::InReg) ||
+          !Attrs.hasAttribute(Attribute::StructRet) || !i->hasStructRetAttr() ||
+          CLI.CB->getArgOperand(ArgIdx) != i) {
         return false;
       }
     }