[llvm] 977f4d6 - [AMDGPU] Employ MemorySSA when rewriting out arguments (NFCI) (#207922)

[via llvm-commits]

Author: Antonio Frighetto
Date: 2026-07-10T09:55:33+02:00
New Revision: 977f4d603ae13c751357352f1256e389a4725e9f

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

LOG: [AMDGPU] Employ MemorySSA when rewriting out arguments (NFCI) (#207922)

While transitioning away from MemoryDependenceAnalysis, replace MDA
`getPointerDependencyFrom` query with a reverse traversal of the memory
accesses of the return basic block (MemoryDefs and MemoryUses), looking
for the must-aliasing store to the out argument. The existing behaviour
is meant to be preserved.

Added: 
    

Modified: 
    llvm/lib/Target/AMDGPU/AMDGPURewriteOutArguments.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Target/AMDGPU/AMDGPURewriteOutArguments.cpp b/llvm/lib/Target/AMDGPU/AMDGPURewriteOutArguments.cpp
index c0802de2763dc..7db637eaa0745 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPURewriteOutArguments.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPURewriteOutArguments.cpp
@@ -44,7 +44,8 @@
 #include "AMDGPU.h"
 #include "Utils/AMDGPUBaseInfo.h"
 #include "llvm/ADT/Statistic.h"
-#include "llvm/Analysis/MemoryDependenceAnalysis.h"
+#include "llvm/Analysis/MemorySSA.h"
+#include "llvm/Analysis/MemorySSAUpdater.h"
 #include "llvm/IR/AttributeMask.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instructions.h"
@@ -81,7 +82,9 @@ namespace {
 class AMDGPURewriteOutArguments : public FunctionPass {
 private:
   const DataLayout *DL = nullptr;
-  MemoryDependenceResults *MDA = nullptr;
+  MemorySSA *MSSA = nullptr;
+  MemorySSAUpdater *MSSAU = nullptr;
+  AAResults *AA = nullptr;
 
   Type *getStoredType(Value &Arg) const;
   Type *getOutArgumentType(Argument &Arg) const;
@@ -92,7 +95,8 @@ class AMDGPURewriteOutArguments : public FunctionPass {
   AMDGPURewriteOutArguments() : FunctionPass(ID) {}
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
-    AU.addRequired<MemoryDependenceWrapperPass>();
+    AU.addRequired<MemorySSAWrapperPass>();
+    AU.addRequired<AAResultsWrapperPass>();
     FunctionPass::getAnalysisUsage(AU);
   }
 
@@ -104,7 +108,7 @@ class AMDGPURewriteOutArguments : public FunctionPass {
 
 INITIALIZE_PASS_BEGIN(AMDGPURewriteOutArguments, DEBUG_TYPE,
                       "AMDGPU Rewrite Out Arguments", false, false)
-INITIALIZE_PASS_DEPENDENCY(MemoryDependenceWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(MemorySSAWrapperPass)
 INITIALIZE_PASS_END(AMDGPURewriteOutArguments, DEBUG_TYPE,
                     "AMDGPU Rewrite Out Arguments", false, false)
 
@@ -170,6 +174,49 @@ bool AMDGPURewriteOutArguments::doInitialization(Module &M) {
   return false;
 }
 
+static StoreInst *findStoreForOutArgument(BasicBlock *BB, Argument *OutArg,
+                                          MemorySSA &MSSA,
+                                          BatchAAResults &BAA) {
+  MemoryLocation ArgLoc = MemoryLocation::getBeforeOrAfter(OutArg);
+  const auto *Accesses = MSSA.getBlockAccesses(BB);
+  if (!Accesses)
+    return nullptr;
+
+  for (const MemoryAccess &Access : reverse(*Accesses)) {
+    const auto *UseOrDef = dyn_cast<MemoryUseOrDef>(&Access);
+    if (!UseOrDef)
+      continue;
+
+    Instruction *I = UseOrDef->getMemoryInst();
+
+    // Return the must-alias store to the out argument.
+    if (auto *Store = dyn_cast<StoreInst>(I))
+      if (Store->getPointerOperand() == OutArg)
+        return Store;
+
+    if (auto *FI = dyn_cast<FenceInst>(I))
+      if (FI->getOrdering() == AtomicOrdering::Release)
+        continue;
+
+    if (auto *LI = dyn_cast<LoadInst>(I)) {
+      if (LI->isAtomic()) {
+        // May-alias reads with monotonic ordering are ignored.
+        if (isStrongerThan(LI->getOrdering(), AtomicOrdering::Monotonic))
+          return nullptr;
+        continue;
+      }
+    }
+
+    // Any other memory access that writes the location prevents the
+    // rewrite.
+    // FIXME: should handle aliasing reads too.
+    if (isModSet(BAA.getModRefInfo(I, ArgLoc)))
+      return nullptr;
+  }
+
+  return nullptr;
+}
+
 bool AMDGPURewriteOutArguments::runOnFunction(Function &F) {
   if (skipFunction(F))
     return false;
@@ -179,8 +226,6 @@ bool AMDGPURewriteOutArguments::runOnFunction(Function &F) {
       AMDGPU::isEntryFunctionCC(F.getCallingConv()))
     return false;
 
-  MDA = &getAnalysis<MemoryDependenceWrapperPass>().getMemDep();
-
   unsigned ReturnNumRegs = 0;
   // Maps an out-argument number to its field index in the return struct.
   // Fields are in processing order, which the retry loop below can reorder
@@ -222,6 +267,13 @@ bool AMDGPURewriteOutArguments::runOnFunction(Function &F) {
   if (Returns.empty())
     return false;
 
+  AA = &getAnalysis<AAResultsWrapperPass>().getAAResults();
+  MSSA = &getAnalysis<MemorySSAWrapperPass>().getMSSA();
+
+  BatchAAResults BatchAA(*AA);
+  MemorySSAUpdater MSSAUpdater(MSSA);
+  MSSAU = &MSSAUpdater;
+
   bool Changing;
 
   do {
@@ -254,17 +306,7 @@ bool AMDGPURewriteOutArguments::runOnFunction(Function &F) {
       for (ReturnInst *RI : Returns) {
         BasicBlock *BB = RI->getParent();
 
-        MemDepResult Q = MDA->getPointerDependencyFrom(
-            MemoryLocation::getBeforeOrAfter(OutArg), true, BB->end(), BB, RI);
-        StoreInst *SI = nullptr;
-        if (Q.isDef())
-          SI = dyn_cast<StoreInst>(Q.getInst());
-
-        // MDA stops at the first may-aliasing store, which need not be to this
-        // argument; only fold a store whose pointer is exactly OutArg.
-        if (SI && SI->getPointerOperand() != OutArg)
-          SI = nullptr;
-
+        StoreInst *SI = findStoreForOutArgument(BB, OutArg, *MSSA, BatchAA);
         if (SI) {
           LLVM_DEBUG(dbgs() << "Found out argument store: " << *SI << '\n');
           ReplaceableStores.emplace_back(RI, SI);
@@ -291,6 +333,7 @@ bool AMDGPURewriteOutArguments::runOnFunction(Function &F) {
         }
 
         ValVec.emplace_back(OutArg, ReplVal);
+        MSSAU->removeMemoryAccess(Store.second);
         Store.second->eraseFromParent();
       }