ColdBlockInfo: consider post-dominating coldness

A block that is post-dominated by a cold exit block is itself cold. The
most basic kind of cold-exit is one that terminates the program.

We can also consider a call to any `Never` returning function as a
cold-exit, though the block(s) leading up to that call may be executed
frequently because of concurrency. For now, I'm ignoring this case.

To help propagate this exit-block coldness, without rewriting the
algorithm to be something more sophisticated, I've simply added an
analysis stage that seeds the ColdBlockInfo with blocks that are
definitely cold due to these cold-exits.

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;
@@ -45,6 +47,9 @@ class ColdBlockInfo {
     RecursionDepthLimit = 3
   };
 
+  // \returns none if there is no known temperature for the block.
+  std::optional<bool> queryIsCold(const SILBasicBlock *BB);
+
   BranchHint getBranchHint(SILValue Cond, int recursionDepth);
 
   bool isSlowPath(const SILBasicBlock *FromBB, const SILBasicBlock *ToBB,
@@ -54,7 +59,11 @@ class ColdBlockInfo {
               int recursionDepth);
 
 public:
-  ColdBlockInfo(DominanceAnalysis *DA): DA(DA) {}
+  ColdBlockInfo(DominanceAnalysis *DA, PostDominanceAnalysis *PDA);
+
+  /// Pre-seed the info with information starting at the exits of the function.
+  /// Recommended prior to querying if blocks are cold in a function.
+  void analyzeExits(SILFunction *Fn);
 
   bool isSlowPath(const SILBasicBlock *FromBB, const SILBasicBlock *ToBB) {
     return isSlowPath(FromBB, ToBB, 0);

include/swift/SILOptimizer/Utils/PerformanceInlinerUtils.h

@@ -429,6 +429,7 @@ class ShortestPathAnalysis {
       return;
 
     BlockInfoStorage.resize(numBlocks);
+    CBI.analyzeExits(F);
 
     // First step: compute the length of the blocks.
     unsigned BlockIdx = 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,11 +14,25 @@
 #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"
 
 using namespace swift;
 
+ColdBlockInfo::ColdBlockInfo(DominanceAnalysis *DA,
+                             PostDominanceAnalysis *PDA) : DA(DA), PDA(PDA) {
+  LLVM_DEBUG(llvm::dbgs() << "ColdBlockInfo: constructed\n");
+}
+
+std::optional<bool> ColdBlockInfo::queryIsCold(const SILBasicBlock *BB) {
+  auto result = ColdBlockMap.find(BB);
+  if (result != ColdBlockMap.end()) {
+    return result->getSecond();
+  }
+  return std::nullopt;
+}
+
 /// Peek through an extract of Bool.value.
 static SILValue getCondition(SILValue C) {
   if (auto *SEI = dyn_cast<StructExtractInst>(C)) {
@@ -29,6 +43,42 @@ 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::ThrowInst:
+  case TermKind::ThrowAddrInst:
+  case TermKind::UnwindInst:
+    return term->getModule().getOptions().EnableThrowsPrediction;
+  case TermKind::UnreachableInst:
+    // For now, assume it's always cold, since it's executed once.
+    // Not all functions in the stdlib are properly annotated as a
+    // "known program termination point", so we don't use
+    // ApplySite::isCalleeKnownProgramTerminationPoint.
+    return term->getModule().getOptions().EnableNoReturnCold;
+  }
+}
+static bool hasColdTerminator(const SILBasicBlock *bb) {
+  assert(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 +185,36 @@ static std::optional<bool> isColdProfiledEdge(const SILBasicBlock *FromBB,
   return takenProbability < WARM_EDGE_MINIMUM;
 }
 
+void ColdBlockInfo::analyzeExits(SILFunction *Fn) {
+  assert(Fn);
+  LLVM_DEBUG(llvm::dbgs()
+    << "ColdBlockInfo: analyzeExits of " << Fn->getName() << "\n");
+
+  PostDominanceInfo *info = PDA->get(Fn);
+  auto *node = info->getRootNode();
+  assert(info->isVirtualRoot(node));
+
+  SmallVector<SILBasicBlock *, 8> scratch;
+  for (auto exitNode : node->children()) {
+    auto exitBB = exitNode->getBlock();
+
+    if (!hasColdTerminator(exitBB))
+      continue;
+
+    ASSERT(exitBB->getNumSuccessors() == 0 && "crucial assumption must hold");
+
+    // Mark all blocks post-dominated by this cold-exit block as cold.
+    info->getDescendants(exitBB, scratch);
+    for (auto dominatedBB : scratch) {
+      ColdBlockMap[dominatedBB] = true;
+      LLVM_DEBUG(llvm::dbgs()
+        << "ColdBlockInfo(" << Fn->getName()
+        << "): analyzeExits says bb"
+        << dominatedBB->getDebugID() << " is cold\n");
+    }
+  }
+}
+
 /// \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.
@@ -147,6 +227,8 @@ bool ColdBlockInfo::isSlowPath(const SILBasicBlock *FromBB,
   (void)DT;
 #endif
 
+  assert(!queryIsCold(ToBB) && "already known temperature?");
+
   // Is this the edge FromBB->ToBB?
   bool isDirectEdge = false;
   for (auto *succ : FromBB->getSuccessorBlocks()) {
@@ -173,20 +255,28 @@ bool ColdBlockInfo::isSlowPath(const SILBasicBlock *FromBB,
         return ToBB == CBI->getTrueBB();
       }
     }
-  }
 
-  // Should we give-up here?
-  if (recursionDepth > RecursionDepthLimit)
-    return false;
+  } else if (recursionDepth <= RecursionDepthLimit) {
+    // If ToBB only has cold predecessors, then the path to it is cold.
+    bool coldPredecessors =
+        llvm::all_of(ToBB->getPredecessorBlocks(), [&](auto *predBB) {
+          return isCold(predBB, recursionDepth+1);
+        });
+
+    if (coldPredecessors)
+      return true;
 
-  for (auto predBB : ToBB->getPredecessorBlocks()) {
-    if (!isCold(predBB, recursionDepth + 1))
-      return false;
+    // If FromBB only has cold successors, then the path from it is cold.
+    bool coldSuccessors =
+        llvm::all_of(FromBB->getSuccessorBlocks(), [&](auto *succBB) {
+          return isCold(succBB, recursionDepth+1);
+        });
+
+    if (coldSuccessors)
+      return true;
   }
 
-  // All predecessors of ToBB are cold, so that implies all paths
-  // FromBB->*->ToBB are "slow".
-  return true;
+  return false; // No reason to believe the path is slow.
 }
 
 /// \return true if the given block is dominated by a _slowPath branch hint.
@@ -212,6 +302,23 @@ bool ColdBlockInfo::isCold(const SILBasicBlock *BB, int recursionDepth) {
   DomChain.push_back(BB);
   bool IsCold = false;
   Node = Node->getIDom();
+
+  // Special case: the entry block
+  if (BB->isEntry()) {
+    assert(!Node && "entry is dominated?");
+    // If the entry block's successors are all cold, then it's cold.
+    if (BB->getNumSuccessors() > 0) {
+      IsCold = llvm::all_of(BB->getSuccessorBlocks(), [&](auto *succBB) {
+        return isCold(succBB, recursionDepth);
+      });
+    } else {
+      // Without any successors, we can't infer the entry block is cold, unless
+      // if it exits in an unusual way.
+      IsCold = hasColdTerminator(BB);
+    }
+  }
+
+  // Check the blocks dominating BB to see if they're cold.
   while (Node) {
     auto fromBB = Node->getBlock();
     auto toBB = DomChain.back();
@@ -220,10 +327,9 @@ bool ColdBlockInfo::isCold(const SILBasicBlock *BB, int recursionDepth) {
       break;
     }
 
-    // If a dominator of BB is cold, then BB is cold.
-    auto I = ColdBlockMap.find(fromBB);
-    if (I != ColdBlockMap.end()) {
-      IsCold = I->second;
+    // If a dominator of BB has a temperature, propagate it.
+    if (auto temp = queryIsCold(fromBB)) {
+      IsCold = *temp;
       break;
     }
 
@@ -234,12 +340,10 @@ bool ColdBlockInfo::isCold(const SILBasicBlock *BB, int recursionDepth) {
 
   for (auto *ChainBB : DomChain) {
     ColdBlockMap[ChainBB] = IsCold;
-
-    if (IsCold) {
-      LLVM_DEBUG(llvm::dbgs()
-        << "ColdBlockInfo(" << BB->getParent()->getName() << "): "
-        << "bb" << ChainBB->getDebugID() << " is cold\n");
-    }
+    LLVM_DEBUG(llvm::dbgs()
+      << "ColdBlockInfo(" << BB->getParent()->getName() << "): "
+      << "bb" << ChainBB->getDebugID()
+      << " is " << (IsCold ? "cold\n" : "warm\n"));
   }
 
   return IsCold;

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,8 @@ void CapturePropagation::run() {
     return;
 
   // Cache cold blocks per function.
-  ColdBlockInfo ColdBlocks(DA);
+  ColdBlockInfo ColdBlocks(DA, PDA);
+  ColdBlocks.analyzeExits(F);
   for (auto &BB : *F) {
     if (ColdBlocks.isCold(&BB))
       continue;

lib/SILOptimizer/LoopTransforms/COWArrayOpt.cpp

@@ -154,10 +154,13 @@ 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) {
+    ColdBlocks.analyzeExits(Function);
+  }
 
   bool run();
 
@@ -1063,6 +1066,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 +1088,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);