[llvm] [AMDGPU] Remove blocks that only branch to other blocks (PR #184908)

[Diana Picus via llvm-commits]

================
@@ -763,6 +765,72 @@ MachineInstrBuilder SIPreEmitPeephole::createUnpackedMI(MachineInstr &I,
   return NewMI;
 }
 
+// Remove blocks that only have an unconditional branch.
+bool SIPreEmitPeephole::removeUnconditionalBranchBlocks(MachineFunction &MF) {
+  SIMachineFunctionInfo *FuncInfo = MF.getInfo<SIMachineFunctionInfo>();
+  SmallVector<MachineBasicBlock *, 2> ToRemoveMBB;
+  bool Changed = false;
+
+  for (MachineBasicBlock &MBB : MF) {
+    MachineBasicBlock *Succ = *MBB.succ_begin();
+    // Cannot remove entry block
+    if (&MBB == &MF.front())
+      continue;
+    // Cannot remove self-loops
+    if (Succ == &MBB)
+      continue;
+
+    // Find blocks that only contain an unconditional branch
+    const MachineInstr *firstMI = nullptr;
+    for (const MachineInstr &MI : MBB) {
+      if (MI.isMetaInstruction())
+        continue;
+      firstMI = &MI;
+      break;
+    }
+    if (firstMI && firstMI->isUnconditionalBranch()) {
+      MachineBasicBlock *FallThrough = nullptr;
+      SmallVector<MachineBasicBlock *, 2> PredecessorInRange;
+      for (MachineBasicBlock *Pred : MBB.predecessors()) {
+        // Ensure that re-directing a branch in Pred from MBB to Succ will
+        // not create a long branch. If a register has been reserved for long
+        // branching and Pred is too many basic blocks away from Succ a long
+        // branch may be created. Using the block difference is a crude but
+        // cheap estimate for the distance.  A very conservative limit has been
+        // chosen.
+        constexpr int BlockNumDiffLimit = 20;
+        if (!FuncInfo->getLongBranchReservedReg() ||
+            std::abs(Pred->getNumber() - Succ->getNumber()) <
+                BlockNumDiffLimit) {
+          if (is_contained(PredecessorInRange, Pred))
+            continue;
+          PredecessorInRange.push_back(Pred);
+          if (Pred->getFallThrough(false) == &MBB)
+            FallThrough = Pred;
+        }
+      }
+      for (MachineBasicBlock *Pred : PredecessorInRange) {
+        Changed = true;
+        Pred->ReplaceUsesOfBlockWith(&MBB, Succ);
+      }
+      if (MBB.predecessors().empty()) {
+        MBB.removeSuccessor(Succ);
+        MBB.clear();
+        ToRemoveMBB.push_back(&MBB);
+      }
+      if (FallThrough && !FallThrough->isLayoutSuccessor(Succ))
----------------
rovka wrote:

If you haven't erased MBB yet, won't Fallthrough's layout successor be MBB? (So isLayoutSuccessor would always return false)

https://github.com/llvm/llvm-project/pull/184908