[AMDGPU] Use UniformityAnalysis in AtomicOptimizer

[Pierre-vh]

Summary

Adds & uses a new isDivergentUse API in UA.
UniformityAnalysis now requires CycleInfo as well as the new temporal divergence API can query it.

---

Original patch that adds isDivergentUse by @sameerds

The user of a temporally divergent value is marked as divergent in the
uniformity analysis. But the same user may also have been marked divergent for
other reasons, thus losing this information about temporal divergence. But some
clients need to specificly check for temporal divergence. This change restores
such an API, that already existed in DivergenceAnalysis.

[foad]

Looks good overall.

llvm/include/llvm/ADT/GenericUniformityImpl.h

466

Was this declaration unused?

1102

Nit: the Cycle->contains check seems like it could get expensive on pathological cases with thousands of nested loops. Would it be more efficient to walk up the cycle tree from getCycle(Def) to nearestCommonAncestor(getCycle(Def), getCycle(Use))?

llvm/lib/Analysis/UniformityAnalysis.cpp

87

I think this should go inside the "if" below, since Def being an Instruction implies that Use is also an Instruction.

[Pierre-vh]

Comments

llvm/include/llvm/ADT/GenericUniformityImpl.h

466

Yes, it had no definition for UA beforehand.

1102

I'm not too familiar with CycleInfo so I'll leave this one to @sameerds

[sameerds]

llvm/include/llvm/ADT/GenericUniformityImpl.h

1102

Every cycle contains a list of all the blocks that it transitively contains, so this function is simply a lookup in a set. This is the same behaviour as LoopInfo.

But now that nested cycles got mentioned, I think this change should have another test with a def inside the inner of two nested cycles, and a use of that def outside the outer cycle.

[foad]

llvm/include/llvm/ADT/GenericUniformityImpl.h

1102

so this function is simply a lookup in a set

It's implemented as a linear search through a std::vector.

[Pierre-vh]

Add nested cycle test + remove unneeded run line

[sameerds]

llvm/include/llvm/ADT/GenericUniformityImpl.h

1102

Ouch, I didn't notice that! We can do that as a separate change. I volunteer.

[sameerds]

LGTM, but wait for @foad to comment.

llvm/include/llvm/ADT/GenericSSAContext.h

public:
   // the smallest complete unit of code in a Module.
   //
   // The compiler produces an error here if this class is implicitly
   // specialized due to an instantiation. An explicit specialization
   // of this template needs to be added before the instantiation point
   // indicated by the compiler.
   using FunctionT = typename _FunctionT::invalidTemplateInstanceError;
 
+  // A UseT represents a data-edge from the defining instruction to the using
+  // instruction.
+  //
+  // using UseT = ...
+
   // Initialize the SSA context with information about the FunctionT being
   // processed.
   //
   // void setFunction(FunctionT &function);
   // FunctionT* getFunction() const;
 
   // Every FunctionT has a unique BlockT marked as its entry.
   //

llvm/include/llvm/ADT/GenericUniformityImpl.h

 /// from sources of divergence to all users. It can be instantiated
 /// for an IR that provides a suitable SSAContext.
 template <typename ContextT> class GenericUniformityAnalysisImpl {
 public:
   using BlockT = typename ContextT::BlockT;
   using FunctionT = typename ContextT::FunctionT;
   using ValueRefT = typename ContextT::ValueRefT;
   using ConstValueRefT = typename ContextT::ConstValueRefT;
+  using UseT = typename ContextT::UseT;
   using InstructionT = typename ContextT::InstructionT;
   using DominatorTreeT = typename ContextT::DominatorTreeT;
 
   using CycleInfoT = GenericCycleInfo<ContextT>;
   using CycleT = typename CycleInfoT::CycleT;
 
   using SyncDependenceAnalysisT = GenericSyncDependenceAnalysis<ContextT>;
   using DivergenceDescriptorT =
if (I.isTerminator()) {
       return DivergentTermBlocks.contains(I.getParent());
     }
     return hasDivergentDefs(I);
   };
 
   /// \brief Whether \p Val is divergent at its definition.
   bool isDivergent(ConstValueRefT V) const { return DivergentValues.count(V); }
 
+  bool isDivergentUse(const UseT &U) const;
+
   bool hasDivergentTerminator(const BlockT &B) const {
     return DivergentTermBlocks.contains(&B);
   }
 
   void print(raw_ostream &out) const;
 
 protected:
   /// \brief Value/block pair representing a single phi input.
void analyzeTemporalDivergence(const InstructionT &I,
                                  const CycleT &OuterDivCycle);
 
   /// \brief Push all users of \p Val (in the region) to the worklist.
   void pushUsers(const InstructionT &I);
   void pushUsers(ConstValueRefT V);
 
   bool usesValueFromCycle(const InstructionT &I, const CycleT &DefCycle) const;
 
-  /// \brief Whether \p Val is divergent when read in \p ObservingBlock.
+  /// \brief Whether \p Def is divergent when read in \p ObservingBlock.
   bool isTemporalDivergent(const BlockT &ObservingBlock,
-                           ConstValueRefT Val) const;
+                           const InstructionT &Def) const;
 };
 
 template <typename ImplT>
 void GenericUniformityAnalysisImplDeleter<ImplT>::operator()(ImplT *Impl) {
   delete Impl;
 }
 
 /// Compute divergence starting with a divergent branch.
getOutermostDivergentCycle(const CycleT *Cycle, const BlockT *DivTermBlock,
   const auto *Int = getIntDivCycle(Cycle, DivTermBlock, JoinBlock, DT, Context);
 
   if (Int)
     return Int;
   return Ext;
 }
 
 template <typename ContextT>
+bool GenericUniformityAnalysisImpl<ContextT>::isTemporalDivergent(
+    const BlockT &ObservingBlock, const InstructionT &Def) const {
+  const BlockT *DefBlock = Def.getParent();
+  for (const CycleT *Cycle = CI.getCycle(DefBlock);
+       Cycle && !Cycle->contains(&ObservingBlock);
+       Cycle = Cycle->getParentCycle()) {
+    if (DivergentExitCycles.contains(Cycle)) {
+      return true;
+    }
+  }
+  return false;
+}
+
+template <typename ContextT>
 void GenericUniformityAnalysisImpl<ContextT>::analyzeControlDivergence(
     const InstructionT &Term) {
   const auto *DivTermBlock = Term.getParent();
   DivergentTermBlocks.insert(DivTermBlock);
   LLVM_DEBUG(dbgs() << "analyzeControlDiv " << Context.print(DivTermBlock)
                     << "\n");
 
   // Don't propagate divergence from unreachable blocks.

 }
 
 template <typename ContextT>
 bool GenericUniformityInfo<ContextT>::isDivergent(const InstructionT *I) const {
   return DA->isDivergent(*I);
 }
 
 template <typename ContextT>
+bool GenericUniformityInfo<ContextT>::isDivergentUse(const UseT &U) const {
+  return DA->isDivergentUse(U);
+}
+
+template <typename ContextT>
 bool GenericUniformityInfo<ContextT>::hasDivergentTerminator(const BlockT &B) {
   return DA->hasDivergentTerminator(B);
 }
 
 /// \brief T helper function for printing.
 template <typename ContextT>
 void GenericUniformityInfo<ContextT>::print(raw_ostream &out) const {
   DA->print(out);

llvm/include/llvm/ADT/GenericUniformityInfo.h

 };
 
 template <typename ContextT> class GenericUniformityInfo {
 public:
   using BlockT = typename ContextT::BlockT;
   using FunctionT = typename ContextT::FunctionT;
   using ValueRefT = typename ContextT::ValueRefT;
   using ConstValueRefT = typename ContextT::ConstValueRefT;
+  using UseT = typename ContextT::UseT;
   using InstructionT = typename ContextT::InstructionT;
   using DominatorTreeT = typename ContextT::DominatorTreeT;
   using ThisT = GenericUniformityInfo<ContextT>;
 
   using CycleInfoT = GenericCycleInfo<ContextT>;
   using CycleT = typename CycleInfoT::CycleT;
 
   GenericUniformityInfo(FunctionT &F, const DominatorTreeT &DT,
public:
 
   // Similar queries for InstructionT. These accept a pointer argument so that
   // in LLVM IR, they overload the equivalent queries for Value*. For example,
   // if querying whether a BranchInst is divergent, it should not be treated as
   // a Value in LLVM IR.
   bool isUniform(const InstructionT *I) const { return !isDivergent(I); };
   bool isDivergent(const InstructionT *I) const;
 
+  /// \brief Whether \p U is divergent. Uses of a uniform value can be
+  /// divergent.
+  bool isDivergentUse(const UseT &U) const;
+
   bool hasDivergentTerminator(const BlockT &B);
 
   void print(raw_ostream &Out) const;
 
 private:
   using ImplT = GenericUniformityAnalysisImpl<ContextT>;
 
   FunctionT *F;

llvm/include/llvm/CodeGen/MachineSSAContext.h

 
 public:
   using BlockT = MachineBasicBlock;
   using FunctionT = MachineFunction;
   using InstructionT = MachineInstr;
   using ValueRefT = Register;
   using ConstValueRefT = Register;
   static const Register ValueRefNull;
+  using UseT = MachineOperand;
   using DominatorTreeT = DominatorTreeBase<BlockT, false>;
 
   void setFunction(MachineFunction &Fn);
   MachineFunction *getFunction() const { return MF; }
 
   static MachineBasicBlock *getEntryBlock(MachineFunction &F);
   static void appendBlockDefs(SmallVectorImpl<Register> &defs,
                               const MachineBasicBlock &block);

llvm/include/llvm/IR/SSAContext.h

 
 public:
   using BlockT = BasicBlock;
   using FunctionT = Function;
   using InstructionT = Instruction;
   using ValueRefT = Value *;
   using ConstValueRefT = const Value *;
   static Value *ValueRefNull;
+  using UseT = Use;
   using DominatorTreeT = DominatorTreeBase<BlockT, false>;
 
   void setFunction(Function &Fn);
   Function *getFunction() const { return F; }
 
   static BasicBlock *getEntryBlock(Function &F);
   static const BasicBlock *getEntryBlock(const Function &F);
 

llvm/lib/Analysis/UniformityAnalysis.cpp

for (const Use &U : I.operands()) {
     if (auto *I = dyn_cast<Instruction>(&U)) {
       if (DefCycle.contains(I->getParent()))
         return true;
     }
   }
   return false;
 }
 
+template <>
+bool llvm::GenericUniformityAnalysisImpl<SSAContext>::isDivergentUse(
+    const Use &U) const {
+  const auto *V = U.get();
+  if (isDivergent(V))
+    return true;
+  if (const auto *DefInstr = dyn_cast<Instruction>(V)) {
+    const auto *UseInstr = cast<Instruction>(U.getUser());
+    return isTemporalDivergent(*UseInstr->getParent(), *DefInstr);
+  }
+  return false;
+}
+
 // This ensures explicit instantiation of
 // GenericUniformityAnalysisImpl::ImplDeleter::operator()
 template class llvm::GenericUniformityInfo<SSAContext>;
 template struct llvm::GenericUniformityAnalysisImplDeleter<
     llvm::GenericUniformityAnalysisImpl<SSAContext>>;
 
 //===----------------------------------------------------------------------===//
 //  UniformityInfoAnalysis and related pass implementations

 
 UniformityInfoWrapperPass::UniformityInfoWrapperPass() : FunctionPass(ID) {
   initializeUniformityInfoWrapperPassPass(*PassRegistry::getPassRegistry());
 }
 
 INITIALIZE_PASS_BEGIN(UniformityInfoWrapperPass, "uniformity",
                       "Uniformity Analysis", true, true)
 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(CycleInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
 INITIALIZE_PASS_END(UniformityInfoWrapperPass, "uniformity",
                     "Uniformity Analysis", true, true)
 
 void UniformityInfoWrapperPass::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.setPreservesAll();
   AU.addRequired<DominatorTreeWrapperPass>();
-  AU.addRequired<CycleInfoWrapperPass>();
+  AU.addRequiredTransitive<CycleInfoWrapperPass>();
   AU.addRequired<TargetTransformInfoWrapperPass>();
 }
 
 bool UniformityInfoWrapperPass::runOnFunction(Function &F) {
   auto &cycleInfo = getAnalysis<CycleInfoWrapperPass>().getResult();
   auto &domTree = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
   auto &targetTransformInfo =
       getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);

llvm/lib/CodeGen/MachineUniformityAnalysis.cpp

for (auto &Op : I.operands()) {
 
     auto *Def = F.getRegInfo().getVRegDef(Reg);
     if (DefCycle.contains(Def->getParent()))
       return true;
   }
   return false;
 }
 
+template <>
+bool llvm::GenericUniformityAnalysisImpl<MachineSSAContext>::isDivergentUse(
+    const MachineOperand &U) const {
+  if (!U.isReg())
+    return false;
+
+  auto Reg = U.getReg();
+  if (isDivergent(Reg))
+    return true;
+
+  const auto &RegInfo = F.getRegInfo();
+  auto *Def = RegInfo.getOneDef(Reg);
+  if (!Def)
+    return true;
+
+  auto *DefInstr = Def->getParent();
+  auto *UseInstr = U.getParent();
+  return isTemporalDivergent(*UseInstr->getParent(), *DefInstr);
+}
+
 // This ensures explicit instantiation of
 // GenericUniformityAnalysisImpl::ImplDeleter::operator()
 template class llvm::GenericUniformityInfo<MachineSSAContext>;
 template struct llvm::GenericUniformityAnalysisImplDeleter<
     llvm::GenericUniformityAnalysisImpl<MachineSSAContext>>;
 
 MachineUniformityInfo
 llvm::computeMachineUniformityInfo(MachineFunction &F,

llvm/lib/Target/AMDGPU/AMDGPUAtomicOptimizer.cpp

 /// This pass optimizes atomic operations by using a single lane of a wavefront
 /// to perform the atomic operation, thus reducing contention on that memory
 /// location.
 //
 //===----------------------------------------------------------------------===//
 
 #include "AMDGPU.h"
 #include "GCNSubtarget.h"
-#include "llvm/Analysis/LegacyDivergenceAnalysis.h"
+#include "llvm/Analysis/UniformityAnalysis.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstVisitor.h"
 #include "llvm/IR/IntrinsicsAMDGPU.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 
struct ReplacementInfo {
   unsigned ValIdx;
   bool ValDivergent;
 };
 
 class AMDGPUAtomicOptimizer : public FunctionPass,
                               public InstVisitor<AMDGPUAtomicOptimizer> {
 private:
   SmallVector<ReplacementInfo, 8> ToReplace;
-  const LegacyDivergenceAnalysis *DA;
+  const UniformityInfo *UA;
   const DataLayout *DL;
   DominatorTree *DT;
   const GCNSubtarget *ST;
   bool IsPixelShader;
 
   Value *buildReduction(IRBuilder<> &B, AtomicRMWInst::BinOp Op, Value *V,
                         Value *const Identity) const;
   Value *buildScan(IRBuilder<> &B, AtomicRMWInst::BinOp Op, Value *V,
                    Value *const Identity) const;
   Value *buildShiftRight(IRBuilder<> &B, Value *V, Value *const Identity) const;
   void optimizeAtomic(Instruction &I, AtomicRMWInst::BinOp Op, unsigned ValIdx,
                       bool ValDivergent) const;
 
 public:
   static char ID;
 
   AMDGPUAtomicOptimizer() : FunctionPass(ID) {}
 
   bool runOnFunction(Function &F) override;
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addPreserved<DominatorTreeWrapperPass>();
-    AU.addRequired<LegacyDivergenceAnalysis>();
+    AU.addRequired<UniformityInfoWrapperPass>();
     AU.addRequired<TargetPassConfig>();
   }
 
   void visitAtomicRMWInst(AtomicRMWInst &I);
   void visitIntrinsicInst(IntrinsicInst &I);
 };
 
 } // namespace
 
 char AMDGPUAtomicOptimizer::ID = 0;
 
 char &llvm::AMDGPUAtomicOptimizerID = AMDGPUAtomicOptimizer::ID;
 
 bool AMDGPUAtomicOptimizer::runOnFunction(Function &F) {
   if (skipFunction(F)) {
     return false;
   }
 
-  DA = &getAnalysis<LegacyDivergenceAnalysis>();
+  UA = &getAnalysis<UniformityInfoWrapperPass>().getUniformityInfo();
   DL = &F.getParent()->getDataLayout();
   DominatorTreeWrapperPass *const DTW =
       getAnalysisIfAvailable<DominatorTreeWrapperPass>();
   DT = DTW ? &DTW->getDomTree() : nullptr;
   const TargetPassConfig &TPC = getAnalysis<TargetPassConfig>();
   const TargetMachine &TM = TPC.getTM<TargetMachine>();
   ST = &TM.getSubtarget<GCNSubtarget>(F);
   IsPixelShader = F.getCallingConv() == CallingConv::AMDGPU_PS;
case AtomicRMWInst::UMin:
     break;
   }
 
   const unsigned PtrIdx = 0;
   const unsigned ValIdx = 1;
 
   // If the pointer operand is divergent, then each lane is doing an atomic
   // operation on a different address, and we cannot optimize that.
-  if (DA->isDivergentUse(&I.getOperandUse(PtrIdx))) {
+  if (UA->isDivergentUse(I.getOperandUse(PtrIdx))) {
     return;
   }
 
-  const bool ValDivergent = DA->isDivergentUse(&I.getOperandUse(ValIdx));
+  const bool ValDivergent = UA->isDivergentUse(I.getOperandUse(ValIdx));
 
   // If the value operand is divergent, each lane is contributing a different
   // value to the atomic calculation. We can only optimize divergent values if
   // we have DPP available on our subtarget, and the atomic operation is 32
   // bits.
   if (ValDivergent &&
       (!ST->hasDPP() || DL->getTypeSizeInBits(I.getType()) != 32)) {
     return;
void AMDGPUAtomicOptimizer::visitIntrinsicInst(IntrinsicInst &I) {
   case Intrinsic::amdgcn_struct_buffer_atomic_umax:
   case Intrinsic::amdgcn_raw_buffer_atomic_umax:
     Op = AtomicRMWInst::UMax;
     break;
   }
 
   const unsigned ValIdx = 0;
 
-  const bool ValDivergent = DA->isDivergentUse(&I.getOperandUse(ValIdx));
+  const bool ValDivergent = UA->isDivergentUse(I.getOperandUse(ValIdx));
 
   // If the value operand is divergent, each lane is contributing a different
   // value to the atomic calculation. We can only optimize divergent values if
   // we have DPP available on our subtarget, and the atomic operation is 32
   // bits.
   if (ValDivergent &&
       (!ST->hasDPP() || DL->getTypeSizeInBits(I.getType()) != 32)) {
     return;
   }
 
   // If any of the other arguments to the intrinsic are divergent, we can't
   // optimize the operation.
   for (unsigned Idx = 1; Idx < I.getNumOperands(); Idx++) {
-    if (DA->isDivergentUse(&I.getOperandUse(Idx))) {
+    if (UA->isDivergentUse(I.getOperandUse(Idx))) {
       return;
     }
   }
 
   // If we get here, we can optimize the atomic using a single wavefront-wide
   // atomic operation to do the calculation for the entire wavefront, so
   // remember the instruction so we can come back to it.
   const ReplacementInfo Info = {&I, Op, ValIdx, ValDivergent};
void AMDGPUAtomicOptimizer::optimizeAtomic(Instruction &I,
   }
 
   // And delete the original.
   I.eraseFromParent();
 }
 
 INITIALIZE_PASS_BEGIN(AMDGPUAtomicOptimizer, DEBUG_TYPE,
                       "AMDGPU atomic optimizations", false, false)
-INITIALIZE_PASS_DEPENDENCY(LegacyDivergenceAnalysis)
+INITIALIZE_PASS_DEPENDENCY(UniformityInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(TargetPassConfig)
 INITIALIZE_PASS_END(AMDGPUAtomicOptimizer, DEBUG_TYPE,
                     "AMDGPU atomic optimizations", false, false)
 
 FunctionPass *llvm::createAMDGPUAtomicOptimizerPass() {
   return new AMDGPUAtomicOptimizer();
 }

llvm/test/CodeGen/AMDGPU/divergence-at-use.ll

 ; RUN: llc -march=amdgcn -mcpu=gfx900 -amdgpu-atomic-optimizations=true < %s | FileCheck %s
-; RUN: llc -march=amdgcn -mcpu=gfx900 -amdgpu-atomic-optimizations=true < %s -use-gpu-divergence-analysis | FileCheck %s
 
 @local = addrspace(3) global i32 undef
 
 define amdgpu_kernel void @reducible() {
 ; CHECK-LABEL: reducible:
 ; CHECK-NOT: dpp
 entry:
   %x = call i32 @llvm.amdgcn.workitem.id.x()
   br label %loop
 loop:
   %i = phi i32 [ 0, %entry ], [ %i1, %loop ]
   %gep = getelementptr i32, ptr addrspace(3) @local, i32 %i
   %cond = icmp ult i32 %i, %x
   %i1 = add i32 %i, 1
   br i1 %cond, label %loop, label %exit
 exit:
   %old = atomicrmw add ptr addrspace(3) %gep, i32 %x acq_rel
   ret void
 }
 
+define amdgpu_kernel void @def_in_nested_cycle() {
+; CHECK-LABEL: def_in_nested_cycle:
+; CHECK-NOT: dpp
+entry:
+  %x = call i32 @llvm.amdgcn.workitem.id.x()
+  br label %loop
+loop:
+  %i = phi i32 [ 0, %entry ], [ 0, %innerloop ], [ %i1, %loop ]
+  %cond = icmp ult i32 %i, %x
+  %i1 = add i32 %i, 1
+  br i1 %cond, label %innerloop, label %loop
+innerloop:
+  %i.inner = phi i32 [ 0, %loop ], [ %i1.inner, %innerloop ]
+  %gep = getelementptr i32, ptr addrspace(3) @local, i32 %i
+  %i1.inner = add i32 %i, 1
+  %cond.inner = icmp ult i32 %i, %x
+  br i1 %cond, label %innerloop, label %loop
+exit:
+  %old = atomicrmw add ptr addrspace(3) %gep, i32 %x acq_rel
+  ret void
+}
+
 declare i32 @llvm.amdgcn.workitem.id.x()