[llvm] GlobalsModRef, ValueTracking: Look through threadlocal.address intrinsic (PR #88418)

[Matthias Braun via llvm-commits]

https://github.com/MatzeB updated https://github.com/llvm/llvm-project/pull/88418

>From 75c583d614063c675edfb18ccc912d4128379ec4 Mon Sep 17 00:00:00 2001
From: Matthias Braun <matze at braunis.de>
Date: Thu, 11 Apr 2024 10:33:00 -0700
Subject: [PATCH] GlobalsModRef, ValueTracking: Look through
 threadlocal.address intrinsic

---
 llvm/include/llvm/Analysis/ValueTracking.h    | 10 +++++-----
 llvm/lib/Analysis/GlobalsModRef.cpp           |  8 ++++++++
 llvm/lib/Analysis/ValueTracking.cpp           |  6 ++++++
 .../GlobalsModRef/nonescaping-noalias.ll      | 19 +++++++++++++++++++
 4 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/llvm/include/llvm/Analysis/ValueTracking.h b/llvm/include/llvm/Analysis/ValueTracking.h
index 3970efba18cc8c..c5b7850990c67e 100644
--- a/llvm/include/llvm/Analysis/ValueTracking.h
+++ b/llvm/include/llvm/Analysis/ValueTracking.h
@@ -694,11 +694,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.
diff --git a/llvm/lib/Analysis/GlobalsModRef.cpp b/llvm/lib/Analysis/GlobalsModRef.cpp
index 527f19b194eeb9..1ceb1b26294183 100644
--- a/llvm/lib/Analysis/GlobalsModRef.cpp
+++ b/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)) {
diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index eb8925c5347537..2c96a75bf94143 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -6282,6 +6282,12 @@ const Value *llvm::getUnderlyingObject(const Value *V, unsigned MaxLookup) {
           continue;
         }
       } else if (auto *Call = dyn_cast<CallBase>(V)) {
+        if (const IntrinsicInst *II = dyn_cast<IntrinsicInst>(V)) {
+          if (II->getIntrinsicID() == Intrinsic::threadlocal_address) {
+            V = II->getArgOperand(0);
+            continue;
+          }
+        }
         // CaptureTracking can know about special capturing properties of some
         // intrinsics like launder.invariant.group, that can't be expressed with
         // the attributes, but have properties like returning aliasing pointer.
diff --git a/llvm/test/Analysis/GlobalsModRef/nonescaping-noalias.ll b/llvm/test/Analysis/GlobalsModRef/nonescaping-noalias.ll
index d109b3b8748ba7..f28e73c48b3e8e 100644
--- a/llvm/test/Analysis/GlobalsModRef/nonescaping-noalias.ll
+++ b/llvm/test/Analysis/GlobalsModRef/nonescaping-noalias.ll
@@ -22,6 +22,25 @@ entry:
   ret i32 %v
 }
 
+ at 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: @test1_tls(
+; CHECK: %p = call ptr @llvm.threadlocal.address.p0(ptr @g1_tls)
+; CHECK: store i32 42, ptr %p
+; CHECK-NOT: load i32
+; CHECK: ret i32 42
+entry:
+  %p = call ptr @llvm.threadlocal.address(ptr @g1_tls)
+  store i32 42, ptr %p
+  store i32 7, ptr %param
+  %v = load i32, ptr %p
+  ret i32 %v
+}
+
 declare ptr @f()
 
 define i32 @test2() {