Expand the test case to prefetch hints.

Author: rlavaee

llvm/lib/CodeGen/BasicBlockSectionsProfileReader.cpp

@@ -328,9 +328,12 @@ Error BasicBlockSectionsProfileReader::ReadV1Profile() {
       // past-the-end element.
       if (FI == ProgramPathAndClusterInfo.end())
         continue;
-      assert(Values.size() == 2);
+      if (Values.size() != 2)
+        return createProfileParseError(Twine("Prefetch hint expected: "+ S));
       SmallVector<StringRef, 2> PrefetchSiteStr;
       Values[0].split(PrefetchSiteStr, '@');
+      if (PrefetchSiteStr.size() != 2)
+        return createProfileParseError(Twine("Prefetch site expected: ") + Values[0]);
       assert(PrefetchSiteStr.size() == 2);
       auto SiteBBID = parseUniqueBBID(PrefetchSiteStr[0]);
       if (!SiteBBID)
@@ -342,7 +345,8 @@ Error BasicBlockSectionsProfileReader::ReadV1Profile() {
 
       SmallVector<StringRef, 3> PrefetchTargetStr;
       Values[1].split(PrefetchTargetStr, '@');
-      assert(PrefetchTargetStr.size() == 3);
+      if (PrefetchTargetStr.size() != 3)
+        return createProfileParseError(Twine("Prefetch target target expected: ") + Values[1]);
       auto TargetBBID = parseUniqueBBID(PrefetchTargetStr[1]);
       if (!TargetBBID)
         return TargetBBID.takeError();
@@ -361,10 +365,13 @@ Error BasicBlockSectionsProfileReader::ReadV1Profile() {
       // past-the-end element.
       if (FI == ProgramPathAndClusterInfo.end())
         continue;
-      assert(Values.size() == 1);
+      if (Values.size() != 1)
+        return createProfileParseError(Twine("Prefetch target expected: ")+ S);
       SmallVector<StringRef, 2> PrefetchTargetStr;
       Values[0].split(PrefetchTargetStr, '@');
       assert(PrefetchTargetStr.size() == 2);
+      if (PrefetchTargetStr.size() != 2)
+        return createProfileParseError(Twine("Prefetch target expected: ")+ Values[0]);
       auto TargetBBID = parseUniqueBBID(PrefetchTargetStr[0]);
       if (!TargetBBID)
         return TargetBBID.takeError();

llvm/lib/CodeGen/InsertCodePrefetch.cpp

@@ -105,8 +105,8 @@ bool InsertCodePrefetch::runOnMachineFunction(MachineFunction &MF) {
   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) {
+  for (auto &[SiteBBID, Hints] : PrefetchHintsBySiteBBID) {
+    llvm::sort(Hints, [](const PrefetchHint &H1, const PrefetchHint &H2) {
       return H1.SiteID.SubblockIndex < H2.SiteID.SubblockIndex;
     });
   }
@@ -122,13 +122,13 @@ bool InsertCodePrefetch::runOnMachineFunction(MachineFunction &MF) {
     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) {
+      while (HintIt != PrefetchHints.end() && 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);
+        TII->insertCodePrefetchInstr(BB, InstrIt, GV);
         ++HintIt;
       }
       if (InstrIt == BB.end())

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

@@ -7,11 +7,14 @@
 ; RUN: echo 't 1@0' >> %t
 ; RUN: echo 't 1@1' >> %t
 ; RUN: echo 't 2@1' >> %t
+; RUN: echo 'i 3@0 _Z3barv@0@0' >> %t
+; RUN: echo 'i 2@1 _Z3foob@1@0' >> %t
 ; RUN: echo 'f _Z3barv' >> %t
 ; RUN: echo 't 0@0' >> %t
+; RUN: echo 'i 0@1 _Z3foob@0@0' >> %t
 ; RUN: echo 't 21@1' >> %t
 ;;
-; RUN: llc < %s -mtriple=x86_64-pc-linux -asm-verbose=false -function-sections -basic-block-sections=%t  | FileCheck %s
+; RUN: llc < %s -O0 -mtriple=x86_64-pc-linux -asm-verbose=false -function-sections -basic-block-sections=%t  | FileCheck %s
 
 define i32 @_Z3foob(i1 zeroext %0) nounwind {
   %2 = alloca i32, align 4
@@ -44,18 +47,27 @@ define i32 @_Z3foob(i1 zeroext %0) nounwind {
 ; CHECK:        callq _Z3bazv@PLT
 ; CHECK-NEXT:   .globl __llvm_prefetch_target__Z3foob_2_1
 ; CHECK-NEXT: __llvm_prefetch_target__Z3foob_2_1:
+; CHECK-NEXT:   prefetchit1	__llvm_prefetch_target__Z3foob_1_0(%rip)
 
 13:                                               ; preds = %11, %9
   %14 = load i32, ptr %2, align 4
   ret i32 %14
+; CHECK:      .LBB0_3:
+; CHECK-NEXT:   prefetchit1	__llvm_prefetch_target__Z3barv_0_0(%rip)
+; CHECK:        retq
+
 }
 
 define weak i32 @_Z3barv() nounwind {
   %1 = call i32 @_Z3bazv()
-  ret i32 %1
+  br label %2
 ; CHECK:      _Z3barv:
 ; CHECK-NEXT:   .weak __llvm_prefetch_target__Z3barv_0_0
 ; CHECK-NEXT: __llvm_prefetch_target__Z3barv_0_0:
+; CHECK:        callq _Z3bazv@PLT
+; CHECK-NEXT:   prefetchit1	__llvm_prefetch_target__Z3foob_0_0(%rip)
+2:
+  ret i32 %1
 }
 
 declare i32 @_Z3bazv() #1