[Dominators][AMDGPU] Don't use virtual exit node in findNearestCommonDominator. Cleanup MachinePostDominators.

[kuhar]

Summary

This patch fixes a bug that originated from passing a virtual exit block (nullptr) to MachinePostDominatorTee::findNearestCommonDominator and resulted in assertion failures inside its callee. It also applies a small cleanup to the class.

The patch introduces a new function in PDT that given a list of MachineBasicBlocks finds their NCD. The new overload of findNearestCommonDominator handles virtual root correctly.

Note that similar handling of virtual root nodes is not necessary in (forward) DominatorTrees, as right now they don't use virtual roots.

[hliao]

llvm/lib/CodeGen/MachinePostDominators.cpp

52

no need to (expensive) check every BB is not null. findNearestCommonDomniator asserts on both of them are non-null.

60

the logic here is weird. if virtual root is the NCD, that means there's no real NCD. but, breaking here returns wrong NCD. need to return nullptr explicitly here

[kuhar]

Small cleanup.

[kuhar]

llvm/lib/CodeGen/MachinePostDominators.cpp

52

This should be a very quick scan over the vector of blocks, not more expensive than the inner check.
The intention was twofold:

1. Document the blocks cannot be nullptr.
2. Fail quickly where the contract is first violated.

60

Virtual root is a real PDT node, but not a real basic block. The basic block value is nullptr and this is was will be returned.
I will update it to return nullptr explicitly, if that's easier to follow.

[hliao]

llvm/lib/CodeGen/MachinePostDominators.cpp

53

the assertion here is still redundant, as findNearestCommonDominator also asserts on both operands must not be null.

[hliao]

llvm/lib/CodeGen/MachinePostDominators.cpp

52

assertions are added to capture abnormal behaviors. we don't need to check abnormal behaviors early and "quick".

[hliao]

LGTM

[kuhar]

Fix a typo.

llvm/include/llvm/CodeGen/MachinePostDominators.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CODEGEN_MACHINEPOSTDOMINATORS_H
 #define LLVM_CODEGEN_MACHINEPOSTDOMINATORS_H
 
 #include "llvm/CodeGen/MachineDominators.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
+#include <memory>
 
 namespace llvm {
 
 ///
-/// PostDominatorTree Class - Concrete subclass of DominatorTree that is used
-/// to compute the post-dominator tree.
+/// MachinePostDominatorTree - an analysis pass wrapper for DominatorTree
+/// used to compute the post-dominator tree for MachineFunctions.
 ///
-struct MachinePostDominatorTree : public MachineFunctionPass {
-private:
- PostDomTreeBase<MachineBasicBlock> *DT;
+class MachinePostDominatorTree : public MachineFunctionPass {
+  using PostDomTreeT = PostDomTreeBase<MachineBasicBlock>;
+  std::unique_ptr<PostDomTreeT> PDT;
 
 public:
   static char ID;
 
   MachinePostDominatorTree();
 
-  ~MachinePostDominatorTree() override;
-
   FunctionPass *createMachinePostDominatorTreePass();
 
   const SmallVectorImpl<MachineBasicBlock *> &getRoots() const {
-    return DT->getRoots();
+    return PDT->getRoots();
   }
 
-  MachineDomTreeNode *getRootNode() const {
-    return DT->getRootNode();
-  }
+  MachineDomTreeNode *getRootNode() const { return PDT->getRootNode(); }
 
   MachineDomTreeNode *operator[](MachineBasicBlock *BB) const {
-    return DT->getNode(BB);
+    return PDT->getNode(BB);
   }
 
   MachineDomTreeNode *getNode(MachineBasicBlock *BB) const {
-    return DT->getNode(BB);
+    return PDT->getNode(BB);
   }
 
   bool dominates(const MachineDomTreeNode *A,
                  const MachineDomTreeNode *B) const {
-    return DT->dominates(A, B);
+    return PDT->dominates(A, B);
   }
 
   bool dominates(const MachineBasicBlock *A, const MachineBasicBlock *B) const {
-    return DT->dominates(A, B);
+    return PDT->dominates(A, B);
   }
 
   bool properlyDominates(const MachineDomTreeNode *A,
                          const MachineDomTreeNode *B) const {
-    return DT->properlyDominates(A, B);
+    return PDT->properlyDominates(A, B);
   }
 
   bool properlyDominates(const MachineBasicBlock *A,
                          const MachineBasicBlock *B) const {
-    return DT->properlyDominates(A, B);
+    return PDT->properlyDominates(A, B);
+  }
+
+  bool isVirtualRoot(const MachineDomTreeNode *BB) const {
+    return PDT->isVirtualRoot(BB);
   }
 
   MachineBasicBlock *findNearestCommonDominator(MachineBasicBlock *A,
-                                                MachineBasicBlock *B) {
-    return DT->findNearestCommonDominator(A, B);
+                                                MachineBasicBlock *B) const {
+    return PDT->findNearestCommonDominator(A, B);
   }
 
+  /// Returns the nearest common dominator of the given blocks.
+  /// If that tree node is a virtual root, a nullptr will be returned.
+  MachineBasicBlock *
+  findNearestCommonDominator(ArrayRef<MachineBasicBlock *> Blocks) const;
+
   bool runOnMachineFunction(MachineFunction &MF) override;
   void getAnalysisUsage(AnalysisUsage &AU) const override;
+  void releaseMemory() override { PDT.reset(nullptr); }
   void print(llvm::raw_ostream &OS, const Module *M = nullptr) const override;
 };
 } //end of namespace llvm
 
 #endif

llvm/include/llvm/CodeGen/MachineRegionInfo.h

 #include "llvm/CodeGen/MachineDominators.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineLoopInfo.h"
 #include <cassert>
 
 namespace llvm {
 
-struct MachinePostDominatorTree;
+class MachinePostDominatorTree;
 class MachineRegion;
 class MachineRegionNode;
 class MachineRegionInfo;
 
 template <> struct RegionTraits<MachineFunction> {
   using FuncT = MachineFunction;
   using BlockT = MachineBasicBlock;
   using RegionT = MachineRegion;

llvm/lib/CodeGen/MachinePostDominators.cpp

 //===- MachinePostDominators.cpp -Machine Post Dominator Calculation ------===//
 //
 // 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 simple dominator construction algorithms for finding
 // post dominators on machine functions.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/CodeGen/MachinePostDominators.h"
 
+#include "llvm/ADT/STLExtras.h"
+
 using namespace llvm;
 
 namespace llvm {
 template class DominatorTreeBase<MachineBasicBlock, true>; // PostDomTreeBase
 }
 
 char MachinePostDominatorTree::ID = 0;
 
 //declare initializeMachinePostDominatorTreePass
 INITIALIZE_PASS(MachinePostDominatorTree, "machinepostdomtree",
                 "MachinePostDominator Tree Construction", true, true)
 
-MachinePostDominatorTree::MachinePostDominatorTree() : MachineFunctionPass(ID) {
+MachinePostDominatorTree::MachinePostDominatorTree()
+    : MachineFunctionPass(ID), PDT(nullptr) {
   initializeMachinePostDominatorTreePass(*PassRegistry::getPassRegistry());
-  DT = new PostDomTreeBase<MachineBasicBlock>();
 }
 
-FunctionPass *
-MachinePostDominatorTree::createMachinePostDominatorTreePass() {
+FunctionPass *MachinePostDominatorTree::createMachinePostDominatorTreePass() {
   return new MachinePostDominatorTree();
 }
 
-bool
-MachinePostDominatorTree::runOnMachineFunction(MachineFunction &F) {
-  DT->recalculate(F);
+bool MachinePostDominatorTree::runOnMachineFunction(MachineFunction &F) {
+  PDT = std::make_unique<PostDomTreeT>();
+  PDT->recalculate(F);
   return false;
 }
 
-MachinePostDominatorTree::~MachinePostDominatorTree() {
-  delete DT;
-}
-
-void
-MachinePostDominatorTree::getAnalysisUsage(AnalysisUsage &AU) const {
+void MachinePostDominatorTree::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.setPreservesAll();
   MachineFunctionPass::getAnalysisUsage(AU);
 }
 
-void
-MachinePostDominatorTree::print(llvm::raw_ostream &OS, const Module *M) const {
-  DT->print(OS);
+MachineBasicBlock *MachinePostDominatorTree::findNearestCommonDominator(
+    ArrayRef<MachineBasicBlock *> Blocks) const {
+  assert(!Blocks.empty());
+  assert(llvm::all_of(Blocks, [](MachineBasicBlock *B) { return B; }) &&
+         "Invalid block found");
+
+  MachineBasicBlock *NCD = Blocks.front();
+  for (MachineBasicBlock *BB : Blocks.drop_front()) {
+    NCD = PDT->findNearestCommonDominator(NCD, BB);
+
+    // Stop when the root is reached.
+    if (PDT->isVirtualRoot(PDT->getNode(NCD)))
+      return nullptr;
+  }
+
+  return NCD;
+}
+
+void MachinePostDominatorTree::print(llvm::raw_ostream &OS,
+                                     const Module *M) const {
+  PDT->print(OS);
 }