[DA][SDA] SyncDependenceAnalysis re-write

[simoll]

Summary

This patch achieves two things:

1. It breaks up the join_blocks interface between the SDA to the DA to return two separate sets for divergent loops exits and divergent, disjoint path joins.
2. It updates the SDA algorithm to run in O(n) time and improves the precision on divergent loop exits.

This fixes https://bugs.llvm.org/show_bug.cgi?id=46372 (by virtue of the improved join_blocks interface) and revealed an imprecise expected result in the Analysis/DivergenceAnalysis/AMDGPU/hidden_loopdiverge.ll test.
This should work as is (feedback on any issues welcome). I'll clean up the code as we proceed with the review.

[arsenm]

llvm/include/llvm/Analysis/SyncDependenceAnalysis.h

41–43

This is just a SetVector?

[foad]

FYI I've tested this patch by running a large chunk of Vulkan CTS with AMDVLK, and I didn't find any problems.

[simoll]

llvm/include/llvm/Analysis/SyncDependenceAnalysis.h

41–43

LoopRPO could be a SetVector. However, this is only modified during construction time and atm we discard the "unique-ing" set once RPO becomes read only.
We still need the RPOIndex map for inverse lookup.

[nhaehnle]

It's interesting that you end up using a modified post-order where loops are kept contiguous, I found the same thing helpful in my work on control flow lowering in AMDGPU.

I haven't looked into the actual divergence propagation, but some early comments on other stuff.

llvm/lib/Analysis/SyncDependenceAnalysis.cpp

141

This is relatively expensive, isn't it? And it gets executed once per loop exit instead of once per loop...

193

"\brief" comments should be above the function, not inside its body.

How is the CFG modified?

223

The explicit initializer for LoopRPO isn't needed.

224–227

This looks like it's computing a post-order, not a reverse post-order. I haven't looked into how this lines up with the rest of the algorithm, but perhaps you can rename LoopRPO to LoopPO?

232

Can this be removed now?

[sameerds]

llvm/include/llvm/Analysis/DivergenceAnalysis.h

134

The comment could be placed on the preceding line to avoid breaking up the declaration.

llvm/lib/Analysis/DivergenceAnalysis.cpp

178

It looks very strange to call a function on the whole Phi inside a loop that is iterating over its incoming values. This is actually a predicate loop that needs to be converted into a predicate function. Something like the following. Like the coding style says, this forces you to pick a name which documents what is happening.

if (!Phi)
  return;
if (!hasInterestingOperands(Phi))
  return;
if (markDivergent(*Phi))
  pushUsers(*Phi);
return;

187

Same as previous comment. The intention of introducing usesLiveOut() as a separate function was to document the loop's intention in the name of the function itself.

267

Loop::Contains is called with the same argument in each iteration. If LoopInfo is available at this point, the function could call LoopInfo::getLoopFor() just once outside the current while-loop and then compare pointers at this point.

llvm/lib/Analysis/SyncDependenceAnalysis.cpp

101

This got reordered. Might be a good idea to clang-format the whole patch.

[simoll]

Thanks for the comments so far.

[simoll]

NFC

• addressed reviewer comments, clean-up, tidy
• rebased

[simoll]

minor cleanup

llvm/lib/Analysis/DivergenceAnalysis.cpp

267

We cannot do that because`getLoopFor` would give us the inner-most loop that contains DivExit. We are looking for the outer-most loop that does not contains DivExit here. I've changed this to use the loop depth instead.

llvm/lib/Analysis/SyncDependenceAnalysis.cpp

141

Actually, this is only called once per loop header.

224–227

Yes, we compute PO and traverse the block list backwards to do the reversal when it's used.

[simoll]

Bug fix: DivLoopExit's loop may be nullptr

[simoll]

• rebased (onto clang-formatting commit for the SDA)
• DCE

[sameerds]

LGTM, to the extent of my high-level understanding of the analysis. I would recommend waiting a few days and give others a chance to finish the review. I do have some nitpicks, but feel free to decide which ones warrant a change.

llvm/lib/Analysis/DivergenceAnalysis.cpp

228

Not necessarily binding on this change, but could this be an error instead? Or perhaps an early exit with a conservative approximation?

279

typo in the debug output: "taintAndPushPhiNodes"

llvm/lib/Analysis/SyncDependenceAnalysis.cpp

116–120

Is this intended to treat each child loop as a single node when traversing a parent loop? Could we just say that here?

132

typo: line needs to be indented.

143

Is LoopHead always the same as Loop->getHeader()? Also, the complete name LoopHeader is preferred.

225

This should either be removed or put in a debug macro

257

Did you mean to start each "if" on a new line? Missing markup for an unordered list?

278

The final return seems to only update the label. How is \top defined?

[simoll]

llvm/lib/Analysis/DivergenceAnalysis.cpp

228

It is non-obvious which threads synchronize in irreducible loops.
The convergence semantics of D85603 should clarify this. Until that patch is in, i'd rather we leave this as an assertion.

llvm/lib/Analysis/SyncDependenceAnalysis.cpp

116–120

We treat child loops as single nodes to speedup the analysis (all loops are reducible and all blocks of child loops will receive the same label as their loop headers).. This is not part of the changes to the CFG that are required for correctness.

[simoll]

• NFC.
• Adressed comments.

I'll let this simmer for a last round of feedback and commit soonish, otw.

llvm/include/llvm/Analysis/DivergenceAnalysis.h

public:
   /// \brief Whether \p BB is part of the region.
   bool inRegion(const BasicBlock &BB) const;
   /// \brief Whether \p I is part of the region.
   bool inRegion(const Instruction &I) const;
 
   /// \brief Mark \p UniVal as a value that is always uniform.
   void addUniformOverride(const Value &UniVal);
 
-  /// \brief Mark \p DivVal as a value that is always divergent.
-  void markDivergent(const Value &DivVal);
+  /// \brief Mark \p DivVal as a value that is always divergent. Will not do so
+  /// if `isAlwaysUniform(DivVal)`.
+  /// \returns Whether the tracked divergence state of \p DivVal changed.
+  bool markDivergent(const Value &DivVal);
 
   /// \brief Propagate divergence to all instructions in the region.
   /// Divergence is seeded by calls to \p markDivergent.
   void compute();
 
   /// \brief Whether any value was marked or analyzed to be divergent.
   bool hasDetectedDivergence() const { return !DivergentValues.empty(); }
 
   /// \brief Whether \p Val will always return a uniform value regardless of its
   /// operands
   bool isAlwaysUniform(const Value &Val) const;
 
   /// \brief Whether \p Val is divergent at its definition.
   bool isDivergent(const Value &Val) const;
 
-  /// \brief Whether \p U is divergent. Uses of a uniform value can be divergent.
+  /// \brief Whether \p U is divergent. Uses of a uniform value can be
+  /// divergent.
   bool isDivergentUse(const Use &U) const;
 
   void print(raw_ostream &OS, const Module *) const;
 
 private:
-  bool updateTerminator(const Instruction &Term) const;
-  bool updatePHINode(const PHINode &Phi) const;
-
-  /// \brief Computes whether \p Inst is divergent based on the
-  /// divergence of its operands.
-  ///
-  /// \returns Whether \p Inst is divergent.
-  ///
-  /// This should only be called for non-phi, non-terminator instructions.
-  bool updateNormalInstruction(const Instruction &Inst) const;
+  /// \brief Mark \p Term as divergent and push all Instructions that become
+  /// divergent as a result on the worklist.
+  void analyzeControlDivergence(const Instruction &Term);
+  /// \brief Mark all phi nodes in \p JoinBlock as divergent and push them on
+  /// the worklist.
+  void taintAndPushPhiNodes(const BasicBlock &JoinBlock);
+
+  /// \brief Identify all Instructions that become divergent because \p DivExit
+  /// is a divergent loop exit of \p DivLoop. Mark those instructions as
+  /// divergent and push them on the worklist.
+  void propagateLoopExitDivergence(const BasicBlock &DivExit,
+                                   const Loop &DivLoop);
+
+  /// \brief Internal implementation function for propagateLoopExitDivergence.
+  void analyzeLoopExitDivergence(const BasicBlock &DivExit,
+                                 const Loop &OuterDivLoop);
+
+  /// \brief Mark all instruction as divergent that use a value defined in \p
+  /// OuterDivLoop. Push their users on the worklist.
+  void analyzeTemporalDivergence(const Instruction &I,
+                                 const Loop &OuterDivLoop);
 
-  /// \brief Mark users of live-out users as divergent.
-  ///
-  /// \param LoopHeader the header of the divergent loop.
-  ///
-  /// Marks all users of live-out values of the loop headed by \p LoopHeader
-  /// as divergent and puts them on the worklist.
-  void taintLoopLiveOuts(const BasicBlock &LoopHeader);
-
-  /// \brief Push all users of \p Val (in the region) to the worklist
+  /// \brief Push all users of \p Val (in the region) to the worklist.
   void pushUsers(const Value &I);
 
-  /// \brief Push all phi nodes in @block to the worklist
-  void pushPHINodes(const BasicBlock &Block);
-
-  /// \brief Mark \p Block as join divergent
-  ///
-  /// A block is join divergent if two threads may reach it from different
-  /// incoming blocks at the same time.
-  void markBlockJoinDivergent(const BasicBlock &Block) {
-    DivergentJoinBlocks.insert(&Block);
-  }
-
   /// \brief Whether \p Val is divergent when read in \p ObservingBlock.
   bool isTemporalDivergent(const BasicBlock &ObservingBlock,
                            const Value &Val) const;
 
   /// \brief Whether \p Block is join divergent
   ///
   /// (see markBlockJoinDivergent).
   bool isJoinDivergent(const BasicBlock &Block) const {
     return DivergentJoinBlocks.find(&Block) != DivergentJoinBlocks.end();
   }
 
-  /// \brief Propagate control-induced divergence to users (phi nodes and
-  /// instructions).
-  //
-  // \param JoinBlock is a divergent loop exit or join point of two disjoint
-  // paths.
-  // \returns Whether \p JoinBlock is a divergent loop exit of \p TermLoop.
-  bool propagateJoinDivergence(const BasicBlock &JoinBlock,
-                               const Loop *TermLoop);
-
-  /// \brief Propagate induced value divergence due to control divergence in \p
-  /// Term.
-  void propagateBranchDivergence(const Instruction &Term);
-
-  /// \brief Propagate divergent caused by a divergent loop exit.
-  ///
-  /// \param ExitingLoop is a divergent loop.
-  void propagateLoopDivergence(const Loop &ExitingLoop);
-
 private:
   const Function &F;
   // If regionLoop != nullptr, analysis is only performed within \p RegionLoop.
   // Otherwise, analyze the whole function
   const Loop *RegionLoop;
 
   const DominatorTree &DT;
   const LoopInfo &LI;
 
   // Recognized divergent loops
-  DenseSet<const Loop *> DivergentLoops;
+  DenseSet<const Loop *>
+      DivergentLoops; // FIXME Deprecated. For statistics only.
 
   // The SDA links divergent branches to divergent control-flow joins.
   SyncDependenceAnalysis &SDA;
 
   // Use simplified code path for LCSSA form.
   bool IsLCSSAForm;
 
   // Set of known-uniform values.
   DenseSet<const Value *> UniformOverrides;
 
   // Blocks with joining divergent control from different predecessors.
-  DenseSet<const BasicBlock *> DivergentJoinBlocks;
+  DenseSet<const BasicBlock *> DivergentJoinBlocks; // FIXME Deprecated
 
   // Detected/marked divergent values.
   DenseSet<const Value *> DivergentValues;
 
   // Internal worklist for divergence propagation.
   std::vector<const Instruction *> Worklist;
 };
 

llvm/include/llvm/Analysis/SyncDependenceAnalysis.h

 #ifndef LLVM_ANALYSIS_SYNC_DEPENDENCE_ANALYSIS_H
 #define LLVM_ANALYSIS_SYNC_DEPENDENCE_ANALYSIS_H
 
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include <memory>
+#include <unordered_map>
 
 namespace llvm {
 
 class BasicBlock;
 class DominatorTree;
 class Loop;
 class PostDominatorTree;
 
 using ConstBlockSet = SmallPtrSet<const BasicBlock *, 4>;
+struct ControlDivergenceDesc {
+  // Join points of divergent disjoint paths.
+  ConstBlockSet JoinDivBlocks;
+  // Divergent loop exits
+  ConstBlockSet LoopDivBlocks;
+};
+
+struct ModifiedRPO {
+  std::vector<const BasicBlock *> LoopRPO;
+  std::unordered_map<const BasicBlock *, unsigned> RPOIndex;
+  void appendBlock(const BasicBlock &BB) {
+    RPOIndex[&BB] = LoopRPO.size();
+    LoopRPO.push_back(&BB);
+  }
+  unsigned getIndexOf(const BasicBlock &BB) const {
+    return RPOIndex.find(&BB)->second;
+  }
+  unsigned size() const { return LoopRPO.size(); }
+  const BasicBlock *getBlockAt(unsigned Idx) const { return LoopRPO[Idx]; }
+};
 
 /// \brief Relates points of divergent control to join points in
 /// reducible CFGs.
 ///
 /// This analysis relates points of divergent control to points of converging
 /// divergent control. The analysis requires all loops to be reducible.
 class SyncDependenceAnalysis {
 public:
   ~SyncDependenceAnalysis();
   SyncDependenceAnalysis(const DominatorTree &DT, const PostDominatorTree &PDT,
                          const LoopInfo &LI);
 
   /// \brief Computes divergent join points and loop exits caused by branch
   /// divergence in \p Term.
   ///
   /// The set of blocks which are reachable by disjoint paths from \p Term.
   /// The set also contains loop exits if there two disjoint paths:
   /// one from \p Term to the loop exit and another from \p Term to the loop
   /// header. Those exit blocks are added to the returned set.
   /// If L is the parent loop of \p Term and an exit of L is in the returned
   /// set then L is a divergent loop.
-  const ConstBlockSet &join_blocks(const Instruction &Term);
-
-  /// \brief Computes divergent join points and loop exits (in the surrounding
-  /// loop) caused by the divergent loop exits of\p Loop.
-  ///
-  /// The set of blocks which are reachable by disjoint paths from the
-  /// loop exits of \p Loop.
-  /// This treats the loop as a single node in \p Loop's parent loop.
-  /// The returned set has the same properties as for join_blocks(TermInst&).
-  const ConstBlockSet &join_blocks(const Loop &Loop);
+  const ControlDivergenceDesc &join_blocks(const Instruction &Term);
 
 private:
-  static ConstBlockSet EmptyBlockSet;
+  static ControlDivergenceDesc EmptyDivergenceDesc;
+
+  ModifiedRPO LoopRPO;
 
-  ReversePostOrderTraversal<const Function *> FuncRPOT;
+  // ReversePostOrderTraversal<const Function *> FuncRPOT;
   const DominatorTree &DT;
   const PostDominatorTree &PDT;
   const LoopInfo &LI;
 
-  std::map<const Loop *, std::unique_ptr<ConstBlockSet>> CachedLoopExitJoins;
-  std::map<const Instruction *, std::unique_ptr<ConstBlockSet>>
-      CachedBranchJoins;
+  std::map<const Instruction *, std::unique_ptr<ControlDivergenceDesc>>
+      CachedControlDivDescs;
 };
 
 } // namespace llvm
 
 #endif // LLVM_ANALYSIS_SYNC_DEPENDENCE_ANALYSIS_H

llvm/lib/Analysis/DivergenceAnalysis.cpp

-//===- DivergenceAnalysis.cpp --------- Divergence Analysis Implementation -==//
+//===---- DivergenceAnalysis.cpp --- Divergence Analysis Implementation ----==//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements a general divergence analysis for loop vectorization

 
 // class DivergenceAnalysis
 DivergenceAnalysis::DivergenceAnalysis(
     const Function &F, const Loop *RegionLoop, const DominatorTree &DT,
     const LoopInfo &LI, SyncDependenceAnalysis &SDA, bool IsLCSSAForm)
     : F(F), RegionLoop(RegionLoop), DT(DT), LI(LI), SDA(SDA),
       IsLCSSAForm(IsLCSSAForm) {}
 
-void DivergenceAnalysis::markDivergent(const Value &DivVal) {
+bool DivergenceAnalysis::markDivergent(const Value &DivVal) {
+  if (isAlwaysUniform(DivVal))
+    return false;
   assert(isa<Instruction>(DivVal) || isa<Argument>(DivVal));
   assert(!isAlwaysUniform(DivVal) && "cannot be a divergent");
-  DivergentValues.insert(&DivVal);
+  return DivergentValues.insert(&DivVal).second;
 }
 
 void DivergenceAnalysis::addUniformOverride(const Value &UniVal) {
   UniformOverrides.insert(&UniVal);
 }
 
-bool DivergenceAnalysis::updateTerminator(const Instruction &Term) const {
-  if (Term.getNumSuccessors() <= 1)
-    return false;
-  if (auto *BranchTerm = dyn_cast<BranchInst>(&Term)) {
-    assert(BranchTerm->isConditional());
-    return isDivergent(*BranchTerm->getCondition());
-  }
-  if (auto *SwitchTerm = dyn_cast<SwitchInst>(&Term)) {
-    return isDivergent(*SwitchTerm->getCondition());
-  }
-  if (isa<InvokeInst>(Term)) {
-    return false; // ignore abnormal executions through landingpad
-  }
-
-  llvm_unreachable("unexpected terminator");
-}
-
-bool DivergenceAnalysis::updateNormalInstruction(const Instruction &I) const {
-  // TODO function calls with side effects, etc
-  for (const auto &Op : I.operands()) {
-    if (isDivergent(*Op))
-      return true;
-  }
-  return false;
-}
-
 bool DivergenceAnalysis::isTemporalDivergent(const BasicBlock &ObservingBlock,
                                              const Value &Val) const {
   const auto *Inst = dyn_cast<const Instruction>(&Val);
   if (!Inst)
     return false;
   // check whether any divergent loop carrying Val terminates before control
   // proceeds to ObservingBlock
   for (const auto *Loop = LI.getLoopFor(Inst->getParent());
        Loop != RegionLoop && !Loop->contains(&ObservingBlock);
        Loop = Loop->getParentLoop()) {
     if (DivergentLoops.find(Loop) != DivergentLoops.end())
       return true;
   }
 
   return false;
 }
 
-bool DivergenceAnalysis::updatePHINode(const PHINode &Phi) const {
-  // joining divergent disjoint path in Phi parent block
-  if (!Phi.hasConstantOrUndefValue() && isJoinDivergent(*Phi.getParent())) {
-    return true;
-  }
+bool DivergenceAnalysis::inRegion(const Instruction &I) const {
+  return I.getParent() && inRegion(*I.getParent());
 
-  // An incoming value could be divergent by itself.
-  // Otherwise, an incoming value could be uniform within the loop
-  // that carries its definition but it may appear divergent
-  // from outside the loop. This happens when divergent loop exits
-  // drop definitions of that uniform value in different iterations.
-  //
-  // for (int i = 0; i < n; ++i) { // 'i' is uniform inside the loop
-  //   if (i % thread_id == 0) break;    // divergent loop exit
-  // }
-  // int divI = i;                 // divI is divergent
-  for (size_t i = 0; i < Phi.getNumIncomingValues(); ++i) {
-    const auto *InVal = Phi.getIncomingValue(i);
-    if (isDivergent(*Phi.getIncomingValue(i)) ||
-        isTemporalDivergent(*Phi.getParent(), *InVal)) {
-      return true;
-    }
-  }
-  return false;
+bool DivergenceAnalysis::inRegion(const BasicBlock &BB) const {
+  return (!RegionLoop && BB.getParent() == &F) || RegionLoop->contains(&BB);
+
+void DivergenceAnalysis::pushUsers(const Value &V) {
+  const auto *I = dyn_cast<const Instruction>(&V);
+
+  if (I && I->isTerminator()) {
+    analyzeControlDivergence(*I);
+    return;
   }
 
-bool DivergenceAnalysis::inRegion(const Instruction &I) const {
-  return I.getParent() && inRegion(*I.getParent());
+  for (const auto *User : V.users()) {
+    const auto *UserInst = dyn_cast<const Instruction>(User);
+    if (!UserInst)
+      continue;
+
+    // only compute divergent inside loop
+    if (!inRegion(*UserInst))
+      continue;
+
+    // All users of divergent values are immediate divergent
+    if (markDivergent(*UserInst))
+      Worklist.push_back(UserInst);
+  }
 }
 
-bool DivergenceAnalysis::inRegion(const BasicBlock &BB) const {
-  return (!RegionLoop && BB.getParent() == &F) || RegionLoop->contains(&BB);
+void DivergenceAnalysis::analyzeTemporalDivergence(const Instruction &I,
+                                                   const Loop &OuterDivLoop) {
+  if (isAlwaysUniform(I))
+    return;
+  if (isDivergent(I))
+    return;
+
+  LLVM_DEBUG(dbgs() << "Analyze temporal divergence: " << I.getName() << "\n");
+  auto *Phi = dyn_cast<const PHINode>(&I);
+  if (Phi) {
+    for (const auto &InVal : Phi->incoming_values()) {
+      const auto *InInst = dyn_cast<const Instruction>(&InVal);
+      if (!InInst)
+        continue;
+      LLVM_DEBUG(dbgs() << "\tinspecting: " << *InInst << "\n");
+      if (!OuterDivLoop.contains(InInst))
+        continue;
+
+      if (markDivergent(*Phi))
+        pushUsers(*Phi);
+      return;
+    }
+    return;
   }
 
-static bool usesLiveOut(const Instruction &I, const Loop *DivLoop) {
-  for (auto &Op : I.operands()) {
-    auto *OpInst = dyn_cast<Instruction>(&Op);
+  assert(!IsLCSSAForm && "T");
+  for (const auto &Op : I.operands()) {
+    const auto *OpInst = dyn_cast<const Instruction>(&Op);
     if (!OpInst)
       continue;
-    if (DivLoop->contains(OpInst->getParent()))
-      return true;
+    if (!OuterDivLoop.contains(OpInst))
+      continue;
+    if (markDivergent(I))
+      pushUsers(I);
+    return;
   }
-  return false;
 }
 
 // marks all users of loop-carried values of the loop headed by LoopHeader as
 // divergent
-void DivergenceAnalysis::taintLoopLiveOuts(const BasicBlock &LoopHeader) {
-  auto *DivLoop = LI.getLoopFor(&LoopHeader);
-  assert(DivLoop && "loopHeader is not actually part of a loop");
+void DivergenceAnalysis::analyzeLoopExitDivergence(const BasicBlock &DivExit,
+                                                   const Loop &OuterDivLoop) {
+  // All users are in immediate exit blocks
+  if (IsLCSSAForm) {
+    for (auto &Phi : DivExit.phis()) {
+      analyzeTemporalDivergence(Phi, OuterDivLoop);
+    }
+    return;
+  }
 
-  SmallVector<BasicBlock *, 8> TaintStack;
-  DivLoop->getExitBlocks(TaintStack);
+  // For non-LCSSA we have to follow all live out edges wherever they may lead.
+  const BasicBlock &LoopHeader = *OuterDivLoop.getHeader();
+  SmallVector<const BasicBlock *, 8> TaintStack;
+  TaintStack.push_back(&DivExit);
 
   // Otherwise potential users of loop-carried values could be anywhere in the
   // dominance region of DivLoop (including its fringes for phi nodes)
   DenseSet<const BasicBlock *> Visited;
-  for (auto *Block : TaintStack) {
-    Visited.insert(Block);
-  }
-  Visited.insert(&LoopHeader);
+  Visited.insert(&DivExit);
 
-  while (!TaintStack.empty()) {
+  do {
     auto *UserBlock = TaintStack.back();
     TaintStack.pop_back();
 
     // don't spread divergence beyond the region
     if (!inRegion(*UserBlock))
       continue;
 
-    assert(!DivLoop->contains(UserBlock) &&
+    assert(!OuterDivLoop.contains(UserBlock) &&
            "irreducible control flow detected");
 
     // phi nodes at the fringes of the dominance region
     if (!DT.dominates(&LoopHeader, UserBlock)) {
       // all PHI nodes of UserBlock become divergent
       for (auto &Phi : UserBlock->phis()) {
-        Worklist.push_back(&Phi);
+        analyzeTemporalDivergence(Phi, OuterDivLoop);
       }
       continue;
     }
 
-    // taint outside users of values carried by DivLoop
+    // Taint outside users of values carried by OuterDivLoop.
     for (auto &I : *UserBlock) {
-      if (isAlwaysUniform(I))
-        continue;
-      if (isDivergent(I))
-        continue;
-      if (!usesLiveOut(I, DivLoop))
-        continue;
-
-      markDivergent(I);
-      if (I.isTerminator()) {
-        propagateBranchDivergence(I);
-      } else {
-        pushUsers(I);
-      }
+      analyzeTemporalDivergence(I, OuterDivLoop);
     }
 
     // visit all blocks in the dominance region
     for (auto *SuccBlock : successors(UserBlock)) {
       if (!Visited.insert(SuccBlock).second) {
         continue;
       }
       TaintStack.push_back(SuccBlock);
     }
-  }
-}
-
-void DivergenceAnalysis::pushPHINodes(const BasicBlock &Block) {
-  for (const auto &Phi : Block.phis()) {
-    if (isDivergent(Phi))
-      continue;
-    Worklist.push_back(&Phi);
-  }
-}
-
-void DivergenceAnalysis::pushUsers(const Value &V) {
-  for (const auto *User : V.users()) {
-    const auto *UserInst = dyn_cast<const Instruction>(User);
-    if (!UserInst)
-      continue;
-
-    if (isDivergent(*UserInst))
-      continue;
-
-    // only compute divergent inside loop
-    if (!inRegion(*UserInst))
-      continue;
-    Worklist.push_back(UserInst);
-  }
+  } while (!TaintStack.empty());
 }
 
-bool DivergenceAnalysis::propagateJoinDivergence(const BasicBlock &JoinBlock,
-                                                 const Loop *BranchLoop) {
-  LLVM_DEBUG(dbgs() << "\tpropJoinDiv " << JoinBlock.getName() << "\n");
+void DivergenceAnalysis::propagateLoopExitDivergence(const BasicBlock &DivExit,
+                                                     const Loop &InnerDivLoop) {
+  LLVM_DEBUG(dbgs() << "\tpropLoopExitDiv " << DivExit.getName() << "\n");
+
+  // Find outer-most loop that does not contain \p DivExit
+  // TODO there should be a way to accelerate this
+  const Loop *DivLoop = &InnerDivLoop;
+  const Loop *OuterDivLoop = DivLoop;
+  while (DivLoop) {
+    DivergentLoops.insert(DivLoop); // all crossed loops are divergent
+    OuterDivLoop = DivLoop;
+    DivLoop = DivLoop->getParentLoop();
+    if (!DivLoop || DivLoop->contains(&DivExit))
+      break;
+  }
+  LLVM_DEBUG(dbgs() << "\tOuter-most left loop: " << OuterDivLoop->getName()
+                    << "\n");
+
+  analyzeLoopExitDivergence(DivExit, *OuterDivLoop);
+}
+
+// this is a divergent join point - mark all phi nodes as divergent and push
+// them onto the stack.
+void DivergenceAnalysis::taintAndPushPhiNodes(const BasicBlock &JoinBlock) {
+  LLVM_DEBUG(dbgs() << "taintAndPhiNodes in " << JoinBlock.getName() << "\n");
 
   // ignore divergence outside the region
   if (!inRegion(JoinBlock)) {
-    return false;
+    return;
   }
 
   // push non-divergent phi nodes in JoinBlock to the worklist
-  pushPHINodes(JoinBlock);
-
-  // disjoint-paths divergent at JoinBlock
-  markBlockJoinDivergent(JoinBlock);
-
-  // JoinBlock is a divergent loop exit
-  return BranchLoop && !BranchLoop->contains(&JoinBlock);
+  for (const auto &Phi : JoinBlock.phis()) {
+    if (isDivergent(Phi))
+      continue;
+    // FIXME Theoretically ,the 'undef' value could be replaced by any other
+    // value causing spurious divergence.
+    if (Phi.hasConstantOrUndefValue())
+      continue;
+    if (markDivergent(Phi))
+      Worklist.push_back(&Phi);
+  }
 }
 
-void DivergenceAnalysis::propagateBranchDivergence(const Instruction &Term) {
-  LLVM_DEBUG(dbgs() << "propBranchDiv " << Term.getParent()->getName() << "\n");
-
-  markDivergent(Term);
+void DivergenceAnalysis::analyzeControlDivergence(const Instruction &Term) {
+  LLVM_DEBUG(dbgs() << "analyzeControlDiv " << Term.getParent()->getName()
+                    << "\n");
 
   // Don't propagate divergence from unreachable blocks.
   if (!DT.isReachableFromEntry(Term.getParent()))
     return;
 
   const auto *BranchLoop = LI.getLoopFor(Term.getParent());
 
-  // whether there is a divergent loop exit from BranchLoop (if any)
-  bool IsBranchLoopDivergent = false;
+  const auto &DivDesc = SDA.join_blocks(Term);
 
-  // iterate over all blocks reachable by disjoint from Term within the loop
-  // also iterates over loop exits that become divergent due to Term.
-  for (const auto *JoinBlock : SDA.join_blocks(Term)) {
-    IsBranchLoopDivergent |= propagateJoinDivergence(*JoinBlock, BranchLoop);
+  // Iterate over all blocks now reachable by a disjoint path join
+  for (const auto *JoinBlock : DivDesc.JoinDivBlocks) {
+    taintAndPushPhiNodes(*JoinBlock);
   }
 
-  // Branch loop is a divergent loop due to the divergent branch in Term
-  if (IsBranchLoopDivergent) {
-    assert(BranchLoop);
-    if (!DivergentLoops.insert(BranchLoop).second) {
-      return;
-    }
-    propagateLoopDivergence(*BranchLoop);
-  }
-}
-
-void DivergenceAnalysis::propagateLoopDivergence(const Loop &ExitingLoop) {
-  LLVM_DEBUG(dbgs() << "propLoopDiv " << ExitingLoop.getName() << "\n");
-
-  // don't propagate beyond region
-  if (!inRegion(*ExitingLoop.getHeader()))
-    return;
-
-  const auto *BranchLoop = ExitingLoop.getParentLoop();
-
-  // Uses of loop-carried values could occur anywhere
-  // within the dominance region of the definition. All loop-carried
-  // definitions are dominated by the loop header (reducible control).
-  // Thus all users have to be in the dominance region of the loop header,
-  // except PHI nodes that can also live at the fringes of the dom region
-  // (incoming defining value).
-  if (!IsLCSSAForm)
-    taintLoopLiveOuts(*ExitingLoop.getHeader());
-
-  // whether there is a divergent loop exit from BranchLoop (if any)
-  bool IsBranchLoopDivergent = false;
-
-  // iterate over all blocks reachable by disjoint paths from exits of
-  // ExitingLoop also iterates over loop exits (of BranchLoop) that in turn
-  // become divergent.
-  for (const auto *JoinBlock : SDA.join_blocks(ExitingLoop)) {
-    IsBranchLoopDivergent |= propagateJoinDivergence(*JoinBlock, BranchLoop);
-  }
-
-  // Branch loop is a divergent due to divergent loop exit in ExitingLoop
-  if (IsBranchLoopDivergent) {
-    assert(BranchLoop);
-    if (!DivergentLoops.insert(BranchLoop).second) {
-      return;
-    }
-    propagateLoopDivergence(*BranchLoop);
+  assert(DivDesc.LoopDivBlocks.empty() || BranchLoop);
+  for (const auto *DivExitBlock : DivDesc.LoopDivBlocks) {
+    propagateLoopExitDivergence(*DivExitBlock, *BranchLoop);
   }
 }
 
 void DivergenceAnalysis::compute() {
-  for (auto *DivVal : DivergentValues) {
+  // Initialize worklist.
+  auto DivValuesCopy = DivergentValues;
+  for (auto *DivVal : DivValuesCopy) {
+    assert(isDivergent(*DivVal) && "Worklist invariant violated!");
     pushUsers(*DivVal);
   }
 
-  // propagate divergence
+  // All values on the Worklist are divergent.
+  // Their users may not have been updated yed.
   while (!Worklist.empty()) {
     const Instruction &I = *Worklist.back();
     Worklist.pop_back();
 
-    // maintain uniformity of overrides
-    if (isAlwaysUniform(I))
-      continue;
-
-    bool WasDivergent = isDivergent(I);
-    if (WasDivergent)
-      continue;
-
-    // propagate divergence caused by terminator
-    if (I.isTerminator()) {
-      if (updateTerminator(I)) {
-        // propagate control divergence to affected instructions
-        propagateBranchDivergence(I);
-        continue;
-      }
-    }
-
-    // update divergence of I due to divergent operands
-    bool DivergentUpd = false;
-    const auto *Phi = dyn_cast<const PHINode>(&I);
-    if (Phi) {
-      DivergentUpd = updatePHINode(*Phi);
-    } else {
-      DivergentUpd = updateNormalInstruction(I);
-    }
-
     // propagate value divergence to users
-    if (DivergentUpd) {
-      markDivergent(I);
-      pushUsers(I);
-    }
-  }
-}
+    assert(isDivergent(I) && "Worklist invariant violated!");
 
 bool DivergenceAnalysis::isAlwaysUniform(const Value &V) const {
   return UniformOverrides.find(&V) != UniformOverrides.end();
 }
 
 bool DivergenceAnalysis::isDivergent(const Value &V) const {
   return DivergentValues.find(&V) != DivergentValues.end();
 }

 }
 
 // class GPUDivergenceAnalysis
 GPUDivergenceAnalysis::GPUDivergenceAnalysis(Function &F,
                                              const DominatorTree &DT,
                                              const PostDominatorTree &PDT,
                                              const LoopInfo &LI,
                                              const TargetTransformInfo &TTI)
-    : SDA(DT, PDT, LI), DA(F, nullptr, DT, LI, SDA, false) {
+    : SDA(DT, PDT, LI), DA(F, nullptr, DT, LI, SDA, /* LCSSA */ false) {
   for (auto &I : instructions(F)) {
     if (TTI.isSourceOfDivergence(&I)) {
       DA.markDivergent(I);
     } else if (TTI.isAlwaysUniform(&I)) {
       DA.addUniformOverride(I);
     }
   }
   for (auto &Arg : F.args()) {

llvm/lib/Analysis/SyncDependenceAnalysis.cpp

-//===- SyncDependenceAnalysis.cpp - Divergent Branch Dependence Calculation
-//--===//
+//===--- SyncDependenceAnalysis.cpp - Compute Control Divergence Effects --===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements an algorithm that returns for a divergent branch

 // by two disjoins paths from A.
 //
 // -- Remarks --
 // In case of loop exits we need to check the disjoint path criterion for loops
 // [2]. To this end, we check whether the definition of x differs between the
 // loop exit and the loop header (_after_ SSA construction).
 //
 //===----------------------------------------------------------------------===//
+#include "llvm/Analysis/SyncDependenceAnalysis.h"
 #include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Analysis/PostDominators.h"
-#include "llvm/Analysis/SyncDependenceAnalysis.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
 
+#include <functional>
 #include <stack>
 #include <unordered_set>
 
 #define DEBUG_TYPE "sync-dependence"
 
+// Custom RPO framework
+namespace {
+using namespace llvm;
+
+using RPOCB = std::function<void(const BasicBlock &)>;
+using VisitedSet = std::set<const BasicBlock *>;
+using BlockStack = std::vector<const BasicBlock *>;
+
+// forward
+static void ComputeLoopPO(const LoopInfo &LI, Loop &Loop, RPOCB CallBack,
+                          VisitedSet &Finalized);
+
+// for a nested region (top-level loop or nested loop)
+static void ComputeStackPO(BlockStack &Stack, const LoopInfo &LI,
+                           const BasicBlock *LoopHead, Loop *Loop,
+                           RPOCB CallBack, VisitedSet &Finalized) {
+  while (!Stack.empty()) {
+    const auto *NextBB = Stack.back();
+
+    auto *NestedLoop = LI.getLoopFor(NextBB);
+    bool IsNestedLoop = NestedLoop != Loop;
+
+    // Treat the loop as a node
+    if (IsNestedLoop) {
+      SmallVector<BasicBlock *, 3> NestedExits;
+      NestedLoop->getUniqueExitBlocks(NestedExits);
+      bool PushedNodes = false;
+      for (const auto *NestedExitBB : NestedExits) {
+        if (NestedExitBB == LoopHead)
+          continue;
+        if (Loop && !Loop->contains(NestedExitBB))
+          continue;
+        if (Finalized.count(NestedExitBB))
+          continue;
+        PushedNodes = true;
+        Stack.push_back(NestedExitBB);
+      }
+      if (!PushedNodes) {
+        // All loop exits finalized -> finish this node
+        Stack.pop_back();
+        ComputeLoopPO(LI, *NestedLoop, CallBack, Finalized);
+      }
+      continue;
+    }
+
+    // DAG-style
+    bool PushedNodes = false;
+    for (const auto *SuccBB : successors(NextBB)) {
+      if (SuccBB == LoopHead)
+        continue;
+      if (Loop && !Loop->contains(SuccBB))
+        continue;
+      if (Finalized.count(SuccBB))
+        continue;
+      PushedNodes = true;
+      Stack.push_back(SuccBB);
+    }
+    if (!PushedNodes) {
+      // Never push nodes twice
+      Stack.pop_back();
+      if (!Finalized.insert(NextBB).second)
+        continue;
+      CallBack(*NextBB);
+    }
+  }
+}
+
+static void ComputeTopLevelPO(Function &F, const LoopInfo &LI, RPOCB CallBack) {
+  VisitedSet Finalized;
+  BlockStack Stack;
+  Stack.reserve(24); // FIXME magic number
+  Stack.push_back(&F.getEntryBlock());
+  ComputeStackPO(Stack, LI, nullptr, nullptr, CallBack, Finalized);
+}
+
+static void ComputeLoopPO(const LoopInfo &LI, Loop &Loop, RPOCB CallBack,
+                          VisitedSet &Finalized) {
+  /// \brief Call CallBack on all loop blocks in the modified CFG
+  std::vector<const BasicBlock *> Stack;
+  const auto *LoopHead = Loop.getHeader();
+
+  // Visit the header last
+  Finalized.insert(LoopHead);
+  CallBack(*LoopHead);
+
+  // Initialize with immediate successors
+  for (const auto *BB : successors(LoopHead)) {
+    if (!Loop.contains(BB))
+      continue;
+    if (BB == LoopHead)
+      continue;
+    Stack.push_back(BB);
+  }
+
+  // Compute PO inside region
+  ComputeStackPO(Stack, LI, LoopHead, &Loop, CallBack, Finalized);
+}
+
+} // namespace
+
 namespace llvm {
 
-ConstBlockSet SyncDependenceAnalysis::EmptyBlockSet;
+ControlDivergenceDesc SyncDependenceAnalysis::EmptyDivergenceDesc;
 
 SyncDependenceAnalysis::SyncDependenceAnalysis(const DominatorTree &DT,
                                                const PostDominatorTree &PDT,
                                                const LoopInfo &LI)
-    : FuncRPOT(DT.getRoot()->getParent()), DT(DT), PDT(PDT), LI(LI) {}
+    : LoopRPO(), DT(DT), PDT(PDT), LI(LI) {
+  ComputeTopLevelPO(*DT.getRoot()->getParent(), LI, [&](const BasicBlock &BB) {
+    // errs() << BB.getName() << "\n";
+    LoopRPO.appendBlock(BB);
+  });
+}
 
 SyncDependenceAnalysis::~SyncDependenceAnalysis() {}
 
 using FunctionRPOT = ReversePostOrderTraversal<const Function *>;
 
 // divergence propagator for reducible CFGs
 struct DivergencePropagator {
-  const FunctionRPOT &FuncRPOT;
+  const ModifiedRPO &LoopRPOT;
   const DominatorTree &DT;
   const PostDominatorTree &PDT;
   const LoopInfo &LI;
+  const BasicBlock &DivTermBlock;
 
-  // identified join points
-  std::unique_ptr<ConstBlockSet> JoinBlocks;
+  // if BlockLabels[IndexOf(B)] == C then C is the dominating definition at
+  // block B if BlockLabels[IndexOf(B)] ~ undef then we haven't seen B yet if
+  // BlockLabels[IndexOf(B)] == B then B is a join point of disjoint paths from
+  // X or B is an immediate successor of X (initial value).
+  using BlockLabelVec = std::vector<const BasicBlock *>;
+  BlockLabelVec BlockLabels;
+  // divergent join and loop exit descriptor.
+  std::unique_ptr<ControlDivergenceDesc> DivDesc;
 
   // reached loop exits (by a path disjoint to a path to the loop header)
-  SmallPtrSet<const BasicBlock *, 4> ReachedLoopExits;
-
-  // if DefMap[B] == C then C is the dominating definition at block B
-  // if DefMap[B] ~ undef then we haven't seen B yet
-  // if DefMap[B] == B then B is a join point of disjoint paths from X or B is
-  // an immediate successor of X (initial value).
-  using DefiningBlockMap = std::map<const BasicBlock *, const BasicBlock *>;
-  DefiningBlockMap DefMap;
+  // SmallPtrSet<const BasicBlock *, 4> ReachedLoopExits;
 
   // all blocks with pending visits
-  std::unordered_set<const BasicBlock *> PendingUpdates;
+  // std::unordered_set<const BasicBlock *> PendingUpdates;
 
-  DivergencePropagator(const FunctionRPOT &FuncRPOT, const DominatorTree &DT,
-                       const PostDominatorTree &PDT, const LoopInfo &LI)
-      : FuncRPOT(FuncRPOT), DT(DT), PDT(PDT), LI(LI),
-        JoinBlocks(new ConstBlockSet) {}
-
-  // set the definition at @block and mark @block as pending for a visit
-  void addPending(const BasicBlock &Block, const BasicBlock &DefBlock) {
-    bool WasAdded = DefMap.emplace(&Block, &DefBlock).second;
-    if (WasAdded)
-      PendingUpdates.insert(&Block);
-  }
+  DivergencePropagator(const ModifiedRPO &LoopRPOT, const DominatorTree &DT,
+                       const PostDominatorTree &PDT, const LoopInfo &LI,
+                       const BasicBlock &DivTermBlock)
+      : LoopRPOT(LoopRPOT), DT(DT), PDT(PDT), LI(LI),
+        DivTermBlock(DivTermBlock), BlockLabels(LoopRPOT.size(), nullptr),
+        DivDesc(new ControlDivergenceDesc) {}
 
   void printDefs(raw_ostream &Out) {
-    Out << "Propagator::DefMap {\n";
-    for (const auto *Block : FuncRPOT) {
-      auto It = DefMap.find(Block);
-      Out << Block->getName() << " : ";
-      if (It == DefMap.end()) {
-        Out << "\n";
+    Out << "Propagator::BlockLabels {\n";
+    for (int i = (int)BlockLabels.size() - 1; i > 0; --i) {
+      const auto *Label = BlockLabels[i];
+      Out << LoopRPOT.getBlockAt(i)->getName().str() << "(" << i << ") : ";
+      if (!Label) {
+        Out << "<null>\n";
       } else {
-        const auto *DefBlock = It->second;
-        Out << (DefBlock ? DefBlock->getName() : "<null>") << "\n";
+        Out << Label->getName() << "\n";
       }
     }
     Out << "}\n";
   }
 
-  // process @succBlock with reaching definition @defBlock
+  // Push a definition (\p DefBlock) to \p SuccBlock and return whether this
+  // raises the definition to '\top'
+  // Updates MaxNextBlock to the highest block index that needs to be visited
+  // (assuming MaxNextBlock already points to an interesting block)
+  bool computeJoin(const BasicBlock &SuccBlock, const BasicBlock &PushedLabel) {
+    auto SuccIdx = LoopRPOT.getIndexOf(SuccBlock);
+
+    // unset or same reaching label
+    const auto *OldLabel = BlockLabels[SuccIdx];
+    if (!OldLabel || (OldLabel == &PushedLabel)) {
+      BlockLabels[SuccIdx] = &PushedLabel;
+      return false;
+    }
+
+    // Update the definition
+    BlockLabels[SuccIdx] = &SuccBlock;
+    return true;
+  }
+
+  // visiting a virtual loop exit edge from the loop header --> temporal
+  // divergence on join
+  bool visitLoopExitEdge(const BasicBlock &ExitBlock,
+                         const BasicBlock &DefBlock, bool FromParentLoop) {
+    // Pushing from a non-parent loop cannot cause temporal divergence.
+    if (!FromParentLoop)
+      return visitEdge(ExitBlock, DefBlock);
+
+    if (!computeJoin(ExitBlock, DefBlock))
+      return false;
+
+    // Identified a divergent loop exit
+    DivDesc->LoopDivBlocks.insert(&ExitBlock);
+    LLVM_DEBUG(dbgs() << "\tDivergent loop exit: " << ExitBlock.getName()
+                      << "\n");
+    return true;
+  }
+
+  // process \p succBlock with reaching definition @defBlock
   // the original divergent branch was in @parentLoop (if any)
-  void visitSuccessor(const BasicBlock &SuccBlock, const Loop *ParentLoop,
-                      const BasicBlock &DefBlock) {
+  bool visitEdge(const BasicBlock &SuccBlock, const BasicBlock &DefBlock) {
+    if (!computeJoin(SuccBlock, DefBlock))
+      return false;
 
-    // @succBlock is a loop exit
-    if (ParentLoop && !ParentLoop->contains(&SuccBlock)) {
-      DefMap.emplace(&SuccBlock, &DefBlock);
-      ReachedLoopExits.insert(&SuccBlock);
-      return;
-    }
-
-    // first reaching def?
-    auto ItLastDef = DefMap.find(&SuccBlock);
-    if (ItLastDef == DefMap.end()) {
-      addPending(SuccBlock, DefBlock);
-      return;
-    }
-
-    // a join of at least two definitions
-    if (ItLastDef->second != &DefBlock) {
-      // do we know this join already?
-      if (!JoinBlocks->insert(&SuccBlock).second)
-        return;
-
-      // update the definition
-      addPending(SuccBlock, SuccBlock);
-    }
+    // Divergent, disjoint paths join.
+    DivDesc->JoinDivBlocks.insert(&SuccBlock);
+    LLVM_DEBUG(dbgs() << "\tDivergent join: " << SuccBlock.getName());
+    return true;
   }
 
-  // find all blocks reachable by two disjoint paths from @rootTerm.
-  // This method works for both divergent terminators and loops with
-  // divergent exits.
-  // @rootBlock is either the block containing the branch or the header of the
-  // divergent loop.
-  // @nodeSuccessors is the set of successors of the node (Loop or Terminator)
-  // headed by @rootBlock.
-  // @parentLoop is the parent loop of the Loop or the loop that contains the
-  // Terminator.
-  template <typename SuccessorIterable>
-  std::unique_ptr<ConstBlockSet>
-  computeJoinPoints(const BasicBlock &RootBlock,
-                    SuccessorIterable NodeSuccessors, const Loop *ParentLoop) {
-    assert(JoinBlocks);
+  std::unique_ptr<ControlDivergenceDesc> computeJoinPoints() {
+    assert(DivDesc);
 
-    LLVM_DEBUG(dbgs() << "SDA:computeJoinPoints. Parent loop: " << (ParentLoop ? ParentLoop->getName() : "<null>") << "\n" );
+    LLVM_DEBUG(dbgs() << "SDA:computeJoinPoints: " << DivTermBlock.getName()
+                      << "\n");
 
-    // bootstrap with branch targets
-    for (const auto *SuccBlock : NodeSuccessors) {
-      DefMap.emplace(SuccBlock, SuccBlock);
+    const auto *DivBlockLoop = LI.getLoopFor(&DivTermBlock);
 
-      if (ParentLoop && !ParentLoop->contains(SuccBlock)) {
-        // immediate loop exit from node.
-        ReachedLoopExits.insert(SuccBlock);
-      } else {
-        // regular successor
-        PendingUpdates.insert(SuccBlock);
-      }
-    }
+    // Early stopping criterion
+    int FloorIdx = LoopRPOT.size() - 1;
+    const BasicBlock *FloorLabel = nullptr;
 
-    LLVM_DEBUG(
-      dbgs() << "SDA: rpo order:\n";
-      for (const auto * RpoBlock : FuncRPOT) {
-        dbgs() << "- " << RpoBlock->getName() << "\n";
-      }
-    );
+    // bootstrap with branch targets
+    int BlockIdx = 0;
 
-    auto ItBeginRPO = FuncRPOT.begin();
-    auto ItEndRPO = FuncRPOT.end();
+    for (const auto *SuccBlock : successors(&DivTermBlock)) {
+      auto SuccIdx = LoopRPOT.getIndexOf(*SuccBlock);
+      BlockLabels[SuccIdx] = SuccBlock;
+
+      // Find the successor with the highest index to start with
+      BlockIdx = std::max<int>(BlockIdx, SuccIdx);
+      FloorIdx = std::min<int>(FloorIdx, SuccIdx);
 
-    // skip until term (TODO RPOT won't let us start at @term directly)
-    for (; *ItBeginRPO != &RootBlock; ++ItBeginRPO) {
-      assert(ItBeginRPO != ItEndRPO && "Unable to find RootBlock");
+      // Identify immediate divergent loop exits
+      if (!DivBlockLoop)
+        continue;
+
+      const auto *BlockLoop = LI.getLoopFor(SuccBlock);
+      if (BlockLoop && DivBlockLoop->contains(BlockLoop))
+        continue;
+      DivDesc->LoopDivBlocks.insert(SuccBlock);
+      LLVM_DEBUG(dbgs() << "\tImmediate divergent loop exit: "
+                        << SuccBlock->getName() << "\n");
     }
 
     // propagate definitions at the immediate successors of the node in RPO
-    auto ItBlockRPO = ItBeginRPO;
-    while ((++ItBlockRPO != ItEndRPO) &&
-           !PendingUpdates.empty()) {
-      const auto *Block = *ItBlockRPO;
-      LLVM_DEBUG(dbgs() << "SDA::joins. visiting " << Block->getName() << "\n");
+    for (; BlockIdx >= FloorIdx; --BlockIdx) {
+      LLVM_DEBUG(dbgs() << "Before next visit:\n"; printDefs(dbgs()));
+
+      // Any label available here
+      const auto *Label = BlockLabels[BlockIdx];
+      if (!Label)
+        continue;
 
-      // skip Block if not pending update
-      auto ItPending = PendingUpdates.find(Block);
-      if (ItPending == PendingUpdates.end())
-        continue;
-      PendingUpdates.erase(ItPending);
-
-      // propagate definition at Block to its successors
-      auto ItDef = DefMap.find(Block);
-      const auto *DefBlock = ItDef->second;
-      assert(DefBlock);
+      // Ok. Get the block
+      const auto *Block = LoopRPOT.getBlockAt(BlockIdx);
+      LLVM_DEBUG(dbgs() << "SDA::joins. visiting " << Block->getName() << "\n");
 
       auto *BlockLoop = LI.getLoopFor(Block);
-      if (ParentLoop &&
-          (ParentLoop != BlockLoop && ParentLoop->contains(BlockLoop))) {
-        // if the successor is the header of a nested loop pretend its a
-        // single node with the loop's exits as successors
+      bool IsLoopHeader = BlockLoop && BlockLoop->getHeader() == Block;
+      bool CausedJoin = false;
+      int LoweredFloorIdx = FloorIdx;
+      if (IsLoopHeader) {
+        // Disconnect from immediate successors and propagate directly to loop
+        // exits.
         SmallVector<BasicBlock *, 4> BlockLoopExits;
         BlockLoop->getExitBlocks(BlockLoopExits);
+
+        bool IsParentLoop = BlockLoop->contains(&DivTermBlock);
         for (const auto *BlockLoopExit : BlockLoopExits) {
-          visitSuccessor(*BlockLoopExit, ParentLoop, *DefBlock);
+          CausedJoin |= visitLoopExitEdge(*BlockLoopExit, *Label, IsParentLoop);
+          LoweredFloorIdx = std::min<int>(LoweredFloorIdx,
+                                          LoopRPOT.getIndexOf(*BlockLoopExit));
         }
-
       } else {
-        // the successors are either on the same loop level or loop exits
+        // Acyclic successor case
         for (const auto *SuccBlock : successors(Block)) {
-          visitSuccessor(*SuccBlock, ParentLoop, *DefBlock);
+          CausedJoin |= visitEdge(*SuccBlock, *Label);
+          LoweredFloorIdx =
+              std::min<int>(LoweredFloorIdx, LoopRPOT.getIndexOf(*SuccBlock));
         }
       }
-    }
-
-    LLVM_DEBUG(dbgs() << "SDA::joins. After propagation:\n"; printDefs(dbgs()));
 
-    // We need to know the definition at the parent loop header to decide
-    // whether the definition at the header is different from the definition at
-    // the loop exits, which would indicate a divergent loop exits.
-    //
-    // A // loop header
-    // |
-    // B // nested loop header
-    // |
-    // C -> X (exit from B loop) -..-> (A latch)
-    // |
-    // D -> back to B (B latch)
-    // |
-    // proper exit from both loops
-    //
-    // analyze reached loop exits
-    if (!ReachedLoopExits.empty()) {
-      const BasicBlock *ParentLoopHeader =
-          ParentLoop ? ParentLoop->getHeader() : nullptr;
-
-      assert(ParentLoop);
-      auto ItHeaderDef = DefMap.find(ParentLoopHeader);
-      const auto *HeaderDefBlock = (ItHeaderDef == DefMap.end()) ? nullptr : ItHeaderDef->second;
-
-      LLVM_DEBUG(printDefs(dbgs()));
-      assert(HeaderDefBlock && "no definition at header of carrying loop");
-
-      for (const auto *ExitBlock : ReachedLoopExits) {
-        auto ItExitDef = DefMap.find(ExitBlock);
-        assert((ItExitDef != DefMap.end()) &&
-               "no reaching def at reachable loop exit");
-        if (ItExitDef->second != HeaderDefBlock) {
-          JoinBlocks->insert(ExitBlock);
-        }
+      // Floor update
+      if (CausedJoin) {
+        // 1. Different labels pushed to successors
+        FloorIdx = LoweredFloorIdx;
+      } else if (FloorLabel != Label) {
+        // 2. No join caused BUT we pushed a label that is different than the
+        // last pushed label
+        FloorIdx = LoweredFloorIdx;
+        FloorLabel = Label;
       }
     }
 
-    return std::move(JoinBlocks);
-  }
-};
+    LLVM_DEBUG(dbgs() << "SDA::joins. After propagation:\n"; printDefs(dbgs()));
 
-const ConstBlockSet &SyncDependenceAnalysis::join_blocks(const Loop &Loop) {
-  using LoopExitVec = SmallVector<BasicBlock *, 4>;
-  LoopExitVec LoopExits;
-  Loop.getExitBlocks(LoopExits);
-  if (LoopExits.size() < 1) {
-    return EmptyBlockSet;
+    return std::move(DivDesc);
   }
+};
 
-  // already available in cache?
-  auto ItCached = CachedLoopExitJoins.find(&Loop);
-  if (ItCached != CachedLoopExitJoins.end()) {
-    return *ItCached->second;
+static void PrintBlockSet(ConstBlockSet &Blocks, raw_ostream &Out) {
+  Out << "[";
+  bool First = true;
+  for (const auto *BB : Blocks) {
+    if (!First)
+      Out << ", ";
+    First = false;
+    Out << BB->getName();
   }
-
-  // compute all join points
-  DivergencePropagator Propagator{FuncRPOT, DT, PDT, LI};
-  auto JoinBlocks = Propagator.computeJoinPoints<const LoopExitVec &>(
-      *Loop.getHeader(), LoopExits, Loop.getParentLoop());
-
-  auto ItInserted = CachedLoopExitJoins.emplace(&Loop, std::move(JoinBlocks));
-  assert(ItInserted.second);
-  return *ItInserted.first->second;
+  Out << "]";
 }
 
-const ConstBlockSet &
+const ControlDivergenceDesc &
 SyncDependenceAnalysis::join_blocks(const Instruction &Term) {
   // trivial case
-  if (Term.getNumSuccessors() < 1) {
-    return EmptyBlockSet;
+  if (Term.getNumSuccessors() <= 1) {
+    return EmptyDivergenceDesc;
   }
 
   // already available in cache?
-  auto ItCached = CachedBranchJoins.find(&Term);
-  if (ItCached != CachedBranchJoins.end())
+  auto ItCached = CachedControlDivDescs.find(&Term);
+  if (ItCached != CachedControlDivDescs.end())
     return *ItCached->second;
 
   // compute all join points
-  DivergencePropagator Propagator{FuncRPOT, DT, PDT, LI};
+  // Special handling of divergent loop exits is not needed for LCSSA
   const auto &TermBlock = *Term.getParent();
-  auto JoinBlocks = Propagator.computeJoinPoints<const_succ_range>(
-      TermBlock, successors(Term.getParent()), LI.getLoopFor(&TermBlock));
+  DivergencePropagator Propagator(LoopRPO, DT, PDT, LI, TermBlock);
+  auto DivDesc = Propagator.computeJoinPoints();
+
+  LLVM_DEBUG(dbgs() << "Result (" << Term.getParent()->getName() << "):\n";
+             dbgs() << "JoinDivBlocks: ";
+             PrintBlockSet(DivDesc->JoinDivBlocks, dbgs());
+             dbgs() << "\nLoopDivBlocks: ";
+             PrintBlockSet(DivDesc->LoopDivBlocks, dbgs()); dbgs() << "\n";);
 
-  auto ItInserted = CachedBranchJoins.emplace(&Term, std::move(JoinBlocks));
+  auto ItInserted = CachedControlDivDescs.emplace(&Term, std::move(DivDesc));
   assert(ItInserted.second);
   return *ItInserted.first->second;
 }
 
 } // namespace llvm

llvm/test/Analysis/DivergenceAnalysis/AMDGPU/hidden_loopdiverge.ll

 G:
   br i1 %div.exitx, label %C, label %L
 ; CHECK: DIVERGENT: br i1 %div.exitx, 
 
 L:
   br i1 %uni.cond, label %D, label %G
 
 X:
-  %div.merge.x = phi i32 [ %a, %entry ], [ %uni.merge.h, %B ] ; temporal divergent phi
+  %uni.merge.x = phi i32 [ %a, %entry ], [ %uni.merge.h, %B ]
   br i1 %uni.cond, label %Y, label %exit
-; CHECK: DIVERGENT: %div.merge.x =
 
 Y:
   %div.merge.y = phi i32 [ 42, %X ], [ %b, %C ]
   br label %exit
 ; CHECK: DIVERGENT: %div.merge.y =
 
 exit:
   %div.merge.exit = phi i32 [ %a, %X ], [ %b, %Y ]

llvm/test/Analysis/DivergenceAnalysis/AMDGPU/trivial-join-at-loop-exit.ll

 ; RUN: opt -mtriple amdgcn-unknown-amdhsa -analyze -divergence -use-gpu-divergence-analysis %s | FileCheck %s
-; XFAIL: *
-
-; https://bugs.llvm.org/show_bug.cgi?id=46372
 
 ; CHECK: bb2:
 ; CHECK-NOT: DIVERGENT:       %Guard.bb2 = phi i1 [ true, %bb1 ], [ false, %bb0 ]
 
 ; Function Attrs: nounwind readnone speculatable
 declare i32 @llvm.amdgcn.workitem.id.x() #0
 
 define protected amdgpu_kernel void @test2(i1 %uni) {