[HotColdSplit] Reflect full cost of parameters in split penalty

[vsk]

Summary

Make the penalty for splitting a region more accurately reflect the cost
of materializing all of the inputs/outputs to/from the region.

This almost entirely eliminates code growth within functions which
undergo splitting in key internal frameworks, and reduces the size of
those frameworks between 2.6% to 3%.

rdar://49167240

[vsk]

• Increase test coverage.

[vsk]

Ping.

[vsk]

Ping.

[hiraditya]

LGTM, sorry for not reviewing this patch. I wasn't actively reviewing that time. Please update the diff.

[hiraditya]

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

89

Should we initialize this to be same as # of argument registers. Any arguments more than supported by argument registers are generally stored on stack and may add to code-size overhead both in caller and callee.

291

Do we have a limit on max # successors outside of the region? In pathological cases this piece of code may have super linear time complexity. If we dont encounter this often then this seems fine though.

[vsk]

I haven't forgotten about this! Hope to have an update by Monday (8/10).

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

89

Ah, do you mean the number of arguments that can be passed by register? If so, yes, that seems like a straightforward improvement.

291

No, we don't have a limit like that, but I'd be open to adding one if a pathological case is found. I've never observed a performance issue here (in practice |SuccsOutsideRegion| is very small).

[hiraditya]

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

89

yes

291

I've never observed a performance issue here (in practice |SuccsOutsideRegion| is very small).

In that case the current approach is fine.

[vsk]

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

89

I had a look at TargetLowering, TargetCallingConv and a few related interfaces, but couldn't find a way to determine the number of gprs available for pointer-sized call arguments. The necessary information is available in .td files, e.g. for AArch64 there's an entry for:

CCIfType<[i32], CCAssignToRegWithShadow<[W0, W1, W2, W3, W4, W5, W6, W7],
                                        [X0, X1, X2, X3, X4, X5, X6, X7]>>,

Do you know whether there's a straightforward way to surface that?

[vsk]

Rebase.

[xbolva00]

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

89

Maybe @ctopper knows?

[rjf]

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

89

Is there a way to get access to a TargetMachine instance here? If so, then maybe then you can use MCRegisterInfo::getNumRegs() which can be accessed via TargetMachine to get the number of registers (Apologies in advance if this isn't helping).

[vsk]

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

89

MCRegisterInfo::getNumRegs() seems to count the number of registers available on the target, but we're just interested in the subset that can be used to pass pointer-sized values given the selected calling convention.

[hiraditya]

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

89

It doesn't look like the calling convention is exposed to the middle end. So it might be tricky to do the plumbing work to get those values. Given that the calling convention do not usually change can we hard-code some commonly used values in this pass?

Briefly looking at CallingConvention files in lib/Target, it seems:

• X86_64 only allows 4 integer arguments to be passed in a register, and 4 floating point.
• Arch64, ARM, RISC-V have 8 argument registers.

Depending on a calling convention these values do change though.

[rjf]

@vsk Hope these following data collected from compiling Firefox helps. The "Delta" values come from my previous patch: https://reviews.llvm.org/D84468. Based on the data, this patch performs very well at reducing the code size.
-Os:

Delta	Size (including dynamic libraries)
D59715	2.184796592 GB
0	2.188262032 GB
5	2.206931464 GB

-O3:

Delta	Size (including dynamic libraries)
-2	2.270277648 GB
0	2.247788640 GB
D59715	2.243288440 GB
2	2.259242024 GB
5	2.270277648 GB

[hiraditya]

Any plans to merge this patch?

[vsk]

@hiraditya @rjf, thanks for sharing those numbers. At the moment I don't have the bandwidth to evaluate setting a different max-params value on x86. Let me know how you'd prefer to proceed - I'm happy to have either of you commandeer the patch, or if you prefer, I can land it in its current state.

[rjf]

@vsk If you're curious about how the max-params value affects code size, I think I can find some time to collect some numbers soon (on firefox). I'll post it here as soon as I have it.

[hiraditya]

@rjf do you have numbers on Firefox. Let us know.

[hiraditya]

or if you prefer, I can land it in its current state.

+1. let's land this.

[hiraditya]

@vsk Let me know if you dont have bandwidth to land. I can land it otherwise.

[hiraditya]

Rebase, fix a typo, change default hotcoldsplit-max-params=4

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

 #include "llvm/Transforms/IPO.h"
 #include "llvm/Transforms/Scalar.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/CodeExtractor.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Transforms/Utils/ValueMapper.h"
 #include <algorithm>
+#include <limits>
 #include <cassert>
 #include <string>
 
 #define DEBUG_TYPE "hotcoldsplit"
 
 STATISTIC(NumColdRegionsFound, "Number of cold regions found.");
 STATISTIC(NumColdRegionsOutlined, "Number of cold regions outlined.");
 
 using namespace llvm;
 
 static cl::opt<bool> EnableStaticAnalysis("hot-cold-static-analysis",
                                           cl::init(true), cl::Hidden);
 
 static cl::opt<int>
     SplittingThreshold("hotcoldsplit-threshold", cl::init(2), cl::Hidden,
                        cl::desc("Base penalty for splitting cold code (as a "
                                 "multiple of TCC_Basic)"));
 
 static cl::opt<bool> EnableColdSection(
     "enable-cold-section", cl::init(false), cl::Hidden,
     cl::desc("Enable placement of extracted cold functions"
              " into a separate section after hot-cold splitting."));
 
 static cl::opt<std::string>
     ColdSectionName("hotcoldsplit-cold-section-name", cl::init("__llvm_cold"),
                     cl::Hidden,
                     cl::desc("Name for the section containing cold functions "
                              "extracted by hot-cold splitting."));
 
+static cl::opt<int> MaxParametersForSplit(
+    "hotcoldsplit-max-params", cl::init(4), cl::Hidden,
+    cl::desc("Maximum number of parameters for a split function"));
+
 namespace {
 // Same as blockEndsInUnreachable in CodeGen/BranchFolding.cpp. Do not modify
 // this function unless you modify the MBB version as well.
 //
 /// A no successor, non-return block probably ends in unreachable and is cold.
 /// Also consider a block that ends in an indirect branch to be a return block,
 /// since many targets use plain indirect branches to return.
 bool blockEndsInUnreachable(const BasicBlock &BB) {
static int getOutliningPenalty(ArrayRef<BasicBlock *> Region,
   int Penalty = SplittingThreshold;
   LLVM_DEBUG(dbgs() << "Applying penalty for splitting: " << Penalty << "\n");
 
   // If the splitting threshold is set at or below zero, skip the usual
   // profitability check.
   if (SplittingThreshold <= 0)
     return Penalty;
 
-  // The typical code size cost for materializing an argument for the outlined
-  // call.
-  LLVM_DEBUG(dbgs() << "Applying penalty for: " << NumInputs << " inputs\n");
-  const int CostForArgMaterialization = TargetTransformInfo::TCC_Basic;
-  Penalty += CostForArgMaterialization * NumInputs;
-
-  // The typical code size cost for an output alloca, its associated store, and
-  // its associated reload.
-  LLVM_DEBUG(dbgs() << "Applying penalty for: " << NumOutputs << " outputs\n");
-  const int CostForRegionOutput = 3 * TargetTransformInfo::TCC_Basic;
-  Penalty += CostForRegionOutput * NumOutputs;
-
   // Find the number of distinct exit blocks for the region. Use a conservative
   // check to determine whether control returns from the region.
   bool NoBlocksReturn = true;
   SmallPtrSet<BasicBlock *, 2> SuccsOutsideRegion;
   for (BasicBlock *BB : Region) {
     // If a block has no successors, only assume it does not return if it's
     // unreachable.
     if (succ_empty(BB)) {
       NoBlocksReturn &= isa<UnreachableInst>(BB->getTerminator());
       continue;
     }
 
     for (BasicBlock *SuccBB : successors(BB)) {
       if (!is_contained(Region, SuccBB)) {
         NoBlocksReturn = false;
         SuccsOutsideRegion.insert(SuccBB);
       }
     }
   }
 
+  // Count the number of phis in exit blocks with >= 2 incoming values from the
+  // outlining region. These phis are split (\ref severSplitPHINodesOfExits),
+  // and new outputs are created to supply the split phis. CodeExtractor can't
+  // report these new outputs until extraction begins, but it's important to
+  // factor the cost of the outputs into the cost calculation.
+  unsigned NumSplitExitPhis = 0;
+  for (BasicBlock *ExitBB : SuccsOutsideRegion) {
+    for (PHINode &PN : ExitBB->phis()) {
+      // Find all incoming values from the outlining region.
+      int NumIncomingVals = 0;
+      for (unsigned i = 0; i < PN.getNumIncomingValues(); ++i)
+        if (find(Region, PN.getIncomingBlock(i)) != Region.end()) {
+          ++NumIncomingVals;
+          if (NumIncomingVals > 1) {
+            ++NumSplitExitPhis;
+            break;
+          }
+        }
+    }
+  }
+
+  // Apply a penalty for calling the split function. Factor in the cost of
+  // materializing all of the parameters.
+  int NumOutputsAndSplitPhis = NumOutputs + NumSplitExitPhis;
+  int NumParams = NumInputs + NumOutputsAndSplitPhis;
+  if (NumParams > MaxParametersForSplit) {
+    LLVM_DEBUG(dbgs() << NumInputs << " inputs and " << NumOutputsAndSplitPhis
+                      << " outputs exceeds parameter limit ("
+                      << MaxParametersForSplit << ")\n");
+    return std::numeric_limits<int>::max();
+  }
+  const int CostForArgMaterialization = 2 * TargetTransformInfo::TCC_Basic;
+  LLVM_DEBUG(dbgs() << "Applying penalty for: " << NumParams << " params\n");
+  Penalty += CostForArgMaterialization * NumParams;
+
+  // Apply the typical code size cost for an output alloca and its associated
+  // reload in the caller. Also penalize the associated store in the callee.
+  LLVM_DEBUG(dbgs() << "Applying penalty for: " << NumOutputsAndSplitPhis
+                    << " outputs/split phis\n");
+  const int CostForRegionOutput = 3 * TargetTransformInfo::TCC_Basic;
+  Penalty += CostForRegionOutput * NumOutputsAndSplitPhis;
+
   // Apply a `noreturn` bonus.
   if (NoBlocksReturn) {
     LLVM_DEBUG(dbgs() << "Applying bonus for: " << Region.size()
                       << " non-returning terminators\n");
     Penalty -= Region.size();
   }
 
   // Apply a penalty for having more than one successor outside of the region.
   // This penalty accounts for the switch needed in the caller.
-  if (!SuccsOutsideRegion.empty()) {
+  if (SuccsOutsideRegion.size() > 1) {
     LLVM_DEBUG(dbgs() << "Applying penalty for: " << SuccsOutsideRegion.size()
                       << " non-region successors\n");
     Penalty += (SuccsOutsideRegion.size() - 1) * TargetTransformInfo::TCC_Basic;
   }
 
   return Penalty;
 }
 

llvm/test/Transforms/CodeExtractor/extract-assume.ll

-; RUN: opt -passes="function(slp-vectorizer),module(hotcoldsplit),function(slp-vectorizer,print<assumptions>)" -disable-output %s 2>&1 | FileCheck %s
+; RUN: opt -passes="function(slp-vectorizer),module(hotcoldsplit),function(slp-vectorizer,print<assumptions>)" -hotcoldsplit-threshold=-1 -disable-output %s 2>&1 | FileCheck %s
 ;
 ; Make sure this compiles. Check that function assumption cache is refreshed
 ; after extracting blocks with assume calls from the function.
 
 ; CHECK: Cached assumptions for function: fun
 ; CHECK-NEXT: Cached assumptions for function: fun.cold
 ; CHECK-NOT: icmp uge
 

llvm/test/Transforms/HotColdSplit/apply-penalty-for-inputs.ll

 ; REQUIRES: asserts
-; RUN: opt -hotcoldsplit -debug-only=hotcoldsplit -S < %s -o /dev/null 2>&1 | FileCheck %s
+; RUN: opt -hotcoldsplit -debug-only=hotcoldsplit -hotcoldsplit-threshold=2 -hotcoldsplit-max-params=2 -S < %s -o /dev/null 2>&1 | FileCheck %s
 
 declare void @sink(i32*, i32, i32) cold
 
 @g = global i32 0
 
 define void @foo(i32 %arg) {
   %local = load i32, i32* @g
   br i1 undef, label %cold, label %exit
 
 cold:
-  ; CHECK: Applying penalty for: 2 inputs
+  ; CHECK: Applying penalty for splitting: 2
+  ; CHECK-NEXT: Applying penalty for: 2 params
+  ; CHECK-NEXT: Applying penalty for: 0 outputs/split phis
+  ; CHECK-NEXT: penalty = 6
   call void @sink(i32* @g, i32 %arg, i32 %local)
   ret void
 
 exit:
   ret void
 }
+
+define void @bar(i32* %p1, i32 %p2, i32 %p3) {
+  br i1 undef, label %cold, label %exit
+
+cold:
+  ; CHECK: Applying penalty for splitting: 2
+  ; CHECK-NEXT: 3 inputs and 0 outputs exceeds parameter limit (2)
+  ; CHECK-NEXT: penalty = 2147483647
+  call void @sink(i32* %p1, i32 %p2, i32 %p3)
+  ret void
+
+exit:
+  ret void
+}

llvm/test/Transforms/HotColdSplit/apply-penalty-for-outputs.ll

 ; REQUIRES: asserts
-; RUN: opt -hotcoldsplit -debug-only=hotcoldsplit -S < %s -o /dev/null 2>&1 | FileCheck %s
+; RUN: opt -hotcoldsplit -debug-only=hotcoldsplit -hotcoldsplit-threshold=2 -S < %s -o /dev/null 2>&1 | FileCheck %s
 
 declare void @sink() cold
 
 @g = global i32 0
 
 define i32 @foo(i32 %arg) {
 entry:
   br i1 undef, label %cold, label %exit
 
 cold:
-  ; CHECK: Applying penalty for: 1 output
-  ; CHECK: Applying penalty for: 1 non-region successors
+  ; CHECK: Applying penalty for splitting: 2
+  ; CHECK-NEXT: Applying penalty for: 1 params
+  ; CHECK-NEXT: Applying penalty for: 1 outputs/split phis
+  ; CHECK-NEXT: penalty = 7
   %local = load i32, i32* @g
   call void @sink()
   br label %exit
 
 exit:
   %p = phi i32 [ %local, %cold ], [ 0, %entry ]
   ret i32 %p
 }

llvm/test/Transforms/HotColdSplit/apply-successor-penalty.ll

 ; REQUIRES: asserts
-; RUN: opt -hotcoldsplit -debug-only=hotcoldsplit -S < %s -o /dev/null 2>&1 | FileCheck %s
+; RUN: opt -hotcoldsplit -debug-only=hotcoldsplit -hotcoldsplit-threshold=2 -S < %s -o /dev/null 2>&1 | FileCheck %s
 
 declare void @sink() cold
 
 ; CHECK-LABEL: Outlining in one_non_region_successor
 define void @one_non_region_successor(i32 %arg) {
 entry:
   br i1 undef, label %cold1, label %exit
 
 cold1:
-  ; CHECK: Applying penalty for: 1 non-region successor
+  ; CHECK: Applying penalty for splitting: 2
+  ; CHECK-NEXT: Applying penalty for: 0 params
+  ; CHECK-NEXT: Applying penalty for: 0 outputs/split phis
+  ; CHECK-NEXT: penalty = 2
   call void @sink()
   br i1 undef, label %cold2, label %cold3
 
 cold2:
   br i1 undef, label %cold4, label %exit
 
 cold3:
   br i1 undef, label %cold4, label %exit
 
 cold4:
   unreachable
 
 exit:
   ret void
 }
 
 ; CHECK-LABEL: Outlining in two_non_region_successor
 define void @two_non_region_successors(i32 %arg) {
 entry:
   br i1 undef, label %cold1, label %exit1
 
 cold1:
-  ; CHECK: Applying penalty for: 2 non-region successors
+  ; CHECK: Applying penalty for splitting: 2
+  ; CHECK-NEXT: Applying penalty for: 0 params
+  ; CHECK-NEXT: Applying penalty for: 0 outputs/split phis
+  ; CHECK-NEXT: Applying penalty for: 2 non-region successors
+  ; CHECK-NEXT: penalty = 3
   call void @sink()
   br i1 undef, label %cold2, label %cold3
 
 cold2:
   br i1 undef, label %cold4, label %exit1
 
 cold3:
   br i1 undef, label %cold4, label %exit2

llvm/test/Transforms/HotColdSplit/assumption-cache-invalidation.ll

 ; REQUIRES: asserts
-; RUN: opt -S -instsimplify -hotcoldsplit -debug < %s 2>&1 | FileCheck %s
+; RUN: opt -S -instsimplify -hotcoldsplit -hotcoldsplit-threshold=-1 -debug < %s 2>&1 | FileCheck %s
 ; RUN: opt -instcombine -hotcoldsplit -instsimplify %s -o /dev/null
 
 target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
 target triple = "aarch64"
 
 %a = type { i64, i64 }
 %b = type { i64 }
 
 ; CHECK: @f
 ; CHECK-LABEL: codeRepl:
 ; CHECK-NOT: @llvm.assume
 ; CHECK: }
 ; CHECK: declare {{.*}}@llvm.assume
-; CHECK: define {{.*}}@f.cold.1(i64 %0)
+; CHECK: define {{.*}}@f.cold.1()
+; CHECK-LABEL: newFuncRoot:
+; CHECK: }
+; CHECK: define {{.*}}@f.cold.2(i64 %0)
 ; CHECK-LABEL: newFuncRoot:
 ; CHECK: %1 = icmp eq i64 %0, 0
 ; CHECK-NOT: call void @llvm.assume
 
 define void @f() {
 entry:
   %0 = getelementptr inbounds %a, %a* null, i64 0, i32 1
   br label %label