[LLVM][Uniformity] Propagate temporal divergence explicitly

[sameerds]

Summary

At a cycle C with divergent exits, UA was using a naive traversal of the exiting
edges to locate blocks that may use values defined inside C. But this traversal
fails when it encounters a cycle. This is now replaced with a much simpler
propagation that iterates over every instruction in C and checks any uses that
are outside C. But such an iteration can be expensive when C is very large; the
original strategy may need to be reconsidered if there is a regression in
compilation times.

Also fixed lit tests that should have originally caught the missed propagation
of temporal divergence.

[foad]

llvm/include/llvm/ADT/GenericUniformityImpl.h

457–458

Comment should be updated. It doesn't mention I.

824

I wonder if there is a quick test you could use to avoid scanning some BBs, e.g. I think if BB does not dominate any edges that exit OuterDivCycle then instructions in BB cannot have any uses outside OuterDivCycle. This should be fairly common because it applies to the body of a simple C "for" loop.

llvm/lib/Analysis/UniformityAnalysis.cpp

72–73

Just iterate over users() instead of uses().

[dstuttard]

Confirmed that this fixes an issue we observed in an app.

[sameerds]

• only examine blocks that dominate an exit.
• FIXME about narrowing this down to live-out values at each divergent exit.

[yassingh]

llvm/lib/CodeGen/MachineUniformityAnalysis.cpp

116–135

Is the case described in the summary also triggering for MIR? If yes maybe you can add a relevant test.

[sameerds]

llvm/lib/CodeGen/MachineUniformityAnalysis.cpp

116–135

I just tried adding an identical MIR test. Unfortunately, it runs into the problem with SI_LOOP that is being fixed in D150438. I'll add the test to that review.

[foad]

Looks fine to me, thanks!

llvm/include/llvm/ADT/GenericUniformityImpl.h

820–821

Conceptually you could unify these two cases by looking for blocks B that dominate an exit edge; but I have no idea whether DomTree has good support for querying dominance on CFG edges.

824

Might be cleaner to use any_of?

[foad]

llvm/include/llvm/ADT/GenericUniformityImpl.h

821

Nit: say "DefCycle" instead of "C" (which you have not defined)?

[sameerds]

llvm/include/llvm/ADT/GenericUniformityImpl.h

820–821

I had a cursory look. The DominatorTree interface has methods to query edges dominating uses, but none to check for defs dominating edges.

llvm/include/llvm/ADT/GenericUniformityImpl.h

private:
   /// is a divergent cycle exit of \p DivCycle. Mark those instructions as
   /// divergent and push them on the worklist.
   void propagateCycleExitDivergence(const BlockT &DivExit,
                                     const CycleT &DivCycle);
 
   /// \brief Internal implementation function for propagateCycleExitDivergence.
   void analyzeCycleExitDivergence(const CycleT &OuterDivCycle);
 
   /// \brief Mark all instruction as divergent that use a value defined in \p
   /// OuterDivCycle. Push their users on the worklist.
-  void analyzeTemporalDivergence(const InstructionT &I,
+  void propagateTemporalDivergence(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 Def is divergent when read in \p ObservingBlock.
   bool isTemporalDivergent(const BlockT &ObservingBlock,
                            const InstructionT &Def) const;
 };
 
 template <typename ImplT>
 void GenericUniformityAnalysisImplDeleter<ImplT>::operator()(ImplT *Impl) {
   delete Impl;

 }
 
 template <typename ContextT>
 void GenericUniformityAnalysisImpl<ContextT>::addUniformOverride(
     const InstructionT &Instr) {
   UniformOverrides.insert(&Instr);
 }
 
-template <typename ContextT>
-void GenericUniformityAnalysisImpl<ContextT>::analyzeTemporalDivergence(
-    const InstructionT &I, const CycleT &OuterDivCycle) {
-  if (isDivergent(I))
-    return;
-
-  LLVM_DEBUG(dbgs() << "Analyze temporal divergence: " << Context.print(&I)
-                    << "\n");
-  if (isAlwaysUniform(I))
-    return;
-
-  if (!usesValueFromCycle(I, OuterDivCycle))
-    return;
-
-  if (markDivergent(I))
-    Worklist.push_back(&I);
-}
-
 // Mark all external users of values defined inside \param
 // OuterDivCycle as divergent.
 //
 // This follows all live out edges wherever they may lead. Potential
 // users of values defined inside DivCycle could be anywhere in the
 // dominance region of DivCycle (including its fringes for phi nodes).
 // A cycle C dominates a block B iff every path from the entry block
 // to B must pass through a block contained in C. If C is a reducible
 // cycle (or natural loop), C dominates B iff the header of C
 // dominates B. But in general, we iteratively examine cycle cycle
 // exits and their successors.
 template <typename ContextT>
 void GenericUniformityAnalysisImpl<ContextT>::analyzeCycleExitDivergence(
     const CycleT &OuterDivCycle) {
-  // Set of blocks that are dominated by the cycle, i.e., each is only
-  // reachable from paths that pass through the cycle.
-  SmallPtrSet<BlockT *, 16> DomRegion;
-
-  // The boundary of DomRegion, formed by blocks that are not
-  // dominated by the cycle.
-  SmallVector<BlockT *> DomFrontier;
-  OuterDivCycle.getExitBlocks(DomFrontier);
-
-  // Returns true if BB is dominated by the cycle.
-  auto isInDomRegion = [&](BlockT *BB) {
-    for (auto *P : predecessors(BB)) {
-      if (OuterDivCycle.contains(P))
-        continue;
-      if (DomRegion.count(P))
-        continue;
-      return false;
-    }
-    return true;
-  };
-
-  // Keep advancing the frontier along successor edges, while
-  // promoting blocks to DomRegion.
-  while (true) {
-    bool Promoted = false;
-    SmallVector<BlockT *> Temp;
-    for (auto *W : DomFrontier) {
-      if (!isInDomRegion(W)) {
-        Temp.push_back(W);
-        continue;
-      }
-      DomRegion.insert(W);
-      Promoted = true;
-      for (auto *Succ : successors(W)) {
-        if (DomRegion.contains(Succ))
-          continue;
-        Temp.push_back(Succ);
-      }
-    }
-    if (!Promoted)
-      break;
-
-    // Restore the set property for the temporary vector
-    llvm::sort(Temp);
-    Temp.erase(std::unique(Temp.begin(), Temp.end()), Temp.end());
-
-    DomFrontier = Temp;
-  }
-
-  // At DomFrontier, only the PHI nodes are affected by temporal
-  // divergence.
-  for (const auto *UserBlock : DomFrontier) {
-    LLVM_DEBUG(dbgs() << "Analyze phis after cycle exit: "
-                      << Context.print(UserBlock) << "\n");
-    for (const auto &Phi : UserBlock->phis()) {
-      LLVM_DEBUG(dbgs() << "  " << Context.print(&Phi) << "\n");
-      analyzeTemporalDivergence(Phi, OuterDivCycle);
-    }
-  }
-
-  // All instructions inside the dominance region are affected by
-  // temporal divergence.
-  for (const auto *UserBlock : DomRegion) {
-    LLVM_DEBUG(dbgs() << "Analyze non-phi users after cycle exit: "
-                      << Context.print(UserBlock) << "\n");
-    for (const auto &I : *UserBlock) {
-      LLVM_DEBUG(dbgs() << "  " << Context.print(&I) << "\n");
-      analyzeTemporalDivergence(I, OuterDivCycle);
+  for (auto *BB : OuterDivCycle.blocks()) {
+    for (auto &II : *BB) {
+      propagateTemporalDivergence(II, OuterDivCycle);
     }
   }
 }
 
 template <typename ContextT>
 void GenericUniformityAnalysisImpl<ContextT>::propagateCycleExitDivergence(
     const BlockT &DivExit, const CycleT &InnerDivCycle) {
   LLVM_DEBUG(dbgs() << "\tpropCycleExitDiv " << Context.print(&DivExit)

llvm/lib/Analysis/UniformityAnalysis.cpp

void llvm::GenericUniformityAnalysisImpl<SSAContext>::pushUsers(
     const Instruction &Instr) {
   assert(!isAlwaysUniform(Instr));
   if (Instr.isTerminator())
     return;
   pushUsers(cast<Value>(&Instr));
 }
 
 template <>
-bool llvm::GenericUniformityAnalysisImpl<SSAContext>::usesValueFromCycle(
-    const Instruction &I, const Cycle &DefCycle) const {
-  assert(!isAlwaysUniform(I));
-  for (const Use &U : I.operands()) {
-    if (auto *I = dyn_cast<Instruction>(&U)) {
-      if (DefCycle.contains(I->getParent()))
-        return true;
-    }
+void llvm::GenericUniformityAnalysisImpl<
+    SSAContext>::propagateTemporalDivergence(const Instruction &I,
+                                             const Cycle &DefCycle) {
+  for (const Use &U : I.uses()) {
+    auto *UserInstr = cast<Instruction>(U.getUser());
+    if (DefCycle.contains(UserInstr->getParent()))
+      continue;
+    if (markDivergent(*UserInstr))
+      Worklist.push_back(UserInstr);
   }
-  return false;
 }
 
 template <>
 bool llvm::GenericUniformityAnalysisImpl<SSAContext>::isDivergentUse(
     const Use &U) const {
   const auto *V = U.get();
   if (isDivergent(V))
     return true;

llvm/lib/CodeGen/MachineUniformityAnalysis.cpp

if (Instr.isTerminator())
     return;
   for (const MachineOperand &op : Instr.operands()) {
     if (op.isReg() && op.isDef() && op.getReg().isVirtual())
       pushUsers(op.getReg());
   }
 }
 
 template <>
-bool llvm::GenericUniformityAnalysisImpl<MachineSSAContext>::usesValueFromCycle(
-    const MachineInstr &I, const MachineCycle &DefCycle) const {
-  assert(!isAlwaysUniform(I));
+void llvm::GenericUniformityAnalysisImpl<MachineSSAContext>::
+    propagateTemporalDivergence(const MachineInstr &I,
+                                const MachineCycle &DefCycle) {
+  const auto &RegInfo = F.getRegInfo();
   for (auto &Op : I.operands()) {
-    if (!Op.isReg() || !Op.readsReg())
+    if (!Op.isReg() || !Op.isDef())
+      continue;
+    if (!Op.getReg().isVirtual())
       continue;
     auto Reg = Op.getReg();
-
-    // FIXME: Physical registers need to be properly checked instead of always
-    // returning true
-    if (Reg.isPhysical())
-      return true;
-
-    auto *Def = F.getRegInfo().getVRegDef(Reg);
-    if (DefCycle.contains(Def->getParent()))
-      return true;
+    for (MachineInstr &UserInstr : RegInfo.use_instructions(Reg)) {
+      if (DefCycle.contains(UserInstr.getParent()))
+        continue;
+      if (markDivergent(UserInstr))
+        Worklist.push_back(&UserInstr);
+    }
   }
-  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,
                                    const MachineCycleInfo &cycleInfo,

llvm/test/Analysis/UniformityAnalysis/AMDGPU/temporal_diverge.ll

X:
   %div.exity = icmp slt i32 %tid, 0
 ; CHECK: DIVERGENT: %div.exity =
   br i1 %div.exity, label %G, label %Y
 ; CHECK: DIVERGENT: br i1 %div.exity,
 
 Y:
   %div.alsouser = add i32 %uni.inc, 5
   ret void
+; CHECK: DIVERGENT: %div.alsouser =
 }
 
 
 ; temporal-divergent use of value carried by divergent loop, user is inside sibling loop
 define amdgpu_kernel void @temporal_diverge_loopuser(i32 %n, i32 %a, i32 %b) #0 {
 ; CHECK-LABEL: for function 'temporal_diverge_loopuser':
 ; CHECK-NOT: DIVERGENT: %uni.
 ; CHECK-NOT: DIVERGENT: br i1 %uni.

 ; CHECK: DIVERGENT: br i1 %div.exitx,
 
 X:
   br label %G
 
 G:
   %div.user = add i32 %uni.inc, 5
   br i1 %uni.cond, label %G, label %Y
+; CHECK: DIVERGENT: %div.user =
 
 Y:
   ret void
 }
 
 ; temporal-divergent use of value carried by divergent loop, user is inside sibling loop, defs and use are carried by a uniform loop
 define amdgpu_kernel void @temporal_diverge_loopuser_nested(i32 %n, i32 %a, i32 %b) #0 {
 ; CHECK-LABEL: for function 'temporal_diverge_loopuser_nested':
 ; CHECK-NOT: DIVERGENT: %uni.
 ; CHECK-NOT: DIVERGENT: br i1 %uni.
 
 entry:
   %tid = call i32 @llvm.amdgcn.workitem.id.x()
   %uni.cond = icmp slt i32 %a, 0
+  br label %G
+
+G:
   br label %H
 
 H:
-  %uni.merge.h = phi i32 [ 0, %entry ], [ %uni.inc, %H ]
+  %uni.merge.h = phi i32 [ 0, %G ], [ %uni.inc, %H ]
   %uni.inc = add i32 %uni.merge.h, 1
   %div.exitx = icmp slt i32 %tid, 0
   br i1 %div.exitx, label %X, label %H ; divergent branch
 ; CHECK: DIVERGENT: %div.exitx =
 ; CHECK: DIVERGENT: br i1 %div.exitx,
 
 X:
-  br label %G
-
-G:
+; CHECK: DIVERGENT: %div.user =
   %div.user = add i32 %uni.inc, 5
-  br i1 %uni.cond, label %G, label %Y
+  br i1 %uni.cond, label %X, label %G
 
 Y:
   ret void
 }
 
 
 declare i32 @llvm.amdgcn.workitem.id.x() #0
 
 attributes #0 = { nounwind readnone }