Skip to content

Conversation

@wdx727
Copy link
Contributor

@wdx727 wdx727 commented Nov 12, 2025

This PR re-submits the previously reverted PR(#165868) and fixes the return type mismatch error.

…nt matching and inference and create clusters.

Co-authored-by: lifengxiang1025 <lifengxiang@kuaishou.com>
Co-authored-by: zcfh <wuminghui03@kuaishou.com>
@llvmbot
Copy link
Member

llvmbot commented Nov 12, 2025

@llvm/pr-subscribers-pgo

@llvm/pr-subscribers-backend-x86

Author: None (wdx727)

Changes

This PR re-submits the previously reverted PR(#165868) and fixes the return type mismatch error.


Patch is 27.93 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/167622.diff

13 Files Affected:

  • (added) llvm/include/llvm/CodeGen/BasicBlockMatchingAndInference.h (+62)
  • (modified) llvm/include/llvm/CodeGen/BasicBlockSectionsProfileReader.h (+7)
  • (modified) llvm/include/llvm/CodeGen/MachineBlockHashInfo.h (+2)
  • (modified) llvm/include/llvm/CodeGen/Passes.h (+4)
  • (modified) llvm/include/llvm/InitializePasses.h (+1)
  • (modified) llvm/include/llvm/Transforms/Utils/SampleProfileInference.h (+16)
  • (added) llvm/lib/CodeGen/BasicBlockMatchingAndInference.cpp (+195)
  • (modified) llvm/lib/CodeGen/BasicBlockSections.cpp (+83-2)
  • (modified) llvm/lib/CodeGen/BasicBlockSectionsProfileReader.cpp (+15)
  • (modified) llvm/lib/CodeGen/CMakeLists.txt (+1)
  • (modified) llvm/lib/CodeGen/TargetPassConfig.cpp (+11-2)
  • (modified) llvm/lib/Transforms/Utils/SampleProfileInference.cpp (-2)
  • (added) llvm/test/CodeGen/X86/basic-block-sections-clusters-bb-hash.ll (+99)
diff --git a/llvm/include/llvm/CodeGen/BasicBlockMatchingAndInference.h b/llvm/include/llvm/CodeGen/BasicBlockMatchingAndInference.h
new file mode 100644
index 0000000000000..6e9bbb969a445
--- /dev/null
+++ b/llvm/include/llvm/CodeGen/BasicBlockMatchingAndInference.h
@@ -0,0 +1,62 @@
+//===- llvm/CodeGen/BasicBlockMatchingAndInference.h ------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Infer weights for all basic blocks using matching and inference.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CODEGEN_BASIC_BLOCK_AND_INFERENCE_H
+#define LLVM_CODEGEN_BASIC_BLOCK_AND_INFERENCE_H
+
+#include "llvm/CodeGen/BasicBlockSectionsProfileReader.h"
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/Transforms/Utils/SampleProfileInference.h"
+
+namespace llvm {
+
+class BasicBlockMatchingAndInference : public MachineFunctionPass {
+private:
+  using Edge = std::pair<const MachineBasicBlock *, const MachineBasicBlock *>;
+  using BlockWeightMap = DenseMap<const MachineBasicBlock *, uint64_t>;
+  using EdgeWeightMap = DenseMap<Edge, uint64_t>;
+  using BlockEdgeMap = DenseMap<const MachineBasicBlock *,
+                                SmallVector<const MachineBasicBlock *, 8>>;
+
+  struct WeightInfo {
+    // Weight of basic blocks.
+    BlockWeightMap BlockWeights;
+    // Weight of edges.
+    EdgeWeightMap EdgeWeights;
+  };
+
+public:
+  static char ID;
+  BasicBlockMatchingAndInference();
+
+  StringRef getPassName() const override {
+    return "Basic Block Matching and Inference";
+  }
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override;
+
+  bool runOnMachineFunction(MachineFunction &F) override;
+
+  std::optional<WeightInfo> getWeightInfo(StringRef FuncName) const;
+
+private:
+  StringMap<WeightInfo> ProgramWeightInfo;
+
+  WeightInfo initWeightInfoByMatching(MachineFunction &MF);
+
+  void generateWeightInfoByInference(MachineFunction &MF,
+                                     WeightInfo &MatchWeight);
+};
+
+} // end namespace llvm
+
+#endif // LLVM_CODEGEN_BASIC_BLOCK_AND_INFERENCE_H
diff --git a/llvm/include/llvm/CodeGen/BasicBlockSectionsProfileReader.h b/llvm/include/llvm/CodeGen/BasicBlockSectionsProfileReader.h
index ee1f28377f7e4..f0d28d863282e 100644
--- a/llvm/include/llvm/CodeGen/BasicBlockSectionsProfileReader.h
+++ b/llvm/include/llvm/CodeGen/BasicBlockSectionsProfileReader.h
@@ -86,6 +86,10 @@ class BasicBlockSectionsProfileReader {
   uint64_t getEdgeCount(StringRef FuncName, const UniqueBBID &SrcBBID,
                         const UniqueBBID &SinkBBID) const;
 
+  // Return the complete function path and cluster info for the given function.
+  std::pair<bool, FunctionPathAndClusterInfo>
+  getFunctionPathAndClusterInfo(StringRef FuncName) const;
+
 private:
   StringRef getAliasName(StringRef FuncName) const {
     auto R = FuncAliasMap.find(FuncName);
@@ -195,6 +199,9 @@ class BasicBlockSectionsProfileReaderWrapperPass : public ImmutablePass {
   uint64_t getEdgeCount(StringRef FuncName, const UniqueBBID &SrcBBID,
                         const UniqueBBID &DestBBID) const;
 
+  std::pair<bool, FunctionPathAndClusterInfo>
+  getFunctionPathAndClusterInfo(StringRef FuncName) const;
+
   // Initializes the FunctionNameToDIFilename map for the current module and
   // then reads the profile for the matching functions.
   bool doInitialization(Module &M) override;
diff --git a/llvm/include/llvm/CodeGen/MachineBlockHashInfo.h b/llvm/include/llvm/CodeGen/MachineBlockHashInfo.h
index d044d5f940b75..6f26819d566ae 100644
--- a/llvm/include/llvm/CodeGen/MachineBlockHashInfo.h
+++ b/llvm/include/llvm/CodeGen/MachineBlockHashInfo.h
@@ -80,6 +80,8 @@ struct BlendedBlockHash {
     return Dist;
   }
 
+  uint16_t getOpcodeHash() const { return OpcodeHash; }
+
 private:
   /// The offset of the basic block from the function start.
   uint16_t Offset{0};
diff --git a/llvm/include/llvm/CodeGen/Passes.h b/llvm/include/llvm/CodeGen/Passes.h
index a8525554b142e..2bf83cfa655b6 100644
--- a/llvm/include/llvm/CodeGen/Passes.h
+++ b/llvm/include/llvm/CodeGen/Passes.h
@@ -69,6 +69,10 @@ LLVM_ABI MachineFunctionPass *createBasicBlockSectionsPass();
 
 LLVM_ABI MachineFunctionPass *createBasicBlockPathCloningPass();
 
+/// createBasicBlockMatchingAndInferencePass - This pass enables matching
+/// and inference when using propeller.
+LLVM_ABI MachineFunctionPass *createBasicBlockMatchingAndInferencePass();
+
 /// createMachineBlockHashInfoPass - This pass computes basic block hashes.
 LLVM_ABI MachineFunctionPass *createMachineBlockHashInfoPass();
 
diff --git a/llvm/include/llvm/InitializePasses.h b/llvm/include/llvm/InitializePasses.h
index 10a4d8525a9e8..18732caf78966 100644
--- a/llvm/include/llvm/InitializePasses.h
+++ b/llvm/include/llvm/InitializePasses.h
@@ -55,6 +55,7 @@ LLVM_ABI void initializeAlwaysInlinerLegacyPassPass(PassRegistry &);
 LLVM_ABI void initializeAssignmentTrackingAnalysisPass(PassRegistry &);
 LLVM_ABI void initializeAssumptionCacheTrackerPass(PassRegistry &);
 LLVM_ABI void initializeAtomicExpandLegacyPass(PassRegistry &);
+LLVM_ABI void initializeBasicBlockMatchingAndInferencePass(PassRegistry &);
 LLVM_ABI void initializeBasicBlockPathCloningPass(PassRegistry &);
 LLVM_ABI void
 initializeBasicBlockSectionsProfileReaderWrapperPassPass(PassRegistry &);
diff --git a/llvm/include/llvm/Transforms/Utils/SampleProfileInference.h b/llvm/include/llvm/Transforms/Utils/SampleProfileInference.h
index 7231e45fe8eb7..e1663d29c1e3c 100644
--- a/llvm/include/llvm/Transforms/Utils/SampleProfileInference.h
+++ b/llvm/include/llvm/Transforms/Utils/SampleProfileInference.h
@@ -130,6 +130,11 @@ template <typename FT> class SampleProfileInference {
   SampleProfileInference(FunctionT &F, BlockEdgeMap &Successors,
                          BlockWeightMap &SampleBlockWeights)
       : F(F), Successors(Successors), SampleBlockWeights(SampleBlockWeights) {}
+  SampleProfileInference(FunctionT &F, BlockEdgeMap &Successors,
+                         BlockWeightMap &SampleBlockWeights,
+                         EdgeWeightMap &SampleEdgeWeights)
+      : F(F), Successors(Successors), SampleBlockWeights(SampleBlockWeights),
+        SampleEdgeWeights(SampleEdgeWeights) {}
 
   /// Apply the profile inference algorithm for a given function
   void apply(BlockWeightMap &BlockWeights, EdgeWeightMap &EdgeWeights);
@@ -157,6 +162,9 @@ template <typename FT> class SampleProfileInference {
 
   /// Map basic blocks to their sampled weights.
   BlockWeightMap &SampleBlockWeights;
+
+  /// Map edges to their sampled weights.
+  EdgeWeightMap SampleEdgeWeights;
 };
 
 template <typename BT>
@@ -266,6 +274,14 @@ FlowFunction SampleProfileInference<BT>::createFlowFunction(
       FlowJump Jump;
       Jump.Source = BlockIndex[BB];
       Jump.Target = BlockIndex[Succ];
+      auto It = SampleEdgeWeights.find(std::make_pair(BB, Succ));
+      if (It != SampleEdgeWeights.end()) {
+        Jump.HasUnknownWeight = false;
+        Jump.Weight = It->second;
+      } else {
+        Jump.HasUnknownWeight = true;
+        Jump.Weight = 0;
+      }
       Func.Jumps.push_back(Jump);
     }
   }
diff --git a/llvm/lib/CodeGen/BasicBlockMatchingAndInference.cpp b/llvm/lib/CodeGen/BasicBlockMatchingAndInference.cpp
new file mode 100644
index 0000000000000..4fa90799f4e10
--- /dev/null
+++ b/llvm/lib/CodeGen/BasicBlockMatchingAndInference.cpp
@@ -0,0 +1,195 @@
+//===- llvm/CodeGen/BasicBlockMatchingAndInference.cpp ----------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// In Propeller's profile, we have already read the hash values of basic blocks,
+// as well as the weights of basic blocks and edges in the CFG. In this file,
+// we first match the basic blocks in the profile with those in the current
+// MachineFunction using the basic block hash, thereby obtaining the weights of
+// some basic blocks and edges. Subsequently, we infer the weights of all basic
+// blocks using an inference algorithm.
+//
+// TODO: Integrate part of the code in this file with BOLT's implementation into
+// the LLVM infrastructure, enabling both BOLT and Propeller to reuse it.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CodeGen/BasicBlockMatchingAndInference.h"
+#include "llvm/CodeGen/BasicBlockSectionsProfileReader.h"
+#include "llvm/CodeGen/MachineBlockHashInfo.h"
+#include "llvm/CodeGen/Passes.h"
+#include "llvm/InitializePasses.h"
+#include <llvm/Support/CommandLine.h>
+
+using namespace llvm;
+
+static cl::opt<float>
+    PropellerInferThreshold("propeller-infer-threshold",
+                            cl::desc("Threshold for infer stale profile"),
+                            cl::init(0.6), cl::Optional);
+
+/// The object is used to identify and match basic blocks given their hashes.
+class StaleMatcher {
+public:
+  /// Initialize stale matcher.
+  void init(const std::vector<MachineBasicBlock *> &Blocks,
+            const std::vector<BlendedBlockHash> &Hashes) {
+    assert(Blocks.size() == Hashes.size() &&
+           "incorrect matcher initialization");
+    for (size_t I = 0; I < Blocks.size(); I++) {
+      MachineBasicBlock *Block = Blocks[I];
+      uint16_t OpHash = Hashes[I].getOpcodeHash();
+      OpHashToBlocks[OpHash].push_back(std::make_pair(Hashes[I], Block));
+    }
+  }
+
+  /// Find the most similar block for a given hash.
+  MachineBasicBlock *matchBlock(BlendedBlockHash BlendedHash) const {
+    auto BlockIt = OpHashToBlocks.find(BlendedHash.getOpcodeHash());
+    if (BlockIt == OpHashToBlocks.end()) {
+      return nullptr;
+    }
+    MachineBasicBlock *BestBlock = nullptr;
+    uint64_t BestDist = std::numeric_limits<uint64_t>::max();
+    for (auto It : BlockIt->second) {
+      MachineBasicBlock *Block = It.second;
+      BlendedBlockHash Hash = It.first;
+      uint64_t Dist = Hash.distance(BlendedHash);
+      if (BestBlock == nullptr || Dist < BestDist) {
+        BestDist = Dist;
+        BestBlock = Block;
+      }
+    }
+    return BestBlock;
+  }
+
+private:
+  using HashBlockPairType = std::pair<BlendedBlockHash, MachineBasicBlock *>;
+  std::unordered_map<uint16_t, std::vector<HashBlockPairType>> OpHashToBlocks;
+};
+
+INITIALIZE_PASS_BEGIN(BasicBlockMatchingAndInference,
+                      "machine-block-match-infer",
+                      "Machine Block Matching and Inference Analysis", true,
+                      true)
+INITIALIZE_PASS_DEPENDENCY(MachineBlockHashInfo)
+INITIALIZE_PASS_DEPENDENCY(BasicBlockSectionsProfileReaderWrapperPass)
+INITIALIZE_PASS_END(BasicBlockMatchingAndInference, "machine-block-match-infer",
+                    "Machine Block Matching and Inference Analysis", true, true)
+
+char BasicBlockMatchingAndInference::ID = 0;
+
+BasicBlockMatchingAndInference::BasicBlockMatchingAndInference()
+    : MachineFunctionPass(ID) {
+  initializeBasicBlockMatchingAndInferencePass(
+      *PassRegistry::getPassRegistry());
+}
+
+void BasicBlockMatchingAndInference::getAnalysisUsage(AnalysisUsage &AU) const {
+  AU.addRequired<MachineBlockHashInfo>();
+  AU.addRequired<BasicBlockSectionsProfileReaderWrapperPass>();
+  AU.setPreservesAll();
+  MachineFunctionPass::getAnalysisUsage(AU);
+}
+
+std::optional<BasicBlockMatchingAndInference::WeightInfo>
+BasicBlockMatchingAndInference::getWeightInfo(StringRef FuncName) const {
+  auto It = ProgramWeightInfo.find(FuncName);
+  if (It == ProgramWeightInfo.end()) {
+    return std::nullopt;
+  }
+  return It->second;
+}
+
+BasicBlockMatchingAndInference::WeightInfo
+BasicBlockMatchingAndInference::initWeightInfoByMatching(MachineFunction &MF) {
+  std::vector<MachineBasicBlock *> Blocks;
+  std::vector<BlendedBlockHash> Hashes;
+  auto BSPR = &getAnalysis<BasicBlockSectionsProfileReaderWrapperPass>();
+  auto MBHI = &getAnalysis<MachineBlockHashInfo>();
+  for (auto &Block : MF) {
+    Blocks.push_back(&Block);
+    Hashes.push_back(BlendedBlockHash(MBHI->getMBBHash(Block)));
+  }
+  StaleMatcher Matcher;
+  Matcher.init(Blocks, Hashes);
+  BasicBlockMatchingAndInference::WeightInfo MatchWeight;
+  auto [IsValid, PathAndClusterInfo] =
+      BSPR->getFunctionPathAndClusterInfo(MF.getName());
+  if (!IsValid)
+    return MatchWeight;
+  for (auto &BlockCount : PathAndClusterInfo.NodeCounts) {
+    if (PathAndClusterInfo.BBHashes.count(BlockCount.first.BaseID)) {
+      auto Hash = PathAndClusterInfo.BBHashes[BlockCount.first.BaseID];
+      MachineBasicBlock *Block = Matcher.matchBlock(BlendedBlockHash(Hash));
+      // When a basic block has clone copies, sum their counts.
+      if (Block != nullptr)
+        MatchWeight.BlockWeights[Block] += BlockCount.second;
+    }
+  }
+  for (auto &PredItem : PathAndClusterInfo.EdgeCounts) {
+    auto PredID = PredItem.first.BaseID;
+    if (!PathAndClusterInfo.BBHashes.count(PredID))
+      continue;
+    auto PredHash = PathAndClusterInfo.BBHashes[PredID];
+    MachineBasicBlock *PredBlock =
+        Matcher.matchBlock(BlendedBlockHash(PredHash));
+    if (PredBlock == nullptr)
+      continue;
+    for (auto &SuccItem : PredItem.second) {
+      auto SuccID = SuccItem.first.BaseID;
+      auto EdgeWeight = SuccItem.second;
+      if (PathAndClusterInfo.BBHashes.count(SuccID)) {
+        auto SuccHash = PathAndClusterInfo.BBHashes[SuccID];
+        MachineBasicBlock *SuccBlock =
+            Matcher.matchBlock(BlendedBlockHash(SuccHash));
+        // When an edge has clone copies, sum their counts.
+        if (SuccBlock != nullptr)
+          MatchWeight.EdgeWeights[std::make_pair(PredBlock, SuccBlock)] +=
+              EdgeWeight;
+      }
+    }
+  }
+  return MatchWeight;
+}
+
+void BasicBlockMatchingAndInference::generateWeightInfoByInference(
+    MachineFunction &MF,
+    BasicBlockMatchingAndInference::WeightInfo &MatchWeight) {
+  BlockEdgeMap Successors;
+  for (auto &Block : MF) {
+    for (auto *Succ : Block.successors())
+      Successors[&Block].push_back(Succ);
+  }
+  SampleProfileInference<MachineFunction> SPI(
+      MF, Successors, MatchWeight.BlockWeights, MatchWeight.EdgeWeights);
+  BlockWeightMap BlockWeights;
+  EdgeWeightMap EdgeWeights;
+  SPI.apply(BlockWeights, EdgeWeights);
+  ProgramWeightInfo.try_emplace(
+      MF.getName(), BasicBlockMatchingAndInference::WeightInfo{
+                        std::move(BlockWeights), std::move(EdgeWeights)});
+}
+
+bool BasicBlockMatchingAndInference::runOnMachineFunction(MachineFunction &MF) {
+  if (MF.empty())
+    return false;
+  auto MatchWeight = initWeightInfoByMatching(MF);
+  // If the ratio of the number of MBBs in matching to the total number of MBBs
+  // in the function is less than the threshold value, the processing should be
+  // abandoned.
+  if (static_cast<float>(MatchWeight.BlockWeights.size()) / MF.size() <
+      PropellerInferThreshold) {
+    return false;
+  }
+  generateWeightInfoByInference(MF, MatchWeight);
+  return false;
+}
+
+MachineFunctionPass *llvm::createBasicBlockMatchingAndInferencePass() {
+  return new BasicBlockMatchingAndInference();
+}
diff --git a/llvm/lib/CodeGen/BasicBlockSections.cpp b/llvm/lib/CodeGen/BasicBlockSections.cpp
index 52e2909bec072..11098d0e1df61 100644
--- a/llvm/lib/CodeGen/BasicBlockSections.cpp
+++ b/llvm/lib/CodeGen/BasicBlockSections.cpp
@@ -70,6 +70,7 @@
 
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/CodeGen/BasicBlockMatchingAndInference.h"
 #include "llvm/CodeGen/BasicBlockSectionUtils.h"
 #include "llvm/CodeGen/BasicBlockSectionsProfileReader.h"
 #include "llvm/CodeGen/MachineDominators.h"
@@ -81,6 +82,7 @@
 #include "llvm/InitializePasses.h"
 #include "llvm/Support/UniqueBBID.h"
 #include "llvm/Target/TargetMachine.h"
+#include "llvm/Transforms/Utils/CodeLayout.h"
 #include <optional>
 
 using namespace llvm;
@@ -175,6 +177,79 @@ updateBranches(MachineFunction &MF,
   }
 }
 
+// This function generates the machine basic block clusters of "hot" blocks.
+// Currently, only support one cluster creation.
+// TODO: Support multi-cluster creation and path cloning.
+static SmallVector<BBClusterInfo>
+createBBClusterInfoForFunction(MachineFunction &MF,
+                               BasicBlockMatchingAndInference *BMI) {
+  unsigned CurrentCluster = 0;
+  SmallVector<BBClusterInfo> BBClusterInfos;
+  auto OptWeightInfo = BMI->getWeightInfo(MF.getName());
+  if (!OptWeightInfo)
+    return BBClusterInfos;
+  auto BlockWeights = OptWeightInfo->BlockWeights;
+  auto EdgeWeights = OptWeightInfo->EdgeWeights;
+
+  SmallVector<const MachineBasicBlock *, 4> HotMBBs;
+  if (MF.size() <= 2) {
+    for (auto &MBB : MF) {
+      if (MBB.isEntryBlock() || BlockWeights[&MBB] > 0) {
+        HotMBBs.push_back(&MBB);
+      }
+    }
+  } else {
+    SmallVector<uint64_t, 0> BlockSizes(MF.size());
+    SmallVector<uint64_t, 0> BlockCounts(MF.size());
+    std::vector<const MachineBasicBlock *> OrigOrder;
+    OrigOrder.reserve(MF.size());
+    SmallVector<codelayout::EdgeCount, 0> JumpCounts;
+
+    // Renumber blocks for running the layout algorithm.
+    MF.RenumberBlocks();
+
+    // Init the MBB size and count.
+    for (auto &MBB : MF) {
+      auto NonDbgInsts =
+          instructionsWithoutDebug(MBB.instr_begin(), MBB.instr_end());
+      int NumInsts = std::distance(NonDbgInsts.begin(), NonDbgInsts.end());
+      BlockSizes[MBB.getNumber()] = 4 * NumInsts;
+      BlockCounts[MBB.getNumber()] = BlockWeights[&MBB];
+      OrigOrder.push_back(&MBB);
+    }
+
+    // Init the edge count.
+    for (auto &MBB : MF) {
+      for (auto *Succ : MBB.successors()) {
+        auto EdgeWeight = EdgeWeights[std::make_pair(&MBB, Succ)];
+        JumpCounts.push_back({static_cast<uint64_t>(MBB.getNumber()),
+                              static_cast<uint64_t>(Succ->getNumber()),
+                              EdgeWeight});
+      }
+    }
+
+    // Run the layout algorithm.
+    auto Result = computeExtTspLayout(BlockSizes, BlockCounts, JumpCounts);
+    for (uint64_t R : Result) {
+      auto Block = OrigOrder[R];
+      if (Block->isEntryBlock() || BlockWeights[Block] > 0)
+        HotMBBs.push_back(Block);
+    }
+  }
+
+  // Generate the "hot" basic block cluster.
+  if (!HotMBBs.empty()) {
+    unsigned CurrentPosition = 0;
+    for (auto &MBB : HotMBBs) {
+      if (MBB->getBBID()) {
+        BBClusterInfos.push_back(
+            {*(MBB->getBBID()), CurrentCluster, CurrentPosition++});
+      }
+    }
+  }
+  return BBClusterInfos;
+}
+
 // This function sorts basic blocks according to the cluster's information.
 // All explicitly specified clusters of basic blocks will be ordered
 // accordingly. All non-specified BBs go into a separate "Cold" section.
@@ -308,8 +383,13 @@ bool BasicBlockSections::handleBBSections(MachineFunction &MF) {
 
   DenseMap<UniqueBBID, BBClusterInfo> FuncClusterInfo;
   if (BBSectionsType == BasicBlockSection::List) {
-    auto ClusterInfo = getAnalysis<BasicBlockSectionsProfileReaderWrapperPass>()
-                           .getClusterInfoForFunction(MF.getName());
+    SmallVector<BBClusterInfo> ClusterInfo;
+    if (auto *BMI = getAnalysisIfAvailable<BasicBlockMatchingAndInference>()) {
+      ClusterInfo = createBBClusterInfoForFunction(MF, BMI);
+    } else {
+      ClusterInfo = getAnalysis<BasicBlockSectionsProfileReaderWrapperPass>()
+                        .getClusterInfoForFunction(MF.getName());
+    }
     if (ClusterInfo.empty())
       return false;
     for (auto &BBClusterInfo : ClusterInfo) {
@@ -399,6 +479,7 @@ bool BasicBlockSections::runOnMachineFunction(MachineFunction &MF) {
 void BasicBlockSections::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.setPreservesAll();
   AU.addRequired<BasicBlockSectionsProfileReaderWrapperPass>();
+  AU.addUsedIfAvailable<BasicBlockMatchingAndInference>();
   AU.addUsedIfAvailable<MachineDominatorTreeWrapperPass>();
   AU.addUsedIfAvailable<MachinePostDominatorTreeWrapperPass>();
  ...
[truncated]

@llvmbot
Copy link
Member

llvmbot commented Nov 12, 2025

@llvm/pr-subscribers-llvm-transforms

Author: None (wdx727)

Changes

This PR re-submits the previously reverted PR(#165868) and fixes the return type mismatch error.


Patch is 27.93 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/167622.diff

13 Files Affected:

  • (added) llvm/include/llvm/CodeGen/BasicBlockMatchingAndInference.h (+62)
  • (modified) llvm/include/llvm/CodeGen/BasicBlockSectionsProfileReader.h (+7)
  • (modified) llvm/include/llvm/CodeGen/MachineBlockHashInfo.h (+2)
  • (modified) llvm/include/llvm/CodeGen/Passes.h (+4)
  • (modified) llvm/include/llvm/InitializePasses.h (+1)
  • (modified) llvm/include/llvm/Transforms/Utils/SampleProfileInference.h (+16)
  • (added) llvm/lib/CodeGen/BasicBlockMatchingAndInference.cpp (+195)
  • (modified) llvm/lib/CodeGen/BasicBlockSections.cpp (+83-2)
  • (modified) llvm/lib/CodeGen/BasicBlockSectionsProfileReader.cpp (+15)
  • (modified) llvm/lib/CodeGen/CMakeLists.txt (+1)
  • (modified) llvm/lib/CodeGen/TargetPassConfig.cpp (+11-2)
  • (modified) llvm/lib/Transforms/Utils/SampleProfileInference.cpp (-2)
  • (added) llvm/test/CodeGen/X86/basic-block-sections-clusters-bb-hash.ll (+99)
diff --git a/llvm/include/llvm/CodeGen/BasicBlockMatchingAndInference.h b/llvm/include/llvm/CodeGen/BasicBlockMatchingAndInference.h
new file mode 100644
index 0000000000000..6e9bbb969a445
--- /dev/null
+++ b/llvm/include/llvm/CodeGen/BasicBlockMatchingAndInference.h
@@ -0,0 +1,62 @@
+//===- llvm/CodeGen/BasicBlockMatchingAndInference.h ------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Infer weights for all basic blocks using matching and inference.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CODEGEN_BASIC_BLOCK_AND_INFERENCE_H
+#define LLVM_CODEGEN_BASIC_BLOCK_AND_INFERENCE_H
+
+#include "llvm/CodeGen/BasicBlockSectionsProfileReader.h"
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/Transforms/Utils/SampleProfileInference.h"
+
+namespace llvm {
+
+class BasicBlockMatchingAndInference : public MachineFunctionPass {
+private:
+  using Edge = std::pair<const MachineBasicBlock *, const MachineBasicBlock *>;
+  using BlockWeightMap = DenseMap<const MachineBasicBlock *, uint64_t>;
+  using EdgeWeightMap = DenseMap<Edge, uint64_t>;
+  using BlockEdgeMap = DenseMap<const MachineBasicBlock *,
+                                SmallVector<const MachineBasicBlock *, 8>>;
+
+  struct WeightInfo {
+    // Weight of basic blocks.
+    BlockWeightMap BlockWeights;
+    // Weight of edges.
+    EdgeWeightMap EdgeWeights;
+  };
+
+public:
+  static char ID;
+  BasicBlockMatchingAndInference();
+
+  StringRef getPassName() const override {
+    return "Basic Block Matching and Inference";
+  }
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override;
+
+  bool runOnMachineFunction(MachineFunction &F) override;
+
+  std::optional<WeightInfo> getWeightInfo(StringRef FuncName) const;
+
+private:
+  StringMap<WeightInfo> ProgramWeightInfo;
+
+  WeightInfo initWeightInfoByMatching(MachineFunction &MF);
+
+  void generateWeightInfoByInference(MachineFunction &MF,
+                                     WeightInfo &MatchWeight);
+};
+
+} // end namespace llvm
+
+#endif // LLVM_CODEGEN_BASIC_BLOCK_AND_INFERENCE_H
diff --git a/llvm/include/llvm/CodeGen/BasicBlockSectionsProfileReader.h b/llvm/include/llvm/CodeGen/BasicBlockSectionsProfileReader.h
index ee1f28377f7e4..f0d28d863282e 100644
--- a/llvm/include/llvm/CodeGen/BasicBlockSectionsProfileReader.h
+++ b/llvm/include/llvm/CodeGen/BasicBlockSectionsProfileReader.h
@@ -86,6 +86,10 @@ class BasicBlockSectionsProfileReader {
   uint64_t getEdgeCount(StringRef FuncName, const UniqueBBID &SrcBBID,
                         const UniqueBBID &SinkBBID) const;
 
+  // Return the complete function path and cluster info for the given function.
+  std::pair<bool, FunctionPathAndClusterInfo>
+  getFunctionPathAndClusterInfo(StringRef FuncName) const;
+
 private:
   StringRef getAliasName(StringRef FuncName) const {
     auto R = FuncAliasMap.find(FuncName);
@@ -195,6 +199,9 @@ class BasicBlockSectionsProfileReaderWrapperPass : public ImmutablePass {
   uint64_t getEdgeCount(StringRef FuncName, const UniqueBBID &SrcBBID,
                         const UniqueBBID &DestBBID) const;
 
+  std::pair<bool, FunctionPathAndClusterInfo>
+  getFunctionPathAndClusterInfo(StringRef FuncName) const;
+
   // Initializes the FunctionNameToDIFilename map for the current module and
   // then reads the profile for the matching functions.
   bool doInitialization(Module &M) override;
diff --git a/llvm/include/llvm/CodeGen/MachineBlockHashInfo.h b/llvm/include/llvm/CodeGen/MachineBlockHashInfo.h
index d044d5f940b75..6f26819d566ae 100644
--- a/llvm/include/llvm/CodeGen/MachineBlockHashInfo.h
+++ b/llvm/include/llvm/CodeGen/MachineBlockHashInfo.h
@@ -80,6 +80,8 @@ struct BlendedBlockHash {
     return Dist;
   }
 
+  uint16_t getOpcodeHash() const { return OpcodeHash; }
+
 private:
   /// The offset of the basic block from the function start.
   uint16_t Offset{0};
diff --git a/llvm/include/llvm/CodeGen/Passes.h b/llvm/include/llvm/CodeGen/Passes.h
index a8525554b142e..2bf83cfa655b6 100644
--- a/llvm/include/llvm/CodeGen/Passes.h
+++ b/llvm/include/llvm/CodeGen/Passes.h
@@ -69,6 +69,10 @@ LLVM_ABI MachineFunctionPass *createBasicBlockSectionsPass();
 
 LLVM_ABI MachineFunctionPass *createBasicBlockPathCloningPass();
 
+/// createBasicBlockMatchingAndInferencePass - This pass enables matching
+/// and inference when using propeller.
+LLVM_ABI MachineFunctionPass *createBasicBlockMatchingAndInferencePass();
+
 /// createMachineBlockHashInfoPass - This pass computes basic block hashes.
 LLVM_ABI MachineFunctionPass *createMachineBlockHashInfoPass();
 
diff --git a/llvm/include/llvm/InitializePasses.h b/llvm/include/llvm/InitializePasses.h
index 10a4d8525a9e8..18732caf78966 100644
--- a/llvm/include/llvm/InitializePasses.h
+++ b/llvm/include/llvm/InitializePasses.h
@@ -55,6 +55,7 @@ LLVM_ABI void initializeAlwaysInlinerLegacyPassPass(PassRegistry &);
 LLVM_ABI void initializeAssignmentTrackingAnalysisPass(PassRegistry &);
 LLVM_ABI void initializeAssumptionCacheTrackerPass(PassRegistry &);
 LLVM_ABI void initializeAtomicExpandLegacyPass(PassRegistry &);
+LLVM_ABI void initializeBasicBlockMatchingAndInferencePass(PassRegistry &);
 LLVM_ABI void initializeBasicBlockPathCloningPass(PassRegistry &);
 LLVM_ABI void
 initializeBasicBlockSectionsProfileReaderWrapperPassPass(PassRegistry &);
diff --git a/llvm/include/llvm/Transforms/Utils/SampleProfileInference.h b/llvm/include/llvm/Transforms/Utils/SampleProfileInference.h
index 7231e45fe8eb7..e1663d29c1e3c 100644
--- a/llvm/include/llvm/Transforms/Utils/SampleProfileInference.h
+++ b/llvm/include/llvm/Transforms/Utils/SampleProfileInference.h
@@ -130,6 +130,11 @@ template <typename FT> class SampleProfileInference {
   SampleProfileInference(FunctionT &F, BlockEdgeMap &Successors,
                          BlockWeightMap &SampleBlockWeights)
       : F(F), Successors(Successors), SampleBlockWeights(SampleBlockWeights) {}
+  SampleProfileInference(FunctionT &F, BlockEdgeMap &Successors,
+                         BlockWeightMap &SampleBlockWeights,
+                         EdgeWeightMap &SampleEdgeWeights)
+      : F(F), Successors(Successors), SampleBlockWeights(SampleBlockWeights),
+        SampleEdgeWeights(SampleEdgeWeights) {}
 
   /// Apply the profile inference algorithm for a given function
   void apply(BlockWeightMap &BlockWeights, EdgeWeightMap &EdgeWeights);
@@ -157,6 +162,9 @@ template <typename FT> class SampleProfileInference {
 
   /// Map basic blocks to their sampled weights.
   BlockWeightMap &SampleBlockWeights;
+
+  /// Map edges to their sampled weights.
+  EdgeWeightMap SampleEdgeWeights;
 };
 
 template <typename BT>
@@ -266,6 +274,14 @@ FlowFunction SampleProfileInference<BT>::createFlowFunction(
       FlowJump Jump;
       Jump.Source = BlockIndex[BB];
       Jump.Target = BlockIndex[Succ];
+      auto It = SampleEdgeWeights.find(std::make_pair(BB, Succ));
+      if (It != SampleEdgeWeights.end()) {
+        Jump.HasUnknownWeight = false;
+        Jump.Weight = It->second;
+      } else {
+        Jump.HasUnknownWeight = true;
+        Jump.Weight = 0;
+      }
       Func.Jumps.push_back(Jump);
     }
   }
diff --git a/llvm/lib/CodeGen/BasicBlockMatchingAndInference.cpp b/llvm/lib/CodeGen/BasicBlockMatchingAndInference.cpp
new file mode 100644
index 0000000000000..4fa90799f4e10
--- /dev/null
+++ b/llvm/lib/CodeGen/BasicBlockMatchingAndInference.cpp
@@ -0,0 +1,195 @@
+//===- llvm/CodeGen/BasicBlockMatchingAndInference.cpp ----------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// In Propeller's profile, we have already read the hash values of basic blocks,
+// as well as the weights of basic blocks and edges in the CFG. In this file,
+// we first match the basic blocks in the profile with those in the current
+// MachineFunction using the basic block hash, thereby obtaining the weights of
+// some basic blocks and edges. Subsequently, we infer the weights of all basic
+// blocks using an inference algorithm.
+//
+// TODO: Integrate part of the code in this file with BOLT's implementation into
+// the LLVM infrastructure, enabling both BOLT and Propeller to reuse it.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CodeGen/BasicBlockMatchingAndInference.h"
+#include "llvm/CodeGen/BasicBlockSectionsProfileReader.h"
+#include "llvm/CodeGen/MachineBlockHashInfo.h"
+#include "llvm/CodeGen/Passes.h"
+#include "llvm/InitializePasses.h"
+#include <llvm/Support/CommandLine.h>
+
+using namespace llvm;
+
+static cl::opt<float>
+    PropellerInferThreshold("propeller-infer-threshold",
+                            cl::desc("Threshold for infer stale profile"),
+                            cl::init(0.6), cl::Optional);
+
+/// The object is used to identify and match basic blocks given their hashes.
+class StaleMatcher {
+public:
+  /// Initialize stale matcher.
+  void init(const std::vector<MachineBasicBlock *> &Blocks,
+            const std::vector<BlendedBlockHash> &Hashes) {
+    assert(Blocks.size() == Hashes.size() &&
+           "incorrect matcher initialization");
+    for (size_t I = 0; I < Blocks.size(); I++) {
+      MachineBasicBlock *Block = Blocks[I];
+      uint16_t OpHash = Hashes[I].getOpcodeHash();
+      OpHashToBlocks[OpHash].push_back(std::make_pair(Hashes[I], Block));
+    }
+  }
+
+  /// Find the most similar block for a given hash.
+  MachineBasicBlock *matchBlock(BlendedBlockHash BlendedHash) const {
+    auto BlockIt = OpHashToBlocks.find(BlendedHash.getOpcodeHash());
+    if (BlockIt == OpHashToBlocks.end()) {
+      return nullptr;
+    }
+    MachineBasicBlock *BestBlock = nullptr;
+    uint64_t BestDist = std::numeric_limits<uint64_t>::max();
+    for (auto It : BlockIt->second) {
+      MachineBasicBlock *Block = It.second;
+      BlendedBlockHash Hash = It.first;
+      uint64_t Dist = Hash.distance(BlendedHash);
+      if (BestBlock == nullptr || Dist < BestDist) {
+        BestDist = Dist;
+        BestBlock = Block;
+      }
+    }
+    return BestBlock;
+  }
+
+private:
+  using HashBlockPairType = std::pair<BlendedBlockHash, MachineBasicBlock *>;
+  std::unordered_map<uint16_t, std::vector<HashBlockPairType>> OpHashToBlocks;
+};
+
+INITIALIZE_PASS_BEGIN(BasicBlockMatchingAndInference,
+                      "machine-block-match-infer",
+                      "Machine Block Matching and Inference Analysis", true,
+                      true)
+INITIALIZE_PASS_DEPENDENCY(MachineBlockHashInfo)
+INITIALIZE_PASS_DEPENDENCY(BasicBlockSectionsProfileReaderWrapperPass)
+INITIALIZE_PASS_END(BasicBlockMatchingAndInference, "machine-block-match-infer",
+                    "Machine Block Matching and Inference Analysis", true, true)
+
+char BasicBlockMatchingAndInference::ID = 0;
+
+BasicBlockMatchingAndInference::BasicBlockMatchingAndInference()
+    : MachineFunctionPass(ID) {
+  initializeBasicBlockMatchingAndInferencePass(
+      *PassRegistry::getPassRegistry());
+}
+
+void BasicBlockMatchingAndInference::getAnalysisUsage(AnalysisUsage &AU) const {
+  AU.addRequired<MachineBlockHashInfo>();
+  AU.addRequired<BasicBlockSectionsProfileReaderWrapperPass>();
+  AU.setPreservesAll();
+  MachineFunctionPass::getAnalysisUsage(AU);
+}
+
+std::optional<BasicBlockMatchingAndInference::WeightInfo>
+BasicBlockMatchingAndInference::getWeightInfo(StringRef FuncName) const {
+  auto It = ProgramWeightInfo.find(FuncName);
+  if (It == ProgramWeightInfo.end()) {
+    return std::nullopt;
+  }
+  return It->second;
+}
+
+BasicBlockMatchingAndInference::WeightInfo
+BasicBlockMatchingAndInference::initWeightInfoByMatching(MachineFunction &MF) {
+  std::vector<MachineBasicBlock *> Blocks;
+  std::vector<BlendedBlockHash> Hashes;
+  auto BSPR = &getAnalysis<BasicBlockSectionsProfileReaderWrapperPass>();
+  auto MBHI = &getAnalysis<MachineBlockHashInfo>();
+  for (auto &Block : MF) {
+    Blocks.push_back(&Block);
+    Hashes.push_back(BlendedBlockHash(MBHI->getMBBHash(Block)));
+  }
+  StaleMatcher Matcher;
+  Matcher.init(Blocks, Hashes);
+  BasicBlockMatchingAndInference::WeightInfo MatchWeight;
+  auto [IsValid, PathAndClusterInfo] =
+      BSPR->getFunctionPathAndClusterInfo(MF.getName());
+  if (!IsValid)
+    return MatchWeight;
+  for (auto &BlockCount : PathAndClusterInfo.NodeCounts) {
+    if (PathAndClusterInfo.BBHashes.count(BlockCount.first.BaseID)) {
+      auto Hash = PathAndClusterInfo.BBHashes[BlockCount.first.BaseID];
+      MachineBasicBlock *Block = Matcher.matchBlock(BlendedBlockHash(Hash));
+      // When a basic block has clone copies, sum their counts.
+      if (Block != nullptr)
+        MatchWeight.BlockWeights[Block] += BlockCount.second;
+    }
+  }
+  for (auto &PredItem : PathAndClusterInfo.EdgeCounts) {
+    auto PredID = PredItem.first.BaseID;
+    if (!PathAndClusterInfo.BBHashes.count(PredID))
+      continue;
+    auto PredHash = PathAndClusterInfo.BBHashes[PredID];
+    MachineBasicBlock *PredBlock =
+        Matcher.matchBlock(BlendedBlockHash(PredHash));
+    if (PredBlock == nullptr)
+      continue;
+    for (auto &SuccItem : PredItem.second) {
+      auto SuccID = SuccItem.first.BaseID;
+      auto EdgeWeight = SuccItem.second;
+      if (PathAndClusterInfo.BBHashes.count(SuccID)) {
+        auto SuccHash = PathAndClusterInfo.BBHashes[SuccID];
+        MachineBasicBlock *SuccBlock =
+            Matcher.matchBlock(BlendedBlockHash(SuccHash));
+        // When an edge has clone copies, sum their counts.
+        if (SuccBlock != nullptr)
+          MatchWeight.EdgeWeights[std::make_pair(PredBlock, SuccBlock)] +=
+              EdgeWeight;
+      }
+    }
+  }
+  return MatchWeight;
+}
+
+void BasicBlockMatchingAndInference::generateWeightInfoByInference(
+    MachineFunction &MF,
+    BasicBlockMatchingAndInference::WeightInfo &MatchWeight) {
+  BlockEdgeMap Successors;
+  for (auto &Block : MF) {
+    for (auto *Succ : Block.successors())
+      Successors[&Block].push_back(Succ);
+  }
+  SampleProfileInference<MachineFunction> SPI(
+      MF, Successors, MatchWeight.BlockWeights, MatchWeight.EdgeWeights);
+  BlockWeightMap BlockWeights;
+  EdgeWeightMap EdgeWeights;
+  SPI.apply(BlockWeights, EdgeWeights);
+  ProgramWeightInfo.try_emplace(
+      MF.getName(), BasicBlockMatchingAndInference::WeightInfo{
+                        std::move(BlockWeights), std::move(EdgeWeights)});
+}
+
+bool BasicBlockMatchingAndInference::runOnMachineFunction(MachineFunction &MF) {
+  if (MF.empty())
+    return false;
+  auto MatchWeight = initWeightInfoByMatching(MF);
+  // If the ratio of the number of MBBs in matching to the total number of MBBs
+  // in the function is less than the threshold value, the processing should be
+  // abandoned.
+  if (static_cast<float>(MatchWeight.BlockWeights.size()) / MF.size() <
+      PropellerInferThreshold) {
+    return false;
+  }
+  generateWeightInfoByInference(MF, MatchWeight);
+  return false;
+}
+
+MachineFunctionPass *llvm::createBasicBlockMatchingAndInferencePass() {
+  return new BasicBlockMatchingAndInference();
+}
diff --git a/llvm/lib/CodeGen/BasicBlockSections.cpp b/llvm/lib/CodeGen/BasicBlockSections.cpp
index 52e2909bec072..11098d0e1df61 100644
--- a/llvm/lib/CodeGen/BasicBlockSections.cpp
+++ b/llvm/lib/CodeGen/BasicBlockSections.cpp
@@ -70,6 +70,7 @@
 
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/CodeGen/BasicBlockMatchingAndInference.h"
 #include "llvm/CodeGen/BasicBlockSectionUtils.h"
 #include "llvm/CodeGen/BasicBlockSectionsProfileReader.h"
 #include "llvm/CodeGen/MachineDominators.h"
@@ -81,6 +82,7 @@
 #include "llvm/InitializePasses.h"
 #include "llvm/Support/UniqueBBID.h"
 #include "llvm/Target/TargetMachine.h"
+#include "llvm/Transforms/Utils/CodeLayout.h"
 #include <optional>
 
 using namespace llvm;
@@ -175,6 +177,79 @@ updateBranches(MachineFunction &MF,
   }
 }
 
+// This function generates the machine basic block clusters of "hot" blocks.
+// Currently, only support one cluster creation.
+// TODO: Support multi-cluster creation and path cloning.
+static SmallVector<BBClusterInfo>
+createBBClusterInfoForFunction(MachineFunction &MF,
+                               BasicBlockMatchingAndInference *BMI) {
+  unsigned CurrentCluster = 0;
+  SmallVector<BBClusterInfo> BBClusterInfos;
+  auto OptWeightInfo = BMI->getWeightInfo(MF.getName());
+  if (!OptWeightInfo)
+    return BBClusterInfos;
+  auto BlockWeights = OptWeightInfo->BlockWeights;
+  auto EdgeWeights = OptWeightInfo->EdgeWeights;
+
+  SmallVector<const MachineBasicBlock *, 4> HotMBBs;
+  if (MF.size() <= 2) {
+    for (auto &MBB : MF) {
+      if (MBB.isEntryBlock() || BlockWeights[&MBB] > 0) {
+        HotMBBs.push_back(&MBB);
+      }
+    }
+  } else {
+    SmallVector<uint64_t, 0> BlockSizes(MF.size());
+    SmallVector<uint64_t, 0> BlockCounts(MF.size());
+    std::vector<const MachineBasicBlock *> OrigOrder;
+    OrigOrder.reserve(MF.size());
+    SmallVector<codelayout::EdgeCount, 0> JumpCounts;
+
+    // Renumber blocks for running the layout algorithm.
+    MF.RenumberBlocks();
+
+    // Init the MBB size and count.
+    for (auto &MBB : MF) {
+      auto NonDbgInsts =
+          instructionsWithoutDebug(MBB.instr_begin(), MBB.instr_end());
+      int NumInsts = std::distance(NonDbgInsts.begin(), NonDbgInsts.end());
+      BlockSizes[MBB.getNumber()] = 4 * NumInsts;
+      BlockCounts[MBB.getNumber()] = BlockWeights[&MBB];
+      OrigOrder.push_back(&MBB);
+    }
+
+    // Init the edge count.
+    for (auto &MBB : MF) {
+      for (auto *Succ : MBB.successors()) {
+        auto EdgeWeight = EdgeWeights[std::make_pair(&MBB, Succ)];
+        JumpCounts.push_back({static_cast<uint64_t>(MBB.getNumber()),
+                              static_cast<uint64_t>(Succ->getNumber()),
+                              EdgeWeight});
+      }
+    }
+
+    // Run the layout algorithm.
+    auto Result = computeExtTspLayout(BlockSizes, BlockCounts, JumpCounts);
+    for (uint64_t R : Result) {
+      auto Block = OrigOrder[R];
+      if (Block->isEntryBlock() || BlockWeights[Block] > 0)
+        HotMBBs.push_back(Block);
+    }
+  }
+
+  // Generate the "hot" basic block cluster.
+  if (!HotMBBs.empty()) {
+    unsigned CurrentPosition = 0;
+    for (auto &MBB : HotMBBs) {
+      if (MBB->getBBID()) {
+        BBClusterInfos.push_back(
+            {*(MBB->getBBID()), CurrentCluster, CurrentPosition++});
+      }
+    }
+  }
+  return BBClusterInfos;
+}
+
 // This function sorts basic blocks according to the cluster's information.
 // All explicitly specified clusters of basic blocks will be ordered
 // accordingly. All non-specified BBs go into a separate "Cold" section.
@@ -308,8 +383,13 @@ bool BasicBlockSections::handleBBSections(MachineFunction &MF) {
 
   DenseMap<UniqueBBID, BBClusterInfo> FuncClusterInfo;
   if (BBSectionsType == BasicBlockSection::List) {
-    auto ClusterInfo = getAnalysis<BasicBlockSectionsProfileReaderWrapperPass>()
-                           .getClusterInfoForFunction(MF.getName());
+    SmallVector<BBClusterInfo> ClusterInfo;
+    if (auto *BMI = getAnalysisIfAvailable<BasicBlockMatchingAndInference>()) {
+      ClusterInfo = createBBClusterInfoForFunction(MF, BMI);
+    } else {
+      ClusterInfo = getAnalysis<BasicBlockSectionsProfileReaderWrapperPass>()
+                        .getClusterInfoForFunction(MF.getName());
+    }
     if (ClusterInfo.empty())
       return false;
     for (auto &BBClusterInfo : ClusterInfo) {
@@ -399,6 +479,7 @@ bool BasicBlockSections::runOnMachineFunction(MachineFunction &MF) {
 void BasicBlockSections::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.setPreservesAll();
   AU.addRequired<BasicBlockSectionsProfileReaderWrapperPass>();
+  AU.addUsedIfAvailable<BasicBlockMatchingAndInference>();
   AU.addUsedIfAvailable<MachineDominatorTreeWrapperPass>();
   AU.addUsedIfAvailable<MachinePostDominatorTreeWrapperPass>();
  ...
[truncated]

@wdx727 wdx727 force-pushed the propeller_bb_hash_fix_PR4 branch from b705818 to 9df098a Compare November 12, 2025 04:33
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants