Implement inserting prefetches into the specified positions.

rlavaee · rlavaee · commit 82b6033548d0 · 2025-11-15T20:09:15.000Z
diff --git a/llvm/include/llvm/CodeGen/BasicBlockSectionsProfileReader.h b/llvm/include/llvm/CodeGen/BasicBlockSectionsProfileReader.h
@@ -52,10 +52,12 @@ struct SubblockID {
   unsigned SubblockIndex;
 };
 
+// This represents a prefetch hint to be injected at site `SiteID`, targetting
+// `TargetID` in function `TargetFunction`.
 struct PrefetchHint {
-  SubblockID SitePosition;
-  StringRef TargetFunctionName;
-  osition TargetPosition;
+  SubblockID SiteID;
+  StringRef TargetFunction;
+  SubblockID TargetID;
 };
 
 // This represents the raw input profile for one function.
@@ -80,27 +82,6 @@ struct FunctionPathAndClusterInfo {
   DenseMap<unsigned, uint64_t> BBHashes;
 };
 
-// Provides DenseMapInfo SubblockID.
-template <> struct DenseMapInfo<SubblockID> {
-  static inline SubblockID getEmptyKey() {
-    return {DenseMapInfo<UniqueBBID>::getEmptyKey(),
-            DenseMapInfo<unsigned>::getEmptyKey()};
-  }
-  static inline SubblockID getTombstoneKey() {
-    return SubblockID{DenseMapInfo<UniqueBBID>::getTombstoneKey(),
-                      DenseMapInfo<unsigned>::getTombstoneKey()};
-  }
-  static unsigned getHashValue(const SubblockID &Val) {
-    std::pair<unsigned, unsigned> PairVal = std::make_pair(
-        DenseMapInfo<UniqueBBID>::getHashValue(Val.BBID), Val.BBOffset);
-    return DenseMapInfo<std::pair<unsigned, unsigned>>::getHashValue(PairVal);
-  }
-  static bool isEqual(const SubblockID &LHS, const SubblockID &RHS) {
-    return DenseMapInfo<UniqueBBID>::isEqual(LHS.BBID, RHS.BBID) &&
-           DenseMapInfo<unsigned>::isEqual(LHS.BBOffset, RHS.BBOffset);
-  }
-};
-
 class BasicBlockSectionsProfileReader {
 public:
   friend class BasicBlockSectionsProfileReaderWrapperPass;
@@ -132,6 +113,7 @@ class BasicBlockSectionsProfileReader {
   SmallVector<SubblockID>
   getPrefetchTargetsForFunction(StringRef FuncName) const;
 
+  // Returns the prefetch hints to be injected in function `FuncName`.
   SmallVector<PrefetchHint>
   getPrefetchHintsForFunction(StringRef FuncName) const;
 
@@ -243,11 +225,10 @@ class BasicBlockSectionsProfileReaderWrapperPass : public ImmutablePass {
 
   uint64_t getEdgeCount(StringRef FuncName, const UniqueBBID &SrcBBID,
                         const UniqueBBID &DestBBID) const;
+
   SmallVector<PrefetchHint>
   getPrefetchHintsForFunction(StringRef FuncName) const;
 
-  DenseSet<SubblockID> getPrefetchTargetsForFunction(StringRef FuncName) const;
-
   SmallVector<SubblockID>
   getPrefetchTargetsForFunction(StringRef FuncName) const;
 
diff --git a/llvm/include/llvm/CodeGen/InsertCodePrefetch.h b/llvm/include/llvm/CodeGen/InsertCodePrefetch.h
@@ -0,0 +1,23 @@
+//===- BasicBlockSectionUtils.h - Utilities for basic block sections     --===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CODEGEN_INSERTCODEPREFETCH_H
+#define LLVM_CODEGEN_INSERTCODEPREFETCH_H
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallString.h"
+#include "llvm/Support/UniqueBBID.h"
+#include "llvm/Support/CommandLine.h"
+
+namespace llvm {
+
+SmallString<128> getPrefetchTargetSymbolName(StringRef FunctionName, const UniqueBBID &BBID, unsigned SubblockIndex);
+
+} // end namespace llvm
+
+#endif // LLVM_CODEGEN_INSERTCODEPREFETCH_H
diff --git a/llvm/include/llvm/CodeGen/MachineBasicBlock.h b/llvm/include/llvm/CodeGen/MachineBasicBlock.h
@@ -1298,12 +1298,6 @@ class MachineBasicBlock
   /// Return the MCSymbol for this basic block.
   LLVM_ABI MCSymbol *getSymbol() const;
 
-  MCSymbol *getCallInstSymbol(unsigned CallInstNumber) const;
-
-  const SmallVector<MCSymbol *, 4> &getCallInstSymbols() const {
-    return CallInstSymbols;
-  }
-
   /// Return the Windows EH Continuation Symbol for this basic block.
   LLVM_ABI MCSymbol *getEHContSymbol() const;
 
diff --git a/llvm/include/llvm/CodeGen/TargetInstrInfo.h b/llvm/include/llvm/CodeGen/TargetInstrInfo.h
@@ -2370,6 +2370,15 @@ class LLVM_ABI TargetInstrInfo : public MCInstrInfo {
     llvm_unreachable("unknown number of operands necessary");
   }
 
+  /// Inserts a code prefetch instruction before `InsertBefore` in block `MBB`
+  /// targetting `GV`.
+  virtual bool insertCodePrefetchInstr(MachineBasicBlock &MBB,
+                                       MachineBasicBlock::iterator InsertBefore,
+                                       const GlobalValue *GV) const {
+    return false;
+  }
+
+
 private:
   mutable std::unique_ptr<MIRFormatter> Formatter;
   unsigned CallFrameSetupOpcode, CallFrameDestroyOpcode;
diff --git a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
@@ -82,6 +82,7 @@
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/LLVMRemarkStreamer.h"
+#include "llvm/CodeGen/InsertCodePrefetch.h"
 #include "llvm/IR/Mangler.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"
@@ -1988,26 +1989,22 @@ void AsmPrinter::emitFunctionBody() {
   FunctionCallGraphInfo FuncCGInfo;
   const auto &CallSitesInfoMap = MF->getCallSitesInfo();
   for (auto &MBB : *MF) {
-    int NextPrefetchTargetIndex = MBB.getPrefetchTargets().empty() ? -1 : 0;
     // Print a label for the basic block.
     emitBasicBlockStart(MBB);
     DenseMap<StringRef, unsigned> MnemonicCounts;
 
     SmallVector<unsigned> PrefetchTargets =
         MBB.getPrefetchTargetSubblockIndexes();
     auto PrefetchTargetIt = PrefetchTargets.begin();
-    unsigned NumCalls = 0;
+    unsigned NumCallsInBB = 0;
     // Helper to emit a symbol for the prefetch target and proceed to the next
     // one.
     auto EmitPrefetchTargetSymbolIfNeeded = [&]() {
       if (PrefetchTargetIt == PrefetchTargets.end())
         return;
-      if (NumCalls < *PrefetchTargetIt)
+      if (NumCallsInBB < *PrefetchTargetIt)
         return;
-      MCSymbol *PrefetchTargetSymbol = OutContext.getOrCreateSymbol(
-          Twine("__llvm_prefetch_target_") + MF->getName() + Twine("_") +
-          utostr(MBB.getBBID()->BaseID) + Twine("_") +
-          utostr(*PrefetchTargetIt));
+      MCSymbol *PrefetchTargetSymbol = OutContext.getOrCreateSymbol(getPrefetchTargetSymbolName(MF->getName(), *MBB.getBBID(), *PrefetchTargetIt));
       // If the function is weak-linkage it may be replaced by a strong version,
       // in which case the prefetch targets should also be replaced.
       OutStreamer->emitSymbolAttribute(
@@ -2159,7 +2156,7 @@ void AsmPrinter::emitFunctionBody() {
       if (MI.isCall()) {
         if (MF->getTarget().Options.BBAddrMap)
           OutStreamer->emitLabel(createCallsiteEndSymbol(MBB));
-        ++NumCalls;
+        ++NumCallsInBB;
       }
 
       if (TM.Options.EmitCallGraphSection && MI.isCall())
diff --git a/llvm/lib/CodeGen/BasicBlockSectionsProfileReader.cpp b/llvm/lib/CodeGen/BasicBlockSectionsProfileReader.cpp
@@ -335,8 +335,8 @@ Error BasicBlockSectionsProfileReader::ReadV1Profile() {
       auto SiteBBID = parseUniqueBBID(PrefetchSiteStr[0]);
       if (!SiteBBID)
         return SiteBBID.takeError();
-      unsigned long long SiteBBOffset;
-      if (getAsUnsignedInteger(PrefetchSiteStr[1], 10, SiteBBOffset))
+      unsigned long long SiteSubblockIndex;
+      if (getAsUnsignedInteger(PrefetchSiteStr[1], 10, SiteSubblockIndex))
         return createProfileParseError(Twine("unsigned integer expected: '") +
                                        PrefetchSiteStr[1]);
 
@@ -346,14 +346,14 @@ Error BasicBlockSectionsProfileReader::ReadV1Profile() {
       auto TargetBBID = parseUniqueBBID(PrefetchTargetStr[1]);
       if (!TargetBBID)
         return TargetBBID.takeError();
-      unsigned long long TargetBBOffset;
-      if (getAsUnsignedInteger(PrefetchTargetStr[2], 10, TargetBBOffset))
+      unsigned long long TargetSubblockIndex;
+      if (getAsUnsignedInteger(PrefetchTargetStr[2], 10, TargetSubblockIndex))
         return createProfileParseError(Twine("unsigned integer expected: '") +
                                        PrefetchTargetStr[2]);
       FI->second.PrefetchHints.push_back(
-          PrefetchHint{{*SiteBBID, static_cast<unsigned>(SiteBBOffset)},
+          PrefetchHint{SubblockID{*SiteBBID, static_cast<unsigned>(SiteSubblockIndex)},
                        PrefetchTargetStr[0],
-                       {*TargetBBID, static_cast<unsigned>(TargetBBOffset)}});
+                       SubblockID{*TargetBBID, static_cast<unsigned>(TargetSubblockIndex)}});
       continue;
     }
     case 't': { // Prefetch target specifier.
diff --git a/llvm/lib/CodeGen/InsertCodePrefetch.cpp b/llvm/lib/CodeGen/InsertCodePrefetch.cpp
@@ -17,6 +17,8 @@
 /// prefetch instruction from any module.
 //===----------------------------------------------------------------------===//
 
+#include "llvm/CodeGen/InsertCodePrefetch.h"
+
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringRef.h"
@@ -26,11 +28,24 @@
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/Passes.h"
+#include "llvm/CodeGen/TargetInstrInfo.h"
 #include "llvm/InitializePasses.h"
 
 using namespace llvm;
 #define DEBUG_TYPE "insert-code-prefetch"
 
+namespace llvm {
+SmallString<128> getPrefetchTargetSymbolName(StringRef FunctionName, const UniqueBBID &BBID, unsigned SubblockIndex) {
+          SmallString<128> R("__llvm_prefetch_target_");
+          R += FunctionName;
+          R += "_";
+          R += utostr(BBID.BaseID);
+          R += "_";
+          R += utostr(SubblockIndex);
+          return R;
+}
+} // namespace llvm
+
 namespace {
 class InsertCodePrefetch : public MachineFunctionPass {
 public:
@@ -82,7 +97,40 @@ bool InsertCodePrefetch::runOnMachineFunction(MachineFunction &MF) {
       continue;
     MBB.setPrefetchTargetSubblockIndexes(R->second);
   }
-  return false;
+  SmallVector<PrefetchHint> PrefetchHints =
+      getAnalysis<BasicBlockSectionsProfileReaderWrapperPass>()
+          .getPrefetchHintsForFunction(MF.getName());
+  DenseMap<UniqueBBID, SmallVector<PrefetchHint>>
+      PrefetchHintsBySiteBBID;
+  for (const auto &H : PrefetchHints)
+    PrefetchHintsBySiteBBID[H.SiteID.BBID].push_back(H);
+  for (auto &[SiteBBID, H]: PrefetchHintsBySiteBBID) {
+    llvm::sort(H, [](const PrefetchHint &H1, const PrefetchHint &H2) {
+      return H1.SiteID.SubblockIndex < H2.SiteID.SubblockIndex;
+    });
+  }
+  auto PtrTy = PointerType::getUnqual(MF.getFunction().getParent()->getContext());
+  const TargetInstrInfo *TII = MF.getSubtarget().getInstrInfo();
+  for (auto &BB : MF) {
+    auto It = PrefetchHintsBySiteBBID.find(*BB.getBBID());
+    if (It == PrefetchHintsBySiteBBID.end())
+      continue;
+    const auto &PrefetchHints = It->second;
+    unsigned NumCallsInBB = 0;
+    auto InstrIt = BB.begin();
+    for(auto HintIt = PrefetchHints.begin() ; HintIt != PrefetchHints.end();) {
+      auto NextInstrIt = InstrIt == BB.end() ? BB.end() : std::next(InstrIt);
+      while (NumCallsInBB >= HintIt->SiteID.SubblockIndex) {
+        auto *GV = MF.getFunction().getParent()->getOrInsertGlobal(getPrefetchTargetSymbolName(HintIt->TargetFunction, HintIt->TargetID.BBID, HintIt->TargetID.SubblockIndex), PtrTy);
+        TII->insertCodePrefetchInstr(BB, NextInstrIt, GV);
+        ++HintIt;
+      }
+      if (InstrIt == BB.end()) break;
+      if (InstrIt->isCall()) ++NumCallsInBB;
+      InstrIt = NextInstrIt;
+    }
+  }
+  return true;
 }
 
 void InsertCodePrefetch::getAnalysisUsage(AnalysisUsage &AU) const {
diff --git a/llvm/lib/CodeGen/MachineBasicBlock.cpp b/llvm/lib/CodeGen/MachineBasicBlock.cpp
@@ -90,19 +90,6 @@ MCSymbol *MachineBasicBlock::getSymbol() const {
   return CachedMCSymbol;
 }
 
-MCSymbol *MachineBasicBlock::getCallInstSymbol(unsigned CallInstNumber) const {
-  if (CallInstSymbols.size() <= CallInstNumber) {
-    const MachineFunction *MF = getParent();
-    MCContext &Ctx = MF->getContext();
-    CallInstSymbols.resize(CallInstNumber + 1);
-    CallInstSymbols[CallInstNumber] = Ctx.createBlockSymbol(
-        "BB" + Twine(MF->getFunctionNumber()) + "_" + Twine(getNumber()) + "_" +
-            Twine(CallInstNumber),
-        /*AlwaysEmit=*/true);
-  }
-  return CallInstSymbols[CallInstNumber];
-}
-
 MCSymbol *MachineBasicBlock::getEHContSymbol() const {
   if (!CachedEHContMCSymbol) {
     const MachineFunction *MF = getParent();
diff --git a/llvm/lib/Target/X86/CMakeLists.txt b/llvm/lib/Target/X86/CMakeLists.txt
@@ -58,7 +58,6 @@ set(sources
   X86IndirectThunks.cpp
   X86InterleavedAccess.cpp
   X86InsertPrefetch.cpp
-  PrefetchInsertion.cpp
   X86InstCombineIntrinsic.cpp
   X86InstrFMA3Info.cpp
   X86InstrFoldTables.cpp
diff --git a/llvm/lib/Target/X86/PrefetchInsertion.cpp b/llvm/lib/Target/X86/PrefetchInsertion.cpp
diff --git a/llvm/lib/Target/X86/X86InstrInfo.cpp b/llvm/lib/Target/X86/X86InstrInfo.cpp
diff --git a/llvm/lib/Target/X86/X86InstrInfo.h b/llvm/lib/Target/X86/X86InstrInfo.h
diff --git a/llvm/lib/Target/X86/X86TargetMachine.cpp b/llvm/lib/Target/X86/X86TargetMachine.cpp
