[llvm] [StatepointLowering] Handle struct return through stack (PR #157251)

[Abhishek Kaushik via llvm-commits]

https://github.com/abhishek-kaushik22 updated https://github.com/llvm/llvm-project/pull/157251

>From 60d53ae26eb3e04b66a1952de291a6194f0599f0 Mon Sep 17 00:00:00 2001
From: Abhishek Kaushik <abhishek.kaushik at intel.com>
Date: Sat, 6 Sep 2025 17:25:44 +0530
Subject: [PATCH 1/2] [Statepoint] Handle struct return through stack

Handle struct return from stack in `lowerCallFromStatepointLoweringInfo`

Patch by @e-kud
Fixes #74612
---
 .../SelectionDAG/StatepointLowering.cpp       |  8 +++-
 llvm/test/CodeGen/statepoint-struct-return.ll | 39 +++++++++++++++++++
 2 files changed, 45 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/CodeGen/statepoint-struct-return.ll

diff --git a/llvm/lib/CodeGen/SelectionDAG/StatepointLowering.cpp b/llvm/lib/CodeGen/SelectionDAG/StatepointLowering.cpp
index 46a5e44374e1c..befb27c2039fd 100644
--- a/llvm/lib/CodeGen/SelectionDAG/StatepointLowering.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/StatepointLowering.cpp
@@ -337,13 +337,17 @@ static std::pair<SDValue, SDNode *> lowerCallFromStatepointLoweringInfo(
   //
   // get_return_value can either be a sequence of CopyFromReg instructions
   // to grab the return value from the return register(s), or it can be a LOAD
-  // to load a value returned by reference via a stack slot.
+  // to load a value returned by reference via a stack slot, or it can be a
+  // struct returned by value through stack.
 
   if (CallEnd->getOpcode() == ISD::EH_LABEL)
     CallEnd = CallEnd->getOperand(0).getNode();
 
-  bool HasDef = !SI.CLI.RetTy->isVoidTy();
+  bool HasDef = !SI.CLI.RetTy->isVoidTy() || !SI.CLI.OutVals.empty();
   if (HasDef) {
+    if (CallEnd->getOpcode() == ISD::TokenFactor)
+      CallEnd = CallEnd->getOperand(0).getNode();
+
     if (CallEnd->getOpcode() == ISD::LOAD)
       CallEnd = CallEnd->getOperand(0).getNode();
     else
diff --git a/llvm/test/CodeGen/statepoint-struct-return.ll b/llvm/test/CodeGen/statepoint-struct-return.ll
new file mode 100644
index 0000000000000..36f5a6ce14473
--- /dev/null
+++ b/llvm/test/CodeGen/statepoint-struct-return.ll
@@ -0,0 +1,39 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc < %s -mtriple=x86_64-- | FileCheck %s
+
+%t = type { i32, i32, i32, i32 }
+
+define %t @foo() gc "statepoint-example" {
+; CHECK-LABEL: foo:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    pushq %rbx
+; CHECK-NEXT:    .cfi_def_cfa_offset 16
+; CHECK-NEXT:    subq $16, %rsp
+; CHECK-NEXT:    .cfi_def_cfa_offset 32
+; CHECK-NEXT:    .cfi_offset %rbx, -16
+; CHECK-NEXT:    movq %rdi, %rbx
+; CHECK-NEXT:    movq %rsp, %rdi
+; CHECK-NEXT:    callq bar at PLT
+; CHECK-NEXT:  .Ltmp0:
+; CHECK-NEXT:    movl {{[0-9]+}}(%rsp), %eax
+; CHECK-NEXT:    movl {{[0-9]+}}(%rsp), %ecx
+; CHECK-NEXT:    movl (%rsp), %edx
+; CHECK-NEXT:    movl {{[0-9]+}}(%rsp), %esi
+; CHECK-NEXT:    movl %edx, (%rbx)
+; CHECK-NEXT:    movl %ecx, 8(%rbx)
+; CHECK-NEXT:    movl %eax, 12(%rbx)
+; CHECK-NEXT:    movl %esi, 4(%rbx)
+; CHECK-NEXT:    movq %rbx, %rax
+; CHECK-NEXT:    addq $16, %rsp
+; CHECK-NEXT:    .cfi_def_cfa_offset 16
+; CHECK-NEXT:    popq %rbx
+; CHECK-NEXT:    .cfi_def_cfa_offset 8
+; CHECK-NEXT:    retq
+  %statepoint_token = call token (i64, i32, ptr, i32, i32, ...) @llvm.experimental.gc.statepoint.p0(i64 2882400000, i32 0, ptr elementtype(%t ()) @bar, i32 0, i32 0, i32 0, i32 0)
+  %res = call %t @llvm.experimental.gc.result.s_zeros(token %statepoint_token)
+  ret %t %res
+}
+
+declare %t @bar()
+declare token @llvm.experimental.gc.statepoint.p0(i64, i32, ptr, i32, i32, ...)
+declare %t @llvm.experimental.gc.result.s_zeros(token)

>From 3bd0d7483105e65cbb65502850b82daa02cf4233 Mon Sep 17 00:00:00 2001
From: Abhishek Kaushik <abhishek.kaushik at intel.com>
Date: Sat, 6 Sep 2025 22:13:17 +0530
Subject: [PATCH 2/2] Add comment and move to correct folder

---
 llvm/test/CodeGen/{ => X86}/statepoint-struct-return.ll | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)
 rename llvm/test/CodeGen/{ => X86}/statepoint-struct-return.ll (79%)

diff --git a/llvm/test/CodeGen/statepoint-struct-return.ll b/llvm/test/CodeGen/X86/statepoint-struct-return.ll
similarity index 79%
rename from llvm/test/CodeGen/statepoint-struct-return.ll
rename to llvm/test/CodeGen/X86/statepoint-struct-return.ll
index 36f5a6ce14473..77e7e6441f8af 100644
--- a/llvm/test/CodeGen/statepoint-struct-return.ll
+++ b/llvm/test/CodeGen/X86/statepoint-struct-return.ll
@@ -1,10 +1,15 @@
 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
 ; RUN: llc < %s -mtriple=x86_64-- | FileCheck %s
 
+; This test checks that statepoint lowering works when the call returns a
+; struct using a stack slot (sret). Previously, this case caused an assertion
+; because the lowering code assumed the chain ended at CALLSEQ_END. The fix
+; handles TokenFactor/LOAD cases so the chain is found correctly.
+
 %t = type { i32, i32, i32, i32 }
 
-define %t @foo() gc "statepoint-example" {
-; CHECK-LABEL: foo:
+define %t @PR74612() gc "statepoint-example" {
+; CHECK-LABEL: PR74612:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    pushq %rbx
 ; CHECK-NEXT:    .cfi_def_cfa_offset 16