[llvm-profgen] Fix index out of bounds error while using ip.advance

[wlei]

Summary

Previously we assume there're some non-executing sections at the bottom of the text section so that we won't hit the array's bound. But on BOLTed binary, it turned out .bolt section is at the bottom of text section which can be profiled, then it crash llvm-profgen. This change try to fix it.

[wlei]

fix lint

[wenlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

358

So this is fixing a separate bug? And ranges end is actually endOffset + 1? If that's the case, some comments for RangesTy Ranges would be helpful.

388–391

Before the loop, RangeBegin <= RangeEnd should be guaranteed, right? Would this look cleaner?

do {

  ...
} while (IP.advance() && IP.Address <= RangeEnd);

[wlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

358

Yeah, a separate bug. Remembering that RangesTy's end is exclusive.

Still kind of related with this patch, because the extra one byte only can cause the out bounds issue.

Will add some comments.

[wlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

388–391

Seems not guaranteed, see the comments:

// Disjoint ranges may have range in the middle of two instr,
// e.g. If Instr1 at Addr1, and Instr2 at Addr2, disjoint range
// can be Addr1+1 to Addr2-1. We should ignore such range.
if (IP.Address > RangeEnd)
  continue;

or change to ?

if (IP.Address > RangeEnd)
     continue;
do {

  ...
} while (IP.advance() && IP.Address <= RangeEnd);

[wenlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

388–391

Ok, I see. Yeah, that works. We could have if (RangeStart > RangeEnd) check before constructing the IP, and together with that comment.

[wlei]

remove code of InstructionPointer.advance and work on offset directly.

[wlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

358

This is indeed error prone. The end range in DWARF is exclusive but the input using by findDisjointRanges is required to be inclusive.

388–391

This is tricky. In this case, RangeStart is actually < RangeEnd but IP.Address is > RangeEnd, because IP.Address is round to next valid instruction.

On a second thought, I'm wondering if we can remove the code with InstructionPointer. you see here it convert Offset --> Address then convert back to Offset. this seems redundant. I'm trying to work on Offset directly.
Please take a look.

[wlei]

remove extra #include<...>

[wenlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

388–391

This is tricky. In this case, RangeStart is actually < RangeEnd but IP.Address is > RangeEnd, because IP.Address is round to next valid instruction.

Ah I see.. thanks for clarification.

On a second thought, I'm wondering if we can remove the code with InstructionPointer. you see here it convert Offset --> Address then convert back to Offset. this seems redundant. I'm trying to work on Offset directly.

I'm still leaning towards keep such indexing details under some abstraction, and in this case, IP. Not too worried about the Addr -> offset conversion. Or if the conversion is a concern, perhaps we can change IP to store Offset by default instead of virtual address, since majority of the use is with Offset.

[wlei]

revert back to use advance to hide index

[wlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

388–391

I see. If we changing to use Offset by default, it seems a lot code need to change in the unwinder where it heavily uses Address . If the conversion is fine, I can just change back to use address like previous version.

[wenlei]

lgtm, thanks.

llvm/tools/llvm-profgen/ProfiledBinary.cpp

679–680

Similarly, would be nice to use abstractions here instead of dealing with index directly.

[hoy]

lgtm, thanks.

llvm/tools/llvm-profgen/ProfiledBinary.h

76

nit: an exclusive

[wlei]

Updating D113238: [llvm-profgen] Fix index out of bounds error while using ip.advance

llvm/test/tools/llvm-profgen/Inputs/out-of-bounds.raw.prof

+3
+0-0:1
+f-fff0:1
+ffff-ffff:1
+0

llvm/test/tools/llvm-profgen/inline-noprobe.test

 ; RUN: llvm-profgen --format=text --perfscript=%S/Inputs/inline-noprobe.perfscript --binary=%S/Inputs/inline-noprobe.perfbin --output=%t --skip-symbolization
 ; RUN: FileCheck %s --input-file %t --check-prefix=CHECK-RAW-PROFILE
 ; RUN: llvm-profgen --format=text --use-dwarf-correlation --perfscript=%S/Inputs/inline-noprobe.perfscript --binary=%S/Inputs/inline-noprobe.perfbin --output=%t --skip-symbolization
 ; RUN: FileCheck %s --input-file %t --check-prefix=CHECK-RAW-PROFILE
 ; RUN: llvm-profgen --format=text --unsymbolized-profile=%t --binary=%S/Inputs/inline-noprobe.perfbin --output=%t1
 ; RUN: FileCheck %s --input-file %t1 --check-prefix=CHECK
 ; RUN: llvm-profgen --format=text --use-dwarf-correlation --perfscript=%S/Inputs/inline-noprobe.perfscript --binary=%S/Inputs/inline-noprobe.perfbin --output=%t
 ; RUN: FileCheck %s --input-file %t --check-prefix=CHECK
 ; RUN: echo -e "0\n0" > %t
 ; RUN: llvm-profgen --format=text --unsymbolized-profile=%t --binary=%S/Inputs/inline-noprobe.perfbin --output=%t1 --fill-zero-for-all-funcs
 ; RUN: FileCheck %s --input-file %t1 --check-prefix=CHECK-ALL-ZERO
+; RUN: llvm-profgen --format=text --unsymbolized-profile=%S/Inputs/out-of-bounds.raw.prof --binary=%S/Inputs/inline-noprobe.perfbin --output=%t1
+; RUN: FileCheck %s --input-file %t1 --check-prefix=CHECK-OB
 
 CHECK: main:188:0
 CHECK:  0: 0
 CHECK:  2: 0
 CHECK:  1: foo:188
 CHECK:   2.1: 42
 CHECK:   3: 62
 CHECK:   3.2: 21

 CHECK-RAW-PROFILE:      3
 CHECK-RAW-PROFILE-NEXT: 650-691:21
 CHECK-RAW-PROFILE-NEXT: 669-677:20
 CHECK-RAW-PROFILE-NEXT: 669-691:21
 CHECK-RAW-PROFILE-NEXT: 2
 CHECK-RAW-PROFILE-NEXT: 677->650:21
 CHECK-RAW-PROFILE-NEXT: 691->669:43
 
+;CHECK-OB: foo:8:0
+;CHECK-OB:  0: 1
+;CHECK-OB:  2.1: 1
+;CHECK-OB:  3: 1
+;CHECK-OB:  3.2: 1
+;CHECK-OB:  4: 1
+;CHECK-OB:  3.1: bar:1
+;CHECK-OB:   1: 1
+;CHECK-OB:  3.2: bar:2
+;CHECK-OB:   1: 1
+;CHECK-OB:   7: 1
+;CHECK-OB: main:8:0
+;CHECK-OB:  0: 1
+;CHECK-OB:  2: 1
+;CHECK-OB:  1: foo:6
+;CHECK-OB:   2.1: 1
+;CHECK-OB:   3: 1
+;CHECK-OB:   3.2: 1
+;CHECK-OB:   4: 1
+;CHECK-OB:   3.1: bar:1
+;CHECK-OB:    1: 1
+;CHECK-OB:   3.2: bar:1
+;CHECK-OB:    1: 1
+;CHECK-OB: bar:2:0
+;CHECK-OB:  1: 1
+;CHECK-OB:  5: 1
+
 ; original code:
 ; clang -O3 -g -fdebug-info-for-profiling test.c -o a.out
 #include <stdio.h>
 
 int bar(int x, int y) {
   if (x % 2) {
     return x - y;
   }

llvm/tools/llvm-profgen/ProfileGenerator.cpp

 }
 
 RangeSample
 ProfileGenerator::preprocessRangeCounter(const RangeSample &RangeCounter) {
   RangeSample Ranges(RangeCounter.begin(), RangeCounter.end());
   if (FillZeroForAllFuncs) {
     for (auto &FuncI : Binary->getAllBinaryFunctions()) {
       for (auto &R : FuncI.second.Ranges) {
-        Ranges[{R.first, R.second}] += 0;
+        Ranges[{R.first, R.second - 1}] += 0;
       }
     }
   } else {
     // For each range, we search for all ranges of the function it belongs to
     // and initialize it with zero count, so it remains zero if doesn't hit any
     // samples. This is to be consistent with compiler that interpret zero count
     // as unexecuted(cold).
     for (auto I : RangeCounter) {
for (auto Range : preprocessRangeCounter(RangeCounter)) {
     uint64_t RangeBegin = Binary->offsetToVirtualAddr(Range.first.first);
     uint64_t RangeEnd = Binary->offsetToVirtualAddr(Range.first.second);
     uint64_t Count = Range.second;
 
     InstructionPointer IP(Binary, RangeBegin, true);
     // Disjoint ranges may have range in the middle of two instr,
     // e.g. If Instr1 at Addr1, and Instr2 at Addr2, disjoint range
     // can be Addr1+1 to Addr2-1. We should ignore such range.
-    while (IP.Address <= RangeEnd) {
+    if (IP.Address > RangeEnd)
+      continue;
+
+    do {
       uint64_t Offset = Binary->virtualAddrToOffset(IP.Address);
       const SampleContextFrameVector &FrameVec =
           Binary->getFrameLocationStack(Offset);
       if (!FrameVec.empty()) {
         FunctionSamples &FunctionProfile = getLeafFrameProfile(FrameVec);
         updateBodySamplesforFunctionProfile(FunctionProfile, FrameVec.back(),
                                             Count);
       }
-      // Move to next IP within the range.
-      IP.advance();
-    }
+    } while (IP.advance() && IP.Address <= RangeEnd);
   }
 }
 
 StringRef ProfileGeneratorBase::getCalleeNameForOffset(uint64_t TargetOffset) {
   // Get the function range by branch target if it's a call branch.
   auto *FRange = Binary->findFuncRangeForStartOffset(TargetOffset);
 
   // We won't accumulate sample count for a range whose start is not the real
for (auto Range : Ranges) {
     // doesn't belong to current context, filter them out.
     if (Count == 0)
       continue;
 
     InstructionPointer IP(Binary, RangeBegin, true);
     // Disjoint ranges may have range in the middle of two instr,
     // e.g. If Instr1 at Addr1, and Instr2 at Addr2, disjoint range
     // can be Addr1+1 to Addr2-1. We should ignore such range.
-    while (IP.Address <= RangeEnd) {
+    if (IP.Address > RangeEnd)
+      continue;
+
+    do {
       uint64_t Offset = Binary->virtualAddrToOffset(IP.Address);
       auto LeafLoc = Binary->getInlineLeafFrameLoc(Offset);
       if (LeafLoc.hasValue()) {
         // Recording body sample for this specific context
         updateBodySamplesforFunctionProfile(FunctionProfile, *LeafLoc, Count);
       }
-
-      // Move to next IP within the range
-      IP.advance();
-    }
+    } while (IP.advance() && IP.Address <= RangeEnd);
   }
 }
 
 void CSProfileGenerator::populateBoundarySamplesForFunction(
     SampleContextFrames ContextId, FunctionSamples &FunctionProfile,
     const BranchSample &BranchCounters) {
 
   for (auto Entry : BranchCounters) {
for (const auto &Range : Ranges) {
     uint64_t RangeEnd = Binary->offsetToVirtualAddr(Range.first.second);
     uint64_t Count = Range.second;
     // Disjoint ranges have introduce zero-filled gap that
     // doesn't belong to current context, filter them out.
     if (Count == 0)
       continue;
 
     InstructionPointer IP(Binary, RangeBegin, true);
-
     // Disjoint ranges may have range in the middle of two instr,
     // e.g. If Instr1 at Addr1, and Instr2 at Addr2, disjoint range
     // can be Addr1+1 to Addr2-1. We should ignore such range.
     if (IP.Address > RangeEnd)
       continue;
 
-    while (IP.Address <= RangeEnd) {
+    do {
       const AddressProbesMap &Address2ProbesMap =
           Binary->getAddress2ProbesMap();
       auto It = Address2ProbesMap.find(IP.Address);
       if (It != Address2ProbesMap.end()) {
         for (const auto &Probe : It->second) {
           if (!Probe.isBlock())
             continue;
           ProbeCounter[&Probe] += Count;
         }
       }
-
-      IP.advance();
-    }
+    } while (IP.advance() && IP.Address <= RangeEnd);
   }
 }
 
 void CSProfileGenerator::populateBodySamplesWithProbes(
     const RangeSample &RangeCounter, SampleContextFrames ContextStack) {
   ProbeCounterMap ProbeCounter;
   // Extract the top frame probes by looking up each address among the range in
   // the Address2ProbeMap

llvm/tools/llvm-profgen/ProfiledBinary.h

union {
     uint64_t Offset = 0;
     // Also used as address in unwinder
     uint64_t Address;
   };
   // Index to the sorted code address array of the binary.
   uint64_t Index = 0;
   InstructionPointer(const ProfiledBinary *Binary, uint64_t Address,
                      bool RoundToNext = false);
-  void advance();
-  void backward();
+  bool advance();
+  bool backward();
   void update(uint64_t Addr);
 };
 
 using RangesTy = std::vector<std::pair<uint64_t, uint64_t>>;
 
 struct BinaryFunction {
   StringRef FuncName;
+  // End of range is an exclusive bound.
   RangesTy Ranges;
 };
 
 // Info about function range. A function can be split into multiple
 // non-continuous ranges, each range corresponds to one FuncRange.
 struct FuncRange {
   uint64_t StartOffset;
-  // EndOffset is a exclusive bound.
+  // EndOffset is an exclusive bound.
   uint64_t EndOffset;
   // Function the range belongs to
   BinaryFunction *Func;
   // Whether the start offset is the real entry of the function.
   bool IsFuncEntry = false;
 
   StringRef getFuncName() { return Func->FuncName; }
 };
 
 // PrologEpilog offset tracker, used to filter out broken stack samples
 // Currently we use a heuristic size (two) to infer prolog and epilog
 // based on the start address and return address. In the future,
 // we will switch to Dwarf CFI based tracker
 struct PrologEpilogTracker {
   // A set of prolog and epilog offsets. Used by virtual unwinding.
   std::unordered_set<uint64_t> PrologEpilogSet;
   ProfiledBinary *Binary;
   PrologEpilogTracker(ProfiledBinary *Bin) : Binary(Bin){};
 
   // Take the two addresses from the start of function as prolog
   void inferPrologOffsets(std::map<uint64_t, FuncRange> &FuncStartOffsetMap) {
     for (auto I : FuncStartOffsetMap) {
       PrologEpilogSet.insert(I.first);
       InstructionPointer IP(Binary, I.first);
-      IP.advance();
+      if (!IP.advance())
+        break;
       PrologEpilogSet.insert(IP.Offset);
     }
   }
 
   // Take the last two addresses before the return address as epilog
   void inferEpilogOffsets(std::unordered_set<uint64_t> &RetAddrs) {
     for (auto Addr : RetAddrs) {
       PrologEpilogSet.insert(Addr);
       InstructionPointer IP(Binary, Addr);
-      IP.backward();
+      if (!IP.backward())
+        break;
       PrologEpilogSet.insert(IP.Offset);
     }
   }
 };
 
 // Track function byte size under different context (outlined version as well as
 // various inlined versions). It also provides query support to get function
 // size with the best matching context, which is used to help pre-inliner use
bool offsetIsTransfer(uint64_t Offset) {
     return BranchOffsets.count(Offset) || RetOffsets.count(Offset) ||
            CallOffsets.count(Offset);
   }
 
   uint64_t getAddressforIndex(uint64_t Index) const {
     return offsetToVirtualAddr(CodeAddrOffsets[Index]);
   }
 
+  size_t getCodeOffsetsSize() const { return CodeAddrOffsets.size(); }
+
   bool usePseudoProbes() const { return UsePseudoProbes; }
   // Get the index in CodeAddrOffsets for the address
   // As we might get an address which is not the code
   // here it would round to the next valid code address by
   // using lower bound operation
   uint32_t getIndexForOffset(uint64_t Offset) const {
     auto Low = llvm::lower_bound(CodeAddrOffsets, Offset);
     return Low - CodeAddrOffsets.begin();

llvm/tools/llvm-profgen/ProfiledBinary.cpp

for (int32_t I = InlineStack.getNumberOfFrames() - 1; I >= 0; I--) {
     CallStack.emplace_back(*It.first, Line);
   }
 
   return CallStack;
 }
 
 void ProfiledBinary::computeInlinedContextSizeForRange(uint64_t StartOffset,
                                                        uint64_t EndOffset) {
-  uint32_t Index = getIndexForOffset(StartOffset);
-  if (CodeAddrOffsets[Index] != StartOffset)
+  uint64_t RangeBegin = offsetToVirtualAddr(StartOffset);
+  uint64_t RangeEnd = offsetToVirtualAddr(EndOffset);
+  InstructionPointer IP(this, RangeBegin, true);
+
+  if (IP.Address != RangeBegin)
     WithColor::warning() << "Invalid start instruction at "
-                         << format("%8" PRIx64, StartOffset) << "\n";
+                         << format("%8" PRIx64, RangeBegin) << "\n";
+
+  if (IP.Address >= RangeEnd)
+    return;
 
-  uint64_t Offset = CodeAddrOffsets[Index];
-  while (Offset < EndOffset) {
+  do {
+    uint64_t Offset = virtualAddrToOffset(IP.Address);
     const SampleContextFrameVector &SymbolizedCallStack =
         getFrameLocationStack(Offset, UsePseudoProbes);
     uint64_t Size = Offset2InstSizeMap[Offset];
 
     // Record instruction size for the corresponding context
     FuncSizeTracker.addInstructionForContext(SymbolizedCallStack, Size);
 
-    Offset = CodeAddrOffsets[++Index];
-  }
+  } while (IP.advance() && IP.Address < RangeEnd);
 }
 
 InstructionPointer::InstructionPointer(const ProfiledBinary *Binary,
                                        uint64_t Address, bool RoundToNext)
     : Binary(Binary), Address(Address) {
   Index = Binary->getIndexForAddr(Address);
   if (RoundToNext) {
     // we might get address which is not the code
     // it should round to the next valid address
+    if (Index >= Binary->getCodeOffsetsSize())
+      this->Address = UINT64_MAX;
+    else
       this->Address = Binary->getAddressforIndex(Index);
   }
 }
 
-void InstructionPointer::advance() {
+bool InstructionPointer::advance() {
   Index++;
+  if (Index >= Binary->getCodeOffsetsSize()) {
+    Address = UINT64_MAX;
+    return false;
+  }
   Address = Binary->getAddressforIndex(Index);
+  return true;
 }
 
-void InstructionPointer::backward() {
+bool InstructionPointer::backward() {
+  if (Index == 0) {
+    Address = 0;
+    return false;
+  }
   Index--;
   Address = Binary->getAddressforIndex(Index);
+  return true;
 }
 
 void InstructionPointer::update(uint64_t Addr) {
   Address = Addr;
   Index = Binary->getIndexForAddr(Address);
 }
 
 } // end namespace sampleprof
 } // end namespace llvm