[llvm] [StackColoring] Do not merge stack slots with `returns_twice` call-sites (PR #196542)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-x86

Author: Antonio Frighetto (antoniofrighetto)

<details>
<summary>Changes</summary>

Do not merge stack slots on disjoint paths if the function may call setjmp/sigsetjmp, as the current algorithm defaults to computing liveness analysis from the actual uses propagated through the CFG, rather than leveraging lifetime markers, thus making it unsound with `returns_twice` calls.

Fixes: https://github.com/llvm/llvm-project/issues/196468.

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


2 Files Affected:

- (modified) llvm/lib/CodeGen/StackColoring.cpp (+6) 
- (added) llvm/test/CodeGen/X86/sjlj-do-not-merge-stack-slots.ll (+50) 


``````````diff
diff --git a/llvm/lib/CodeGen/StackColoring.cpp b/llvm/lib/CodeGen/StackColoring.cpp
index 65f9ae5433195..2d708f72f15ed 100644
--- a/llvm/lib/CodeGen/StackColoring.cpp
+++ b/llvm/lib/CodeGen/StackColoring.cpp
@@ -1221,6 +1221,12 @@ bool StackColoring::run(MachineFunction &Func, bool OnlyRemoveMarkers) {
   if (!NumSlots)
     return false;
 
+  // Conservatively do not merge stack slots on disjoint paths if we have calls
+  // to setjmp/sigsetjmp, as longjmp may re-enter the function on a different
+  // path.
+  if (MF->exposesReturnsTwice())
+    return false;
+
   SmallVector<int, 8> SortedSlots;
   SortedSlots.reserve(NumSlots);
   Intervals.reserve(NumSlots);
diff --git a/llvm/test/CodeGen/X86/sjlj-do-not-merge-stack-slots.ll b/llvm/test/CodeGen/X86/sjlj-do-not-merge-stack-slots.ll
new file mode 100644
index 0000000000000..ef2da73bbbf5b
--- /dev/null
+++ b/llvm/test/CodeGen/X86/sjlj-do-not-merge-stack-slots.ll
@@ -0,0 +1,50 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 6
+; RUN: llc -mtriple=x86_64-- < %s | FileCheck %s
+
+define void @sjlj_do_not_merge_stack_slots() nounwind {
+; CHECK-LABEL: sjlj_do_not_merge_stack_slots:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    subq $40, %rsp
+; CHECK-NEXT:    leaq {{[0-9]+}}(%rsp), %rdi
+; CHECK-NEXT:    callq setjmp at PLT
+; CHECK-NEXT:    testl %eax, %eax
+; CHECK-NEXT:    je .LBB0_1
+; CHECK-NEXT:  # %bb.2: # %else
+; CHECK-NEXT:    leaq {{[0-9]+}}(%rsp), %rdi
+; CHECK-NEXT:    callq opaque at PLT
+; CHECK-NEXT:    addq $40, %rsp
+; CHECK-NEXT:    retq
+; CHECK-NEXT:  .LBB0_1: # %then
+; CHECK-NEXT:    leaq {{[0-9]+}}(%rsp), %rdi
+; CHECK-NEXT:    callq escape at PLT
+; CHECK-NEXT:    addq $40, %rsp
+; CHECK-NEXT:    retq
+entry:
+  %obj = alloca [8 x i8], align 4
+  %buf = alloca [12 x i8], align 4
+  %obj.1 = alloca [8 x i8], align 4
+  call void @llvm.lifetime.start.p0(ptr %obj)
+  call void @llvm.lifetime.start.p0(ptr %buf)
+  %rv = call i32 @setjmp(ptr %buf)
+  %c = icmp eq i32 %rv, 0
+  br i1 %c, label %then, label %else
+
+then:
+  call void @escape(ptr %obj)
+  br label %exit
+
+else:
+  call void @llvm.lifetime.start.p0(ptr %obj.1)
+  call void @opaque(ptr %obj.1)
+  call void @llvm.lifetime.end.p0(ptr %obj.1)
+  br label %exit
+
+exit:
+  call void @llvm.lifetime.end.p0(ptr %buf)
+  call void @llvm.lifetime.end.p0(ptr %obj)
+  ret void
+}
+
+declare i32 @setjmp(ptr) returns_twice
+declare void @escape(ptr)
+declare void @opaque(ptr)

``````````

</details>


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