[llvm] [StackColoring] Add test for stack-coloring and setjmp (PR #199959)

[Mikołaj Piróg via llvm-commits]

https://github.com/mikolaj-pirog updated https://github.com/llvm/llvm-project/pull/199959

>From 45f9007a388dc9252f3f76c3fe623163e5ed35a0 Mon Sep 17 00:00:00 2001
From: "Pirog, Mikolaj Maciej" <mikolaj.maciej.pirog at intel.com>
Date: Wed, 27 May 2026 13:04:33 +0200
Subject: [PATCH 1/2] Add test for stack-coloring and setjmp

---
 .../test/CodeGen/X86/stack-coloring-setjmp.ll | 76 +++++++++++++++++++
 1 file changed, 76 insertions(+)
 create mode 100644 llvm/test/CodeGen/X86/stack-coloring-setjmp.ll

diff --git a/llvm/test/CodeGen/X86/stack-coloring-setjmp.ll b/llvm/test/CodeGen/X86/stack-coloring-setjmp.ll
new file mode 100644
index 0000000000000..bd7a13900b22b
--- /dev/null
+++ b/llvm/test/CodeGen/X86/stack-coloring-setjmp.ll
@@ -0,0 +1,76 @@
+; RUN: llc -mtriple=x86_64-linux -no-stack-coloring=false -debug-only=stack-coloring < %s -o /dev/null 2>&1 | FileCheck %s
+; REQUIRES: asserts
+
+; Test that volatile stack slots accessed after setjmp are not merged.
+; Volatile variables must retain their values across longjmp, so their
+; stack slots cannot be reused even if their lifetimes don't overlap.
+
+declare i32 @setjmp(ptr) returns_twice
+declare void @baz(ptr)
+declare void @llvm.lifetime.start.p0(i64 immarg, ptr nocapture)
+declare void @llvm.lifetime.end.p0(i64 immarg, ptr nocapture)
+declare dso_local void @stash(ptr noundef, ptr noundef) local_unnamed_addr 
+declare dso_local void @use(ptr noundef) local_unnamed_addr 
+
+; CHECK-LABEL: setjmp_test
+; CHECK: Conservative slots : { 1 1 }
+; CHECK: Merge 0 slots
+
+define void @setjmp_test(ptr %jump_buffer) {
+entry:
+  %foo = alloca i32, align 4
+  %bar = alloca [100 x i32], align 4
+
+  call void @llvm.lifetime.start.p0(i64 4, ptr %foo)
+  call void @llvm.lifetime.start.p0(i64 400, ptr %bar)
+
+  %setjmp_result = call i32 @setjmp(ptr %jump_buffer)
+  %cmp = icmp eq i32 %setjmp_result, 0
+  br i1 %cmp, label %after_setjmp, label %continue
+
+after_setjmp:
+  store volatile i32 100, ptr %foo, align 4
+  br label %exit
+
+continue:
+  store i32 100, ptr %bar, align 4
+  call void @baz(ptr %bar)
+  call void @llvm.lifetime.end.p0(i64 4, ptr %foo)
+  call void @llvm.lifetime.end.p0(i64 400, ptr %bar)
+  br label %exit
+
+exit:
+  ret void
+}
+
+; CHECK-LABEL: setjmp_test_2
+; CHECK: Conservative slots : { 1 1 }
+; CHECK: Merge 0 slots
+
+
+%struct.T = type { [100 x i32] }
+
+define dso_local void @setjmp_test_2(ptr %jump_buffer) local_unnamed_addr  {
+entry:
+  %test1 = alloca %struct.T, align 4
+  %test2 = alloca %struct.T, align 4
+  call void @llvm.lifetime.start.p0(i64 1, ptr %test1)
+  %call = call i32 @setjmp(ptr %jump_buffer)
+  %cmp = icmp eq i32 %call, 0
+  br i1 %cmp, label %if.then, label %if.else
+
+if.then:
+  call void @stash(ptr %jump_buffer, ptr %test1) 
+  br label %if.end
+
+if.else:
+  call void @llvm.lifetime.start.p0(i64 2, ptr %test2) 
+  call void @use(ptr %test2) 
+  call void @llvm.lifetime.end.p0(i64 2, ptr %test2) 
+  br label %if.end
+
+if.end:
+  call void @llvm.lifetime.end.p0(i64 1, ptr %test1) 
+  ret void
+}
+

>From e5ec9879c29b076d8c2cd46f142de432d1291c18 Mon Sep 17 00:00:00 2001
From: "Pirog, Mikolaj Maciej" <mikolaj.maciej.pirog at intel.com>
Date: Tue, 9 Jun 2026 15:44:03 +0200
Subject: [PATCH 2/2] Review suggestions

---
 .../test/CodeGen/X86/stack-coloring-setjmp.ll | 20 +++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/llvm/test/CodeGen/X86/stack-coloring-setjmp.ll b/llvm/test/CodeGen/X86/stack-coloring-setjmp.ll
index bd7a13900b22b..226afe7827fe4 100644
--- a/llvm/test/CodeGen/X86/stack-coloring-setjmp.ll
+++ b/llvm/test/CodeGen/X86/stack-coloring-setjmp.ll
@@ -7,8 +7,8 @@
 
 declare i32 @setjmp(ptr) returns_twice
 declare void @baz(ptr)
-declare void @llvm.lifetime.start.p0(i64 immarg, ptr nocapture)
-declare void @llvm.lifetime.end.p0(i64 immarg, ptr nocapture)
+declare void @llvm.lifetime.start.p0(ptr nocapture)
+declare void @llvm.lifetime.end.p0(ptr nocapture)
 declare dso_local void @stash(ptr noundef, ptr noundef) local_unnamed_addr 
 declare dso_local void @use(ptr noundef) local_unnamed_addr 
 
@@ -21,8 +21,8 @@ entry:
   %foo = alloca i32, align 4
   %bar = alloca [100 x i32], align 4
 
-  call void @llvm.lifetime.start.p0(i64 4, ptr %foo)
-  call void @llvm.lifetime.start.p0(i64 400, ptr %bar)
+  call void @llvm.lifetime.start.p0(ptr %foo)
+  call void @llvm.lifetime.start.p0(ptr %bar)
 
   %setjmp_result = call i32 @setjmp(ptr %jump_buffer)
   %cmp = icmp eq i32 %setjmp_result, 0
@@ -35,8 +35,8 @@ after_setjmp:
 continue:
   store i32 100, ptr %bar, align 4
   call void @baz(ptr %bar)
-  call void @llvm.lifetime.end.p0(i64 4, ptr %foo)
-  call void @llvm.lifetime.end.p0(i64 400, ptr %bar)
+  call void @llvm.lifetime.end.p0(ptr %foo)
+  call void @llvm.lifetime.end.p0(ptr %bar)
   br label %exit
 
 exit:
@@ -54,7 +54,7 @@ define dso_local void @setjmp_test_2(ptr %jump_buffer) local_unnamed_addr  {
 entry:
   %test1 = alloca %struct.T, align 4
   %test2 = alloca %struct.T, align 4
-  call void @llvm.lifetime.start.p0(i64 1, ptr %test1)
+  call void @llvm.lifetime.start.p0(ptr %test1)
   %call = call i32 @setjmp(ptr %jump_buffer)
   %cmp = icmp eq i32 %call, 0
   br i1 %cmp, label %if.then, label %if.else
@@ -64,13 +64,13 @@ if.then:
   br label %if.end
 
 if.else:
-  call void @llvm.lifetime.start.p0(i64 2, ptr %test2) 
+  call void @llvm.lifetime.start.p0(ptr %test2) 
   call void @use(ptr %test2) 
-  call void @llvm.lifetime.end.p0(i64 2, ptr %test2) 
+  call void @llvm.lifetime.end.p0(ptr %test2) 
   br label %if.end
 
 if.end:
-  call void @llvm.lifetime.end.p0(i64 1, ptr %test1) 
+  call void @llvm.lifetime.end.p0(ptr %test1) 
   ret void
 }