[LowerSwitch] Delete all new dead blocks accurately

[Nuullll]

Summary

https://llvm.org/pr52383

LowerSwitch should be responsible for removing all newly introduced dead blocks.
The original implementation only deleted the unreachable default block itself, leaving
its dominatee-successors (if any) unreachable.

This patch uses DominatorTree so that we know the accurate set of blocks that need to
be removed.

[Nuullll]

Preserve analysis for new pass manager.

[Nuullll]

Fix AMDGPU llc-pipeline test

[ChuanqiXu]

If the newly added blocks are dead, they should be deleted by later DCE passes in one shot, am I right?
I think it is good to eliminate dead blocks if we could make it cheaply. But if we need to calculate reachable definitions by DT in a loop, it looks a little bit expensive.
I mean this patch may cause the compile time increases but wouldn't change the code generated finally.

[joerg]

I fully agree that constructing DT can be expensive especially if it is potentially done multiple times. I don't mind using it when available, but the gain seems limited for forcing it.

[aeubanks]

+1 to what the others have said

[Nuullll]

Thanks all for leaving the comments. I agree with you.

llvm/lib/Transforms/Utils/LowerSwitch.cpp

 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/LazyValueInfo.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/ConstantRange.h"
 #include "llvm/IR/Constants.h"
+#include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/IR/Value.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/Casting.h"
unsigned Clusterify(CaseVector &Cases, SwitchInst *SI) {
 
   return NumSimpleCases;
 }
 
 /// Replace the specified switch instruction with a sequence of chained if-then
 /// insts in a balanced binary search.
 void ProcessSwitchInst(SwitchInst *SI,
                        SmallPtrSetImpl<BasicBlock *> &DeleteList,
-                       AssumptionCache *AC, LazyValueInfo *LVI) {
+                       AssumptionCache *AC, LazyValueInfo *LVI,
+                       DominatorTree *DT) {
   BasicBlock *OrigBlock = SI->getParent();
   Function *F = OrigBlock->getParent();
   Value *Val = SI->getCondition();  // The value we are switching on...
   BasicBlock* Default = SI->getDefaultDest();
+  // The basic blocks dominated by the original default block won't change.
+  // So we query the descendants of the default block in the dominator tree
+  // before processing the switch instruction.
+  SmallVector<BasicBlock *, 2> BBDominatedByOrigDefault;
+  DT->getDescendants(Default, BBDominatedByOrigDefault);
+
+  // This helper function is supposed to be called when switch processing is
+  // done.
+  // 1. The SwitchInst SI will be erased.
+  // 2. DT will be rebuilt to check whether the original default block become
+  //    unreachable from entry. If so, add the OrigDefault and its dominatees
+  //    to the DeleteList.
+  auto EraseSwitchInstAndDeleteOrigDefaultIfUnreachable = [&]() {
+    BasicBlock *OrigDefault = SI->getDefaultDest();
+    SI->eraseFromParent();
+    DT->recalculate(*OrigBlock->getParent());
+    if (!DT->isReachableFromEntry(OrigDefault))
+      DeleteList.insert(BBDominatedByOrigDefault.begin(),
+                        BBDominatedByOrigDefault.end());
+  };
 
   // Don't handle unreachable blocks. If there are successors with phis, this
   // would leave them behind with missing predecessors.
-  if ((OrigBlock != &F->getEntryBlock() && pred_empty(OrigBlock)) ||
-      OrigBlock->getSinglePredecessor() == OrigBlock) {
+  if (!DT->isReachableFromEntry(OrigBlock)) {
+    LLVM_DEBUG(dbgs() << "Unreachable block " << OrigBlock->getName()
+                      << " contains a switch inst, removing\n");
     DeleteList.insert(OrigBlock);
     return;
   }
 
   // Prepare cases vector.
   CaseVector Cases;
   const unsigned NumSimpleCases = Clusterify(Cases, SI);
   LLVM_DEBUG(dbgs() << "Clusterify finished. Total clusters: " << Cases.size()
#endif
     assert(MaxPop > 0 && PopSucc);
     Default = PopSucc;
     llvm::erase_if(Cases,
                    [PopSucc](const CaseRange &R) { return R.BB == PopSucc; });
 
     // If there are no cases left, just branch.
     if (Cases.empty()) {
       BranchInst::Create(Default, OrigBlock);
-      SI->eraseFromParent();
       // As all the cases have been replaced with a single branch, only keep
       // one entry in the PHI nodes.
       for (unsigned I = 0 ; I < (MaxPop - 1) ; ++I)
         PopSucc->removePredecessor(OrigBlock);
+      EraseSwitchInstAndDeleteOrigDefaultIfUnreachable();
       return;
     }
 
     // If the condition was a PHI node with the switch block as a predecessor
     // removing predecessors may have caused the condition to be erased.
     // Getting the condition value again here protects against that.
     Val = SI->getCondition();
   }
#endif
   // If there are entries in any PHI nodes for the default edge, make sure
   // to update them as well.
   FixPhis(Default, OrigBlock, NewDefault);
 
   // Branch to our shiny new if-then stuff...
   BranchInst::Create(SwitchBlock, OrigBlock);
 
   // We are now done with the switch instruction, delete it.
-  BasicBlock *OldDefault = SI->getDefaultDest();
-  OrigBlock->getInstList().erase(SI);
-
-  // If the Default block has no more predecessors just add it to DeleteList.
-  if (pred_empty(OldDefault))
-    DeleteList.insert(OldDefault);
+  EraseSwitchInstAndDeleteOrigDefaultIfUnreachable();
 }
 
-bool LowerSwitch(Function &F, LazyValueInfo *LVI, AssumptionCache *AC) {
+bool LowerSwitch(Function &F, LazyValueInfo *LVI, AssumptionCache *AC,
+                 DominatorTree *DT) {
   bool Changed = false;
   SmallPtrSet<BasicBlock *, 8> DeleteList;
 
   // We use make_early_inc_range here so that we don't traverse new blocks.
   for (BasicBlock &Cur : llvm::make_early_inc_range(F)) {
     // If the block is a dead Default block that will be deleted later, don't
     // waste time processing it.
     if (DeleteList.count(&Cur))
       continue;
 
     if (SwitchInst *SI = dyn_cast<SwitchInst>(Cur.getTerminator())) {
       Changed = true;
-      ProcessSwitchInst(SI, DeleteList, AC, LVI);
+      ProcessSwitchInst(SI, DeleteList, AC, LVI, DT);
     }
   }
 
-  for (BasicBlock *BB : DeleteList) {
+  // Keep LazyValueInfo up-to-date.
+  for (auto *BB : DeleteList)
     LVI->eraseBlock(BB);
-    DeleteDeadBlock(BB);
-  }
+  SmallVector<BasicBlock *, 8> BBVec(DeleteList.begin(), DeleteList.end());
+  DeleteDeadBlocks(BBVec);
 
   return Changed;
 }
 
 /// Replace all SwitchInst instructions with chained branch instructions.
 class LowerSwitchLegacyPass : public FunctionPass {
 public:
   // Pass identification, replacement for typeid
   static char ID;
 
   LowerSwitchLegacyPass() : FunctionPass(ID) {
     initializeLowerSwitchLegacyPassPass(*PassRegistry::getPassRegistry());
   }
 
   bool runOnFunction(Function &F) override;
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<DominatorTreeWrapperPass>();
     AU.addRequired<LazyValueInfoWrapperPass>();
+    AU.addPreserved<DominatorTreeWrapperPass>();
+    AU.addPreserved<LazyValueInfoWrapperPass>();
   }
 };
 
 } // end anonymous namespace
 
 char LowerSwitchLegacyPass::ID = 0;
 
 // Publicly exposed interface to pass...
 char &llvm::LowerSwitchID = LowerSwitchLegacyPass::ID;
 
 INITIALIZE_PASS_BEGIN(LowerSwitchLegacyPass, "lowerswitch",
                       "Lower SwitchInst's to branches", false, false)
 INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
+INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(LazyValueInfoWrapperPass)
 INITIALIZE_PASS_END(LowerSwitchLegacyPass, "lowerswitch",
                     "Lower SwitchInst's to branches", false, false)
 
 // createLowerSwitchPass - Interface to this file...
 FunctionPass *llvm::createLowerSwitchPass() {
   return new LowerSwitchLegacyPass();
 }
 
 bool LowerSwitchLegacyPass::runOnFunction(Function &F) {
   LazyValueInfo *LVI = &getAnalysis<LazyValueInfoWrapperPass>().getLVI();
   auto *ACT = getAnalysisIfAvailable<AssumptionCacheTracker>();
   AssumptionCache *AC = ACT ? &ACT->getAssumptionCache(F) : nullptr;
-  return LowerSwitch(F, LVI, AC);
+  DominatorTree *DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
+  return LowerSwitch(F, LVI, AC, DT);
 }
 
 PreservedAnalyses LowerSwitchPass::run(Function &F,
                                        FunctionAnalysisManager &AM) {
   LazyValueInfo *LVI = &AM.getResult<LazyValueAnalysis>(F);
   AssumptionCache *AC = AM.getCachedResult<AssumptionAnalysis>(F);
-  return LowerSwitch(F, LVI, AC) ? PreservedAnalyses::none()
-                                 : PreservedAnalyses::all();
+  DominatorTree *DT = &AM.getResult<DominatorTreeAnalysis>(F);
+  if (!LowerSwitch(F, LVI, AC, DT))
+    return PreservedAnalyses::all();
+  PreservedAnalyses PA;
+  PA.preserve<DominatorTreeAnalysis>();
+  PA.preserve<LazyValueAnalysis>();
+  return PA;
 }

llvm/test/CodeGen/AMDGPU/llc-pipeline.ll

 ; GCN-O0-NEXT:      Scalarize Masked Memory Intrinsics
 ; GCN-O0-NEXT:      Expand reduction intrinsics
 ; GCN-O0-NEXT:    AMDGPU Attributor
 ; GCN-O0-NEXT:    CallGraph Construction
 ; GCN-O0-NEXT:    Call Graph SCC Pass Manager
 ; GCN-O0-NEXT:      AMDGPU Annotate Kernel Features
 ; GCN-O0-NEXT:      FunctionPass Manager
 ; GCN-O0-NEXT:        AMDGPU Lower Kernel Arguments
+; GCN-O0-NEXT:        Dominator Tree Construction
 ; GCN-O0-NEXT:        Lazy Value Information Analysis
 ; GCN-O0-NEXT:        Lower SwitchInst's to branches
 ; GCN-O0-NEXT:        Lower invoke and unwind, for unwindless code generators
 ; GCN-O0-NEXT:        Remove unreachable blocks from the CFG
 ; GCN-O0-NEXT:        Post-Dominator Tree Construction
 ; GCN-O0-NEXT:        Dominator Tree Construction
 ; GCN-O0-NEXT:        Natural Loop Information
 ; GCN-O0-NEXT:        Legacy Divergence Analysis
 ; GCN-O0-NEXT:        Unify divergent function exit nodes
+; GCN-O0-NEXT:        Dominator Tree Construction
 ; GCN-O0-NEXT:        Lazy Value Information Analysis
 ; GCN-O0-NEXT:        Lower SwitchInst's to branches
-; GCN-O0-NEXT:        Dominator Tree Construction
 ; GCN-O0-NEXT:        Natural Loop Information
 ; GCN-O0-NEXT:        Convert irreducible control-flow into natural loops
 ; GCN-O0-NEXT:        Fixup each natural loop to have a single exit block
 ; GCN-O0-NEXT:        Post-Dominator Tree Construction
 ; GCN-O0-NEXT:        Dominance Frontier Construction
 ; GCN-O0-NEXT:        Detect single entry single exit regions
 ; GCN-O0-NEXT:        Region Pass Manager
 ; GCN-O0-NEXT:          Structurize control flow

 ; GCN-O1-NEXT:    CallGraph Construction
 ; GCN-O1-NEXT:    Call Graph SCC Pass Manager
 ; GCN-O1-NEXT:      AMDGPU Annotate Kernel Features
 ; GCN-O1-NEXT:      FunctionPass Manager
 ; GCN-O1-NEXT:        AMDGPU Lower Kernel Arguments
 ; GCN-O1-NEXT:        Dominator Tree Construction
 ; GCN-O1-NEXT:        Natural Loop Information
 ; GCN-O1-NEXT:        CodeGen Prepare
+; GCN-O1-NEXT:        Dominator Tree Construction
 ; GCN-O1-NEXT:        Lazy Value Information Analysis
 ; GCN-O1-NEXT:        Lower SwitchInst's to branches
 ; GCN-O1-NEXT:        Lower invoke and unwind, for unwindless code generators
 ; GCN-O1-NEXT:        Remove unreachable blocks from the CFG
 ; GCN-O1-NEXT:        Dominator Tree Construction
 ; GCN-O1-NEXT:        Basic Alias Analysis (stateless AA impl)
 ; GCN-O1-NEXT:        Function Alias Analysis Results
 ; GCN-O1-NEXT:        Flatten the CFG
 ; GCN-O1-NEXT:        Dominator Tree Construction
 ; GCN-O1-NEXT:        Post-Dominator Tree Construction
 ; GCN-O1-NEXT:        Natural Loop Information
 ; GCN-O1-NEXT:        Legacy Divergence Analysis
 ; GCN-O1-NEXT:        AMDGPU IR late optimizations
 ; GCN-O1-NEXT:        Basic Alias Analysis (stateless AA impl)
 ; GCN-O1-NEXT:        Function Alias Analysis Results
 ; GCN-O1-NEXT:        Code sinking
 ; GCN-O1-NEXT:        Legacy Divergence Analysis
 ; GCN-O1-NEXT:        Unify divergent function exit nodes
+; GCN-O1-NEXT:        Dominator Tree Construction
 ; GCN-O1-NEXT:        Lazy Value Information Analysis
 ; GCN-O1-NEXT:        Lower SwitchInst's to branches
-; GCN-O1-NEXT:        Dominator Tree Construction
 ; GCN-O1-NEXT:        Natural Loop Information
 ; GCN-O1-NEXT:        Convert irreducible control-flow into natural loops
 ; GCN-O1-NEXT:        Fixup each natural loop to have a single exit block
 ; GCN-O1-NEXT:        Post-Dominator Tree Construction
 ; GCN-O1-NEXT:        Dominance Frontier Construction
 ; GCN-O1-NEXT:        Detect single entry single exit regions
 ; GCN-O1-NEXT:        Region Pass Manager
 ; GCN-O1-NEXT:          Structurize control flow

 ; GCN-O1-OPTS-NEXT:        Natural Loop Information
 ; GCN-O1-OPTS-NEXT:        Legacy Divergence Analysis
 ; GCN-O1-OPTS-NEXT:        AMDGPU IR late optimizations
 ; GCN-O1-OPTS-NEXT:        Basic Alias Analysis (stateless AA impl)
 ; GCN-O1-OPTS-NEXT:        Function Alias Analysis Results
 ; GCN-O1-OPTS-NEXT:        Code sinking
 ; GCN-O1-OPTS-NEXT:        Legacy Divergence Analysis
 ; GCN-O1-OPTS-NEXT:        Unify divergent function exit nodes
+; GCN-O1-OPTS-NEXT:        Dominator Tree Construction
 ; GCN-O1-OPTS-NEXT:        Lazy Value Information Analysis
 ; GCN-O1-OPTS-NEXT:        Lower SwitchInst's to branches
-; GCN-O1-OPTS-NEXT:        Dominator Tree Construction
 ; GCN-O1-OPTS-NEXT:        Natural Loop Information
 ; GCN-O1-OPTS-NEXT:        Convert irreducible control-flow into natural loops
 ; GCN-O1-OPTS-NEXT:        Fixup each natural loop to have a single exit block
 ; GCN-O1-OPTS-NEXT:        Post-Dominator Tree Construction
 ; GCN-O1-OPTS-NEXT:        Dominance Frontier Construction
 ; GCN-O1-OPTS-NEXT:        Detect single entry single exit regions
 ; GCN-O1-OPTS-NEXT:        Region Pass Manager
 ; GCN-O1-OPTS-NEXT:          Structurize control flow

 ; GCN-O2-NEXT:        Natural Loop Information
 ; GCN-O2-NEXT:        Legacy Divergence Analysis
 ; GCN-O2-NEXT:        AMDGPU IR late optimizations
 ; GCN-O2-NEXT:        Basic Alias Analysis (stateless AA impl)
 ; GCN-O2-NEXT:        Function Alias Analysis Results
 ; GCN-O2-NEXT:        Code sinking
 ; GCN-O2-NEXT:        Legacy Divergence Analysis
 ; GCN-O2-NEXT:        Unify divergent function exit nodes
+; GCN-O2-NEXT:        Dominator Tree Construction
 ; GCN-O2-NEXT:        Lazy Value Information Analysis
 ; GCN-O2-NEXT:        Lower SwitchInst's to branches
-; GCN-O2-NEXT:        Dominator Tree Construction
 ; GCN-O2-NEXT:        Natural Loop Information
 ; GCN-O2-NEXT:        Convert irreducible control-flow into natural loops
 ; GCN-O2-NEXT:        Fixup each natural loop to have a single exit block
 ; GCN-O2-NEXT:        Post-Dominator Tree Construction
 ; GCN-O2-NEXT:        Dominance Frontier Construction
 ; GCN-O2-NEXT:        Detect single entry single exit regions
 ; GCN-O2-NEXT:        Region Pass Manager
 ; GCN-O2-NEXT:          Structurize control flow

 ; GCN-O3-NEXT:        Natural Loop Information
 ; GCN-O3-NEXT:        Legacy Divergence Analysis
 ; GCN-O3-NEXT:        AMDGPU IR late optimizations
 ; GCN-O3-NEXT:        Basic Alias Analysis (stateless AA impl)
 ; GCN-O3-NEXT:        Function Alias Analysis Results
 ; GCN-O3-NEXT:        Code sinking
 ; GCN-O3-NEXT:        Legacy Divergence Analysis
 ; GCN-O3-NEXT:        Unify divergent function exit nodes
+; GCN-O3-NEXT:        Dominator Tree Construction
 ; GCN-O3-NEXT:        Lazy Value Information Analysis
 ; GCN-O3-NEXT:        Lower SwitchInst's to branches
-; GCN-O3-NEXT:        Dominator Tree Construction
 ; GCN-O3-NEXT:        Natural Loop Information
 ; GCN-O3-NEXT:        Convert irreducible control-flow into natural loops
 ; GCN-O3-NEXT:        Fixup each natural loop to have a single exit block
 ; GCN-O3-NEXT:        Post-Dominator Tree Construction
 ; GCN-O3-NEXT:        Dominance Frontier Construction
 ; GCN-O3-NEXT:        Detect single entry single exit regions
 ; GCN-O3-NEXT:        Region Pass Manager
 ; GCN-O3-NEXT:          Structurize control flow

llvm/test/Transforms/LowerSwitch/delete-all-new-dead-blocks.ll

+; RUN: opt -S -lowerswitch %s | FileCheck %s
+; RUN: opt %s -passes=lowerswitch -S | FileCheck %s
+
+; Check that the new dead blocks introduced by LowerSwitch pass (they were alive before the pass), are deleted.
+
+; CHECK-LABEL: define i32 @f
+; CHECK-NOT: out_of_switch_range
+; CHECK-NOT: successor_of_unreachable
+
+define i32 @f(i32 %c) {
+chk65:
+  %cmp = icmp sgt i32 %c, 65
+  br i1 %cmp, label %return, label %chk0
+
+chk0:                                             ; preds = %chk65
+  %cmp1 = icmp slt i32 %c, 0
+  br i1 %cmp, label %return, label %bb_if
+
+bb_if:                                            ; preds = %chk0
+  %ashr.val = ashr exact i32 %c, 2
+  %cmp2 = icmp sgt i32 %ashr.val, 14
+  br i1 %cmp2, label %bb_switch, label %return
+
+bb_switch:                                        ; preds = %bb_if
+; %ashr.val must be 15 or 16 here (analyzed by LazyValueInfo)
+  switch i32 %ashr.val, label %out_of_switch_range [
+    i32 15, label %return
+    i32 16, label %another_ret
+  ]
+
+; unreachable default case
+out_of_switch_range:                              ; preds = %bb_switch
+  br label %successor_of_unreachable
+
+; unreachable successors
+successor_of_unreachable:                         ; preds = %out_of_switch_range
+  br label %return
+
+another_ret:                                      ; preds = %bb_switch
+  br label %return
+
+return:                                           ; preds = %another_ret, %successor_of_unreachable, %bb_switch, %bb_if, %chk0, %chk65
+  %retval = phi i32 [ -1, %chk0 ], [ -1, %chk65 ], [ -1, %bb_if ], [ 42, %bb_switch ], [ 42, %another_ret ], [ -42, %successor_of_unreachable ]
+  ret i32 %retval
+}
+
+; CHECK-LABEL: define i32 @g
+; CHECK: bb_switch:
+; CHECK-NEXT: br label %return
+; CHECK-NOT: out_of_switch_range
+define i32 @g(i32 %c) {
+bb_switch:
+  switch i32 %c, label %out_of_switch_range [
+    i32 16, label %return
+  ]
+
+out_of_switch_range:                              ; preds = %bb_switch
+  unreachable
+
+return:                                           ; preds = %bb_switch
+  ret i32 %c
+}
+
+; CHECK-LABEL: define i32 @h
+; CHECK-NOT: out_of_switch_range
+; CHECK-NOT: successor_of_unreachable
+define i32 @h(i32 %c) {
+chk65:
+  %cmp = icmp sgt i32 %c, 65
+  br i1 %cmp, label %return, label %chk0
+
+chk0:                                             ; preds = %chk65
+  %cmp1 = icmp slt i32 %c, 0
+  br i1 %cmp, label %return, label %bb_if
+
+bb_if:                                            ; preds = %chk0
+  %ashr.val = ashr exact i32 %c, 2
+  %cmp2 = icmp sgt i32 %ashr.val, 14
+  br i1 %cmp2, label %bb_switch, label %return
+
+bb_switch:                                        ; preds = %bb_if
+; %ashr.val must be 15 or 16 here (analyzed by LazyValueInfo)
+  switch i32 %ashr.val, label %out_of_switch_range [
+    i32 15, label %return
+    i32 16, label %another_ret
+  ]
+
+; unreachable default case
+out_of_switch_range:                              ; preds = %bb_switch, %successor_of_unreachable
+  br label %successor_of_unreachable
+
+; unreachable successors
+successor_of_unreachable:                         ; preds = %out_of_switch_range
+  br label %out_of_switch_range
+
+another_ret:                                      ; preds = %bb_switch
+  br label %return
+
+return:                                           ; preds = %another_ret, %bb_switch, %bb_if, %chk0, %chk65
+  %retval = phi i32 [ -1, %chk0 ], [ -1, %chk65 ], [ -1, %bb_if ], [ 42, %bb_switch ], [ 42, %another_ret ]
+  ret i32 %retval
+}

llvm/test/Transforms/Util/lowerswitch.ll

 ; CHECK-LABEL: @test4
 entry:
   switch i32 %mode, label %cleanup17 [
     i32 0, label %return
     i32 9, label %return
   ]
 
 cleanup17:
-; CHECK: cleanup17:
+; CHECK-NOT: cleanup17:
 ; CHECK-NOT: phi i16 [ undef, %entry ]
 ; CHECK: return:
 
   %retval.4 = phi i16 [ undef, %entry ]
   unreachable
 
 return:
   ret void