ColdBlockInfo: overhaul analysis pass

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/Basic/ProfileCounter.h

@@ -27,8 +27,8 @@ class ProfileCounter {
   uint64_t count;
 
 public:
-  explicit ProfileCounter() : count(UINT64_MAX) {}
-  ProfileCounter(uint64_t Count) : count(Count) {
+  explicit constexpr ProfileCounter() : count(UINT64_MAX) {}
+  constexpr ProfileCounter(uint64_t Count) : count(Count) {
     if (Count == UINT64_MAX) {
       count = UINT64_MAX - 1;
     }
@@ -40,6 +40,22 @@ class ProfileCounter {
     return count;
   }
   explicit operator bool() const { return hasValue(); }
+
+  /// Saturating addition of another counter to this one, meaning that overflow
+  /// is avoided. If overflow would have happened, this function returns true
+  /// and the maximum representable value will be set in this counter.
+  bool add_saturating(ProfileCounter other) {
+    assert(hasValue() && other.hasValue());
+
+    // Will we go over the max representable value by adding other?
+    if (count > ((UINT64_MAX-1) - other.count)) {
+      count = UINT64_MAX - 1;
+      return true;
+    }
+
+    count += other.count;
+    return false;
+  }
 };
 } // end namespace swift
 

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

@@ -15,52 +15,74 @@
 
 #include "llvm/ADT/DenseMap.h"
 #include "swift/SIL/SILValue.h"
+#include "swift/Basic/FixedBitSet.h"
 
 namespace swift {
 class DominanceAnalysis;
+class PostDominanceAnalysis;
 class SILBasicBlock;
+class CondBranchInst;
 
 /// Cache a set of basic blocks that have been determined to be cold or hot.
 ///
 /// This does not inherit from SILAnalysis because it is not worth preserving
 /// across passes.
 class ColdBlockInfo {
+public:
+  // Represents the temperatures of edges flowing into a block.
+  //
+  //         T = "top" -- both warm and cold edges
+  //        /  \
+  //     Warm  Cold
+  //        \  /
+  //         B = "bottom" -- neither warm or cold edges
+  struct State {
+    // Each state kind, as an integer, is its position in any bit vectors.
+    enum Temperature {
+      Warm = 0,
+      Cold = 1
+    };
+
+    // Number of states, excluding Top or Bottom, in this flow problem.
+    static constexpr unsigned NumStates = 2;
+  };
+
+  using Energy = FixedBitSet<State::NumStates, State::Temperature>;
+
+private:
   DominanceAnalysis *DA;
+  PostDominanceAnalysis *PDA;
 
-  /// Each block in this map has been determined to be either cold or hot.
-  llvm::DenseMap<const SILBasicBlock*, bool> ColdBlockMap;
+  /// Each block in this map has been determined to be cold and/or warm.
+  /// Make sure to always use the set/resetToCold methods to update this!
+  llvm::DenseMap<const SILBasicBlock*, Energy> EnergyMap;
 
   // This is a cache and shouldn't be copied around.
   ColdBlockInfo(const ColdBlockInfo &) = delete;
   ColdBlockInfo &operator=(const ColdBlockInfo &) = delete;
+  LLVM_DUMP_METHOD void dump() const;
 
-  /// Tri-value return code for checking branch hints.
-  enum BranchHint : unsigned {
-    None,
-    LikelyTrue,
-    LikelyFalse
-  };
+  /// Tracks whether any information was changed in the energy map.
+  bool changedMap = false;
+  void resetToCold(const SILBasicBlock *BB);
+  void set(const SILBasicBlock *BB, State::Temperature temp);
 
-  enum {
-    RecursionDepthLimit = 3
-  };
+  using ExpectedValue = std::optional<bool>;
 
-  BranchHint getBranchHint(SILValue Cond, int recursionDepth);
+  void setExpectedCondition(CondBranchInst *CondBranch, ExpectedValue value);
 
-  bool isSlowPath(const SILBasicBlock *FromBB, const SILBasicBlock *ToBB,
-                  int recursionDepth);
+  ExpectedValue searchForExpectedValue(SILValue Cond,
+                                       unsigned recursionDepth = 0);
+  void searchForExpectedValue(CondBranchInst *CondBranch);
 
-  bool isCold(const SILBasicBlock *BB,
-              int recursionDepth);
+  bool inferFromEdgeProfile(SILBasicBlock *BB);
 
 public:
-  ColdBlockInfo(DominanceAnalysis *DA): DA(DA) {}
+  ColdBlockInfo(DominanceAnalysis *DA, PostDominanceAnalysis *PDA);
 
-  bool isSlowPath(const SILBasicBlock *FromBB, const SILBasicBlock *ToBB) {
-    return isSlowPath(FromBB, ToBB, 0);
-  }
+  void analyze(SILFunction *F);
 
-  bool isCold(const SILBasicBlock *BB) { return isCold(BB, 0); }
+  bool isCold(const SILBasicBlock *BB) const;
 };
 } // end namespace swift
 

include/swift/SILOptimizer/Utils/PerformanceInlinerUtils.h

@@ -429,6 +429,7 @@ class ShortestPathAnalysis {
       return;
 
     BlockInfoStorage.resize(numBlocks);
+    CBI.analyze(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);