[llvm] Avoid BlockFrequency overflow problems (PR #66280)

[Matthias Braun via llvm-commits]

https://github.com/MatzeB updated https://github.com/llvm/llvm-project/pull/66280:

>From 0bd06f14228c5db4b7de62c90601ccfc3e1b2fb5 Mon Sep 17 00:00:00 2001
From: Matthias Braun <matze at braunis.de>
Date: Mon, 11 Sep 2023 19:35:23 -0700
Subject: [PATCH] Avoid BlockFrequency overflow problems

Multiplying raw block frequency with an integer carries a high risk
of overflow.

- Introduce a new `BlockFrequency::mul` function returning a `bool`
  indicating overflow.
- Mark function with `__attribute__((warn_unused_result))` to avoid
  users accidentally ignoring the indicator.
- Fix two instances where overflow were leading to wrong results for me.
---
 llvm/include/llvm/Support/BlockFrequency.h |  4 ++++
 llvm/lib/Analysis/InlineCost.cpp           | 11 ++++++-----
 llvm/lib/CodeGen/CodeGenPrepare.cpp        |  6 +++---
 llvm/lib/Support/BlockFrequency.cpp        |  9 +++++++++
 4 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/llvm/include/llvm/Support/BlockFrequency.h b/llvm/include/llvm/Support/BlockFrequency.h
index 6c624d7dad7d801..12a753301b36aba 100644
--- a/llvm/include/llvm/Support/BlockFrequency.h
+++ b/llvm/include/llvm/Support/BlockFrequency.h
@@ -15,6 +15,7 @@
 
 #include <cassert>
 #include <cstdint>
+#include <optional>
 
 namespace llvm {
 
@@ -76,6 +77,9 @@ class BlockFrequency {
     return NewFreq;
   }
 
+  /// Multiplies frequency with `Factor`. Returns `nullopt` in case of overflow.
+  std::optional<BlockFrequency> mul(uint64_t Factor) const;
+
   /// Shift block frequency to the right by count digits saturating to 1.
   BlockFrequency &operator>>=(const unsigned count) {
     // Frequency can never be 0 by design.
diff --git a/llvm/lib/Analysis/InlineCost.cpp b/llvm/lib/Analysis/InlineCost.cpp
index a9de1dde7c7f717..84dc412a3ab6e18 100644
--- a/llvm/lib/Analysis/InlineCost.cpp
+++ b/llvm/lib/Analysis/InlineCost.cpp
@@ -118,7 +118,7 @@ static cl::opt<int> ColdCallSiteRelFreq(
              "entry frequency, for a callsite to be cold in the absence of "
              "profile information."));
 
-static cl::opt<int> HotCallSiteRelFreq(
+static cl::opt<uint64_t> HotCallSiteRelFreq(
     "hot-callsite-rel-freq", cl::Hidden, cl::init(60),
     cl::desc("Minimum block frequency, expressed as a multiple of caller's "
              "entry frequency, for a callsite to be hot in the absence of "
@@ -1820,10 +1820,11 @@ InlineCostCallAnalyzer::getHotCallSiteThreshold(CallBase &Call,
   // potentially cache the computation of scaled entry frequency, but the added
   // complexity is not worth it unless this scaling shows up high in the
   // profiles.
-  auto CallSiteBB = Call.getParent();
-  auto CallSiteFreq = CallerBFI->getBlockFreq(CallSiteBB).getFrequency();
-  auto CallerEntryFreq = CallerBFI->getEntryFreq();
-  if (CallSiteFreq >= CallerEntryFreq * HotCallSiteRelFreq)
+  const BasicBlock *CallSiteBB = Call.getParent();
+  BlockFrequency CallSiteFreq = CallerBFI->getBlockFreq(CallSiteBB);
+  BlockFrequency CallerEntryFreq = CallerBFI->getEntryFreq();
+  std::optional<BlockFrequency> Limit = CallerEntryFreq.mul(HotCallSiteRelFreq);
+  if (Limit && CallSiteFreq >= *Limit)
     return Params.LocallyHotCallSiteThreshold;
 
   // Otherwise treat it normally.
diff --git a/llvm/lib/CodeGen/CodeGenPrepare.cpp b/llvm/lib/CodeGen/CodeGenPrepare.cpp
index f07fc4fc52bffba..faee623d7c62fba 100644
--- a/llvm/lib/CodeGen/CodeGenPrepare.cpp
+++ b/llvm/lib/CodeGen/CodeGenPrepare.cpp
@@ -198,7 +198,7 @@ static cl::opt<bool> BBSectionsGuidedSectionPrefix(
              "impacted, i.e., their prefixes will be decided by FDO/sampleFDO "
              "profiles."));
 
-static cl::opt<unsigned> FreqRatioToSkipMerge(
+static cl::opt<uint64_t> FreqRatioToSkipMerge(
     "cgp-freq-ratio-to-skip-merge", cl::Hidden, cl::init(2),
     cl::desc("Skip merging empty blocks if (frequency of empty block) / "
              "(frequency of destination block) is greater than this ratio"));
@@ -986,8 +986,8 @@ bool CodeGenPrepare::isMergingEmptyBlockProfitable(BasicBlock *BB,
         DestBB == findDestBlockOfMergeableEmptyBlock(SameValueBB))
       BBFreq += BFI->getBlockFreq(SameValueBB);
 
-  return PredFreq.getFrequency() <=
-         BBFreq.getFrequency() * FreqRatioToSkipMerge;
+  std::optional<BlockFrequency> Limit = BBFreq.mul(FreqRatioToSkipMerge);
+  return !Limit || PredFreq <= *Limit;
 }
 
 /// Return true if we can merge BB into DestBB if there is a single
diff --git a/llvm/lib/Support/BlockFrequency.cpp b/llvm/lib/Support/BlockFrequency.cpp
index a4a1e477d9403f7..329f1e12cdc29b2 100644
--- a/llvm/lib/Support/BlockFrequency.cpp
+++ b/llvm/lib/Support/BlockFrequency.cpp
@@ -12,6 +12,7 @@
 
 #include "llvm/Support/BlockFrequency.h"
 #include "llvm/Support/BranchProbability.h"
+#include "llvm/Support/MathExtras.h"
 
 using namespace llvm;
 
@@ -36,3 +37,11 @@ BlockFrequency BlockFrequency::operator/(BranchProbability Prob) const {
   Freq /= Prob;
   return Freq;
 }
+
+std::optional<BlockFrequency> BlockFrequency::mul(uint64_t Factor) const {
+  bool Overflow;
+  uint64_t ResultFrequency = SaturatingMultiply(Frequency, Factor, &Overflow);
+  if (Overflow)
+    return {};
+  return BlockFrequency(ResultFrequency);
+}