[mlir][llvm] Add thread-local address intrinsic

[LukasSommerTu]

Summary

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

Signed-off-by: Lukas Sommer <[e-mail]>

[zero9178]

mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td

316

What's the behaviour of the op if the pointer operand is null or not a thread local? The reference sadly does not state anything https://llvm.org/docs/LangRef.html#llvm-threadlocal-address-intrinsic but guessing by your use of nonnull below it is undefined for null?

If undefined, it should be marked NoMemoryEffects rather than Pure since it can't always be speculatively executed for all inputs.

[LukasSommerTu]

@zero9178 Thanks for your review! I addressed your comment inline.

mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td

316

The LLVM definition of the intrinsic has the property IntrSpeculatable, which is described as:

// This property indicates that the intrinsic is safe to speculate.

Given that, I assume it can be executed speculatively and it is OK to mark this as Pure. In LLVM itself, there's test-cases with and without nonnull.

[zero9178]

LGTM

mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td

315

Tiny nit

316

I see, thank you for double checking!

[LukasSommerTu]

Add "Op" suffix to intrinsic.

[LukasSommerTu]

@zero9178 Thanks for the review, I addressed the last nit. If you're happy with the update, could you please commit this for me, I don't have commit access.

[gysit]

LGTM!

mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td

def LLVM_ExpectWithProbabilityOp
   }];
   string mlirBuilder = [{
     $res = $_builder.create<LLVM::ExpectWithProbabilityOp>(
       $_location, $val, $expected, $_float_attr($prob));
   }];
   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.
 //
 
 def LLVM_CoroIdOp : LLVM_IntrOp<"coro.id", [], [], [], 1> {
   let arguments = (ins I32:$align,
                        LLVM_i8Ptr:$promise,
                        LLVM_i8Ptr:$coroaddr,

mlir/test/Target/LLVMIR/Import/intrinsic.ll

 ; CHECK-SAME:  %[[VAL:[a-zA-Z0-9]+]]
 define void @expect_with_probability(i16 %0) {
   ; CHECK:  %[[EXP:.+]] = llvm.mlir.constant(42 : i16) : i16
   ; CHECK:  llvm.intr.expect.with.probability %[[VAL]], %[[EXP]], 5.000000e-01 : i16
   %2 = call i16 @llvm.expect.with.probability.i16(i16 %0, i16 42, double 0.5)
   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
   %3 = call token @llvm.coro.id(i32 %0, ptr %1, ptr %1, ptr null)
   ret void
 }
 
 ; CHECK-LABEL:  llvm.func @coro_begin

 declare <8 x i32> @llvm.usub.sat.v8i32(<8 x i32>, <8 x i32>)
 declare i32 @llvm.sshl.sat.i32(i32, i32)
 declare <8 x i32> @llvm.sshl.sat.v8i32(<8 x i32>, <8 x i32>)
 declare i32 @llvm.ushl.sat.i32(i32, i32)
 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()
 declare i32 @llvm.coro.size.i32()
 declare i64 @llvm.coro.align.i64()
 declare i32 @llvm.coro.align.i32()
 declare token @llvm.coro.save(ptr)
 declare i8 @llvm.coro.suspend(token, i1)

mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir

 // CHECK-LABEL: @expect_with_probability
 llvm.func @expect_with_probability(%arg0: i16) {
   %0 = llvm.mlir.constant(42 : i16) : i16
   // CHECK: call i16 @llvm.expect.with.probability.i16(i16 %{{.*}}, i16 42, double 5.000000e-01)
   %1 = llvm.intr.expect.with.probability %arg0, %0, 5.000000e-01 : 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
   "llvm.intr.sadd.sat"(%arg0, %arg1) : (i32, i32) -> i32
   // CHECK: call <8 x i32> @llvm.sadd.sat.v8i32
   "llvm.intr.sadd.sat"(%arg2, %arg3) : (vector<8xi32>, vector<8xi32>) -> vector<8xi32>
   llvm.return
 }

 // CHECK-DAG: declare <8 x i32> @llvm.usub.sat.v8i32(<8 x i32>, <8 x i32>)
 // CHECK-DAG: declare i32 @llvm.sshl.sat.i32(i32, i32)
 // CHECK-DAG: declare <8 x i32> @llvm.sshl.sat.v8i32(<8 x i32>, <8 x i32>)
 // CHECK-DAG: declare i32 @llvm.ushl.sat.i32(i32, i32)
 // CHECK-DAG: declare <8 x i32> @llvm.ushl.sat.v8i32(<8 x i32>, <8 x i32>)
 // 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()
 // CHECK-DAG: declare i32 @llvm.coro.size.i32()
 // CHECK-DAG: declare token @llvm.coro.save(ptr)
 // CHECK-DAG: declare i8 @llvm.coro.suspend(token, i1)
 // CHECK-DAG: declare i1 @llvm.coro.end(ptr, i1)
 // CHECK-DAG: declare ptr @llvm.coro.free(token, ptr nocapture readonly)