[Inline][Cloning] Reliably remove unreachable blocks during cloning (PR53206)

[nikic]

Summary

The pruning cloner already tries to remove unreachable blocks. The original cloning process will simplify instructions and constant terminators, and only clone blocks that are reachable at that point. However, phi nodes can only be simplified after everything has been cloned. For that reason, additional blocks may become unreachable after phi simplification.

The code does try to handle this as well, but only removes blocks that don't have predecessors. It misses unreachable cycles. This can cause issues if SEH exception handling code is part of an unreachable cycle, as the inliner is not prepared to deal with that.

This patch instead performs an explicit scan for reachable blocks, and drops everything else.

Fixes https://github.com/llvm/llvm-project/issues/53206.

[lebedev.ri]

Have you considered using removeUnreachableBlocks()?

[nikic]

In D118449#3278938, @lebedev.ri wrote:

Have you considered using removeUnreachableBlocks()?

I did consider that, but the implementation seems to be doing significantly more than just removing unreachable blocks, and I was not entirely certain that it is appropriate here.

If desired, I can switch to that though (with an interface change to accept a start block rather than a function).

[rnk]

removeUnreachableBlocks operates on the whole function and preserves extra analyses. I think it makes sense to have a mini-DFS here to consider only a single inlined function body.

lgtm

llvm/lib/Transforms/Utils/CloneFunction.cpp

for (unsigned Idx = 0; Idx != Worklist.size(); ++Idx) {
 
     // If the original instruction had no side effects, remove it.
     if (isInstructionTriviallyDead(I))
       I->eraseFromParent();
     else
       VMap[OrigV] = I;
   }
 
+  // Simplify conditional branches and switches with a constant operand. We try
+  // to prune these out when cloning, but if the simplification required
+  // looking through PHI nodes, those are only available after forming the full
+  // basic block. That may leave some here, and we still want to prune the dead
+  // code as early as possible.
+  Function::iterator Begin = cast<BasicBlock>(VMap[StartingBB])->getIterator();
+  for (BasicBlock &BB : make_range(Begin, NewFunc->end()))
+    ConstantFoldTerminator(&BB);
+
+  // Some blocks may have become unreachable as a result. Find and delete them.
+  {
+    SmallPtrSet<BasicBlock *, 16> ReachableBlocks;
+    SmallVector<BasicBlock *, 16> Worklist;
+    Worklist.push_back(&*Begin);
+    while (!Worklist.empty()) {
+      BasicBlock *BB = Worklist.pop_back_val();
+      if (ReachableBlocks.insert(BB).second)
+        append_range(Worklist, successors(BB));
+    }
+
+    SmallVector<BasicBlock *, 16> UnreachableBlocks;
+    for (BasicBlock &BB : make_range(Begin, NewFunc->end()))
+      if (!ReachableBlocks.contains(&BB))
+        UnreachableBlocks.push_back(&BB);
+    DeleteDeadBlocks(UnreachableBlocks);
+  }
+
   // Now that the inlined function body has been fully constructed, go through
   // and zap unconditional fall-through branches. This happens all the time when
   // specializing code: code specialization turns conditional branches into
   // uncond branches, and this code folds them.
-  Function::iterator Begin = cast<BasicBlock>(VMap[StartingBB])->getIterator();
   Function::iterator I = Begin;
   while (I != NewFunc->end()) {
-    // We need to simplify conditional branches and switches with a constant
-    // operand. We try to prune these out when cloning, but if the
-    // simplification required looking through PHI nodes, those are only
-    // available after forming the full basic block. That may leave some here,
-    // and we still want to prune the dead code as early as possible.
-    //
-    // Do the folding before we check if the block is dead since we want code
-    // like
-    //  bb:
-    //    br i1 undef, label %bb, label %bb
-    // to be simplified to
-    //  bb:
-    //    br label %bb
-    // before we call I->getSinglePredecessor().
-    ConstantFoldTerminator(&*I);
-
-    // Check if this block has become dead during inlining or other
-    // simplifications. Note that the first block will appear dead, as it has
-    // not yet been wired up properly.
-    if (I != Begin && (pred_empty(&*I) || I->getSinglePredecessor() == &*I)) {
-      BasicBlock *DeadBB = &*I++;
-      DeleteDeadBlock(DeadBB);
-      continue;
-    }
-
     BranchInst *BI = dyn_cast<BranchInst>(I->getTerminator());
     if (!BI || BI->isConditional()) {
       ++I;
       continue;
     }
 
     BasicBlock *Dest = BI->getSuccessor(0);
     if (!Dest->getSinglePredecessor()) {

llvm/test/Transforms/Inline/pr53206.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -inline < %s | FileCheck %s
+
+; Check that the exception handling code is fully pruned, and does not
+; leave behind invalid IR.
+
+define internal void @foo() personality i32 (...)* undef {
+entry:
+  br i1 false, label %join, label %split
+
+split:
+  br label %join
+
+join:
+  %phi = phi i64 [ 1, %split ], [ 0, %entry ]
+  %cmp = icmp ugt i64 1, %phi
+  br i1 %cmp, label %invoke1, label %exit
+
+invoke1:
+  invoke void undef()
+  to label %exit unwind label %cleanup1
+
+cleanup1:
+  %pad1 = cleanuppad within none []
+  br label %cleanup1.cont
+
+cleanup1.cont:
+  br i1 undef, label %cleanupret, label %invoke2
+
+invoke2:
+  invoke void undef() [ "funclet"(token %pad1) ]
+  to label %cleanup1.cont unwind label %cleanup2
+
+cleanup2:
+  %pad2 = cleanuppad within %pad1 []
+  unreachable
+
+cleanupret:
+  unreachable
+
+exit:
+  ret void
+}
+
+define void @test() personality i32 (...)* undef {
+; CHECK-LABEL: @test(
+; CHECK-NEXT:    ret void
+;
+  call void @foo()
+  ret void
+}