[PartialInliner]: Handle code regions in a switch stmt cases

[etiotto]

Summary

Currently the following C example fails to be partially inlined. The issue is that computeOutliningColdRegionsInfo checks for candidate blocks that have a single entry, but blocks that dominate only themselves fail the check. Example (this is also added as a LLVM IR test):

int callee(int c1, int c2) {
  int rc = 0;

  switch(c1) {
  case 0: // cold
    rc = 1;
    break;
  case 1: // warm
    rc = 2;
    break;
  case 2: // cold
    rc = 4;
    break;
  default: //hot
    rc = c2;
  }

  return rc;
}

int caller(int c1) {
  int rc = callee(c1, c1);
  return rc;
}

With this patch the code in the 2 cold switch cases will be outlined and the remaining function inlined in its caller.

[Whitney]

llvm/lib/Transforms/IPO/PartialInlining.cpp

418

This change allows the single basic block to have multiple predecessor, which is what this function want to avoid.
How about changing BlockList.size() > 1 && to BlockList.size() >= 1 &&?

[etiotto]

Address comments

[fhahn]

llvm/test/Transforms/PartialInlining/switch_stmt.ll

3

run line for old PM, as it is still the default?

[etiotto]

Address comments.

[etiotto]

Add LLVM_DEBUG stmts.

[etiotto]

Fix formatting.

[etiotto]

ping

[Whitney]

LGTM. I will approve tomorrow if no one object..

[fhahn]

llvm/lib/Transforms/IPO/PartialInlining.cpp

509

It seems like it would be simpler to get rid of the isSingleEntry lambda all together and inline the checks.

The check BlockList.size() >= 1 in particular seems redundant. The only case where DominateVector may be empty is if unreachable I think. SI is a successor of a node we found during DFS traversal from the entry block, so all SIs should be reachable? So could this just be something like

  assert(!DominateVector.empty() && "SI should be reachable and have at least itself as descendant");
  // We can only outline single entry regions (for now).
  if (!DominateVector.front()->hasNPredecessors(1)) {
...

511

nit: isSingleEntry checks the predecessors in the dom tree, not the direct predecessors of the node. Might be good to clarify

533

I am not sure why we say ABORT: here but not at the other, similar continues?

[etiotto]

llvm/lib/Transforms/IPO/PartialInlining.cpp

533

I'll add similar debug messages in the other 'continue' conditions

[etiotto]

Address code review comments.

[fhahn]

Lgtm, thanks

llvm/lib/Transforms/IPO/PartialInlining.cpp

PartialInlinerImpl::computeOutliningColdRegionsInfo(
   // Return if we don't have profiling information.
   if (!PSI.hasInstrumentationProfile())
     return std::unique_ptr<FunctionOutliningMultiRegionInfo>();
 
   std::unique_ptr<FunctionOutliningMultiRegionInfo> OutliningInfo =
       std::make_unique<FunctionOutliningMultiRegionInfo>();
 
   auto IsSingleEntry = [](SmallVectorImpl<BasicBlock *> &BlockList) {
+    if (BlockList.size() == 1)
+      return true;
     BasicBlock *Dom = BlockList.front();
     return BlockList.size() > 1 && Dom->hasNPredecessors(1);
   };
 
   auto IsSingleExit =
       [&ORE](SmallVectorImpl<BasicBlock *> &BlockList) -> BasicBlock * {
     BasicBlock *ExitBlock = nullptr;
     for (auto *Block : BlockList) {
for (auto SI = succ_begin(ThisBB); SI != succ_end(ThisBB); ++SI) {
       LLVM_DEBUG(dbgs() << "Found cold edge: " << ThisBB->getName() << "->"
                         << SI->getName()
                         << "\nBranch Probability = " << SuccProb << "\n";);
 
       SmallVector<BasicBlock *, 8> DominateVector;
       DT.getDescendants(*SI, DominateVector);
 
       // We can only outline single entry regions (for now).
       if (!IsSingleEntry(DominateVector))
         continue;
 
       BasicBlock *ExitBlock = nullptr;
       // We can only outline single exit regions (for now).
       if (!(ExitBlock = IsSingleExit(DominateVector)))
         continue;
 
       int OutlineRegionCost = 0;
       for (auto *BB : DominateVector)
         OutlineRegionCost += computeBBInlineCost(BB, &GetTTI(*BB->getParent()));
 
       LLVM_DEBUG(dbgs() << "OutlineRegionCost = " << OutlineRegionCost
                         << "\n";);
 
       if (OutlineRegionCost < MinOutlineRegionCost) {
         ORE.emit([&]() {
           return OptimizationRemarkAnalysis(DEBUG_TYPE, "TooCostly",
                                             &SI->front())
                  << ore::NV("Callee", &F)
                  << " inline cost-savings smaller than "
                  << ore::NV("Cost", MinOutlineRegionCost);
         });
         continue;
       }
       // For now, ignore blocks that belong to a SISE region that is a
       // candidate for outlining.  In the future, we may want to look
       // at inner regions because the outer region may have live-exit
       // variables.
       for (auto *BB : DominateVector)
         VisitedMap[BB] = true;
 
       // ReturnBlock here means the block after the outline call

llvm/test/Transforms/PartialInlining/switch_stmt.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes="partial-inliner" -skip-partial-inlining-cost-analysis -S < %s | FileCheck %s
+
+define dso_local signext i32 @callee(i32 signext %c1, i32 signext %c2) !prof !31 {
+; CHECK-LABEL: @callee(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[RC:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    store i32 0, i32* [[RC]], align 4
+; CHECK-NEXT:    switch i32 [[C1:%.*]], label [[SW_DEFAULT:%.*]] [
+; CHECK-NEXT:    i32 0, label [[SW_BB:%.*]]
+; CHECK-NEXT:    i32 1, label [[SW_BB1:%.*]]
+; CHECK-NEXT:    i32 2, label [[SW_BB2:%.*]]
+; CHECK-NEXT:    ], !prof !31
+; CHECK:       sw.bb:
+; CHECK-NEXT:    store i32 1, i32* [[RC]], align 4
+; CHECK-NEXT:    br label [[SW_EPILOG:%.*]]
+; CHECK:       sw.bb1:
+; CHECK-NEXT:    store i32 2, i32* [[RC]], align 4
+; CHECK-NEXT:    br label [[SW_EPILOG]]
+; CHECK:       sw.bb2:
+; CHECK-NEXT:    store i32 4, i32* [[RC]], align 4
+; CHECK-NEXT:    br label [[SW_EPILOG]]
+; CHECK:       sw.default:
+; CHECK-NEXT:    store i32 [[C2:%.*]], i32* [[RC]], align 4
+; CHECK-NEXT:    br label [[SW_EPILOG]]
+; CHECK:       sw.epilog:
+; CHECK-NEXT:    [[TMP0:%.*]] = load i32, i32* [[RC]], align 4
+; CHECK-NEXT:    ret i32 [[TMP0]]
+;
+entry:
+  %rc = alloca i32, align 4
+  store i32 0, i32* %rc, align 4
+  switch i32 %c1, label %sw.default [
+  i32 0, label %sw.bb
+  i32 1, label %sw.bb1
+  i32 2, label %sw.bb2
+  ], !prof !32
+
+sw.bb: ;; cold
+  store i32 1, i32* %rc, align 4
+  br label %sw.epilog
+
+sw.bb1:
+  store i32 2, i32* %rc, align 4
+  br label %sw.epilog
+
+sw.bb2: ;; cold
+  store i32 4, i32* %rc, align 4
+  br label %sw.epilog
+
+sw.default:
+  store i32 %c2, i32* %rc, align 4
+  br label %sw.epilog
+
+sw.epilog:
+  %0 = load i32, i32* %rc, align 4
+  ret i32 %0
+}
+
+define dso_local signext i32 @caller(i32 signext %c) !prof !31 {
+; CHECK-LABEL: @caller(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[RC_I:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[TMP0:%.*]] = bitcast i32* [[RC_I]] to i8*
+; CHECK-NEXT:    call void @llvm.lifetime.start.p0i8(i64 4, i8* [[TMP0]])
+; CHECK-NEXT:    store i32 0, i32* [[RC_I]], align 4
+; CHECK-NEXT:    switch i32 [[C:%.*]], label [[SW_DEFAULT_I:%.*]] [
+; CHECK-NEXT:    i32 0, label [[CODEREPL_I:%.*]]
+; CHECK-NEXT:    i32 1, label [[SW_BB1_I:%.*]]
+; CHECK-NEXT:    i32 2, label [[CODEREPL1_I:%.*]]
+; CHECK-NEXT:    ], !prof !31
+; CHECK:       codeRepl.i:
+; CHECK-NEXT:    call void @callee.1.sw.bb(i32* [[RC_I]])
+; CHECK-NEXT:    br label [[CALLEE_1_EXIT:%.*]]
+; CHECK:       sw.bb1.i:
+; CHECK-NEXT:    store i32 2, i32* [[RC_I]], align 4
+; CHECK-NEXT:    br label [[CALLEE_1_EXIT]]
+; CHECK:       codeRepl1.i:
+; CHECK-NEXT:    call void @callee.1.sw.bb2(i32* [[RC_I]])
+; CHECK-NEXT:    br label [[CALLEE_1_EXIT]]
+; CHECK:       sw.default.i:
+; CHECK-NEXT:    store i32 [[C]], i32* [[RC_I]], align 4
+; CHECK-NEXT:    br label [[CALLEE_1_EXIT]]
+; CHECK:       callee.1.exit:
+; CHECK-NEXT:    [[TMP1:%.*]] = load i32, i32* [[RC_I]], align 4
+; CHECK-NEXT:    [[TMP2:%.*]] = bitcast i32* [[RC_I]] to i8*
+; CHECK-NEXT:    call void @llvm.lifetime.end.p0i8(i64 4, i8* [[TMP2]])
+;
+entry:
+  %0 = call signext i32 @callee(i32 signext %c, i32 signext %c)
+  ret i32 %0 
+}
+
+!llvm.module.flags = !{!0, !1}
+
+!0 = !{i32 1, !"wchar_size", i32 4}
+!1 = !{i32 1, !"ProfileSummary", !2}
+!2 = !{!3, !4, !5, !6, !7, !8, !9, !10, !11, !12}
+!3 = !{!"ProfileFormat", !"InstrProf"}
+!4 = !{!"TotalCount", i64 2}
+!5 = !{!"MaxCount", i64 1}
+!6 = !{!"MaxInternalCount", i64 1}
+!7 = !{!"MaxFunctionCount", i64 1}
+!8 = !{!"NumCounts", i64 4}
+!9 = !{!"NumFunctions", i64 2}
+!10 = !{!"IsPartialProfile", i64 0}
+!11 = !{!"PartialProfileRatio", double 0.000000e+00}
+!12 = !{!"DetailedSummary", !13}
+!13 = !{!14, !15, !16, !17, !18, !19, !20, !21, !22, !23, !24, !25, !26, !27, !28, !29}
+!14 = !{i32 10000, i64 0, i32 0}
+!15 = !{i32 100000, i64 0, i32 0}
+!16 = !{i32 200000, i64 0, i32 0}
+!17 = !{i32 300000, i64 0, i32 0}
+!18 = !{i32 400000, i64 0, i32 0}
+!19 = !{i32 500000, i64 1, i32 2}
+!20 = !{i32 600000, i64 1, i32 2}
+!21 = !{i32 700000, i64 1, i32 2}
+!22 = !{i32 800000, i64 1, i32 2}
+!23 = !{i32 900000, i64 1, i32 2}
+!24 = !{i32 950000, i64 1, i32 2}
+!25 = !{i32 990000, i64 1, i32 2}
+!26 = !{i32 999000, i64 1, i32 2}
+!27 = !{i32 999900, i64 1, i32 2}
+!28 = !{i32 999990, i64 1, i32 2}
+!29 = !{i32 999999, i64 1, i32 2}
+!31 = !{!"function_entry_count", i64 1000}
+!32 = !{!"branch_weights", i32 500, i32 10, i32 150, i32 40}