llvm-reduce: Fix block reduction with unreachable blocks

[arsenm]

Summary

Previously this would produce many invalid reductions with
"Instruction does not dominate uses" verifier errors.

This fixes issues in cases where the incoming IR
has unreachable blocks, and the resulting reduction
introduced new reachable blocks.

Have basic-blocks skip functions that have unreachable
blocks, Introduce a separate reduction which only
deletes unreachable blocks. Cleanup any newly unreachable
blocks after trimming out the requested deletions.

Includes a variety of meta-reduced tests for llvm-reduce
itself with -abort-on-invalid-reduction that were failing
on different iterations of this patch.

Bugpoint's implementation is much simpler (but currently I don't
understand how it avoids disconnecting interesting blocks from the CFG).

[aeubanks]

nice!

llvm/tools/llvm-reduce/deltas/ReduceBasicBlocks.cpp

137

remove parens

199

probably not super important, but keeping phis limits the instructions that the basic block pass messes with

[arsenm]

llvm/tools/llvm-reduce/deltas/ReduceBasicBlocks.cpp

199

Doesn't seem to do anything in the tests

[arsenm]

45a91c1521b243b728d4ce69f8d19a449627d251

llvm/test/tools/llvm-reduce/reduce-bb-unreachable-does-not-dominate-error0.ll

+; RUN: llvm-reduce --abort-on-invalid-reduction --delta-passes=basic-blocks --test FileCheck --test-arg --check-prefixes=CHECK-INTERESTINGNESS --test-arg %s --test-arg --input-file %s -o %t
+; RUN: FileCheck %s < %t
+
+
+; Check that no invalid reduction is produced. Remaining unreachable
+; blocks can leave instructions that dominate uses. When %bb3 was made
+; unreachable it produced this verifier error:
+; Instruction does not dominate all uses!
+;   %i4 = icmp eq i32 0, 0
+;   %i6 = select i1 %i4, i1 false, i1 false
+
+
+; CHECK: define void @func() {
+; CHECK-NEXT: bb:
+; CHECK-NEXT: br label %bb1
+
+; CHECK: bb1:
+; CHECK-NEXT: label %bb3
+
+; CHECK: bb2:
+; CHECK-NEXT: br i1 false, label %bb1, label %bb2
+
+; CHECK: bb3:
+; CHECK-NEXT: %i = phi i32 [ 0, %bb1 ]
+; CHECK: %i4 = icmp eq i32 0, 0
+; CHECK-NEXT: br label %bb5
+
+; CHECK: bb5:
+; CHECK-NEXT: %i6 = select i1 %i4, i1 false, i1 false
+; CHECK-NEXT: store i32 0
+; CHECK-NEXT: ret void
+define void @func() {
+bb:
+  br label %bb1
+
+bb1:                                              ; preds = %bb2, %bb
+  br label %bb3
+
+bb2:                                              ; preds = %bb2
+  br i1 false, label %bb1, label %bb2
+
+bb3:                                              ; preds = %bb1
+  %i = phi i32 [ 0, %bb1 ]
+  %i4 = icmp eq i32 0, 0
+  br label %bb5
+
+bb5:                                              ; preds = %bb3
+  %i6 = select i1 %i4, i1 false, i1 false
+  ; CHECK-INTERESTINGNESS: store
+  store i32 0, ptr undef, align 4
+  ret void
+}

llvm/test/tools/llvm-reduce/reduce-bb-unreachable-does-not-dominate-error1.ll

+; RUN: llvm-reduce --abort-on-invalid-reduction --delta-passes=unreachable-basic-blocks,basic-blocks --test FileCheck --test-arg --check-prefixes=CHECK-INTERESTINGNESS --test-arg %s --test-arg --input-file %s -o %t
+; RUN: FileCheck %s < %t
+
+
+; CHECK-INTERESTINGNESS: @func(
+
+; CHECK-INTERESTINGNESS: store
+; CHECK-INTERESTINGNESS: store
+; CHECK-INTERESTINGNESS: store
+
+
+; CHECK: bb:
+; CHECK-NEXT: br i1 %arg1, label %bb3, label %bb7
+
+; CHECK: bb3: ; preds = %bb
+; CHECK-NEXT: br i1 %arg2, label %bb4, label %bb7
+
+; CHECK: bb4: ; preds = %bb3
+; CHECK-NEXT: store i32 0, ptr addrspace(1) null, align 4
+; CHECK-NEXT: br label %bb10
+
+; CHECK: bb7: ; preds = %bb3, %bb
+; CHECK-NEXT: %i = phi i1 [ false, %bb ], [ true, %bb3 ]
+; CHECK-NEXT: store i32 1, ptr addrspace(1) null, align 4
+; CHECK-NEXT: br label %bb10
+
+; CHECK: bb10: ; preds = %bb7, %bb4
+; CHECK-NEXT: store i32 2, ptr addrspace(1) null, align 4
+; CHECK-NEXT: unreachable
+define amdgpu_kernel void @func(i1 %arg, i1 %arg1, i1 %arg2) {
+bb:
+  br i1 %arg1, label %bb3, label %bb7
+
+bb3:                                              ; preds = %bb
+  br i1 %arg2, label %bb4, label %bb7
+
+bb4:                                              ; preds = %bb3
+  store i32 0, ptr addrspace(1) null
+  br label %bb5
+
+bb5:                                              ; preds = %bb4
+  unreachable
+
+bb6:                                              ; No predecessors!
+  unreachable
+
+bb7:                                              ; preds = %bb3, %bb
+  %i = phi i1 [ false, %bb ], [ true, %bb3 ]
+  store i32 1, ptr addrspace(1) null
+  br i1 %arg, label %bb10, label %bb8
+
+bb8:                                              ; preds = %bb7
+  br i1 %i, label %bb9, label %bb9
+
+bb9:                                              ; preds = %bb8, %bb8
+  unreachable
+
+bb10:                                             ; preds = %bb7
+  store i32 2, ptr addrspace(1) null
+  unreachable
+}

llvm/test/tools/llvm-reduce/reduce-bb-unreachable-does-not-dominate-error2.ll

+; RUN: llvm-reduce --abort-on-invalid-reduction --delta-passes=basic-blocks --test FileCheck --test-arg --check-prefixes=CHECK-INTERESTINGNESS --test-arg %s --test-arg --input-file %s -o %t
+; RUN: FileCheck %s < %t
+
+; Make sure an invalid reduction isn't produced due to leaving behind
+; invalid code in %bb8 after it becomes unreachable.
+
+; CHECK-INTERESTINGNESS: store i32 0,
+; CHECK-INTERESTINGNESS: store i32 1,
+; CHECK-INTERESTINGNESS: store i32 2,
+
+
+; CHECK: bb:
+; CHECK-NEXT: store i32 0, ptr addrspace(3) null, align 4
+
+; CHECK: bb6: ; preds = %bb8, %bb
+; CHECK-NEXT: store i32 1, ptr addrspace(3) null, align 4
+
+; CHECK: bb8: ; preds = %bb6
+; CHECK-NEXT: %tmp = phi ptr addrspace(5) [ null, %bb6 ]
+define amdgpu_kernel void @foo(i32 %arg) {
+bb:
+  store i32 0, ptr addrspace(3) null
+  br label %bb6
+
+bb6:                                              ; preds = %bb10, %bb9, %bb8, %bb
+  store i32 1, ptr addrspace(3) null
+  switch i32 0, label %bb7 [
+    i32 0, label %bb8
+  ]
+
+bb7:                                              ; preds = %bb6
+  unreachable
+
+bb8:                                              ; preds = %bb6
+  %tmp = phi ptr addrspace(5) [ null, %bb6 ]
+  store i32 2, ptr addrspace(5) %tmp
+  switch i32 %arg, label %bb6 [
+    i32 0, label %bb10
+    i32 1, label %bb9
+  ]
+
+bb9:                                              ; preds = %bb8
+  br label %bb6
+
+bb10:                                             ; preds = %bb8
+  br label %bb6
+}

llvm/test/tools/llvm-reduce/reduce-bb-unreachable-does-not-dominate-error3.ll

+; RUN: llvm-reduce --abort-on-invalid-reduction --delta-passes=unreachable-basic-blocks,basic-blocks --test FileCheck --test-arg --check-prefixes=CHECK-INTERESTINGNESS --test-arg %s --test-arg --input-file %s -o %t
+; RUN: FileCheck %s < %t
+
+
+; CHECK-INTERESTINGNESS: store i32 0,
+; CHECK-INTERESTINGNESS: store i32 1,
+
+
+; CHECK: bb:
+; CHECK-NEXT: %tmp = icmp eq i8 0, 0
+; CHECK-NEXT: %tmp12 = load i32, ptr addrspace(4) inttoptr (i64 3016 to ptr addrspace(4)), align 8
+; CHECK-NEXT: %tmp13 = load i32, ptr addrspace(4) null, align 8
+; CHECK-NEXT: br label %bb20
+
+; CHECK: bb20:
+; CHECK-NEXT: store i32 0, ptr addrspace(3) null, align 4
+; CHECK-NEXT: br label %bb21
+
+; CHECK: bb21:
+; CHECK-NEXT: store i32 1, ptr addrspace(3) null, align 4
+; CHECK-NEXT: ret void
+
+define void @snork() {
+bb:
+  %tmp = icmp eq i8 0, 0
+  %tmp12 = load i32, ptr addrspace(4) inttoptr (i64 3016 to ptr addrspace(4)), align 8
+  %tmp13 = load i32, ptr addrspace(4) null, align 8
+  br label %bb14
+
+bb14:                                             ; preds = %bb21, %bb20, %bb19, %bb14, %bb
+  switch i32 %tmp12, label %bb22 [
+    i32 2, label %bb14
+    i32 1, label %bb19
+  ]
+
+bb15:                                             ; preds = %bb17
+  %tmp16 = fadd contract double %tmp18, 1.000000e+00
+  unreachable
+
+bb17:                                             ; preds = %bb17
+  %tmp18 = fadd contract double 0.000000e+00, 0.000000e+00
+  br i1 false, label %bb15, label %bb17
+
+bb19:                                             ; preds = %bb14
+  switch i32 %tmp13, label %bb14 [
+    i32 2, label %bb21
+    i32 1, label %bb20
+  ]
+
+bb20:                                             ; preds = %bb19
+  store i32 0, ptr addrspace(3) null, align 4
+  br label %bb14
+
+bb21:                                             ; preds = %bb19
+  store i32 1, ptr addrspace(3) null, align 4
+  br label %bb14
+
+bb22:                                             ; preds = %bb14
+  unreachable
+}

llvm/test/tools/llvm-reduce/reduce-blocks-only-phi-nodes-may-reference-own-value.ll

+; RUN: llvm-reduce --abort-on-invalid-reduction --delta-passes=basic-blocks --test=FileCheck --test-arg=--check-prefix=CHECK-INTERESTINGNESS --test-arg=%s --test-arg=--input-file %s -o %t
+; RUN: FileCheck %s < %t
+
+; Make sure an invalid reduction isn't tried when deleting %bb5,
+; causing the or in %bb6 to use its own output value.
+
+
+; CHECK-INTERESTINGNESS: store i32 0
+; CHECK-INTERESTINGNESS: store i32 1
+; CHECK-INTERESTINGNESS: store i32 2
+
+; CHECK: store i32 0
+; CHECK-NEXT: br label %bb5
+
+; CHECK: bb5:
+; CHECK-NEXT: switch
+
+; CHECK: bb6:
+; CHECK-NEXT: %tmp = phi i32 [ %tmp7, %bb6 ]
+; CHECK-NEXT: store i32 1
+; CHECK-NEXT: %tmp7 = or i32 %tmp, 0
+; CHECK-NEXT: br label %bb6
+
+; CHECK-NOT: bb7
+; CHECK: bb8:
+; CHECK-NEXT: store i32 2,
+define amdgpu_kernel void @snork(i32 %arg, i1 %arg1) {
+bb:
+  store i32 0, ptr addrspace(3) null
+  br i1 %arg1, label %bb5, label %bb7
+
+bb5:                                              ; preds = %bb5, %bb
+  switch i32 %arg, label %bb5 [
+    i32 0, label %bb8
+    i32 1, label %bb6
+  ]
+
+bb6:                                              ; preds = %bb6, %bb5
+  %tmp = phi i32 [ %tmp7, %bb6 ], [ 0, %bb5 ]
+  store i32 1, ptr addrspace(3) null
+  %tmp7 = or i32 %tmp, 0
+  br label %bb6
+
+bb7:
+  store i32 3, ptr addrspace(3) null
+  br label %bb8
+
+bb8:                                              ; preds = %bb5
+  store i32 2, ptr addrspace(3) null
+  unreachable
+}

llvm/test/tools/llvm-reduce/remove-bb-switch-default.ll

 ; RUN: llvm-reduce -abort-on-invalid-reduction --delta-passes=basic-blocks --test FileCheck --test-arg --check-prefix=CHECK-INTERESTINGNESS --test-arg %s --test-arg --input-file %s -o %t
 ; RUN: FileCheck %s < %t
 
 ; CHECK-INTERESTINGNESS: store i32 1,
 ; CHECK-INTERESTINGNESS: store i32 2,
 
 
 ; CHECK: bb:
 ; CHECK-NEXT: %bb.load = load i32, ptr null, align 4
 ; CHECK-NEXT: store i32 0, ptr null, align 4
 ; CHECK-NEXT: br i1 %arg0, label %bb1, label %bb2
 
 ; CHECK: bb1:
-; CHECK-NEXT: %bb1.phi = phi i32 [ %bb.load, %bb ], [ %bb2.phi, %bb2 ], [ %bb2.phi, %bb2 ]
 ; CHECK-NEXT: store i32 1, ptr null, align 4
 ; CHECK-NEXT: ret void
 
 ; CHECK: bb2: ; preds = %bb
-; CHECK-NEXT: %bb2.phi = phi i32 [ %bb.load, %bb ]
 ; CHECK-NEXT: store i32 2, ptr null, align 4
-; CHECK-NEXT: switch i32 %bb2.phi, label %bb1 [
+; CHECK-NEXT: switch i32 %bb.load, label %bb1 [
 ; CHECK-NEXT: i32 0, label %bb1
 ; CHECK-NEXT: ]
 define void @main(i1 %arg0) {
 bb:
   %bb.load = load i32, ptr null
   store i32 0, ptr null
   br i1 %arg0, label %bb1, label %bb2
 

llvm/test/tools/llvm-reduce/remove-bbs-unreachable.ll

 ; Check that verification doesn't fail when reducing a function with
 ; unreachable blocks.
 ;
+; RUN: llvm-reduce --abort-on-invalid-reduction --delta-passes=unreachable-basic-blocks --test FileCheck --test-arg --check-prefixes=CHECK-INTERESTINGNESS --test-arg %s --test-arg --input-file %s -o %t
+; RUN: FileCheck -check-prefix=UNREACHABLE %s < %t
+
 ; RUN: llvm-reduce --abort-on-invalid-reduction --delta-passes=basic-blocks --test FileCheck --test-arg --check-prefixes=CHECK-INTERESTINGNESS --test-arg %s --test-arg --input-file %s -o %t
-; RUN: FileCheck %s < %t
+; RUN: FileCheck -check-prefix=REACHABLE %s < %t
+
+; CHECK-INTERESTINGNESS: test0
+; CHECK-INTERESTINGNESS: test1
+
+; UNREACHABLE: define void @test0() {
+; UNREACHABLE-NEXT: entry:
+; UNREACHABLE-NEXT:   br label %exit
+
+; UNREACHABLE-NOT: unreachable
+; UNREACHABLE: exit:
+; UNREACHABLE-NEXT: ret void
 
-; CHECK-INTERESTINGNESS: test
 
-; CHECK: define void @test() {
-; CHECK-NEXT:   unreachable:
-; CHECK-NEXT:     ret void
+; basic-blocks cannot deal with unreachable blocks, leave it behind
+; REACHABLE: define void @test0() {
+; REACHABLE: entry:
+; REACHABLE: unreachable:
+; REACHABLE: exit:
 
-define void @test() {
+define void @test0() {
 entry:
   br label %exit
 
 unreachable:                                        ; No predecessors!
   br label %exit
 
 exit:
   ret void
 }
+
+; UNREACHABLE: define void @test1() {
+; UNREACHABLE-NEXT: entry:
+; UNREACHABLE-NEXT:   br label %exit
+
+; REACHABLE: define void @test1() {
+; REACHABLE: entry:
+; REACHABLE: unreachable:
+; REACHABLE: exit:
+define void @test1() {
+entry:
+  br label %exit
+
+unreachable:
+  br label %unreachable
+
+exit:
+  ret void
+}

llvm/tools/llvm-reduce/DeltaManager.cpp

 #define DELTA_PASSES                                                           \
   do {                                                                         \
     DELTA_PASS("functions", reduceFunctionsDeltaPass)                          \
     DELTA_PASS("function-bodies", reduceFunctionBodiesDeltaPass)               \
     DELTA_PASS("special-globals", reduceSpecialGlobalsDeltaPass)               \
     DELTA_PASS("aliases", reduceAliasesDeltaPass)                              \
     DELTA_PASS("simplify-conditionals-true", reduceConditionalsTrueDeltaPass)  \
     DELTA_PASS("simplify-conditionals-false", reduceConditionalsFalseDeltaPass)\
+    DELTA_PASS("unreachable-basic-blocks", reduceUnreachableBasicBlocksDeltaPass) \
     DELTA_PASS("basic-blocks", reduceBasicBlocksDeltaPass)                     \
+    DELTA_PASS("simplify-cfg", reduceUsingSimplifyCFGDeltaPass)                \
     DELTA_PASS("global-values", reduceGlobalValuesDeltaPass)                   \
     DELTA_PASS("global-objects", reduceGlobalObjectsDeltaPass)                 \
     DELTA_PASS("global-initializers", reduceGlobalsInitializersDeltaPass)      \
     DELTA_PASS("global-variables", reduceGlobalsDeltaPass)                     \
     DELTA_PASS("di-metadata", reduceDIMetadataDeltaPass)                       \
     DELTA_PASS("metadata", reduceMetadataDeltaPass)                            \
     DELTA_PASS("arguments", reduceArgumentsDeltaPass)                          \
     DELTA_PASS("instructions", reduceInstructionsDeltaPass)                    \
     DELTA_PASS("simplify-instructions", simplifyInstructionsDeltaPass)         \
     DELTA_PASS("ir-passes", runIRPassesDeltaPass)                              \
     DELTA_PASS("operands-zero", reduceOperandsZeroDeltaPass)                   \
     DELTA_PASS("operands-one", reduceOperandsOneDeltaPass)                     \
     DELTA_PASS("operands-nan", reduceOperandsNaNDeltaPass)                     \
     DELTA_PASS("operands-to-args", reduceOperandsToArgsDeltaPass)              \
     DELTA_PASS("operands-skip", reduceOperandsSkipDeltaPass)                   \
     DELTA_PASS("operand-bundles", reduceOperandBundesDeltaPass)                \
-    DELTA_PASS("simplify-cfg", reduceUsingSimplifyCFGDeltaPass)                \
     DELTA_PASS("attributes", reduceAttributesDeltaPass)                        \
     DELTA_PASS("module-data", reduceModuleDataDeltaPass)                       \
     DELTA_PASS("opcodes", reduceOpcodesDeltaPass)                              \
   } while (false)
 
 #define DELTA_PASSES_MIR                                                       \
   do {                                                                         \
     DELTA_PASS("instructions", reduceInstructionsMIRDeltaPass)                 \

llvm/tools/llvm-reduce/deltas/ReduceBasicBlocks.h

 //===- ReduceArguments.h - Specialized Delta Pass ---------------*- C++ -*-===//
 //
 // 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 a function which calls the Generic Delta pass in order
-// to reduce uninteresting Arguments from defined functions.
+// to reduce uninteresting BasicBlocks from defined functions.
 //
 //===----------------------------------------------------------------------===//
 #ifndef LLVM_TOOLS_LLVM_REDUCE_DELTAS_REDUCEBASICBLOCKS_H
 #define LLVM_TOOLS_LLVM_REDUCE_DELTAS_REDUCEBASICBLOCKS_H
 
 #include "Delta.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 
 namespace llvm {
 void reduceBasicBlocksDeltaPass(TestRunner &Test);
+void reduceUnreachableBasicBlocksDeltaPass(TestRunner &Test);
 } // namespace llvm
 
 #endif

llvm/tools/llvm-reduce/deltas/ReduceBasicBlocks.cpp

 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Value.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include "llvm/Transforms/Utils/Local.h"
+
 #include <vector>
 
+#define DEBUG_TYPE "llvm-reduce"
+
 using namespace llvm;
 
 /// Replaces BB Terminator with one that only contains Chunk BBs
 static void replaceBranchTerminator(BasicBlock &BB,
-                                    const DenseSet<BasicBlock *> &BBsToKeep) {
+                                    const DenseSet<BasicBlock *> &BBsToDelete) {
   auto *Term = BB.getTerminator();
   std::vector<BasicBlock *> ChunkSuccessors;
-  for (auto *Succ : successors(&BB))
-    if (BBsToKeep.count(Succ))
+  for (auto *Succ : successors(&BB)) {
+    if (!BBsToDelete.count(Succ))
       ChunkSuccessors.push_back(Succ);
+  }
 
   // BB only references Chunk BBs
   if (ChunkSuccessors.size() == Term->getNumSuccessors())
     return;
 
   bool IsBranch = isa<BranchInst>(Term) || isa<InvokeInst>(Term);
   Value *Address = nullptr;
   if (auto *IndBI = dyn_cast<IndirectBrInst>(Term))
     Address = IndBI->getAddress();
 
   Term->replaceAllUsesWith(getDefaultValue(Term->getType()));
   Term->eraseFromParent();
 
   if (ChunkSuccessors.empty()) {
     // Scan forward in BB list to try find a block that is kept.
     Function &F = *BB.getParent();
     Function::iterator FI = BB.getIterator();
     FI++;
     while (FI != F.end()) {
       auto &FIB = *FI;
-      if (BBsToKeep.count(&FIB) && !isa<PHINode>(FIB.begin())) {
+      if (!BBsToDelete.count(&FIB) && !isa<PHINode>(FIB.begin())) {
         BranchInst::Create(&FIB, &BB);
         return;
       }
       FI++;
     }
     // If that fails then resort to replacing with a ret.
     auto *FnRetTy = BB.getParent()->getReturnType();
     ReturnInst::Create(BB.getContext(),
static void replaceBranchTerminator(BasicBlock &BB,
   }
 }
 
 /// Removes uninteresting BBs from switch, if the default case ends up being
 /// uninteresting, the switch is replaced with a void return (since it has to be
 /// replace with something)
 static void
 removeUninterestingBBsFromSwitch(SwitchInst &SwInst,
-                                 const DenseSet<BasicBlock *> &BBsToKeep) {
+                                 const DenseSet<BasicBlock *> &BBsToDelete) {
   for (int I = 0, E = SwInst.getNumCases(); I != E; ++I) {
     auto Case = SwInst.case_begin() + I;
-    if (!BBsToKeep.count(Case->getCaseSuccessor())) {
+    if (BBsToDelete.count(Case->getCaseSuccessor())) {
       SwInst.removeCase(Case);
       --I;
       --E;
     }
   }
 
-  if (!BBsToKeep.count(SwInst.getDefaultDest())) {
+  if (BBsToDelete.count(SwInst.getDefaultDest())) {
     if (SwInst.getNumCases() == 0) {
       auto *FnRetTy = SwInst.getParent()->getParent()->getReturnType();
       Value *RetValue =
           FnRetTy->isVoidTy() ? nullptr : getDefaultValue(FnRetTy);
       ReturnInst::Create(SwInst.getContext(), RetValue, SwInst.getParent());
       SwInst.eraseFromParent();
       return;
     }
 
     // Replace the default dest with one of the other cases
     auto Case = SwInst.case_begin();
 
     BasicBlock *NewDefault = Case->getCaseSuccessor();
     SwInst.setDefaultDest(NewDefault);
 
     for (PHINode &SuccPHI : NewDefault->phis()) {
       SuccPHI.addIncoming(SuccPHI.getIncomingValueForBlock(SwInst.getParent()),
                           SwInst.getParent());
     }
   }
 }
 
-/// It's OK to add a block to the set of removed blocks if the first
-/// basic block in the function that survives all of the deletions is
-/// a legal entry block. Keep at least one block in a function.
-static bool okToRemove(BasicBlock &Candidate, Function &F,
-                       const DenseSet<BasicBlock *> &BBsToDelete) {
-  size_t NumBlocksDeleted = 0;
-  bool FoundNewEntryBlock = false;
-  for (auto &B : F) {
-    if (&B == &Candidate)
-      continue;
-    if (BBsToDelete.count(&B)) {
-      ++NumBlocksDeleted;
-      continue;
-    }
-    if (!FoundNewEntryBlock) {
-      /// Ok we've found the first block that's not going to be deleted,
-      /// it will be the new entry block -- that's only legal if this
-      /// block has no predecessors among blocks that survive the
-      /// deletions
-      for (BasicBlock *Pred : predecessors(&B)) {
-        if (!BBsToDelete.contains(Pred))
-          return false;
-      }
-      FoundNewEntryBlock = true;
-    }
-  }
-  // Don't delete the last block.
-  return NumBlocksDeleted + 1 < F.size();
-}
-
 /// Removes out-of-chunk arguments from functions, and modifies their calls
 /// accordingly. It also removes allocations of out-of-chunk arguments.
 static void extractBasicBlocksFromModule(Oracle &O, Module &Program) {
-  DenseSet<BasicBlock *> BBsToKeep, BBsToDelete;
+  DenseSet<BasicBlock *> BBsToDelete;
+  df_iterator_default_set<BasicBlock *> Reachable;
 
   for (auto &F : Program) {
-    for (auto &BB : F) {
-      if (!okToRemove(BB, F, BBsToDelete) || O.shouldKeep())
-        BBsToKeep.insert(&BB);
-      else
-        BBsToDelete.insert(&BB);
+    if (F.empty())
+      continue;
+
+    BasicBlock &Entry = F.getEntryBlock();
+    for (auto *BB : depth_first_ext(&Entry, Reachable)) {
+      (void)BB;
     }
+
+    // Skip any function with unreachable blocks. It's somewhat difficult to
+    // avoid producing invalid IR without deleting them.
+    //
+    // We also do not want to unconditionally delete them, as doing so would
+    // break the invariant of changing the number of chunks during counting.
+
+    const bool HasUnreachableBlocks = Reachable.size() != F.size();
+    Reachable.clear();
+
+    if (HasUnreachableBlocks) {
+      LLVM_DEBUG(dbgs() << "Skipping function with unreachable blocks\n");
+      continue;
+    }
+
+    for (BasicBlock &BB : F) {
+      if (&BB != &Entry && !O.shouldKeep())
+        BBsToDelete.insert(&BB);
     }
 
     // Replace terminators that reference out-of-chunk BBs
-  for (auto &F : Program)
-    for (auto &BB : F) {
+    for (BasicBlock &BB : F) {
       if (auto *SwInst = dyn_cast<SwitchInst>(BB.getTerminator()))
-        removeUninterestingBBsFromSwitch(*SwInst, BBsToKeep);
+        removeUninterestingBBsFromSwitch(*SwInst, BBsToDelete);
       else
-        replaceBranchTerminator(BB, BBsToKeep);
+        replaceBranchTerminator(BB, BBsToDelete);
     }
 
-  // Remove out-of-chunk BB from successor phi nodes
-  for (auto &BB : BBsToDelete) {
-    for (auto *Succ : successors(BB))
-      Succ->removePredecessor(BB, /*KeepOneInputPHIs=*/true);
-  }
-
-  // Replace out-of-chunk switch uses
-  for (auto &BB : BBsToDelete) {
-    // Instructions might be referenced in other BBs
-    for (auto &I : *BB)
-      I.replaceAllUsesWith(getDefaultValue(I.getType()));
-    // Should not be completely removing the body of a function.
-    assert(BB->getParent()->size() > 1);
-    BB->eraseFromParent();
+    // Cleanup any blocks that are now dead after eliminating this set. This
+    // will likely be larger than the number of blocks the oracle told us to
+    // delete.
+    EliminateUnreachableBlocks(F);
+    BBsToDelete.clear();
   }
 }
 
 void llvm::reduceBasicBlocksDeltaPass(TestRunner &Test) {
   runDeltaPass(Test, extractBasicBlocksFromModule, "Reducing Basic Blocks");
 }
+
+static void removeUnreachableBasicBlocksFromModule(Oracle &O, Module &M) {
+  std::vector<BasicBlock *> DeadBlocks;
+  df_iterator_default_set<BasicBlock *> Reachable;
+
+  for (Function &F : M) {
+    if (F.empty())
+      continue;
+
+    // Mark all reachable blocks.
+    for (BasicBlock *BB : depth_first_ext(&F, Reachable))
+      (void)BB;
+
+    if (Reachable.size() != F.size() && !O.shouldKeep()) {
+      for (BasicBlock &BB : F) {
+        if (!Reachable.count(&BB))
+          DeadBlocks.push_back(&BB);
+      }
+
+      // Delete the dead blocks.
+      DeleteDeadBlocks(DeadBlocks, nullptr, /*KeepOneInputPHIs*/ false);
+      DeadBlocks.clear();
+    }
+
+    Reachable.clear();
+  }
+}
+
+void llvm::reduceUnreachableBasicBlocksDeltaPass(TestRunner &Test) {
+  runDeltaPass(Test, removeUnreachableBasicBlocksFromModule,
+               "Removing Unreachable Basic Blocks");
+}