Revert "[AMDGPU][UnifyDivergentExitNodes][StructurizeCFG] Add support for callbr instruction with inline-asm" (llvm#166186)

Reverts llvm#152161

Need to revert to fix changed logic for the expensive checks.

Author: ro-i

llvm/lib/Target/AMDGPU/AMDGPUUnifyDivergentExitNodes.cpp

@@ -181,52 +181,14 @@ BasicBlock *AMDGPUUnifyDivergentExitNodesImpl::unifyReturnBlockSet(
   return NewRetBlock;
 }
 
-static BasicBlock *
-createDummyReturnBlock(Function &F,
-                       SmallVector<BasicBlock *, 4> &ReturningBlocks) {
-  BasicBlock *DummyReturnBB =
-      BasicBlock::Create(F.getContext(), "DummyReturnBlock", &F);
-  Type *RetTy = F.getReturnType();
-  Value *RetVal = RetTy->isVoidTy() ? nullptr : PoisonValue::get(RetTy);
-  ReturnInst::Create(F.getContext(), RetVal, DummyReturnBB);
-  ReturningBlocks.push_back(DummyReturnBB);
-  return DummyReturnBB;
-}
-
-/// Handle conditional branch instructions (-> 2 targets) and callbr
-/// instructions with N targets.
-static void handleNBranch(Function &F, BasicBlock *BB, Instruction *BI,
-                          BasicBlock *DummyReturnBB,
-                          std::vector<DominatorTree::UpdateType> &Updates) {
-  SmallVector<BasicBlock *, 2> Successors(successors(BB));
-
-  // Create a new transition block to hold the conditional branch.
-  BasicBlock *TransitionBB = BB->splitBasicBlock(BI, "TransitionBlock");
-
-  Updates.reserve(Updates.size() + 2 * Successors.size() + 2);
-
-  // 'Successors' become successors of TransitionBB instead of BB,
-  // and TransitionBB becomes a single successor of BB.
-  Updates.emplace_back(DominatorTree::Insert, BB, TransitionBB);
-  for (BasicBlock *Successor : Successors) {
-    Updates.emplace_back(DominatorTree::Insert, TransitionBB, Successor);
-    Updates.emplace_back(DominatorTree::Delete, BB, Successor);
-  }
-
-  // Create a branch that will always branch to the transition block and
-  // references DummyReturnBB.
-  BB->getTerminator()->eraseFromParent();
-  BranchInst::Create(TransitionBB, DummyReturnBB,
-                     ConstantInt::getTrue(F.getContext()), BB);
-  Updates.emplace_back(DominatorTree::Insert, BB, DummyReturnBB);
-}
-
 bool AMDGPUUnifyDivergentExitNodesImpl::run(Function &F, DominatorTree *DT,
                                             const PostDominatorTree &PDT,
                                             const UniformityInfo &UA) {
+  assert(hasOnlySimpleTerminator(F) && "Unsupported block terminator.");
+
   if (PDT.root_size() == 0 ||
       (PDT.root_size() == 1 &&
-       !isa<BranchInst, CallBrInst>(PDT.getRoot()->getTerminator())))
+       !isa<BranchInst>(PDT.getRoot()->getTerminator())))
     return false;
 
   // Loop over all of the blocks in a function, tracking all of the blocks that
@@ -260,27 +222,46 @@ bool AMDGPUUnifyDivergentExitNodesImpl::run(Function &F, DominatorTree *DT,
       if (HasDivergentExitBlock)
         UnreachableBlocks.push_back(BB);
     } else if (BranchInst *BI = dyn_cast<BranchInst>(BB->getTerminator())) {
-      if (!DummyReturnBB)
-        DummyReturnBB = createDummyReturnBlock(F, ReturningBlocks);
+
+      ConstantInt *BoolTrue = ConstantInt::getTrue(F.getContext());
+      if (DummyReturnBB == nullptr) {
+        DummyReturnBB =
+            BasicBlock::Create(F.getContext(), "DummyReturnBlock", &F);
+        Type *RetTy = F.getReturnType();
+        Value *RetVal = RetTy->isVoidTy() ? nullptr : PoisonValue::get(RetTy);
+        ReturnInst::Create(F.getContext(), RetVal, DummyReturnBB);
+        ReturningBlocks.push_back(DummyReturnBB);
+      }
 
       if (BI->isUnconditional()) {
         BasicBlock *LoopHeaderBB = BI->getSuccessor(0);
         BI->eraseFromParent(); // Delete the unconditional branch.
         // Add a new conditional branch with a dummy edge to the return block.
-        BranchInst::Create(LoopHeaderBB, DummyReturnBB,
-                           ConstantInt::getTrue(F.getContext()), BB);
+        BranchInst::Create(LoopHeaderBB, DummyReturnBB, BoolTrue, BB);
+        Updates.emplace_back(DominatorTree::Insert, BB, DummyReturnBB);
+      } else { // Conditional branch.
+        SmallVector<BasicBlock *, 2> Successors(successors(BB));
+
+        // Create a new transition block to hold the conditional branch.
+        BasicBlock *TransitionBB = BB->splitBasicBlock(BI, "TransitionBlock");
+
+        Updates.reserve(Updates.size() + 2 * Successors.size() + 2);
+
+        // 'Successors' become successors of TransitionBB instead of BB,
+        // and TransitionBB becomes a single successor of BB.
+        Updates.emplace_back(DominatorTree::Insert, BB, TransitionBB);
+        for (BasicBlock *Successor : Successors) {
+          Updates.emplace_back(DominatorTree::Insert, TransitionBB, Successor);
+          Updates.emplace_back(DominatorTree::Delete, BB, Successor);
+        }
+
+        // Create a branch that will always branch to the transition block and
+        // references DummyReturnBB.
+        BB->getTerminator()->eraseFromParent();
+        BranchInst::Create(TransitionBB, DummyReturnBB, BoolTrue, BB);
         Updates.emplace_back(DominatorTree::Insert, BB, DummyReturnBB);
-      } else {
-        handleNBranch(F, BB, BI, DummyReturnBB, Updates);
       }
       Changed = true;
-    } else if (CallBrInst *CBI = dyn_cast<CallBrInst>(BB->getTerminator())) {
-      if (!DummyReturnBB)
-        DummyReturnBB = createDummyReturnBlock(F, ReturningBlocks);
-
-      handleNBranch(F, BB, CBI, DummyReturnBB, Updates);
-    } else {
-      llvm_unreachable("unsupported block terminator");
     }
   }
 

llvm/lib/Transforms/Scalar/StructurizeCFG.cpp

@@ -558,10 +558,11 @@ void StructurizeCFG::analyzeLoops(RegionNode *N) {
   } else {
     // Test for successors as back edge
     BasicBlock *BB = N->getNodeAs<BasicBlock>();
-    if (BranchInst *Term = dyn_cast<BranchInst>(BB->getTerminator()))
-      for (BasicBlock *Succ : Term->successors())
-        if (Visited.count(Succ))
-          Loops[Succ] = BB;
+    BranchInst *Term = cast<BranchInst>(BB->getTerminator());
+
+    for (BasicBlock *Succ : Term->successors())
+      if (Visited.count(Succ))
+        Loops[Succ] = BB;
   }
 }
 
@@ -593,7 +594,7 @@ void StructurizeCFG::gatherPredicates(RegionNode *N) {
 
   for (BasicBlock *P : predecessors(BB)) {
     // Ignore it if it's a branch from outside into our region entry
-    if (!ParentRegion->contains(P) || !dyn_cast<BranchInst>(P->getTerminator()))
+    if (!ParentRegion->contains(P))
       continue;
 
     Region *R = RI->getRegionFor(P);
@@ -1401,17 +1402,13 @@ bool StructurizeCFG::makeUniformRegion(Region *R, UniformityInfo &UA) {
 /// Run the transformation for each region found
 bool StructurizeCFG::run(Region *R, DominatorTree *DT,
                          const TargetTransformInfo *TTI) {
-  // CallBr and its corresponding direct target blocks are for now ignored by
-  // this pass. This is not a limitation for the currently intended uses cases
-  // of callbr in the AMDGPU backend.
-  // Parent and child regions are not affected by this (current) restriction.
-  // See `llvm/test/Transforms/StructurizeCFG/callbr.ll` for details.
-  if (R->isTopLevelRegion() || isa<CallBrInst>(R->getEntry()->getTerminator()))
+  if (R->isTopLevelRegion())
     return false;
 
   this->DT = DT;
   this->TTI = TTI;
   Func = R->getEntry()->getParent();
+  assert(hasOnlySimpleTerminator(*Func) && "Unsupported block terminator.");
 
   ParentRegion = R;
 

llvm/test/CodeGen/AMDGPU/callbr.ll

@@ -3,7 +3,6 @@
 
 declare void @foo(ptr)
 declare i1 @bar(ptr)
-declare i32 @bar32(ptr)
 
 define void @musttail_call_without_return_value(ptr %p) {
 ; CHECK-LABEL: define void @musttail_call_without_return_value(
@@ -29,31 +28,6 @@ bb.1:
   ret void
 }
 
-define void @musttail_call_without_return_value_callbr(ptr %p) {
-; CHECK-LABEL: define void @musttail_call_without_return_value_callbr(
-; CHECK-SAME: ptr [[P:%.*]]) #[[ATTR0]] {
-; CHECK-NEXT:  [[ENTRY:.*:]]
-; CHECK-NEXT:    [[LOAD:%.*]] = load i32, ptr [[P]], align 1
-; CHECK-NEXT:    callbr void asm "", "r,!i"(i32 [[LOAD]])
-; CHECK-NEXT:            to label %[[BB_0:.*]] [label %bb.1]
-; CHECK:       [[BB_0]]:
-; CHECK-NEXT:    musttail call void @foo(ptr [[P]])
-; CHECK-NEXT:    ret void
-; CHECK:       [[BB_1:.*:]]
-; CHECK-NEXT:    ret void
-;
-entry:
-  %load = load i32, ptr %p, align 1
-  callbr void asm "", "r,!i"(i32 %load) to label %bb.0 [label %bb.1]
-
-bb.0:
-  musttail call void @foo(ptr %p)
-  ret void
-
-bb.1:
-  ret void
-}
-
 define i1 @musttail_call_with_return_value(ptr %p) {
 ; CHECK-LABEL: define i1 @musttail_call_with_return_value(
 ; CHECK-SAME: ptr [[P:%.*]]) #[[ATTR0]] {
@@ -77,28 +51,3 @@ bb.0:
 bb.1:
   ret i1 %load
 }
-
-define i32 @musttail_call_with_return_value_callbr(ptr %p) {
-; CHECK-LABEL: define i32 @musttail_call_with_return_value_callbr(
-; CHECK-SAME: ptr [[P:%.*]]) #[[ATTR0]] {
-; CHECK-NEXT:  [[ENTRY:.*:]]
-; CHECK-NEXT:    [[LOAD:%.*]] = load i32, ptr [[P]], align 1
-; CHECK-NEXT:    callbr void asm "", "r,!i"(i32 [[LOAD]])
-; CHECK-NEXT:            to label %[[BB_0:.*]] [label %bb.1]
-; CHECK:       [[BB_0]]:
-; CHECK-NEXT:    [[RET:%.*]] = musttail call i32 @bar32(ptr [[P]])
-; CHECK-NEXT:    ret i32 [[RET]]
-; CHECK:       [[BB_1:.*:]]
-; CHECK-NEXT:    ret i32 [[LOAD]]
-;
-entry:
-  %load = load i32, ptr %p, align 1
-  callbr void asm "", "r,!i"(i32 %load) to label %bb.0 [label %bb.1]
-
-bb.0:
-  %ret = musttail call i32 @bar32(ptr %p)
-  ret i32 %ret
-
-bb.1:
-  ret i32 %load
-}