[llvm] InferAddressSpaces: Improve handling of instructions with multiple pointer uses (PR #101922)

via llvm-commits llvm-commits at lists.llvm.org
Sun Aug 4 22:49:57 PDT 2024


llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-amdgpu

Author: Matt Arsenault (arsenm)

<details>
<summary>Changes</summary>

The use list iteration worked correctly for the load and store case. The atomic
instructions happen to have the pointer value as the last visited operand, but we
rejected the instruction as simple after the first encountered use.

Ignore the use list for the recognized load/store/atomic instructions, and just
try to directly replace the known pointer use.

---
Full diff: https://github.com/llvm/llvm-project/pull/101922.diff


2 Files Affected:

- (modified) llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp (+34-18) 
- (modified) llvm/test/Transforms/InferAddressSpaces/AMDGPU/store-pointer-to-self.ll (+39-3) 


``````````diff
diff --git a/llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp b/llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp
index 830c15249582c..943b68efc05b9 100644
--- a/llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp
+++ b/llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp
@@ -1002,31 +1002,48 @@ bool InferAddressSpacesImpl::updateAddressSpace(
   return true;
 }
 
+/// Replace operand \p OpIdx in \p Inst, if the value is the same as \p OldVal
+/// with \p NewVal.
+static void replaceOperandIfSame(Instruction *Inst, unsigned OpIdx,
+                                 Value *OldVal, Value *NewVal) {
+  Use &U = Inst->getOperandUse(OpIdx);
+  if (U.get() == OldVal)
+    U.set(NewVal);
+}
+
 /// \p returns true if \p U is the pointer operand of a memory instruction with
 /// a single pointer operand that can have its address space changed by simply
 /// mutating the use to a new value. If the memory instruction is volatile,
 /// return true only if the target allows the memory instruction to be volatile
 /// in the new address space.
 static bool isSimplePointerUseValidToReplace(const TargetTransformInfo &TTI,
-                                             Use &U, unsigned AddrSpace) {
-  User *Inst = U.getUser();
-  unsigned OpNo = U.getOperandNo();
-
-  if (auto *LI = dyn_cast<LoadInst>(Inst))
-    return OpNo == LoadInst::getPointerOperandIndex() &&
-           (!LI->isVolatile() || TTI.hasVolatileVariant(LI, AddrSpace));
+                                             User *Inst, unsigned AddrSpace,
+                                             Value *OldV, Value *NewV) {
+  if (auto *LI = dyn_cast<LoadInst>(Inst)) {
+    if (!LI->isVolatile() || TTI.hasVolatileVariant(LI, AddrSpace))
+      replaceOperandIfSame(LI, LoadInst::getPointerOperandIndex(), OldV, NewV);
+    return true;
+  }
 
-  if (auto *SI = dyn_cast<StoreInst>(Inst))
-    return OpNo == StoreInst::getPointerOperandIndex() &&
-           (!SI->isVolatile() || TTI.hasVolatileVariant(SI, AddrSpace));
+  if (auto *SI = dyn_cast<StoreInst>(Inst)) {
+    if (!SI->isVolatile() || TTI.hasVolatileVariant(SI, AddrSpace))
+      replaceOperandIfSame(SI, StoreInst::getPointerOperandIndex(), OldV, NewV);
+    return true;
+  }
 
-  if (auto *RMW = dyn_cast<AtomicRMWInst>(Inst))
-    return OpNo == AtomicRMWInst::getPointerOperandIndex() &&
-           (!RMW->isVolatile() || TTI.hasVolatileVariant(RMW, AddrSpace));
+  if (auto *RMW = dyn_cast<AtomicRMWInst>(Inst)) {
+    if (!RMW->isVolatile() || TTI.hasVolatileVariant(RMW, AddrSpace))
+      replaceOperandIfSame(RMW, AtomicRMWInst::getPointerOperandIndex(), OldV,
+                           NewV);
+    return true;
+  }
 
-  if (auto *CmpX = dyn_cast<AtomicCmpXchgInst>(Inst))
-    return OpNo == AtomicCmpXchgInst::getPointerOperandIndex() &&
-           (!CmpX->isVolatile() || TTI.hasVolatileVariant(CmpX, AddrSpace));
+  if (auto *CmpX = dyn_cast<AtomicCmpXchgInst>(Inst)) {
+    if (!CmpX->isVolatile() || TTI.hasVolatileVariant(CmpX, AddrSpace))
+      replaceOperandIfSame(CmpX, AtomicCmpXchgInst::getPointerOperandIndex(),
+                           OldV, NewV);
+    return true;
+  }
 
   return false;
 }
@@ -1226,11 +1243,10 @@ bool InferAddressSpacesImpl::rewriteWithNewAddressSpaces(
       I = skipToNextUser(I, E);
 
       if (isSimplePointerUseValidToReplace(
-              *TTI, U, V->getType()->getPointerAddressSpace())) {
+              *TTI, CurUser, V->getType()->getPointerAddressSpace(), V, NewV)) {
         // If V is used as the pointer operand of a compatible memory operation,
         // sets the pointer operand to NewV. This replacement does not change
         // the element type, so the resultant load/store is still valid.
-        U.set(NewV);
         continue;
       }
 
diff --git a/llvm/test/Transforms/InferAddressSpaces/AMDGPU/store-pointer-to-self.ll b/llvm/test/Transforms/InferAddressSpaces/AMDGPU/store-pointer-to-self.ll
index bce0e4ec1fe16..e1cb2b307986d 100644
--- a/llvm/test/Transforms/InferAddressSpaces/AMDGPU/store-pointer-to-self.ll
+++ b/llvm/test/Transforms/InferAddressSpaces/AMDGPU/store-pointer-to-self.ll
@@ -27,7 +27,7 @@ define ptr @atomicrmw_xchg_flat_pointer_to_self() {
 ; CHECK-LABEL: define ptr @atomicrmw_xchg_flat_pointer_to_self() {
 ; CHECK-NEXT:    [[ALLOCA:%.*]] = alloca ptr, align 8, addrspace(5)
 ; CHECK-NEXT:    [[FLAT:%.*]] = addrspacecast ptr addrspace(5) [[ALLOCA]] to ptr
-; CHECK-NEXT:    [[XCHG:%.*]] = atomicrmw xchg ptr [[FLAT]], ptr [[FLAT]] seq_cst, align 8
+; CHECK-NEXT:    [[XCHG:%.*]] = atomicrmw xchg ptr addrspace(5) [[ALLOCA]], ptr [[FLAT]] seq_cst, align 8
 ; CHECK-NEXT:    call void @user(ptr [[FLAT]])
 ; CHECK-NEXT:    ret ptr [[XCHG]]
 ;
@@ -43,7 +43,7 @@ define { ptr, i1 } @cmpxchg_flat_pointer_new_to_self(ptr %cmp) {
 ; CHECK-SAME: ptr [[CMP:%.*]]) {
 ; CHECK-NEXT:    [[ALLOCA:%.*]] = alloca ptr, align 8, addrspace(5)
 ; CHECK-NEXT:    [[FLAT:%.*]] = addrspacecast ptr addrspace(5) [[ALLOCA]] to ptr
-; CHECK-NEXT:    [[CMPX:%.*]] = cmpxchg ptr [[FLAT]], ptr [[CMP]], ptr [[FLAT]] seq_cst seq_cst, align 8
+; CHECK-NEXT:    [[CMPX:%.*]] = cmpxchg ptr addrspace(5) [[ALLOCA]], ptr [[CMP]], ptr [[FLAT]] seq_cst seq_cst, align 8
 ; CHECK-NEXT:    call void @user(ptr [[FLAT]])
 ; CHECK-NEXT:    ret { ptr, i1 } [[CMPX]]
 ;
@@ -59,7 +59,7 @@ define { ptr, i1 } @cmpxchg_flat_pointer_cmp_to_self(ptr %new) {
 ; CHECK-SAME: ptr [[NEW:%.*]]) {
 ; CHECK-NEXT:    [[ALLOCA:%.*]] = alloca ptr, align 8, addrspace(5)
 ; CHECK-NEXT:    [[FLAT:%.*]] = addrspacecast ptr addrspace(5) [[ALLOCA]] to ptr
-; CHECK-NEXT:    [[CMPX:%.*]] = cmpxchg ptr [[FLAT]], ptr [[FLAT]], ptr [[NEW]] seq_cst seq_cst, align 8
+; CHECK-NEXT:    [[CMPX:%.*]] = cmpxchg ptr addrspace(5) [[ALLOCA]], ptr [[FLAT]], ptr [[NEW]] seq_cst seq_cst, align 8
 ; CHECK-NEXT:    call void @user(ptr [[FLAT]])
 ; CHECK-NEXT:    ret { ptr, i1 } [[CMPX]]
 ;
@@ -69,3 +69,39 @@ define { ptr, i1 } @cmpxchg_flat_pointer_cmp_to_self(ptr %new) {
   call void @user(ptr %flat)
   ret { ptr, i1 } %cmpx
 }
+
+define { ptr, i1 } @cmpxchg_flat_pointer_cmp_new_self() {
+; CHECK-LABEL: define { ptr, i1 } @cmpxchg_flat_pointer_cmp_new_self() {
+; CHECK-NEXT:    [[ALLOCA:%.*]] = alloca ptr, align 8, addrspace(5)
+; CHECK-NEXT:    [[FLAT:%.*]] = addrspacecast ptr addrspace(5) [[ALLOCA]] to ptr
+; CHECK-NEXT:    [[CMPX:%.*]] = cmpxchg ptr addrspace(5) [[ALLOCA]], ptr [[FLAT]], ptr [[FLAT]] seq_cst seq_cst, align 8
+; CHECK-NEXT:    call void @user(ptr [[FLAT]])
+; CHECK-NEXT:    ret { ptr, i1 } [[CMPX]]
+;
+  %alloca = alloca ptr, align 8, addrspace(5)
+  %flat = addrspacecast ptr addrspace(5) %alloca to ptr
+  %cmpx = cmpxchg ptr %flat, ptr %flat, ptr %flat seq_cst seq_cst, align 8
+  call void @user(ptr %flat)
+  ret { ptr, i1 } %cmpx
+}
+
+define void @multi_store_flat_pointer_to_self() {
+; CHECK-LABEL: define void @multi_store_flat_pointer_to_self() {
+; CHECK-NEXT:    [[ALLOCA:%.*]] = alloca ptr, align 8, addrspace(5)
+; CHECK-NEXT:    [[FLAT:%.*]] = addrspacecast ptr addrspace(5) [[ALLOCA]] to ptr
+; CHECK-NEXT:    store ptr [[FLAT]], ptr addrspace(5) [[ALLOCA]], align 8
+; CHECK-NEXT:    store ptr [[FLAT]], ptr addrspace(5) [[ALLOCA]], align 8
+; CHECK-NEXT:    call void @user(ptr [[FLAT]])
+; CHECK-NEXT:    store ptr [[FLAT]], ptr addrspace(5) [[ALLOCA]], align 8
+; CHECK-NEXT:    store ptr addrspace(5) [[ALLOCA]], ptr addrspace(5) [[ALLOCA]], align 8
+; CHECK-NEXT:    ret void
+;
+  %alloca = alloca ptr, align 8, addrspace(5)
+  %flat = addrspacecast ptr addrspace(5) %alloca to ptr
+  store ptr %flat, ptr %flat, align 8
+  store ptr %flat, ptr %flat, align 8
+  call void @user(ptr %flat)
+  store ptr %flat, ptr addrspace(5) %alloca, align 8
+  store ptr addrspace(5) %alloca, ptr %flat, align 8
+  ret void
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/101922


More information about the llvm-commits mailing list