[llvm] MachineBlockPlacement: Avoid overflow problems in scaleThreshold() (PR #68272)

[Matthias Braun via llvm-commits]

https://github.com/MatzeB created https://github.com/llvm/llvm-project/pull/68272

Multiplying block frequency values with integers is dangerous and can produce overflows. Change some block placement code to use the `BlockFrequency::mul` API which returns `std::nullopt` in case of overflow and return the maximum number which will fail the checks in the placement code.

>From aa2c539f8d71f7957f8ff1e5988311833574169a Mon Sep 17 00:00:00 2001
From: Matthias Braun <matze at braunis.de>
Date: Wed, 4 Oct 2023 13:34:32 -0700
Subject: [PATCH] MachineBlockPlacement: Avoid overflow problems in
 scaleThreshold()

---
 llvm/include/llvm/Support/BlockFrequency.h |  5 ++++-
 llvm/lib/CodeGen/MachineBlockPlacement.cpp | 21 ++++++++++++++-------
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/llvm/include/llvm/Support/BlockFrequency.h b/llvm/include/llvm/Support/BlockFrequency.h
index 12a753301b36aba..3523b13eed97643 100644
--- a/llvm/include/llvm/Support/BlockFrequency.h
+++ b/llvm/include/llvm/Support/BlockFrequency.h
@@ -28,9 +28,12 @@ class BlockFrequency {
 public:
   BlockFrequency(uint64_t Freq = 0) : Frequency(Freq) { }
 
-  /// Returns the maximum possible frequency, the saturation value.
+  /// Returns the maximum possible frequency, the saturation value integer.
   static uint64_t getMaxFrequency() { return UINT64_MAX; }
 
+  /// Returns the maximum possible frequency, the saturation value.
+  static BlockFrequency max() { return BlockFrequency(UINT64_MAX); }
+
   /// Returns the frequency as a fixpoint number scaled by the entry
   /// frequency.
   uint64_t getFrequency() const { return Frequency; }
diff --git a/llvm/lib/CodeGen/MachineBlockPlacement.cpp b/llvm/lib/CodeGen/MachineBlockPlacement.cpp
index 603c0e9600afc32..5a7305c81f60ae4 100644
--- a/llvm/lib/CodeGen/MachineBlockPlacement.cpp
+++ b/llvm/lib/CodeGen/MachineBlockPlacement.cpp
@@ -452,7 +452,7 @@ class MachineBlockPlacement : public MachineFunctionPass {
   }
 
   /// Scale the DupThreshold according to basic block size.
-  BlockFrequency scaleThreshold(MachineBasicBlock *BB);
+  BlockFrequency scaleThreshold(const MachineBasicBlock &BB) const;
   void initDupThreshold();
 
   /// Decrease the UnscheduledPredecessors count for all blocks in chain, and
@@ -3146,9 +3146,9 @@ bool MachineBlockPlacement::maybeTailDuplicateBlock(
 }
 
 // Count the number of actual machine instructions.
-static uint64_t countMBBInstruction(MachineBasicBlock *MBB) {
+static uint64_t countMBBInstructions(const MachineBasicBlock &MBB) {
   uint64_t InstrCount = 0;
-  for (MachineInstr &MI : *MBB) {
+  for (const MachineInstr &MI : MBB) {
     if (!MI.isPHI() && !MI.isMetaInstruction())
       InstrCount += 1;
   }
@@ -3158,8 +3158,15 @@ static uint64_t countMBBInstruction(MachineBasicBlock *MBB) {
 // The size cost of duplication is the instruction size of the duplicated block.
 // So we should scale the threshold accordingly. But the instruction size is not
 // available on all targets, so we use the number of instructions instead.
-BlockFrequency MachineBlockPlacement::scaleThreshold(MachineBasicBlock *BB) {
-  return DupThreshold.getFrequency() * countMBBInstruction(BB);
+BlockFrequency
+MachineBlockPlacement::scaleThreshold(const MachineBasicBlock &MBB) const {
+  std::optional<BlockFrequency> threshold =
+      DupThreshold.mul(countMBBInstructions(MBB));
+  // Return maximum value in case of overflow so the `Gains > Threshold`
+  // in callers do not succeed.
+  if (!threshold)
+    return BlockFrequency::max();
+  return *threshold;
 }
 
 // Returns true if BB is Pred's best successor.
@@ -3196,7 +3203,7 @@ bool MachineBlockPlacement::isBestSuccessor(MachineBasicBlock *BB,
   // instead of another successor. Then compare it with threshold.
   BlockFrequency PredFreq = getBlockCountOrFrequency(Pred);
   BlockFrequency Gain = PredFreq * (BBProb - BestProb);
-  return Gain > scaleThreshold(BB);
+  return Gain > scaleThreshold(*BB);
 }
 
 // Find out the predecessors of BB and BB can be beneficially duplicated into
@@ -3207,7 +3214,7 @@ void MachineBlockPlacement::findDuplicateCandidates(
     BlockFilterSet *BlockFilter) {
   MachineBasicBlock *Fallthrough = nullptr;
   BranchProbability DefaultBranchProb = BranchProbability::getZero();
-  BlockFrequency BBDupThreshold(scaleThreshold(BB));
+  BlockFrequency BBDupThreshold(scaleThreshold(*BB));
   SmallVector<MachineBasicBlock *, 8> Preds(BB->predecessors());
   SmallVector<MachineBasicBlock *, 8> Succs(BB->successors());