From df37e2211ea34f128a53912835bcbc96147a1408 Mon Sep 17 00:00:00 2001
From: Lukas Sommer <lukas.sommer@codeplay.com>
Date: Tue, 30 May 2023 06:44:31 +0000
Subject: [PATCH] [mlir][llvm] Add thread-local address intrinsic

Add the `llvm.threadlocal.address` intrinsic to the LLVM dialect.

Signed-off-by: Lukas Sommer <lukas.sommer@codeplay.com>

Reviewed By: gysit, zero9178

Differential Revision: https://reviews.llvm.org/D151566
---
 mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td | 5 +++++
 mlir/test/Target/LLVMIR/Import/intrinsic.ll          | 8 ++++++++
 mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir       | 8 ++++++++
 3 files changed, 21 insertions(+)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td
index a409223ade15..e032bcc47b37 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td
@@ -324,6 +324,11 @@ def LLVM_ExpectWithProbabilityOp
   let assemblyFormat = "$val `,` $expected `,` $prob attr-dict `:` type($val)";
 }
 
+def LLVM_ThreadlocalAddressOp : LLVM_OneResultIntrOp<"threadlocal.address", [],
+                                [0], [Pure]> {
+  let arguments = (ins LLVM_AnyPointer:$global);
+}
+
 //
 // Coroutine intrinsics.
 //
diff --git a/mlir/test/Target/LLVMIR/Import/intrinsic.ll b/mlir/test/Target/LLVMIR/Import/intrinsic.ll
index e9b361509d03..3cc8b2f6fb78 100644
--- a/mlir/test/Target/LLVMIR/Import/intrinsic.ll
+++ b/mlir/test/Target/LLVMIR/Import/intrinsic.ll
@@ -613,6 +613,13 @@ define void @expect_with_probability(i16 %0) {
   ret void
 }
 
+; CHECK-LABEL: llvm.func @threadlocal_test
+define void @threadlocal_test(ptr %0) {
+  ; CHECK: "llvm.intr.threadlocal.address"(%{{.*}}) : (!llvm.ptr) -> !llvm.ptr
+  %local = call ptr @llvm.threadlocal.address.p0(ptr %0)
+  ret void
+}
+
 ; CHECK-LABEL:  llvm.func @coro_id
 define void @coro_id(i32 %0, ptr %1) {
   ; CHECK: llvm.intr.coro.id %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : (i32, !llvm.ptr, !llvm.ptr, !llvm.ptr) -> !llvm.token
@@ -955,6 +962,7 @@ declare <8 x i32> @llvm.ushl.sat.v8i32(<8 x i32>, <8 x i32>)
 declare i1 @llvm.is.constant.i32(i32)
 declare i32 @llvm.expect.i32(i32, i32)
 declare i16 @llvm.expect.with.probability.i16(i16, i16, double immarg)
+declare nonnull ptr @llvm.threadlocal.address.p0(ptr nonnull)
 declare token @llvm.coro.id(i32, ptr readnone, ptr nocapture readonly, ptr)
 declare ptr @llvm.coro.begin(token, ptr writeonly)
 declare i64 @llvm.coro.size.i64()
diff --git a/mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir b/mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir
index ec619b9a9d36..3e07ac90784c 100644
--- a/mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir
+++ b/mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir
@@ -559,6 +559,13 @@ llvm.func @expect_with_probability(%arg0: i16) {
   llvm.return
 }
 
+// CHECK-LABEL: @threadlocal_test
+llvm.func @threadlocal_test(%arg0 : !llvm.ptr) {
+  // CHECK: call ptr @llvm.threadlocal.address.p0(ptr %{{.*}})
+  "llvm.intr.threadlocal.address"(%arg0) : (!llvm.ptr) -> !llvm.ptr
+  llvm.return
+}
+
 // CHECK-LABEL: @sadd_sat_test
 llvm.func @sadd_sat_test(%arg0: i32, %arg1: i32, %arg2: vector<8xi32>, %arg3: vector<8xi32>) {
   // CHECK: call i32 @llvm.sadd.sat.i32
@@ -996,6 +1003,7 @@ llvm.func @lifetime(%p: !llvm.ptr) {
 // CHECK-DAG: declare i1 @llvm.is.constant.i32(i32)
 // CHECK-DAG: declare i32 @llvm.expect.i32(i32, i32)
 // CHECK-DAG: declare i16 @llvm.expect.with.probability.i16(i16, i16, double immarg)
+// CHECK-DAG: declare nonnull ptr @llvm.threadlocal.address.p0(ptr nonnull)
 // CHECK-DAG: declare token @llvm.coro.id(i32, ptr readnone, ptr nocapture readonly, ptr)
 // CHECK-DAG: declare ptr @llvm.coro.begin(token, ptr writeonly)
 // CHECK-DAG: declare i64 @llvm.coro.size.i64()