Implement inserting prefetches into the specified positions.

Author: rlavaee

llvm/include/llvm/CodeGen/BasicBlockSectionsProfileReader.h

@@ -50,10 +50,12 @@ struct CallsiteID {
   unsigned CallsiteIndex;
 };
 
+// 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.
@@ -78,27 +80,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;
@@ -130,6 +111,7 @@ class BasicBlockSectionsProfileReader {
   SmallVector<CallsiteID>
   getPrefetchTargetsForFunction(StringRef FuncName) const;
 
+  // Returns the prefetch hints to be injected in function `FuncName`.
   SmallVector<PrefetchHint>
   getPrefetchHintsForFunction(StringRef FuncName) const;
 
@@ -241,11 +223,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<CallsiteID>
   getPrefetchTargetsForFunction(StringRef FuncName) const;
 

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

llvm/include/llvm/CodeGen/MachineBasicBlock.h

@@ -1297,12 +1297,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;
 

llvm/include/llvm/CodeGen/TargetInstrInfo.h

@@ -2381,6 +2381,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;

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,7 +1989,6 @@ 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;

llvm/lib/CodeGen/BasicBlockSectionsProfileReader.cpp

@@ -365,8 +365,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]);
 
@@ -376,14 +376,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.

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:
@@ -87,7 +102,40 @@ bool InsertCodePrefetch::runOnMachineFunction(MachineFunction &MF) {
       continue;
     MBB.setPrefetchTargetCallsiteIndexes(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 {

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