Refine and polish.

Author: rlavaee

llvm/include/llvm/CodeGen/BasicBlockSectionsProfileReader.h

@@ -42,14 +42,11 @@ struct BBClusterInfo {
   unsigned PositionInCluster;
 };
 
-// Assuming a block is split into subblocks across its callsites, this struct
-// uniquely identifies the subblock in block `BBID` which starts from right
-// after call number `SubblockIndex` (or the beginning of the block if
-// `SubblockIndex` is zero) to the call number `SubblockIndex+1` (or the end of
-// the block if `SubblockIndex` is the last call in the block).
-struct SubblockID {
+// The prefetch symbol is emitted immediately after the call of the given index
+// in block `BBID` (or at the beginning of the block if CallsiteIndex is -1).
+struct CallsiteID {
   UniqueBBID BBID;
-  unsigned SubblockIndex;
+  int CallsiteIndex;
 };
 
 // This represents the raw input profile for one function.
@@ -60,9 +57,9 @@ struct FunctionPathAndClusterInfo {
   // the edge a -> b (a is not cloned). The index of the path in this vector
   // determines the `UniqueBBID::CloneID` of the cloned blocks in that path.
   SmallVector<SmallVector<unsigned>> ClonePaths;
-  // Code prefetch targets, specified by the subblock ID of which beginning must
-  // be targetted for prefetching.
-  SmallVector<SubblockID> PrefetchTargets;
+  // Code prefetch targets, specified by the callsite ID immediately after
+  // which beginning must be targetted for prefetching.
+  SmallVector<CallsiteID> PrefetchTargets;
   // Node counts for each basic block.
   DenseMap<UniqueBBID, uint64_t> NodeCounts;
   // Edge counts for each edge.
@@ -99,9 +96,9 @@ class BasicBlockSectionsProfileReader {
   uint64_t getEdgeCount(StringRef FuncName, const UniqueBBID &SrcBBID,
                         const UniqueBBID &SinkBBID) const;
 
-  // Returns the prefetch targets (identified by their containing subblocks) for
-  // function `FuncName`.
-  SmallVector<SubblockID>
+  // Returns the prefetch targets (identified by their containing callsite IDs)
+  // for function `FuncName`.
+  SmallVector<CallsiteID>
   getPrefetchTargetsForFunction(StringRef FuncName) const;
 
 private:
@@ -213,7 +210,7 @@ class BasicBlockSectionsProfileReaderWrapperPass : public ImmutablePass {
   uint64_t getEdgeCount(StringRef FuncName, const UniqueBBID &SrcBBID,
                         const UniqueBBID &DestBBID) const;
 
-  SmallVector<SubblockID>
+  SmallVector<CallsiteID>
   getPrefetchTargetsForFunction(StringRef FuncName) const;
 
   // Initializes the FunctionNameToDIFilename map for the current module and

llvm/include/llvm/CodeGen/MachineBasicBlock.h

@@ -229,12 +229,10 @@ class MachineBasicBlock
   /// is only computed once and is cached.
   mutable MCSymbol *CachedMCSymbol = nullptr;
 
-  /// Contains the subblock indices in this block that are targets of code
-  /// prefetching. The subblock indexed `i` specifies that region after the
-  /// `i`th call (or the beginning of the block if `i==0`) to before the`i+1`th
-  /// callsite (or the end of the block). The prefetch target is always the
-  /// beginning of the subblock.
-  SmallVector<unsigned> PrefetchTargetSubblockIndexes;
+  /// Contains the callsite indices in this block that are targets of code
+  /// prefetching. The index `i` specifies the `i`th call, with `-1`
+  /// representing the beginning of the block.
+  SmallVector<int> PrefetchTargetCallsiteIndexes;
 
   /// Cached MCSymbol for this block (used if IsEHContTarget).
   mutable MCSymbol *CachedEHContMCSymbol = nullptr;
@@ -717,12 +715,12 @@ class MachineBasicBlock
 
   std::optional<UniqueBBID> getBBID() const { return BBID; }
 
-  const SmallVector<unsigned> &getPrefetchTargetSubblockIndexes() const {
-    return PrefetchTargetSubblockIndexes;
+  const SmallVector<int> &getPrefetchTargetCallsiteIndexes() const {
+    return PrefetchTargetCallsiteIndexes;
   }
 
-  void setPrefetchTargetSubblockIndexes(const SmallVector<unsigned> &V) {
-    PrefetchTargetSubblockIndexes = V;
+  void setPrefetchTargetCallsiteIndexes(const SmallVector<int> &V) {
+    PrefetchTargetCallsiteIndexes = V;
   }
 
   /// Returns the section ID of this basic block.

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

@@ -1986,32 +1986,31 @@ void AsmPrinter::emitFunctionBody() {
     emitBasicBlockStart(MBB);
     DenseMap<StringRef, unsigned> MnemonicCounts;
 
-    SmallVector<unsigned> PrefetchTargets =
-        MBB.getPrefetchTargetSubblockIndexes();
+    SmallVector<int> PrefetchTargets =
+        MBB.getPrefetchTargetCallsiteIndexes();
     auto PrefetchTargetIt = PrefetchTargets.begin();
-    unsigned NumCalls = 0;
+    int CurrentCallsiteIndex = -1;
     // 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)
-        return;
-      MCSymbol *PrefetchTargetSymbol = OutContext.getOrCreateSymbol(
-          Twine("__llvm_prefetch_target_") + MF->getName() + Twine("_") +
-          utostr(MBB.getBBID()->BaseID) + Twine("_") +
-          utostr(*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(
-          PrefetchTargetSymbol,
-          MF->getFunction().isWeakForLinker() ? MCSA_Weak : MCSA_Global);
-      OutStreamer->emitLabel(PrefetchTargetSymbol);
-      ++PrefetchTargetIt;
+      if (PrefetchTargetIt != PrefetchTargets.end() &&
+             *PrefetchTargetIt == CurrentCallsiteIndex) {
+        MCSymbol *PrefetchTargetSymbol = OutContext.getOrCreateSymbol(
+            Twine("__llvm_prefetch_target_") + MF->getName() + Twine("_") +
+                      utostr(MBB.getBBID()->BaseID) + Twine("_") + utostr(static_cast<unsigned>(*PrefetchTargetIt + 1)));
+        // 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(
+            PrefetchTargetSymbol,
+            MF->getFunction().isWeakForLinker() ? MCSA_Weak : MCSA_Global);
+        OutStreamer->emitLabel(PrefetchTargetSymbol);
+        ++PrefetchTargetIt;
+      }
     };
 
     for (auto &MI : MBB) {
       EmitPrefetchTargetSymbolIfNeeded();
+
       // Print the assembly for the instruction.
       if (!MI.isPosition() && !MI.isImplicitDef() && !MI.isKill() &&
           !MI.isDebugInstr()) {
@@ -2152,7 +2151,7 @@ void AsmPrinter::emitFunctionBody() {
       if (MI.isCall()) {
         if (MF->getTarget().Options.BBAddrMap)
           OutStreamer->emitLabel(createCallsiteEndSymbol(MBB));
-        ++NumCalls;
+        CurrentCallsiteIndex++;
       }
 
       if (TM.Options.EmitCallGraphSection && MI.isCall())
@@ -2165,8 +2164,7 @@ void AsmPrinter::emitFunctionBody() {
       for (auto &Handler : Handlers)
         Handler->endInstruction();
     }
-    // If the block ends with a call, we may need to emit a prefetch target
-    // at the end.
+    // Emit the last prefetch target in case the last instruction was a call.
     EmitPrefetchTargetSymbolIfNeeded();
 
     // We must emit temporary symbol for the end of this basic block, if either

llvm/lib/CodeGen/BasicBlockSectionsProfileReader.cpp

@@ -93,7 +93,7 @@ uint64_t BasicBlockSectionsProfileReader::getEdgeCount(
   return EdgeIt->second;
 }
 
-SmallVector<SubblockID>
+SmallVector<CallsiteID>
 BasicBlockSectionsProfileReader::getPrefetchTargetsForFunction(
     StringRef FuncName) const {
   return ProgramPathAndClusterInfo.lookup(getAliasName(FuncName))
@@ -155,6 +155,35 @@ BasicBlockSectionsProfileReader::getPrefetchTargetsForFunction(
 //                            +-->: 5 :
 //                                ....
 // ****************************************************************************
+// This profile can also specify prefetch targets (starting with 't') which
+// instruct the compiler to emit a prefetch symbol for the given target.
+// A prefetch target is specified by a pair "<bbid>,<subblock_index>" where
+// bbid specifies the target basic block and subblock_index is a zero-based
+// index. Subblock 0 refers to the region at the beginning of the block up to
+// the first callsite. Subblock `i > 0` refers to the region immediately after
+// the `i`-th callsite up to the `i+1`-th callsite (or the end of the block).
+// The prefetch target is always emitted at the beginning of the subblock.
+// This is the beginning of the basic block for `i = 0` and immediately after
+// the `i`-th call for every `i > 0`.
+//
+// Example: A basic block in function "foo" with BBID 10 and two call instructions (call_A, call_B).
+// This block is conceptually split into subblocks, with the prefetch target
+// symbol emitted at the beginning of each subblock.
+//
+// +----------------------------------+
+// | __llvm_prefetch_target_foo_10_0: | <- Subblock 0 (before call_A)
+// |  Instruction 1                   |
+// |  Instruction 2                   |
+// |  call_A (Callsite 0)             |
+// | __llvm_prefetch_target_foo_10_1: | <--- Subblock 1 (after call_A,
+// |                                  |                  before call_B)
+// |  Instruction 3                   |
+// |  call_B (Callsite 1)             |
+// | __llvm_prefetch_target_foo_10_2: | <--- Subblock 2 (after call_B,
+// |                                  |                  before call_C)
+// |  Instruction 4                   |
+// +----------------------------------+
+//
 Error BasicBlockSectionsProfileReader::ReadV1Profile() {
   auto FI = ProgramPathAndClusterInfo.end();
 
@@ -315,7 +344,7 @@ Error BasicBlockSectionsProfileReader::ReadV1Profile() {
       }
       continue;
     }
-    case 't': { // Prefetch target specifier.
+    case 't': { // Callsite target specifier.
       // Skip the profile when we the profile iterator (FI) refers to the
       // past-the-end element.
       if (FI == ProgramPathAndClusterInfo.end())
@@ -324,16 +353,16 @@ Error BasicBlockSectionsProfileReader::ReadV1Profile() {
       Values[0].split(PrefetchTargetStr, ',');
       if (PrefetchTargetStr.size() != 2)
         return createProfileParseError(
-            Twine("Prefetch target target expected: ") + Values[0]);
+            Twine("Callsite target expected: ") + Values[0]);
       auto TargetBBID = parseUniqueBBID(PrefetchTargetStr[0]);
       if (!TargetBBID)
         return TargetBBID.takeError();
-      unsigned long long TargetCallsiteIndex;
-      if (getAsUnsignedInteger(PrefetchTargetStr[1], 10, TargetCallsiteIndex))
-        return createProfileParseError(Twine("unsigned integer expected: '") +
+      long long CallsiteIndex;
+      if (getAsSignedInteger(PrefetchTargetStr[1], 10, CallsiteIndex))
+        return createProfileParseError(Twine("signed integer expected: '") +
                                        PrefetchTargetStr[1]);
       FI->second.PrefetchTargets.push_back(
-          SubblockID{*TargetBBID, static_cast<unsigned>(TargetCallsiteIndex)});
+          CallsiteID{*TargetBBID, static_cast<int>(CallsiteIndex - 1)});
       continue;
     }
     default:
@@ -542,7 +571,7 @@ uint64_t BasicBlockSectionsProfileReaderWrapperPass::getEdgeCount(
   return BBSPR.getEdgeCount(FuncName, SrcBBID, SinkBBID);
 }
 
-SmallVector<SubblockID>
+SmallVector<CallsiteID>
 BasicBlockSectionsProfileReaderWrapperPass::getPrefetchTargetsForFunction(
     StringRef FuncName) const {
   return BBSPR.getPrefetchTargetsForFunction(FuncName);

llvm/lib/CodeGen/InsertCodePrefetch.cpp

@@ -70,17 +70,17 @@ bool InsertCodePrefetch::runOnMachineFunction(MachineFunction &MF) {
     return false;
   // Set each block's prefetch targets so AsmPrinter can emit a special symbol
   // there.
-  SmallVector<SubblockID> PrefetchTargets =
+  SmallVector<CallsiteID> PrefetchTargets =
       getAnalysis<BasicBlockSectionsProfileReaderWrapperPass>()
           .getPrefetchTargetsForFunction(MF.getName());
-  DenseMap<UniqueBBID, SmallVector<unsigned>> PrefetchTargetsByBBID;
+  DenseMap<UniqueBBID, SmallVector<int>> PrefetchTargetsByBBID;
   for (const auto &Target : PrefetchTargets)
-    PrefetchTargetsByBBID[Target.BBID].push_back(Target.SubblockIndex);
+    PrefetchTargetsByBBID[Target.BBID].push_back(Target.CallsiteIndex);
   for (auto &MBB : MF) {
     auto R = PrefetchTargetsByBBID.find(*MBB.getBBID());
     if (R == PrefetchTargetsByBBID.end())
       continue;
-    MBB.setPrefetchTargetSubblockIndexes(R->second);
+    MBB.setPrefetchTargetCallsiteIndexes(R->second);
   }
   return false;
 }

llvm/test/CodeGen/X86/basic-block-sections-code-prefetch.ll

@@ -7,6 +7,7 @@
 ; RUN: echo 't 1,0' >> %t
 ; RUN: echo 't 1,1' >> %t
 ; RUN: echo 't 2,1' >> %t
+; RUN: echo 't 4,0' >> %t
 ; RUN: echo 'f _Z3barv' >> %t
 ; RUN: echo 't 0,0' >> %t
 ; RUN: echo 't 21,1' >> %t
@@ -48,6 +49,10 @@ define i32 @_Z3foob(i1 zeroext %0) nounwind {
 13:                                               ; preds = %11, %9
   %14 = load i32, ptr %2, align 4
   ret i32 %14
+; CHECK:      .LBB0_3:
+; CHECK-NEXT:   .globl	__llvm_prefetch_target__Z3foob_4_0
+; CHECK-NEXT: __llvm_prefetch_target__Z3foob_4_0:
+
 }
 
 define weak i32 @_Z3barv() nounwind {