[StatepointLowering] Handle struct return through stack

[abhishek-kaushik22]

Handle struct return from stack in lowerCallFromStatepointLoweringInfo

Patch by @e-kud
Fixes #74612

[llvmbot]

@llvm/pr-subscribers-backend-x86

@llvm/pr-subscribers-llvm-selectiondag

Author: Abhishek Kaushik (abhishek-kaushik22)

Changes

Handle struct return from stack in lowerCallFromStatepointLoweringInfo

Patch by @e-kud
Fixes #74612

---

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

2 Files Affected:

• (modified) llvm/lib/CodeGen/SelectionDAG/StatepointLowering.cpp (+6-2)
• (added) llvm/test/CodeGen/statepoint-struct-return.ll (+39)

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@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)

[arsenm]

Suggested change

	bool HasDef = !SI.CLI.RetTy->isVoidTy() || !SI.CLI.OutVals.empty();
	bool HasDef = !SI.CLI.OutVals.empty();

Should be redundant?

[arsenm]

This is making an assumption about the ordering of the TokenFactor inputs, don't do that? You should have an explicit reference to the call end direct from lowerInvokable?