[llvm] [AMDGPU] Invalidate stale UniformityInfo entries on instruction erase (PR #215455)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-amdgpu

Author: Arseniy Obolenskiy (aobolensk)

<details>
<summary>Changes</summary>

AMDGPUUniformIntrinsicCombine holds one UniformityInfo for the whole function while erasing and creating instructions in a loop, so a freed instruction address could be reused and misclassified as uniform

Found as a result of an audit after finding the same class of bug fixed for Reassociate in #<!-- -->214719

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


1 Files Affected:

- (modified) llvm/lib/Target/AMDGPU/AMDGPUUniformIntrinsicCombine.cpp (+10-7) 


``````````diff
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUUniformIntrinsicCombine.cpp b/llvm/lib/Target/AMDGPU/AMDGPUUniformIntrinsicCombine.cpp
index 6b68c25825d7b..3d26d108952da 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUUniformIntrinsicCombine.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUUniformIntrinsicCombine.cpp
@@ -44,7 +44,7 @@ using namespace llvm::PatternMatch;
 /// Wrapper for querying uniformity info that first checks locally tracked
 /// instructions.
 static bool
-isDivergentUseWithNew(const Use &U, const UniformityInfo &UI,
+isDivergentUseWithNew(const Use &U, UniformityInfo &UI,
                       const ValueMap<const Value *, bool> &Tracker) {
   Value *V = U.get();
   if (auto It = Tracker.find(V); It != Tracker.end())
@@ -53,8 +53,7 @@ isDivergentUseWithNew(const Use &U, const UniformityInfo &UI,
 }
 
 /// Optimizes uniform intrinsics calls if their operand can be proven uniform.
-static bool optimizeUniformIntrinsic(IntrinsicInst &II,
-                                     const UniformityInfo &UI,
+static bool optimizeUniformIntrinsic(IntrinsicInst &II, UniformityInfo &UI,
                                      ValueMap<const Value *, bool> &Tracker) {
   llvm::Intrinsic::ID IID = II.getIntrinsicID();
   /// We deliberately do not simplify readfirstlane with a uniform argument, so
@@ -68,6 +67,7 @@ static bool optimizeUniformIntrinsic(IntrinsicInst &II,
       return false;
     LLVM_DEBUG(dbgs() << "Replacing " << II << " with " << *Src << '\n');
     II.replaceAllUsesWith(Src);
+    UI.forgetValue(&II);
     II.eraseFromParent();
     return true;
   }
@@ -103,8 +103,10 @@ static bool optimizeUniformIntrinsic(IntrinsicInst &II,
       }
     }
     // Erase the intrinsic if it has no remaining uses.
-    if (II.use_empty())
+    if (II.use_empty()) {
+      UI.forgetValue(&II);
       II.eraseFromParent();
+    }
     return Changed;
   }
   case Intrinsic::amdgcn_wave_shuffle: {
@@ -114,6 +116,7 @@ static bool optimizeUniformIntrinsic(IntrinsicInst &II,
     // Like with readlane, if Value is uniform then just propagate it
     if (!isDivergentUseWithNew(Val, UI, Tracker)) {
       II.replaceAllUsesWith(Val);
+      UI.forgetValue(&II);
       II.eraseFromParent();
       return true;
     }
@@ -136,7 +139,7 @@ static bool optimizeUniformIntrinsic(IntrinsicInst &II,
 }
 
 /// Iterates over intrinsic calls in the Function to optimize.
-static bool runUniformIntrinsicCombine(Function &F, const UniformityInfo &UI) {
+static bool runUniformIntrinsicCombine(Function &F, UniformityInfo &UI) {
   bool IsChanged = false;
   ValueMap<const Value *, bool> Tracker;
 
@@ -152,7 +155,7 @@ static bool runUniformIntrinsicCombine(Function &F, const UniformityInfo &UI) {
 PreservedAnalyses
 AMDGPUUniformIntrinsicCombinePass::run(Function &F,
                                        FunctionAnalysisManager &AM) {
-  const auto &UI = AM.getResult<UniformityInfoAnalysis>(F);
+  auto &UI = AM.getResult<UniformityInfoAnalysis>(F);
   if (!runUniformIntrinsicCombine(F, UI))
     return PreservedAnalyses::all();
 
@@ -184,7 +187,7 @@ char &llvm::AMDGPUUniformIntrinsicCombineLegacyPassID =
 bool AMDGPUUniformIntrinsicCombineLegacy::runOnFunction(Function &F) {
   if (skipFunction(F))
     return false;
-  const UniformityInfo &UI =
+  UniformityInfo &UI =
       getAnalysis<UniformityInfoWrapperPass>().getUniformityInfo();
   return runUniformIntrinsicCombine(F, UI);
 }

``````````

</details>


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