GlobalsModRef, ValueTracking: Look through threadlocal.address intrinsic

Author: MatzeB

llvm/include/llvm/Analysis/ValueTracking.h

@@ -690,11 +690,11 @@ inline Value *getArgumentAliasingToReturnedPointer(CallBase *Call,
 bool isIntrinsicReturningPointerAliasingArgumentWithoutCapturing(
     const CallBase *Call, bool MustPreserveNullness);
 
-/// This method strips off any GEP address adjustments and pointer casts from
-/// the specified value, returning the original object being addressed. Note
-/// that the returned value has pointer type if the specified value does. If
-/// the MaxLookup value is non-zero, it limits the number of instructions to
-/// be stripped off.
+/// This method strips off any GEP address adjustments, pointer casts
+/// or `llvm.threadlocal.address` from the specified value \p V, returning the
+/// original object being addressed. Note that the returned value has pointer
+/// type if the specified value does. If the \p MaxLookup value is non-zero, it
+/// limits the number of instructions to be stripped off.
 const Value *getUnderlyingObject(const Value *V, unsigned MaxLookup = 6);
 inline Value *getUnderlyingObject(Value *V, unsigned MaxLookup = 6) {
   // Force const to avoid infinite recursion.

llvm/lib/Analysis/GlobalsModRef.cpp

@@ -344,6 +344,14 @@ bool GlobalsAAResult::AnalyzeUsesOfPointer(Value *V,
       if (AnalyzeUsesOfPointer(I, Readers, Writers, OkayStoreDest))
         return true;
     } else if (auto *Call = dyn_cast<CallBase>(I)) {
+      if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(I)) {
+        if (II->getIntrinsicID() == Intrinsic::threadlocal_address &&
+            V == II->getArgOperand(0)) {
+          if (AnalyzeUsesOfPointer(II, Readers, Writers))
+            return true;
+          continue;
+        }
+      }
       // Make sure that this is just the function being called, not that it is
       // passing into the function.
       if (Call->isDataOperand(&U)) {

llvm/lib/Analysis/ValueTracking.cpp

@@ -6255,6 +6255,10 @@ bool llvm::isIntrinsicReturningPointerAliasingArgumentWithoutCapturing(
     return true;
   case Intrinsic::ptrmask:
     return !MustPreserveNullness;
+  case Intrinsic::threadlocal_address:
+    // The underlying variable changes with thread ID. The Thread ID may change
+    // at coroutine suspend points.
+    return !Call->getParent()->getParent()->isPresplitCoroutine();
   default:
     return false;
   }

llvm/test/Analysis/GlobalsModRef/nonescaping-noalias.ll

@@ -22,6 +22,67 @@ entry:
   ret i32 %v
 }
 
+@g1_tls = internal thread_local global i32 0
+
+define i32 @test1_tls(ptr %param) {
+; Ensure that we can fold a store to a load of a global across a store to
+; a parameter when the global is non-escaping.
+;
+; CHECK-LABEL: define i32 @test1_tls(
+; CHECK-SAME: ptr [[PARAM:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[P:%.*]] = call ptr @llvm.threadlocal.address.p0(ptr @g1_tls)
+; CHECK-NEXT:    store i32 42, ptr [[P]], align 4
+; CHECK-NEXT:    store i32 7, ptr [[PARAM]], align 4
+; CHECK-NEXT:    ret i32 42
+;
+entry:
+  %p = call ptr @llvm.threadlocal.address(ptr @g1_tls)
+  store i32 42, ptr %p
+  store i32 7, ptr %param
+  %p2 = call ptr @llvm.threadlocal.address(ptr @g1_tls)
+  %v = load i32, ptr %p2
+  ret i32 %v
+}
+
+define ptr @test1_tls_noopt(ptr %coro, ptr %param) presplitcoroutine {
+; CHECK-LABEL: define ptr @test1_tls_noopt(
+; CHECK-SAME: ptr [[CORO:%.*]], ptr [[PARAM:%.*]]) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[P:%.*]] = call ptr @llvm.threadlocal.address.p0(ptr @g1_tls)
+; CHECK-NEXT:    store i32 42, ptr [[P]], align 4
+; CHECK-NEXT:    [[TMP0:%.*]] = call i8 @llvm.coro.suspend(token none, i1 false)
+; CHECK-NEXT:    switch i8 [[TMP0]], label [[SUSPEND:%.*]] [
+; CHECK-NEXT:      i8 0, label [[RESUME:%.*]]
+; CHECK-NEXT:      i8 1, label [[SUSPEND]]
+; CHECK-NEXT:    ]
+; CHECK:       resume:
+; CHECK-NEXT:    [[P2:%.*]] = call ptr @llvm.threadlocal.address.p0(ptr @g1_tls)
+; CHECK-NEXT:    [[V:%.*]] = load i32, ptr [[P2]], align 4
+; CHECK-NEXT:    store i32 [[V]], ptr [[PARAM]], align 4
+; CHECK-NEXT:    ret ptr [[CORO]]
+; CHECK:       suspend:
+; CHECK-NEXT:    [[TMP1:%.*]] = call i1 @llvm.coro.end(ptr [[CORO]], i1 false, token none)
+; CHECK-NEXT:    ret ptr [[CORO]]
+;
+entry:
+  %p = call ptr @llvm.threadlocal.address(ptr @g1_tls)
+  store i32 42, ptr %p
+
+  %0 = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %0, label %suspend [i8 0, label %resume
+  i8 1, label %suspend]
+resume:
+  %p2 = call ptr @llvm.threadlocal.address(ptr @g1_tls)
+  %v = load i32, ptr %p2
+  store i32 %v, ptr %param, align 4
+  ret ptr %coro
+
+suspend:
+  call i1 @llvm.coro.end(ptr %coro, i1 0)
+  ret ptr %coro
+}
+
 declare ptr @f()
 
 define i32 @test2() {

llvm/unittests/Analysis/ValueTrackingTest.cpp

@@ -1273,19 +1273,35 @@ TEST_F(ValueTrackingTest, getUnderlyingObjectCastsAliases) {
 
 TEST_F(ValueTrackingTest, getUnderlyingObjectIntrinsics) {
   parseAssembly(R"IR(
+    @tlsvar = thread_local global i32 0
     define void @test(ptr %arg) {
       ; intrinsic with Return<> arg attribute
       %A = call ptr @llvm.objc.retain(ptr %arg)
       %A2 = call ptr @llvm.ssa.copy(ptr %arg)
       ; special cased intrinsics
       %A3 = call ptr @llvm.launder.invariant.group(ptr %arg)
+      %A4 = call ptr @llvm.threadlocal.address(ptr @tlsvar)
       ret void
     }
 )IR");
   Value *arg = F->getArg(0);
   EXPECT_EQ(getUnderlyingObject(A), arg);
   EXPECT_EQ(getUnderlyingObject(A2), arg);
   EXPECT_EQ(getUnderlyingObject(A3), arg);
+  Value *tlsvar = M->getNamedGlobal("tlsvar");
+  EXPECT_EQ(getUnderlyingObject(A4), tlsvar);
+}
+
+TEST_F(ValueTrackingTest, getUnderlyingObjectNoSkipTLS) {
+  parseAssembly(R"IR(
+    @tlsvar = thread_local global i32 0
+    define void @test() presplitcoroutine {
+      %A = call ptr @llvm.threadlocal.address(ptr @tlsvar)
+      ret void
+    }
+)IR");
+  // Should not skip `threadlocal.address` when functions can switch thread ids.
+  EXPECT_EQ(getUnderlyingObject(A), A);
 }
 
 TEST_F(ValueTrackingTest, getUnderlyingObjectPtrInt) {