Adding Matching and Inference Functionality to Propeller-PR4: Implement matching and inference and create clusters (#167622)

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

Co-authored-by: lifengxiang1025 <[email protected]>
Co-authored-by: zcfh <[email protected]>

Author: 3 people

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

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;

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};

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();
 

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 &);

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);
     }
   }

llvm/lib/CodeGen/BasicBlockMatchingAndInference.cpp

@@ -0,0 +1,196 @@
+//===- 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>
+#include <unordered_map>
+
+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();
+}