[PartialInlining] Fix incorrect costing when IR has unreachable BBs

Partial Inlining identifies basic blocks that can be outlined into a
function. It is possible that an unreachable basic block is marked for
outlining. During costing of the outlined region, such unreachable basic
blocks are included as well. However, the CodeExtractor eliminates such
unreachable basic blocks and emits outlined function without them.

Thus, during costing of the outlined function, it is possible that the
cost of the outlined function comes out to be lesser than the cost of
outlined region, which triggers an assert.

Assertion `OutlinedFunctionCost >= Cloner.OutlinedRegionCost && "Outlined
function cost should be no less than the outlined region"' failed.

This patch adds code to eliminate unreachable blocks from the function
body before passing it on to be inlined. It also adds a test that checks
for behaviour of costing in case of unreachable basic blocks.

Discussion: https://discourse.llvm.org/t/incorrect-costing-in-partialinliner-if-ir-has-unreachable-basic-blocks/70163

Reviewed By: fhahn

Differential Revision: https://reviews.llvm.org/D149130

llvm/lib/Transforms/IPO/PartialInlining.cpp

@@ -14,6 +14,7 @@
 #include "llvm/Transforms/IPO/PartialInlining.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/DepthFirstIterator.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
@@ -1178,14 +1179,14 @@ PartialInlinerImpl::FunctionCloner::doSingleRegionFunctionOutlining() {
   ToExtract.push_back(ClonedOI->NonReturnBlock);
   OutlinedRegionCost += PartialInlinerImpl::computeBBInlineCost(
       ClonedOI->NonReturnBlock, ClonedFuncTTI);
-  for (BasicBlock &BB : *ClonedFunc)
-    if (!ToBeInlined(&BB) && &BB != ClonedOI->NonReturnBlock) {
-      ToExtract.push_back(&BB);
+  for (BasicBlock *BB : depth_first(&ClonedFunc->getEntryBlock()))
+    if (!ToBeInlined(BB) && BB != ClonedOI->NonReturnBlock) {
+      ToExtract.push_back(BB);
       // FIXME: the code extractor may hoist/sink more code
       // into the outlined function which may make the outlining
       // overhead (the difference of the outlined function cost
       // and OutliningRegionCost) look larger.
-      OutlinedRegionCost += computeBBInlineCost(&BB, ClonedFuncTTI);
+      OutlinedRegionCost += computeBBInlineCost(BB, ClonedFuncTTI);
     }
 
   // Extract the body of the if.

llvm/test/Transforms/PartialInlining/unreachable_basic_block.ll

@@ -0,0 +1,29 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes=partial-inliner -S < %s | FileCheck %s
+declare i1 @llvm.public.type.test(ptr, metadata)
+
+define void @dummy() {
+; CHECK-LABEL: @dummy(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br i1 false, label [[WHILE_END:%.*]], label [[WHILE_BODY:%.*]]
+; CHECK:       while.body:
+; CHECK-NEXT:    call void @dummy.1.while.body()
+; CHECK-NEXT:    br label [[WHILE_END]]
+; CHECK:       while.end:
+; CHECK-NEXT:    ret void
+;
+entry:
+  br i1 false, label %while.end, label %while.body
+
+while.body:                                       ; preds = %entry
+  call void @dummy()
+  br label %while.end
+
+unreachable.block:               ; No predecessors!
+  %0 = tail call i1 @llvm.public.type.test(ptr null, metadata !"test-function")
+  %result = getelementptr ptr, ptr null, i64 1
+  br label %while.end
+
+while.end:                                        ; preds = %unreachable.block, %while.body, %entry
+  ret void
+}