ColdBlockInfo: consider post-dominating coldness

A block that is post-dominated by a cold block is itself cold. The most
basic kind of cold block is one that ends with a fatalError or some
other non-returning call, as that block is only executed at most once
per task.

To help propagate this exit-block coldness, without rewriting the
algorithm to be something that iterates to a fixed-point, I've added
some fast checks for blocks that reach "obviously cold" blocks.

Author: kavon

include/swift/AST/SILOptions.h

@@ -122,6 +122,9 @@ class SILOptions {
   /// within the optimizer. This influences static branch prediction.
   bool EnableThrowsPrediction = false;
 
+  /// Controls whether to say that blocks ending in an 'unreachable' are cold.
+  bool EnableNoReturnCold = false;
+
   /// Should we run any SIL performance optimizations
   ///
   /// Useful when you want to enable -O LLVM opts but not -O SIL opts.

include/swift/Option/FrontendOptions.td

@@ -545,6 +545,11 @@ def enable_throws_prediction : Flag<["-"], "enable-throws-prediction">,
 def disable_throws_prediction : Flag<["-"], "disable-throws-prediction">,
   HelpText<"Disables optimization assumption that functions rarely throw errors.">;
 
+def enable_noreturn_prediction : Flag<["-"], "enable-noreturn-prediction">,
+  HelpText<"Enables optimization assumption that calls to no-return functions are cold.">;
+def disable_noreturn_prediction : Flag<["-"], "disable-noreturn-prediction">,
+  HelpText<"Disables optimization assumption that calls to no-return functions are cold.">;
+
 def disable_access_control : Flag<["-"], "disable-access-control">,
   HelpText<"Don't respect access control restrictions">;
 def enable_access_control : Flag<["-"], "enable-access-control">,

include/swift/SILOptimizer/Analysis/ColdBlockInfo.h

@@ -18,6 +18,7 @@
 
 namespace swift {
 class DominanceAnalysis;
+class PostDominanceAnalysis;
 class SILBasicBlock;
 
 /// Cache a set of basic blocks that have been determined to be cold or hot.
@@ -26,6 +27,7 @@ class SILBasicBlock;
 /// across passes.
 class ColdBlockInfo {
   DominanceAnalysis *DA;
+  PostDominanceAnalysis *PDA;
 
   /// Each block in this map has been determined to be either cold or hot.
   llvm::DenseMap<const SILBasicBlock*, bool> ColdBlockMap;
@@ -53,8 +55,11 @@ class ColdBlockInfo {
   bool isCold(const SILBasicBlock *BB,
               int recursionDepth);
 
+  bool reachesColdBlock(const SILBasicBlock *BB);
+
 public:
-  ColdBlockInfo(DominanceAnalysis *DA): DA(DA) {}
+  ColdBlockInfo(DominanceAnalysis *DA, PostDominanceAnalysis *PDA)
+      : DA(DA), PDA(PDA) {}
 
   bool isSlowPath(const SILBasicBlock *FromBB, const SILBasicBlock *ToBB) {
     return isSlowPath(FromBB, ToBB, 0);

lib/DriverTool/sil_opt_main.cpp

@@ -293,6 +293,10 @@ struct SILOptOptions {
   EnableThrowsPrediction = llvm::cl::opt<bool>("enable-throws-prediction",
                      llvm::cl::desc("Enables optimization assumption that functions rarely throw errors."));
 
+  llvm::cl::opt<bool>
+  EnableNoReturnCold = llvm::cl::opt<bool>("enable-noreturn-prediction",
+                     llvm::cl::desc("Enables optimization assumption that calls to no-return functions are cold."));
+
   llvm::cl::opt<bool>
   EnableMoveInoutStackProtection = llvm::cl::opt<bool>("enable-move-inout-stack-protector",
                     llvm::cl::desc("Enable the stack protector by moving values to temporaries."));
@@ -852,6 +856,7 @@ int sil_opt_main(ArrayRef<const char *> argv, void *MainAddr) {
   SILOpts.EnableSpeculativeDevirtualization = options.EnableSpeculativeDevirtualization;
   SILOpts.EnableAsyncDemotion = options.EnableAsyncDemotion;
   SILOpts.EnableThrowsPrediction = options.EnableThrowsPrediction;
+  SILOpts.EnableNoReturnCold = options.EnableNoReturnCold;
   SILOpts.IgnoreAlwaysInline = options.IgnoreAlwaysInline;
   SILOpts.EnableOSSAModules = options.EnableOSSAModules;
   SILOpts.EnableSILOpaqueValues = options.EnableSILOpaqueValues;

lib/Frontend/CompilerInvocation.cpp

@@ -2661,6 +2661,9 @@ static bool ParseSILArgs(SILOptions &Opts, ArgList &Args,
   Opts.EnableThrowsPrediction = Args.hasFlag(
       OPT_enable_throws_prediction, OPT_disable_throws_prediction,
       Opts.EnableThrowsPrediction);
+  Opts.EnableNoReturnCold = Args.hasFlag(
+      OPT_enable_noreturn_prediction, OPT_disable_noreturn_prediction,
+      Opts.EnableNoReturnCold);
   Opts.EnableActorDataRaceChecks |= Args.hasFlag(
       OPT_enable_actor_data_race_checks,
       OPT_disable_actor_data_race_checks, /*default=*/false);

lib/SILOptimizer/Analysis/ColdBlockInfo.cpp

@@ -14,6 +14,7 @@
 #include "swift/SILOptimizer/Analysis/DominanceAnalysis.h"
 #include "swift/SIL/SILArgument.h"
 #include "swift/AST/SemanticAttrs.h"
+#include "swift/SIL/SILModule.h"
 
 #define DEBUG_TYPE "cold-block-info"
 
@@ -29,6 +30,37 @@ static SILValue getCondition(SILValue C) {
   return C;
 }
 
+/// A cold terminator is one where it's unlikely to be reached, which are
+/// function exits that are less-common. A cold terminator implies a cold block.
+///   - 'unreachable', as it's never executed.
+///   - 'throw', if throws prediction is enabled.
+static bool isColdTerminator(const TermInst *term) {
+  switch (term->getTermKind()) {
+  case TermKind::AwaitAsyncContinuationInst:
+  case TermKind::BranchInst:
+  case TermKind::CondBranchInst:
+  case TermKind::SwitchValueInst:
+  case TermKind::SwitchEnumInst:
+  case TermKind::SwitchEnumAddrInst:
+  case TermKind::DynamicMethodBranchInst:
+  case TermKind::CheckedCastBranchInst:
+  case TermKind::CheckedCastAddrBranchInst:
+  case TermKind::TryApplyInst:
+  case TermKind::YieldInst:
+  case TermKind::ReturnInst:
+    return false;
+  case TermKind::UnreachableInst:
+    return term->getModule().getOptions().EnableNoReturnCold;
+  case TermKind::ThrowInst:
+  case TermKind::ThrowAddrInst:
+  case TermKind::UnwindInst:
+    return term->getModule().getOptions().EnableThrowsPrediction;
+  }
+}
+static bool hasColdTerminator(const SILBasicBlock *bb) {
+  return isColdTerminator(bb->getTerminator());
+}
+
 /// \return a BranchHint if this conditional's likely value can be inferred.
 ColdBlockInfo::BranchHint ColdBlockInfo::getBranchHint(SILValue Cond,
                                                        int recursionDepth) {
@@ -135,6 +167,57 @@ static std::optional<bool> isColdProfiledEdge(const SILBasicBlock *FromBB,
   return takenProbability < WARM_EDGE_MINIMUM;
 }
 
+// Does this block always reach a cold block?
+//
+// The `isCold` function is not used here, to avoid duplicated work and
+// keep this function fast and non-recursive.
+bool ColdBlockInfo::reachesColdBlock(const SILBasicBlock *BB) {
+  assert(ColdBlockMap.find(BB) == ColdBlockMap.end() && "already known temp?");
+  auto *fn = const_cast<SILFunction *>(BB->getParent());
+  PostDominanceInfo *info = PDA->get(fn);
+  auto *node = info->getNode(BB);
+
+  // Advance to the block closest to the root that post-dominates BB.
+  // Such a block is the only kind that might exit the function.
+  SmallVector<const SILBasicBlock *, 8> chain;
+  std::optional<bool> lastBBColdness;
+  while (node && !info->isVirtualRoot(node)) {
+    auto postDomBB = node->getBlock();
+
+    // If the block has a known temperature, we can stop early.
+    auto result = ColdBlockMap.find(postDomBB);
+    if (result != ColdBlockMap.end()) {
+      lastBBColdness = result->getSecond();
+      break;
+    }
+
+    chain.push_back(node->getBlock());
+    node = node->getIDom();
+  }
+
+  // If we don't already know the last block's temperature, do a fast check.
+  if (!lastBBColdness && hasColdTerminator(chain.back()))
+    lastBBColdness = true;
+
+  // We don't have a definitive answer, so say it does not reach a cold one.
+  if (!lastBBColdness)
+    return false;
+
+  // Mark all blocks in the dominance chain with the given temperature.
+  for (auto chainBB : chain) {
+    ColdBlockMap[chainBB] = *lastBBColdness;
+
+    if (*lastBBColdness) {
+      LLVM_DEBUG(llvm::dbgs()
+                     << "ColdBlockInfo(" << BB->getParent()->getName() << "): "
+                     << "bb" << chainBB->getDebugID() << " is cold\n");
+    }
+  }
+
+  assert(ColdBlockMap.find(BB) != ColdBlockMap.end() && "didn't cache?");
+  return *lastBBColdness;
+}
+
 /// \return true if the CFG path FromBB->*->ToBB is either
 ///   1. directly gated by a _slowPath branch hint.
 ///   2. has a low probability of being taken according to profile information.
@@ -175,6 +258,10 @@ bool ColdBlockInfo::isSlowPath(const SILBasicBlock *FromBB,
     }
   }
 
+  // If the target BB is a cold function exit, then the path is cold.
+  if (hasColdTerminator(ToBB))
+    return true;
+
   // Should we give-up here?
   if (recursionDepth > RecursionDepthLimit)
     return false;
@@ -208,6 +295,10 @@ bool ColdBlockInfo::isCold(const SILBasicBlock *BB, int recursionDepth) {
     return true;
   }
 
+  // If this block always reaches a cold exit, then it's cold.
+  if (reachesColdBlock(BB))
+    return true;
+
   std::vector<const SILBasicBlock*> DomChain;
   DomChain.push_back(BB);
   bool IsCold = false;
@@ -220,7 +311,7 @@ bool ColdBlockInfo::isCold(const SILBasicBlock *BB, int recursionDepth) {
       break;
     }
 
-    // If a dominator of BB is cold, then BB is cold.
+    // If a dominator of BB has a temperature, propagate it.
     auto I = ColdBlockMap.find(fromBB);
     if (I != ColdBlockMap.end()) {
       IsCold = I->second;

lib/SILOptimizer/IPO/CapturePropagation.cpp

@@ -590,6 +590,7 @@ bool CapturePropagation::optimizePartialApply(PartialApplyInst *PAI) {
 
 void CapturePropagation::run() {
   DominanceAnalysis *DA = PM->getAnalysis<DominanceAnalysis>();
+  PostDominanceAnalysis *PDA = PM->getAnalysis<PostDominanceAnalysis>();
   auto *F = getFunction();
   bool HasChanged = false;
 
@@ -598,7 +599,7 @@ void CapturePropagation::run() {
     return;
 
   // Cache cold blocks per function.
-  ColdBlockInfo ColdBlocks(DA);
+  ColdBlockInfo ColdBlocks(DA, PDA);
   for (auto &BB : *F) {
     if (ColdBlocks.isCold(&BB))
       continue;

lib/SILOptimizer/LoopTransforms/COWArrayOpt.cpp

@@ -154,10 +154,11 @@ class COWArrayOpt {
   // analyzing.
   SILValue CurrentArrayAddr;
 public:
-  COWArrayOpt(RCIdentityFunctionInfo *RCIA, SILLoop *L, DominanceAnalysis *DA)
+  COWArrayOpt(RCIdentityFunctionInfo *RCIA, SILLoop *L, DominanceAnalysis *DA,
+              PostDominanceAnalysis *PDA)
       : RCIA(RCIA), Function(L->getHeader()->getParent()), Loop(L),
         Preheader(L->getLoopPreheader()), DomTree(DA->get(Function)),
-        ColdBlocks(DA), CachedSafeLoop(false, false), ReachingBlocks(Function) {}
+        ColdBlocks(DA, PDA), CachedSafeLoop(false, false), ReachingBlocks(Function) {}
 
   bool run();
 
@@ -1063,6 +1064,7 @@ class COWArrayOptPass : public SILFunctionTransform {
                             << getFunction()->getName() << "\n");
 
     auto *DA = PM->getAnalysis<DominanceAnalysis>();
+    auto *PDA = PM->getAnalysis<PostDominanceAnalysis>();
     auto *LA = PM->getAnalysis<SILLoopAnalysis>();
     auto *RCIA =
       PM->getAnalysis<RCIdentityAnalysis>()->get(getFunction());
@@ -1084,7 +1086,7 @@ class COWArrayOptPass : public SILFunctionTransform {
 
     bool HasChanged = false;
     for (auto *L : Loops)
-      HasChanged |= COWArrayOpt(RCIA, L, DA).run();
+      HasChanged |= COWArrayOpt(RCIA, L, DA, PDA).run();
 
     if (HasChanged)
       invalidateAnalysis(SILAnalysis::InvalidationKind::CallsAndInstructions);

lib/SILOptimizer/Transforms/PerformanceInliner.cpp

@@ -239,10 +239,11 @@ class SILPerformanceInliner {
   SILPerformanceInliner(StringRef PassName, SILOptFunctionBuilder &FuncBuilder,
                         InlineSelection WhatToInline,
                         SILPassManager *pm, DominanceAnalysis *DA,
+                        PostDominanceAnalysis *PDA,
                         SILLoopAnalysis *LA, BasicCalleeAnalysis *BCA,
                         OptimizationMode OptMode, OptRemark::Emitter &ORE)
       : PassName(PassName), FuncBuilder(FuncBuilder),
-        WhatToInline(WhatToInline), pm(pm), DA(DA), LA(LA), BCA(BCA), CBI(DA), ORE(ORE),
+        WhatToInline(WhatToInline), pm(pm), DA(DA), LA(LA), BCA(BCA), CBI(DA, PDA), ORE(ORE),
         OptMode(OptMode) {}
 
   bool inlineCallsIntoFunction(SILFunction *F);
@@ -525,9 +526,15 @@ bool SILPerformanceInliner::isProfitableToInline(
     return false;
   }
 
-  SILLoopInfo *LI = LA->get(Callee);
-  ShortestPathAnalysis *SPA = getSPA(Callee, LI);
-  assert(SPA->isValid());
+  SILLoopInfo *CalleeLI = LA->get(Callee);
+  ShortestPathAnalysis *CalleeSPA = getSPA(Callee, CalleeLI);
+  if (!CalleeSPA->isValid()) {
+    CalleeSPA->analyze(CBI, [](FullApplySite FAS) {
+      // We don't compute SPA for another call-level. Functions called from
+      // the callee are assumed to have DefaultApplyLength.
+      return DefaultApplyLength;
+    });
+  }
 
   ConstantTracker constTracker(Callee, &callerTracker, AI);
   DominanceInfo *DT = DA->get(Callee);
@@ -560,7 +567,7 @@ bool SILPerformanceInliner::isProfitableToInline(
   // benefits.
   while (SILBasicBlock *block = domOrder.getNext()) {
     constTracker.beginBlock();
-    Weight BlockW = SPA->getWeight(block, CallerWeight);
+    Weight BlockW = CalleeSPA->getWeight(block, CallerWeight);
 
     for (SILInstruction &I : *block) {
       constTracker.trackInst(&I);
@@ -773,7 +780,7 @@ bool SILPerformanceInliner::isProfitableToInline(
   LLVM_DEBUG(dumpCaller(AI.getFunction());
              llvm::dbgs() << "    decision {c=" << CalleeCost
                           << ", b=" << Benefit
-                          << ", l=" << SPA->getScopeLength(CalleeEntry, 0)
+                          << ", l=" << CalleeSPA->getScopeLength(CalleeEntry, 0)
                           << ", c-w=" << CallerWeight
                           << ", bb=" << Callee->size()
                           << ", c-bb=" << NumCallerBlocks
@@ -1331,6 +1338,7 @@ class SILPerformanceInlinerPass : public SILFunctionTransform {
 
   void run() override {
     DominanceAnalysis *DA = PM->getAnalysis<DominanceAnalysis>();
+    PostDominanceAnalysis *PDA = PM->getAnalysis<PostDominanceAnalysis>();
     SILLoopAnalysis *LA = PM->getAnalysis<SILLoopAnalysis>();
     BasicCalleeAnalysis *BCA = PM->getAnalysis<BasicCalleeAnalysis>();
     OptRemark::Emitter ORE(DEBUG_TYPE, *getFunction());
@@ -1344,7 +1352,7 @@ class SILPerformanceInlinerPass : public SILFunctionTransform {
     SILOptFunctionBuilder FuncBuilder(*this);
 
     SILPerformanceInliner Inliner(getID(), FuncBuilder, WhatToInline,
-                                  getPassManager(), DA, LA, BCA, OptMode, ORE);
+                              getPassManager(), DA, PDA, LA, BCA, OptMode, ORE);
 
     assert(getFunction()->isDefinition() &&
            "Expected only functions with bodies!");