[llvm] [Attributor] Indicate optimistic fixed point if an instruction already has non-zero address space (PR #101589)

[Shilei Tian via llvm-commits]

https://github.com/shiltian created https://github.com/llvm/llvm-project/pull/101589

None

>From e83e23266035121cd79b7c8c2e5a86009898f644 Mon Sep 17 00:00:00 2001
From: Shilei Tian <i at tianshilei.me>
Date: Thu, 1 Aug 2024 20:20:13 -0400
Subject: [PATCH] [Attributor] Indicate optimistic fixed point if an
 instruction already has non-zero address space

---
 llvm/lib/Transforms/IPO/AttributorAttributes.cpp        | 3 +++
 llvm/test/Transforms/Attributor/memory_locations_gpu.ll | 6 ++++--
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/IPO/AttributorAttributes.cpp b/llvm/lib/Transforms/IPO/AttributorAttributes.cpp
index 2816a85743faa..f69054334c739 100644
--- a/llvm/lib/Transforms/IPO/AttributorAttributes.cpp
+++ b/llvm/lib/Transforms/IPO/AttributorAttributes.cpp
@@ -12491,6 +12491,9 @@ struct AAAddressSpaceImpl : public AAAddressSpace {
   void initialize(Attributor &A) override {
     assert(getAssociatedType()->isPtrOrPtrVectorTy() &&
            "Associated value is not a pointer");
+    auto *PtrTy = cast<PointerType>(getAssociatedType());
+    if (PtrTy->getAddressSpace())
+      indicateOptimisticFixpoint();
   }
 
   ChangeStatus updateImpl(Attributor &A) override {
diff --git a/llvm/test/Transforms/Attributor/memory_locations_gpu.ll b/llvm/test/Transforms/Attributor/memory_locations_gpu.ll
index 2c7a98a41f86f..c10883b54ad59 100644
--- a/llvm/test/Transforms/Attributor/memory_locations_gpu.ll
+++ b/llvm/test/Transforms/Attributor/memory_locations_gpu.ll
@@ -55,7 +55,8 @@ define i32 @test_const_as_call2() {
 ; CHECK-LABEL: define {{[^@]+}}@test_const_as_call2
 ; CHECK-SAME: () #[[ATTR3:[0-9]+]] {
 ; CHECK-NEXT:    [[P2:%.*]] = call ptr @ptr() #[[ATTR4]]
-; CHECK-NEXT:    [[L2:%.*]] = load i32, ptr [[P2]], align 4
+; CHECK-NEXT:    [[C2:%.*]] = addrspacecast ptr [[P2]] to ptr addrspace(4)
+; CHECK-NEXT:    [[L2:%.*]] = load i32, ptr addrspace(4) [[C2]], align 4
 ; CHECK-NEXT:    ret i32 [[L2]]
 ;
   %p2 = call ptr @ptr()
@@ -84,7 +85,8 @@ define i32 @test_shared_as_call2() {
 ; CHECK-LABEL: define {{[^@]+}}@test_shared_as_call2
 ; CHECK-SAME: () #[[ATTR2]] {
 ; CHECK-NEXT:    [[P2:%.*]] = call ptr @ptr() #[[ATTR4]]
-; CHECK-NEXT:    [[L2:%.*]] = load i32, ptr [[P2]], align 4
+; CHECK-NEXT:    [[C2:%.*]] = addrspacecast ptr [[P2]] to ptr addrspace(3)
+; CHECK-NEXT:    [[L2:%.*]] = load i32, ptr addrspace(3) [[C2]], align 4
 ; CHECK-NEXT:    ret i32 [[L2]]
 ;
   %p2 = call ptr @ptr()