From 5abeaa821f53403a88cfb5724583d688e648a2e3 Mon Sep 17 00:00:00 2001
From: joboet <jonasboettiger@icloud.com>
Date: Fri, 4 Apr 2025 19:47:44 +0200
Subject: [PATCH] rustc_codegen_llvm: use `threadlocal.address` intrinsic to
 access TLS

---
 compiler/rustc_codegen_llvm/src/builder.rs |  9 +++++++--
 compiler/rustc_codegen_llvm/src/context.rs |  1 +
 tests/codegen/thread-local.rs              | 16 ++++++++++------
 3 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/compiler/rustc_codegen_llvm/src/builder.rs b/compiler/rustc_codegen_llvm/src/builder.rs
index 297f104d124..caf8e044e8b 100644
--- a/compiler/rustc_codegen_llvm/src/builder.rs
+++ b/compiler/rustc_codegen_llvm/src/builder.rs
@@ -1454,9 +1454,14 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
 impl<'ll> StaticBuilderMethods for Builder<'_, 'll, '_> {
     fn get_static(&mut self, def_id: DefId) -> &'ll Value {
         // Forward to the `get_static` method of `CodegenCx`
-        let s = self.cx().get_static(def_id);
+        let g = self.cx().get_static(def_id);
         // Cast to default address space if globals are in a different addrspace
-        self.cx().const_pointercast(s, self.type_ptr())
+        let g = self.cx().const_pointercast(g, self.type_ptr());
+        if self.cx().tcx.is_thread_local_static(def_id) {
+            self.call_intrinsic("llvm.threadlocal.address", &[g])
+        } else {
+            g
+        }
     }
 }
 
diff --git a/compiler/rustc_codegen_llvm/src/context.rs b/compiler/rustc_codegen_llvm/src/context.rs
index f7b096ff976..867d48f41ce 100644
--- a/compiler/rustc_codegen_llvm/src/context.rs
+++ b/compiler/rustc_codegen_llvm/src/context.rs
@@ -1198,6 +1198,7 @@ impl<'ll> CodegenCx<'ll, '_> {
         }
 
         ifn!("llvm.ptrmask", fn(ptr, t_isize) -> ptr);
+        ifn!("llvm.threadlocal.address", fn(ptr) -> ptr);
 
         None
     }
diff --git a/tests/codegen/thread-local.rs b/tests/codegen/thread-local.rs
index 9ce34473b91..41df8c9be1b 100644
--- a/tests/codegen/thread-local.rs
+++ b/tests/codegen/thread-local.rs
@@ -14,13 +14,14 @@ use std::cell::Cell;
 
 thread_local!(static A: Cell<u32> = const { Cell::new(1) });
 
-// CHECK: [[TLS_AUX:@.+]] = external thread_local local_unnamed_addr global i64
-// CHECK: [[TLS:@.+]] = internal thread_local unnamed_addr global
+// CHECK: [[TLS_AUX:@.+]] = external thread_local{{.*}} global i64
+// CHECK: [[TLS:@.+]] = internal thread_local{{.*}} global
 
 // CHECK-LABEL: @get
 #[no_mangle]
 fn get() -> u32 {
-    // CHECK: [[RET_0:%.+]] = load i32, {{.*}}[[TLS]]{{.*}}
+    // CHECK: [[PTR:%.+]] = tail call {{.*}} ptr @llvm.threadlocal.address.p0(ptr [[TLS]])
+    // CHECK-NEXT: [[RET_0:%.+]] = load i32, ptr [[PTR]]
     // CHECK-NEXT: ret i32 [[RET_0]]
     A.with(|a| a.get())
 }
@@ -28,7 +29,8 @@ fn get() -> u32 {
 // CHECK-LABEL: @set
 #[no_mangle]
 fn set(v: u32) {
-    // CHECK: store i32 %0, {{.*}}[[TLS]]{{.*}}
+    // CHECK: [[PTR:%.+]] = tail call {{.*}} ptr @llvm.threadlocal.address.p0(ptr [[TLS]])
+    // CHECK-NEXT: store i32 %0, ptr [[PTR]]
     // CHECK-NEXT: ret void
     A.with(|a| a.set(v))
 }
@@ -36,7 +38,8 @@ fn set(v: u32) {
 // CHECK-LABEL: @get_aux
 #[no_mangle]
 fn get_aux() -> u64 {
-    // CHECK: [[RET_1:%.+]] = load i64, {{.*}}[[TLS_AUX]]
+    // CHECK: [[PTR:%.+]] = tail call {{.*}} ptr @llvm.threadlocal.address.p0(ptr [[TLS_AUX]])
+    // CHECK-NEXT: [[RET_1:%.+]] = load i64, ptr [[PTR]]
     // CHECK-NEXT: ret i64 [[RET_1]]
     aux::A.with(|a| a.get())
 }
@@ -44,7 +47,8 @@ fn get_aux() -> u64 {
 // CHECK-LABEL: @set_aux
 #[no_mangle]
 fn set_aux(v: u64) {
-    // CHECK: store i64 %0, {{.*}}[[TLS_AUX]]
+    // CHECK: [[PTR:%.+]] = tail call {{.*}} ptr @llvm.threadlocal.address.p0(ptr [[TLS_AUX]])
+    // CHECK-NEXT: store i64 %0, ptr [[PTR]]
     // CHECK-NEXT: ret void
     aux::A.with(|a| a.set(v))
 }