[AMDGPU] Improve clobbering checks in the kernel argument promotion

[rampitec]

Summary

Use same MSSA clobbering checks as in the AMDGPUAnnotateUniformValues.
Kernel argument promotion needs exactly the same information so factor
out utility function isClobberedInFunction.

[vpykhtin]

llvm/lib/Target/AMDGPU/Utils/AMDGPUMemoryUtils.cpp

34

This is usefull routine outside of the eclosing function, I would move it on the file scope and replace Load argument with its pointer, I don't see other uses of Load there.

[rampitec]

Factored out isReallyAClobber().

[vpykhtin]

LGTM.

[vpykhtin]

llvm/lib/Target/AMDGPU/Utils/AMDGPUMemoryUtils.cpp

29

Nit: move the debug output back to isClobberedInFunction

[rampitec]

Moved debug output.

llvm/lib/Target/AMDGPU/AMDGPUAnnotateUniformValues.cpp

 //===-- AMDGPUAnnotateUniformValues.cpp - ---------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 /// \file
 /// This pass adds amdgpu.uniform metadata to IR values so this information
 /// can be used during instruction selection.
 //
 //===----------------------------------------------------------------------===//
 
 #include "AMDGPU.h"
 #include "Utils/AMDGPUBaseInfo.h"
-#include "llvm/ADT/SmallSet.h"
+#include "Utils/AMDGPUMemoryUtils.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/LegacyDivergenceAnalysis.h"
 #include "llvm/Analysis/MemorySSA.h"
 #include "llvm/IR/InstVisitor.h"
-#include "llvm/IR/IntrinsicsAMDGPU.h"
 #include "llvm/InitializePasses.h"
 
 #define DEBUG_TYPE "amdgpu-annotate-uniform"
 
 using namespace llvm;
 
 namespace {
 
void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<LegacyDivergenceAnalysis>();
     AU.addRequired<MemorySSAWrapperPass>();
     AU.addRequired<AAResultsWrapperPass>();
     AU.setPreservesAll();
  }
 
   void visitBranchInst(BranchInst &I);
   void visitLoadInst(LoadInst &I);
-  bool isClobberedInFunction(LoadInst * Load);
 };
 
 } // End anonymous namespace
 
 INITIALIZE_PASS_BEGIN(AMDGPUAnnotateUniformValues, DEBUG_TYPE,
                       "Add AMDGPU uniform metadata", false, false)
 INITIALIZE_PASS_DEPENDENCY(LegacyDivergenceAnalysis)
 INITIALIZE_PASS_DEPENDENCY(MemorySSAWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(AAResultsWrapperPass)
 INITIALIZE_PASS_END(AMDGPUAnnotateUniformValues, DEBUG_TYPE,
                     "Add AMDGPU uniform metadata", false, false)
 
 char AMDGPUAnnotateUniformValues::ID = 0;
 
 static void setUniformMetadata(Instruction *I) {
   I->setMetadata("amdgpu.uniform", MDNode::get(I->getContext(), {}));
 }
 static void setNoClobberMetadata(Instruction *I) {
   I->setMetadata("amdgpu.noclobber", MDNode::get(I->getContext(), {}));
 }
 
-bool AMDGPUAnnotateUniformValues::isClobberedInFunction(LoadInst *Load) {
-  MemorySSAWalker *Walker = MSSA->getWalker();
-  SmallVector<MemoryAccess *> WorkList{Walker->getClobberingMemoryAccess(Load)};
-  SmallSet<MemoryAccess *, 8> Visited;
-  MemoryLocation Loc(MemoryLocation::get(Load));
-
-  const auto isReallyAClobber = [this, Load](MemoryDef *Def) -> bool {
-    Instruction *DefInst = Def->getMemoryInst();
-    LLVM_DEBUG(dbgs() << "  Def: " << *DefInst << '\n');
-
-    if (isa<FenceInst>(DefInst))
-      return false;
-
-    if (const IntrinsicInst *II = dyn_cast<IntrinsicInst>(DefInst)) {
-      switch (II->getIntrinsicID()) {
-      case Intrinsic::amdgcn_s_barrier:
-      case Intrinsic::amdgcn_wave_barrier:
-        return false;
-      default:
-        break;
-      }
-    }
-
-    // Ignore atomics not aliasing with the original load, any atomic is a
-    // universal MemoryDef from MSSA's point of view too, just like a fence.
-    const auto checkNoAlias = [this, Load](auto I) -> bool {
-      return I && AA->isNoAlias(I->getPointerOperand(),
-                                Load->getPointerOperand());
-    };
-
-    if (checkNoAlias(dyn_cast<AtomicCmpXchgInst>(DefInst)) ||
-        checkNoAlias(dyn_cast<AtomicRMWInst>(DefInst)))
-      return false;
-
-    return true;
-  };
-
-  LLVM_DEBUG(dbgs() << "Checking clobbering of: " << *Load << '\n');
-
-  // Start with a nearest dominating clobbering access, it will be either
-  // live on entry (nothing to do, load is not clobbered), MemoryDef, or
-  // MemoryPhi if several MemoryDefs can define this memory state. In that
-  // case add all Defs to WorkList and continue going up and checking all
-  // the definitions of this memory location until the root. When all the
-  // defs are exhausted and came to the entry state we have no clobber.
-  // Along the scan ignore barriers and fences which are considered clobbers
-  // by the MemorySSA, but not really writing anything into the memory.
-  while (!WorkList.empty()) {
-    MemoryAccess *MA = WorkList.pop_back_val();
-    if (!Visited.insert(MA).second)
-      continue;
-
-    if (MSSA->isLiveOnEntryDef(MA))
-      continue;
-
-    if (MemoryDef *Def = dyn_cast<MemoryDef>(MA)) {
-      if (isReallyAClobber(Def)) {
-        LLVM_DEBUG(dbgs() << "      -> load is clobbered\n");
-        return true;
-      }
-
-      WorkList.push_back(
-          Walker->getClobberingMemoryAccess(Def->getDefiningAccess(), Loc));
-      continue;
-    }
-
-    const MemoryPhi *Phi = cast<MemoryPhi>(MA);
-    for (auto &Use : Phi->incoming_values())
-      WorkList.push_back(cast<MemoryAccess>(&Use));
-  }
-
-  LLVM_DEBUG(dbgs() << "      -> no clobber\n");
-  return false;
-}
-
 void AMDGPUAnnotateUniformValues::visitBranchInst(BranchInst &I) {
   if (DA->isUniform(&I))
     setUniformMetadata(&I);
 }
 
 void AMDGPUAnnotateUniformValues::visitLoadInst(LoadInst &I) {
   Value *Ptr = I.getPointerOperand();
   if (!DA->isUniform(Ptr))
     return;
   Instruction *PtrI = dyn_cast<Instruction>(Ptr);
   if (PtrI)
     setUniformMetadata(PtrI);
 
   // We're tracking up to the Function boundaries, and cannot go beyond because
   // of FunctionPass restrictions. We can ensure that is memory not clobbered
   // for memory operations that are live in to entry points only.
   if (!isEntryFunc)
     return;
   bool GlobalLoad = I.getPointerAddressSpace() == AMDGPUAS::GLOBAL_ADDRESS;
-  bool NotClobbered = GlobalLoad && !isClobberedInFunction(&I);
-  if (NotClobbered)
+  if (GlobalLoad && !AMDGPU::isClobberedInFunction(&I, MSSA, AA))
     setNoClobberMetadata(&I);
 }
 
 bool AMDGPUAnnotateUniformValues::doInitialization(Module &M) {
   return false;
 }
 
 bool AMDGPUAnnotateUniformValues::runOnFunction(Function &F) {

llvm/lib/Target/AMDGPU/AMDGPUPromoteKernelArguments.cpp

 ///
 /// The pass walks kernel's pointer arguments, then loads from them. If a loaded
 /// value is a pointer and loaded pointer is unmodified in the kernel before the
 /// load, then promote loaded pointer to global. Then recursively continue.
 //
 //===----------------------------------------------------------------------===//
 
 #include "AMDGPU.h"
+#include "Utils/AMDGPUMemoryUtils.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/MemorySSA.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/InitializePasses.h"
 
 #define DEBUG_TYPE "amdgpu-promote-kernel-arguments"
 
 using namespace llvm;
 
 namespace {
 
 class AMDGPUPromoteKernelArguments : public FunctionPass {
   MemorySSA *MSSA;
 
+  AliasAnalysis *AA;
+
   Instruction *ArgCastInsertPt;
 
   SmallVector<Value *> Ptrs;
 
   void enqueueUsers(Value *Ptr);
 
   bool promotePointer(Value *Ptr);
 
 public:
   static char ID;
 
   AMDGPUPromoteKernelArguments() : FunctionPass(ID) {}
 
-  bool run(Function &F, MemorySSA &MSSA);
+  bool run(Function &F, MemorySSA &MSSA, AliasAnalysis &AA);
 
   bool runOnFunction(Function &F) override;
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<AAResultsWrapperPass>();
     AU.addRequired<MemorySSAWrapperPass>();
     AU.setPreservesAll();
   }
 };
 
 } // end anonymous namespace
 
 void AMDGPUPromoteKernelArguments::enqueueUsers(Value *Ptr) {
case Instruction::Load: {
       LoadInst *LD = cast<LoadInst>(U);
       PointerType *PT = dyn_cast<PointerType>(LD->getType());
       if (!PT ||
           (PT->getAddressSpace() != AMDGPUAS::FLAT_ADDRESS &&
            PT->getAddressSpace() != AMDGPUAS::GLOBAL_ADDRESS &&
            PT->getAddressSpace() != AMDGPUAS::CONSTANT_ADDRESS) ||
           LD->getPointerOperand()->stripInBoundsOffsets() != Ptr)
         break;
-      const MemoryAccess *MA = MSSA->getWalker()->getClobberingMemoryAccess(LD);
       // TODO: This load poprobably can be promoted to constant address space.
-      if (MSSA->isLiveOnEntryDef(MA))
+      if (!AMDGPU::isClobberedInFunction(LD, MSSA, AA))
         Ptrs.push_back(LD);
       break;
     }
     case Instruction::GetElementPtr:
     case Instruction::AddrSpaceCast:
     case Instruction::BitCast:
       if (U->getOperand(0)->stripInBoundsOffsets() == Ptr)
         PtrUsers.append(U->user_begin(), U->user_end());
for (BasicBlock::iterator E = BB.end(); InsPt != E; ++InsPt) {
     // so loads will need to be inserted before it.
     if (!AI || !AI->isStaticAlloca())
       break;
   }
 
   return InsPt;
 }
 
-bool AMDGPUPromoteKernelArguments::run(Function &F, MemorySSA &MSSA) {
+bool AMDGPUPromoteKernelArguments::run(Function &F, MemorySSA &MSSA,
+                                       AliasAnalysis &AA) {
   if (skipFunction(F))
     return false;
 
   CallingConv::ID CC = F.getCallingConv();
   if (CC != CallingConv::AMDGPU_KERNEL || F.arg_empty())
     return false;
 
   ArgCastInsertPt = &*getInsertPt(*F.begin());
   this->MSSA = &MSSA;
+  this->AA = &AA;
 
   for (Argument &Arg : F.args()) {
     if (Arg.use_empty())
       continue;
 
     PointerType *PT = dyn_cast<PointerType>(Arg.getType());
     if (!PT || (PT->getAddressSpace() != AMDGPUAS::FLAT_ADDRESS &&
                 PT->getAddressSpace() != AMDGPUAS::GLOBAL_ADDRESS &&
while (!Ptrs.empty()) {
     Changed |= promotePointer(Ptr);
   }
 
   return Changed;
 }
 
 bool AMDGPUPromoteKernelArguments::runOnFunction(Function &F) {
   MemorySSA &MSSA = getAnalysis<MemorySSAWrapperPass>().getMSSA();
-  return run(F, MSSA);
+  AliasAnalysis &AA = getAnalysis<AAResultsWrapperPass>().getAAResults();
+  return run(F, MSSA, AA);
 }
 
 INITIALIZE_PASS_BEGIN(AMDGPUPromoteKernelArguments, DEBUG_TYPE,
                       "AMDGPU Promote Kernel Arguments", false, false)
+INITIALIZE_PASS_DEPENDENCY(AAResultsWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(MemorySSAWrapperPass)
 INITIALIZE_PASS_END(AMDGPUPromoteKernelArguments, DEBUG_TYPE,
                     "AMDGPU Promote Kernel Arguments", false, false)
 
 char AMDGPUPromoteKernelArguments::ID = 0;
 
 FunctionPass *llvm::createAMDGPUPromoteKernelArgumentsPass() {
   return new AMDGPUPromoteKernelArguments();
 }
 
 PreservedAnalyses
 AMDGPUPromoteKernelArgumentsPass::run(Function &F,
                                       FunctionAnalysisManager &AM) {
   MemorySSA &MSSA = AM.getResult<MemorySSAAnalysis>(F).getMSSA();
-  if (AMDGPUPromoteKernelArguments().run(F, MSSA)) {
+  AliasAnalysis &AA = AM.getResult<AAManager>(F);
+  if (AMDGPUPromoteKernelArguments().run(F, MSSA, AA)) {
     PreservedAnalyses PA;
     PA.preserveSet<CFGAnalyses>();
     PA.preserve<MemorySSAAnalysis>();
     return PA;
   }
   return PreservedAnalyses::all();
 }

llvm/lib/Target/AMDGPU/Utils/AMDGPUMemoryUtils.h

+//===- AMDGPUMemoryUtils.h - Memory related helper functions -*- C++ -*----===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_LIB_TARGET_AMDGPU_UTILS_AMDGPUMEMORYUTILS_H
+#define LLVM_LIB_TARGET_AMDGPU_UTILS_AMDGPUMEMORYUTILS_H
+
+namespace llvm {
+
+class AAResults;
+class LoadInst;
+class MemoryDef;
+class MemorySSA;
+class Value;
+
+namespace AMDGPU {
+
+/// Given a \p Def clobbering a load from \p Ptr accroding to the MSSA check
+/// if this is actually a memory update or an artifical clobber to facilitate
+/// ordering constraints.
+bool isReallyAClobber(const Value *Ptr, MemoryDef *Def, AAResults *AA);
+
+/// Check is a \p Load is clobbered in its function.
+bool isClobberedInFunction(const LoadInst *Load, MemorySSA *MSSA,
+                           AAResults *AA);
+
+} // end namespace AMDGPU
+
+} // end namespace llvm
+
+#endif // LLVM_LIB_TARGET_AMDGPU_UTILS_AMDGPUMEMORYUTILS_H

llvm/lib/Target/AMDGPU/Utils/AMDGPUMemoryUtils.cpp

+//===-- AMDGPUMemoryUtils.cpp - -------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "AMDGPUMemoryUtils.h"
+#include "AMDGPU.h"
+#include "llvm/ADT/SmallSet.h"
+#include "llvm/Analysis/AliasAnalysis.h"
+#include "llvm/Analysis/MemorySSA.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/IntrinsicsAMDGPU.h"
+#include "llvm/IR/IntrinsicInst.h"
+
+#define DEBUG_TYPE "amdgpu-memory-utils"
+
+using namespace llvm;
+
+namespace llvm {
+
+namespace AMDGPU {
+
+bool isReallyAClobber(const Value *Ptr, MemoryDef *Def, AAResults *AA) {
+  Instruction *DefInst = Def->getMemoryInst();
+
+  if (isa<FenceInst>(DefInst))
+    return false;
+
+  if (const IntrinsicInst *II = dyn_cast<IntrinsicInst>(DefInst)) {
+    switch (II->getIntrinsicID()) {
+    case Intrinsic::amdgcn_s_barrier:
+    case Intrinsic::amdgcn_wave_barrier:
+      return false;
+    default:
+      break;
+    }
+  }
+
+  // Ignore atomics not aliasing with the original load, any atomic is a
+  // universal MemoryDef from MSSA's point of view too, just like a fence.
+  const auto checkNoAlias = [AA, Ptr](auto I) -> bool {
+    return I && AA->isNoAlias(I->getPointerOperand(), Ptr);
+  };
+
+  if (checkNoAlias(dyn_cast<AtomicCmpXchgInst>(DefInst)) ||
+      checkNoAlias(dyn_cast<AtomicRMWInst>(DefInst)))
+    return false;
+
+  return true;
+}
+
+bool isClobberedInFunction(const LoadInst *Load, MemorySSA *MSSA,
+                           AAResults *AA) {
+  MemorySSAWalker *Walker = MSSA->getWalker();
+  SmallVector<MemoryAccess *> WorkList{Walker->getClobberingMemoryAccess(Load)};
+  SmallSet<MemoryAccess *, 8> Visited;
+  MemoryLocation Loc(MemoryLocation::get(Load));
+
+  LLVM_DEBUG(dbgs() << "Checking clobbering of: " << *Load << '\n');
+
+  // Start with a nearest dominating clobbering access, it will be either
+  // live on entry (nothing to do, load is not clobbered), MemoryDef, or
+  // MemoryPhi if several MemoryDefs can define this memory state. In that
+  // case add all Defs to WorkList and continue going up and checking all
+  // the definitions of this memory location until the root. When all the
+  // defs are exhausted and came to the entry state we have no clobber.
+  // Along the scan ignore barriers and fences which are considered clobbers
+  // by the MemorySSA, but not really writing anything into the memory.
+  while (!WorkList.empty()) {
+    MemoryAccess *MA = WorkList.pop_back_val();
+    if (!Visited.insert(MA).second)
+      continue;
+
+    if (MSSA->isLiveOnEntryDef(MA))
+      continue;
+
+    if (MemoryDef *Def = dyn_cast<MemoryDef>(MA)) {
+      LLVM_DEBUG(dbgs() << "  Def: " << *Def->getMemoryInst() << '\n');
+
+      if (isReallyAClobber(Load->getPointerOperand(), Def, AA)) {
+        LLVM_DEBUG(dbgs() << "      -> load is clobbered\n");
+        return true;
+      }
+
+      WorkList.push_back(
+          Walker->getClobberingMemoryAccess(Def->getDefiningAccess(), Loc));
+      continue;
+    }
+
+    const MemoryPhi *Phi = cast<MemoryPhi>(MA);
+    for (auto &Use : Phi->incoming_values())
+      WorkList.push_back(cast<MemoryAccess>(&Use));
+  }
+
+  LLVM_DEBUG(dbgs() << "      -> no clobber\n");
+  return false;
+}
+
+} // end namespace AMDGPU
+
+} // end namespace llvm

llvm/lib/Target/AMDGPU/Utils/CMakeLists.txt

 add_llvm_component_library(LLVMAMDGPUUtils
   AMDGPUAsmUtils.cpp
   AMDGPUBaseInfo.cpp
   AMDGPULDSUtils.cpp
+  AMDGPUMemoryUtils.cpp
   AMDGPUPALMetadata.cpp
   AMDKernelCodeTUtils.cpp
 
   LINK_COMPONENTS
   Analysis
   Core
   MC
   BinaryFormat
   Support
 
   ADD_TO_COMPONENT
   AMDGPU
   )

llvm/test/CodeGen/AMDGPU/promote-kernel-arguments.ll

entry:
   %arrayidx711 = getelementptr inbounds [4 x float], [4 x float] addrspace(3)* @LDS, i32 0, i32 %sub
   %i6 = load float, float addrspace(3)* %arrayidx711, align 4
   %idxprom8 = sext i32 %X to i64
   %arrayidx9 = getelementptr inbounds float, float* %i1, i64 %idxprom8
   store float %i6, float* %arrayidx9, align 4
   ret void
 }
 
+; GCN-LABEL: ptr_nest_3_barrier:
+; GCN-COUNT-2: global_load_dwordx2
+; GCN:         global_store_dword
+define amdgpu_kernel void @ptr_nest_3_barrier(float** addrspace(1)* nocapture readonly %Arg) {
+; CHECK-LABEL: @ptr_nest_3_barrier(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[I:%.*]] = tail call i32 @llvm.amdgcn.workitem.id.x()
+; CHECK-NEXT:    [[P1:%.*]] = getelementptr inbounds float**, float** addrspace(1)* [[ARG:%.*]], i32 [[I]]
+; CHECK-NEXT:    tail call void @llvm.amdgcn.s.barrier()
+; CHECK-NEXT:    [[P2:%.*]] = load float**, float** addrspace(1)* [[P1]], align 8
+; CHECK-NEXT:    [[P2_GLOBAL:%.*]] = addrspacecast float** [[P2]] to float* addrspace(1)*
+; CHECK-NEXT:    [[P3:%.*]] = load float*, float* addrspace(1)* [[P2_GLOBAL]], align 8
+; CHECK-NEXT:    [[P3_GLOBAL:%.*]] = addrspacecast float* [[P3]] to float addrspace(1)*
+; CHECK-NEXT:    store float 0.000000e+00, float addrspace(1)* [[P3_GLOBAL]], align 4
+; CHECK-NEXT:    ret void
+;
+entry:
+  %i = tail call i32 @llvm.amdgcn.workitem.id.x()
+  %p1 = getelementptr inbounds float**, float** addrspace(1)* %Arg, i32 %i
+  tail call void @llvm.amdgcn.s.barrier()
+  %p2 = load float**, float** addrspace(1)* %p1, align 8
+  %p3 = load float*, float** %p2, align 8
+  store float 0.000000e+00, float* %p3, align 4
+  ret void
+}
+
 declare i32 @llvm.amdgcn.workitem.id.x()
+declare void @llvm.amdgcn.s.barrier()