[llvm] fe146e9 - [InferAddressSpaces] Fix constant replace to avoid modifying other functions (#70611)

[via llvm-commits]

Author: Wenju He
Date: 2023-11-13T13:28:56+08:00
New Revision: fe146e9b59891e50bddc0e61037773999335e1fe

URL: https://github.com/llvm/llvm-project/commit/fe146e9b59891e50bddc0e61037773999335e1fe
DIFF: https://github.com/llvm/llvm-project/commit/fe146e9b59891e50bddc0e61037773999335e1fe.diff

LOG: [InferAddressSpaces] Fix constant replace to avoid modifying other functions (#70611)

A constant value is unique in llvm context. InferAddressSpaces was
replacing its users in other functions as well. This leads to unexpected
behavior in our downstream use case after the pass.

InferAddressSpaces is a function passe, so it shall not modify functions
other than currently processed one.

Co-authored-by: Abhinav Gaba <abhinav.gaba at intel.com>

---------

Co-authored-by: Abhinav Gaba <abhinav.gaba at intel.com>

Added: 
    llvm/test/Transforms/InferAddressSpaces/optnone-func-unchanged.ll

Modified: 
    llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp
    llvm/test/Transforms/InferAddressSpaces/AMDGPU/noop-ptrint-pair.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp b/llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp
index 28fe1b5e75327e6..1bf50d79e5331e9 100644
--- a/llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp
+++ b/llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp
@@ -477,7 +477,7 @@ void InferAddressSpacesImpl::appendsFlatAddressExpressionToPostorderStack(
 }
 
 // Returns all flat address expressions in function F. The elements are ordered
-// ordered in postorder.
+// in postorder.
 std::vector<WeakTrackingVH>
 InferAddressSpacesImpl::collectFlatAddressExpressions(Function &F) const {
   // This function implements a non-recursive postorder traversal of a partial
@@ -1170,6 +1170,8 @@ bool InferAddressSpacesImpl::rewriteWithNewAddressSpaces(
   }
 
   SmallVector<Instruction *, 16> DeadInstructions;
+  ValueToValueMapTy VMap;
+  ValueMapper VMapper(VMap, RF_NoModuleLevelChanges | RF_IgnoreMissingLocals);
 
   // Replaces the uses of the old address expressions with the new ones.
   for (const WeakTrackingVH &WVH : Postorder) {
@@ -1188,7 +1190,30 @@ bool InferAddressSpacesImpl::rewriteWithNewAddressSpaces(
       if (C != Replace) {
         LLVM_DEBUG(dbgs() << "Inserting replacement const cast: " << Replace
                           << ": " << *Replace << '\n');
-        C->replaceAllUsesWith(Replace);
+        SmallVector<User *, 16> WorkList;
+        for (User *U : make_early_inc_range(C->users())) {
+          if (auto *I = dyn_cast<Instruction>(U)) {
+            if (I->getFunction() == F)
+              I->replaceUsesOfWith(C, Replace);
+          } else {
+            WorkList.append(U->user_begin(), U->user_end());
+          }
+        }
+        if (!WorkList.empty()) {
+          VMap[C] = Replace;
+          DenseSet<User *> Visited{WorkList.begin(), WorkList.end()};
+          while (!WorkList.empty()) {
+            User *U = WorkList.pop_back_val();
+            if (auto *I = dyn_cast<Instruction>(U)) {
+              if (I->getFunction() == F)
+                VMapper.remapInstruction(*I);
+              continue;
+            }
+            for (User *U2 : U->users())
+              if (Visited.insert(U2).second)
+                WorkList.push_back(U2);
+          }
+        }
         V = Replace;
       }
     }
@@ -1214,6 +1239,11 @@ bool InferAddressSpacesImpl::rewriteWithNewAddressSpaces(
       // Skip if the current user is the new value itself.
       if (CurUser == NewV)
         continue;
+
+      if (auto *CurUserI = dyn_cast<Instruction>(CurUser);
+          CurUserI && CurUserI->getFunction() != F)
+        continue;
+
       // Handle more complex cases like intrinsic that need to be remangled.
       if (auto *MI = dyn_cast<MemIntrinsic>(CurUser)) {
         if (!MI->isVolatile() && handleMemIntrinsicPtrUse(MI, V, NewV))

diff  --git a/llvm/test/Transforms/InferAddressSpaces/AMDGPU/noop-ptrint-pair.ll b/llvm/test/Transforms/InferAddressSpaces/AMDGPU/noop-ptrint-pair.ll
index b6713e96bed3ff5..422ac0dfd2cd470 100644
--- a/llvm/test/Transforms/InferAddressSpaces/AMDGPU/noop-ptrint-pair.ll
+++ b/llvm/test/Transforms/InferAddressSpaces/AMDGPU/noop-ptrint-pair.ll
@@ -70,7 +70,7 @@ define ptr @noop_ptrint_pair_ce2() {
 }
 
 ; COMMON-LABEL: @noop_ptrint_pair_ce2_vec(
-; AMDGCN-NEXT: ret <2 x ptr> <ptr addrspacecast (ptr addrspace(1) @g to ptr), ptr inttoptr (i64 ptrtoint (ptr addrspace(3) @l to i64) to ptr)>
+; AMDGCN-NEXT: ret <2 x ptr> <ptr inttoptr (i64 ptrtoint (ptr addrspace(1) @g to i64) to ptr), ptr inttoptr (i64 ptrtoint (ptr addrspace(3) @l to i64) to ptr)>
 ; NOTTI-NEXT: ret <2 x ptr> <ptr inttoptr (i64 ptrtoint (ptr addrspace(1) @g to i64) to ptr), ptr inttoptr (i64 ptrtoint (ptr addrspace(3) @l to i64) to ptr)>
 define <2 x ptr> @noop_ptrint_pair_ce2_vec() {
   ret <2 x ptr> <ptr inttoptr (i64 ptrtoint (ptr addrspace(1) @g to i64) to ptr), ptr inttoptr (i64 ptrtoint (ptr addrspace(3) @l to i64) to ptr)>

diff  --git a/llvm/test/Transforms/InferAddressSpaces/optnone-func-unchanged.ll b/llvm/test/Transforms/InferAddressSpaces/optnone-func-unchanged.ll
new file mode 100644
index 000000000000000..b0213f7486928ae
--- /dev/null
+++ b/llvm/test/Transforms/InferAddressSpaces/optnone-func-unchanged.ll
@@ -0,0 +1,32 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 3
+; RUN: opt -assume-default-is-flat-addrspace -S -passes=infer-address-spaces < %s 2>&1 | FileCheck %s
+
+ at g = addrspace(1) global i32 0, align 4
+
+define ptr @f2() {
+; CHECK-LABEL: define ptr @f2() {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[X2:%.*]] = addrspacecast ptr addrspace(1) @g to ptr
+; CHECK-NEXT:    ret ptr [[X2]]
+;
+entry:
+  %x1 = addrspacecast ptr addrspacecast (ptr addrspace(1) @g to ptr) to ptr addrspace(1)
+  %x2 = addrspacecast ptr addrspace(1) %x1 to ptr
+  ret ptr %x2
+}
+
+define ptr @f3() #0 {
+; CHECK-LABEL: define ptr @f3(
+; CHECK-SAME: ) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[X1:%.*]] = addrspacecast ptr addrspacecast (ptr addrspace(1) @g to ptr) to ptr addrspace(1)
+; CHECK-NEXT:    [[X2:%.*]] = addrspacecast ptr addrspace(1) [[X1]] to ptr
+; CHECK-NEXT:    ret ptr [[X2]]
+;
+entry:
+  %x1 = addrspacecast ptr addrspacecast (ptr addrspace(1) @g to ptr) to ptr addrspace(1)
+  %x2 = addrspacecast ptr addrspace(1) %x1 to ptr
+  ret ptr %x2
+}
+
+attributes #0 = { noinline optnone }