[AMDGPU][FixIrreducible][UnifyLoopExits] Support callbr with inline-asm (#149308)

First batch of changes to add support for inline-asm callbr for the AMDGPU backend.

Author: ro-i

llvm/include/llvm/ADT/GenericCycleImpl.h

@@ -561,6 +561,17 @@ auto GenericCycleInfo<ContextT>::getSmallestCommonCycle(CycleT *A,
  return A;
 }
 
+/// \brief Find the innermost cycle containing both given blocks.
+///
+/// \returns the innermost cycle containing both \p A and \p B
+/// or nullptr if there is no such cycle.
+template <typename ContextT>
+auto GenericCycleInfo<ContextT>::getSmallestCommonCycle(BlockT *A,
+ BlockT *B) const
+ -> CycleT * {
+ return getSmallestCommonCycle(getCycle(A), getCycle(B));
+}
+
 /// \brief get the depth for the cycle which containing a given block.
 ///
 /// \returns the depth for the innermost cycle containing \p Block or 0 if it is

llvm/include/llvm/ADT/GenericCycleInfo.h

@@ -298,6 +298,7 @@ template <typename ContextT> class GenericCycleInfo {
 
  CycleT *getCycle(const BlockT *Block) const;
  CycleT *getSmallestCommonCycle(CycleT *A, CycleT *B) const;
+ CycleT *getSmallestCommonCycle(BlockT *A, BlockT *B) const;
  unsigned getCycleDepth(const BlockT *Block) const;
  CycleT *getTopLevelParentCycle(BlockT *Block);
 

llvm/include/llvm/Support/GenericLoopInfo.h

@@ -615,6 +615,17 @@ template <class BlockT, class LoopT> class LoopInfoBase {
  return L ? L->getLoopDepth() : 0;
  }
 
+ /// \brief Find the innermost loop containing both given loops.
+ ///
+ /// \returns the innermost loop containing both \p A and \p B
+ /// or nullptr if there is no such loop.
+ LoopT *getSmallestCommonLoop(LoopT *A, LoopT *B) const;
+ /// \brief Find the innermost loop containing both given blocks.
+ ///
+ /// \returns the innermost loop containing both \p A and \p B
+ /// or nullptr if there is no such loop.
+ LoopT *getSmallestCommonLoop(BlockT *A, BlockT *B) const;
+
  // True if the block is a loop header node
  bool isLoopHeader(const BlockT *BB) const {
  const LoopT *L = getLoopFor(BB);

llvm/include/llvm/Support/GenericLoopInfoImpl.h

@@ -355,7 +355,7 @@ void LoopBase<BlockT, LoopT>::verifyLoop() const {
  if (BB == getHeader()) {
  assert(!OutsideLoopPreds.empty() && "Loop is unreachable!");
  } else if (!OutsideLoopPreds.empty()) {
- // A non-header loop shouldn't be reachable from outside the loop,
+ // A non-header loop block shouldn't be reachable from outside the loop,
  // though it is permitted if the predecessor is not itself actually
  // reachable.
  BlockT *EntryBB = &BB->getParent()->front();
@@ -645,6 +645,36 @@ LoopInfoBase<BlockT, LoopT>::getLoopsInReverseSiblingPreorder() const {
  return PreOrderLoops;
 }
 
+template <class BlockT, class LoopT>
+LoopT *LoopInfoBase<BlockT, LoopT>::getSmallestCommonLoop(LoopT *A,
+ LoopT *B) const {
+ if (!A || !B)
+ return nullptr;
+
+ // If lops A and B have different depth replace them with parent loop
+ // until they have the same depth.
+ while (A->getLoopDepth() > B->getLoopDepth())
+ A = A->getParentLoop();
+ while (B->getLoopDepth() > A->getLoopDepth())
+ B = B->getParentLoop();
+
+ // Loops A and B are at same depth but may be disjoint, replace them with
+ // parent loops until we find loop that contains both or we run out of
+ // parent loops.
+ while (A != B) {
+ A = A->getParentLoop();
+ B = B->getParentLoop();
+ }
+
+ return A;
+}
+
+template <class BlockT, class LoopT>
+LoopT *LoopInfoBase<BlockT, LoopT>::getSmallestCommonLoop(BlockT *A,
+ BlockT *B) const {
+ return getSmallestCommonLoop(getLoopFor(A), getLoopFor(B));
+}
+
 // Debugging
 template <class BlockT, class LoopT>
 void LoopInfoBase<BlockT, LoopT>::print(raw_ostream &OS) const {

llvm/include/llvm/Transforms/Utils/BasicBlockUtils.h

@@ -19,6 +19,7 @@
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/IR/BasicBlock.h"
+#include "llvm/IR/CycleInfo.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/Printable.h"
@@ -262,6 +263,34 @@ LLVM_ABI BasicBlock *SplitEdge(BasicBlock *From, BasicBlock *To,
  MemorySSAUpdater *MSSAU = nullptr,
  const Twine &BBName = "");
 
+/// \brief Create a new intermediate target block for a callbr edge.
+///
+/// Create a new basic block between a callbr instruction and one of its
+/// successors. The new block replaces the original successor in the callbr
+/// instruction and unconditionally branches to the original successor. This
+/// is useful for normalizing control flow, e.g., when transforming
+/// irreducible loops.
+///
+/// \param CallBrBlock block containing the callbr instruction
+/// \param Succ original successor block
+/// \param SuccIdx index of the original successor in the callbr
+/// instruction
+/// \param DTU optional \p DomTreeUpdater for updating the
+/// dominator tree
+/// \param CI optional \p CycleInfo for updating cycle membership
+/// \param LI optional \p LoopInfo for updating loop membership
+/// \param UpdatedLI optional output flag indicating if \p LoopInfo has
+/// been updated
+///
+/// \returns newly created intermediate target block
+///
+/// \note This function updates PHI nodes, dominator tree, loop info, and
+/// cycle info as needed.
+LLVM_ABI BasicBlock *
+SplitCallBrEdge(BasicBlock *CallBrBlock, BasicBlock *Succ, unsigned SuccIdx,
+ DomTreeUpdater *DTU = nullptr, CycleInfo *CI = nullptr,
+ LoopInfo *LI = nullptr, bool *UpdatedLI = nullptr);
+
 /// Sets the unwind edge of an instruction to a particular successor.
 LLVM_ABI void setUnwindEdgeTo(Instruction *TI, BasicBlock *Succ);
 

llvm/include/llvm/Transforms/Utils/ControlFlowUtils.h

@@ -15,10 +15,13 @@
 
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/IR/CycleInfo.h"
 
 namespace llvm {
 
 class BasicBlock;
+class CallBrInst;
+class LoopInfo;
 class DomTreeUpdater;
 
 /// Given a set of branch descriptors [BB, Succ0, Succ1], create a "hub" such
@@ -104,7 +107,8 @@ struct ControlFlowHub {
  : BB(BB), Succ0(Succ0), Succ1(Succ1) {}
  };
 
- void addBranch(BasicBlock *BB, BasicBlock *Succ0, BasicBlock *Succ1) {
+ void addBranch(BasicBlock *BB, BasicBlock *Succ0,
+ BasicBlock *Succ1 = nullptr) {
  assert(BB);
  assert(Succ0 || Succ1);
  Branches.emplace_back(BB, Succ0, Succ1);

llvm/lib/Transforms/Utils/BasicBlockUtils.cpp

@@ -674,6 +674,79 @@ BasicBlock *llvm::SplitEdge(BasicBlock *BB, BasicBlock *Succ, DominatorTree *DT,
  return SplitBlock(BB, BB->getTerminator(), DT, LI, MSSAU, BBName);
 }
 
+/// Helper function to update the cycle or loop information after inserting a
+/// new block between a callbr instruction and one of its target blocks. Adds
+/// the new block to the innermost cycle or loop that the callbr instruction and
+/// the original target block share.
+/// \p LCI cycle or loop information to update
+/// \p CallBrBlock block containing the callbr instruction
+/// \p CallBrTarget new target block of the callbr instruction
+/// \p Succ original target block of the callbr instruction
+template <typename TI, typename T>
+static bool updateCycleLoopInfo(TI *LCI, BasicBlock *CallBrBlock,
+ BasicBlock *CallBrTarget, BasicBlock *Succ) {
+ static_assert(std::is_same_v<TI, CycleInfo> || std::is_same_v<TI, LoopInfo>,
+ "type must be CycleInfo or LoopInfo");
+ if (!LCI)
+ return false;
+
+ T *LC;
+ if constexpr (std::is_same_v<TI, CycleInfo>)
+ LC = LCI->getSmallestCommonCycle(CallBrBlock, Succ);
+ else
+ LC = LCI->getSmallestCommonLoop(CallBrBlock, Succ);
+ if (!LC)
+ return false;
+
+ if constexpr (std::is_same_v<TI, CycleInfo>)
+ LCI->addBlockToCycle(CallBrTarget, LC);
+ else
+ LC->addBasicBlockToLoop(CallBrTarget, *LCI);
+
+ return true;
+}
+
+BasicBlock *llvm::SplitCallBrEdge(BasicBlock *CallBrBlock, BasicBlock *Succ,
+ unsigned SuccIdx, DomTreeUpdater *DTU,
+ CycleInfo *CI, LoopInfo *LI,
+ bool *UpdatedLI) {
+ CallBrInst *CallBr = dyn_cast<CallBrInst>(CallBrBlock->getTerminator());
+ assert(CallBr && "expected callbr terminator");
+ assert(SuccIdx < CallBr->getNumSuccessors() &&
+ Succ == CallBr->getSuccessor(SuccIdx) && "invalid successor index");
+
+ // Create a new block between callbr and the specified successor.
+ // splitBlockBefore cannot be re-used here since it cannot split if the split
+ // point is a PHI node (because BasicBlock::splitBasicBlockBefore cannot
+ // handle that). But we don't need to rewire every part of a potential PHI
+ // node. We only care about the edge between CallBrBlock and the original
+ // successor.
+ BasicBlock *CallBrTarget =
+ BasicBlock::Create(CallBrBlock->getContext(),
+ CallBrBlock->getName() + ".target." + Succ->getName(),
+ CallBrBlock->getParent());
+ // Rewire control flow from the new target block to the original successor.
+ Succ->replacePhiUsesWith(CallBrBlock, CallBrTarget);
+ // Rewire control flow from callbr to the new target block.
+ CallBr->setSuccessor(SuccIdx, CallBrTarget);
+ // Jump from the new target block to the original successor.
+ BranchInst::Create(Succ, CallBrTarget);
+
+ bool Updated =
+ updateCycleLoopInfo<LoopInfo, Loop>(LI, CallBrBlock, CallBrTarget, Succ);
+ if (UpdatedLI)
+ *UpdatedLI = Updated;
+ updateCycleLoopInfo<CycleInfo, Cycle>(CI, CallBrBlock, CallBrTarget, Succ);
+ if (DTU) {
+ DTU->applyUpdates({{DominatorTree::Insert, CallBrBlock, CallBrTarget}});
+ if (DTU->getDomTree().dominates(CallBrBlock, Succ))
+ DTU->applyUpdates({{DominatorTree::Delete, CallBrBlock, Succ},
+ {DominatorTree::Insert, CallBrTarget, Succ}});
+ }
+
+ return CallBrTarget;
+}
+
 void llvm::setUnwindEdgeTo(Instruction *TI, BasicBlock *Succ) {
  if (auto *II = dyn_cast<InvokeInst>(TI))
  II->setUnwindDest(Succ);

llvm/lib/Transforms/Utils/ControlFlowUtils.cpp

@@ -13,6 +13,7 @@
 #include "llvm/Transforms/Utils/ControlFlowUtils.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/Analysis/DomTreeUpdater.h"
+#include "llvm/Analysis/LoopInfo.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/ValueHandle.h"
@@ -281,7 +282,9 @@ std::pair<BasicBlock *, bool> ControlFlowHub::finalize(
 
  for (auto [BB, Succ0, Succ1] : Branches) {
 #ifndef NDEBUG
- assert(Incoming.insert(BB).second && "Duplicate entry for incoming block.");
+ assert(
+ (Incoming.insert(BB).second || isa<CallBrInst>(BB->getTerminator())) &&
+ "Duplicate entry for incoming block.");
 #endif
  if (Succ0)
  Outgoing.insert(Succ0);