[llvm-profgen][CSSPGO] On-demand function size computation for preinliner

[wlei]

Summary

Similar to https://reviews.llvm.org/D110465, we can compute function size on-demand for the functions that's hit by samples.

Here we leverage the raw range samples' address to compute a set of sample hit function. Then BinarySizeContextTracker just works on those function range for the size.

[hoy]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

430

Can this be done when preinliner is enabled?

447

Does this include ranges not executed? Those ranges should also be counted for size calculation.

[wenlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

443

Some comment here would be helpful. Specifically what's in the map, and what computeFuncSizeForRange does.. I think probably computeInlinedContextSizeForRange is probably better name for computeFuncSizeForRange. (saying this because if the map contains start and end offset for a function, then compute size from given start and end offset seems weird..)

llvm/tools/llvm-profgen/ProfileGenerator.h

231

nit: computeSizeForProfiledFunctions

llvm/tools/llvm-profgen/ProfiledBinary.cpp

557

Is the EndAddr inclusive for function end? What is we have a 1-byte instr at EndAddr?

[wenlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

430

See ctor of ProfiledBinary, TrackFuncContextSize = EnableCSPreInliner && UseContextCostForPreInliner,

447

Yes, not executed range is included by findFuncOffsetRange - that gives the function start, end for the containing top level inliner.

[hoy]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

430

Sounds good, thanks for pointing it out.

447

I see. This converts each lbr range to its corresponding function range, so should cover the whole function.

llvm/tools/llvm-profgen/ProfiledBinary.cpp

557

EndAddr should be the start of the next function, as computed in ProfiledBinary::dissassembleSymbol.

[wlei]

Update :

1. Rename and add comments
2. Refactor IP.Address related code to avoid redundant "offset -> address -> offset"

[wlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

443

comments added.

computeInlinedContextSizeForRange

That's better name, thanks for the suggestion!

llvm/tools/llvm-profgen/ProfiledBinary.cpp

557

Yeah, in dissassembleSymbol the logic is like

uint64_t Offset = StartOffset;
while (Offset < EndOffset) { 
  ...
  Offset += Size;
 }

so it guarantees that the end Offset in this function(the last IP.Address) won't be EndOffset.

[hoy]

lgtm

llvm/tools/llvm-profgen/ProfileGenerator.cpp

446

nit: range -> ranges

llvm/tools/llvm-profgen/ProfiledBinary.cpp

555

nit: "Invalid function start" -> "Invalid start"?

[wlei]

refine comments

[wlei]

Updating D110466: [llvm-profgen][CSSPGO] On-demand function size computation for preinliner

[wenlei]

llvm/tools/llvm-profgen/ProfiledBinary.cpp

553–554

Is this the same as addressIsCode? or perhaps we need addrOffsetIsCode?

557

ok, thanks for clarification. Then the name EndOffset could be confusing. We could do StartOffset + Size (for the map in ProfiledBinary too) or rename them to avoid confusion.

[wlei]

llvm/tools/llvm-profgen/ProfiledBinary.cpp

553–554

The input of addressIsCode is the address. I guess the confusing part is CodeAddrs so I changed to codeAddrOffsets.

557

Yeah, I changed some code in dissassembleSymbol to ensure the EndOffset is the last valid offset of the function and the previous EndOffset is renamed to NextStartOffset

[wlei]

Updating D110466: [llvm-profgen][CSSPGO] On-demand function size computation for preinliner

[wenlei]

llvm/tools/llvm-profgen/ProfiledBinary.cpp

553–554

What I meant is we could have something like this:

bool addrOffsetIsCode(uint64_t Offset) const {
  return Offset2LocStackMap.find(Offset) != Offset2LocStackMap.end();
}

which is similar to addressIsCode, except that we use offset now.

bool addressIsCode(uint64_t Address) const {
  uint64_t Offset = virtualAddrToOffset(Address);
  return Offset2LocStackMap.find(Offset) != Offset2LocStackMap.end();
}

I think it's cleaner to use API instead of dealing with internals like codeAddrOffsets directly when possible. Also Offset2LocStackMap is unordered_map while codeAddrOffsets is sorted vector. So look up with Offset2LocStackMap is faster.

[wlei]

llvm/tools/llvm-profgen/ProfiledBinary.cpp

553–554

I see. Here I used uint32_t Index = getIndexForOffset(StartOffset); because I need to use the Index to iterate all the invalid instructions which are stored in codeAddrOffsets. Offset2LocStackMap doesn't give the Index.

Index =  getIndexForOffset(StartOffset);
while(codeAddrOffsets[Index] <= End) {
    // use codeAddrOffsets[Index]  to do something.
    Index++;
}

To avoid using codeAddrOffsets directly, previously we use IP, but IP only working on Address and this introduces many redundancy. i. e. we need to change offset --> address and then change address back --> offset.

[wenlei]

llvm/tools/llvm-profgen/ProfiledBinary.cpp

553–554

Ok, I missed the use of that Index later.. However why do we care about invalid instructions? I assume invalid instructions won't be in the middle of a function, then for those from padding at the end, should we ignore them when computing size?

[wlei]

llvm/tools/llvm-profgen/ProfiledBinary.cpp

553–554

Yeah, we actually only care about valid instructions here.

uint32_t Index = getIndexForOffset(StartOffset);
uint64_t Offset = CodeAddrOffsets[Index];

If StartOffset is an invalid offset(give a warning), it will round to next valid Index. for other instructions, they're all got from the CodeAddrOffsets by increasing the index, so they're also valid.

[wenlei]

lgtm, thanks.

llvm/tools/llvm-profgen/ProfiledBinary.cpp

553–554

I see. Though with current callers, StartOffset should always be valid offset, right? The check and warning is more about making the function robust for other use cases.

[wlei]

llvm/tools/llvm-profgen/ProfiledBinary.cpp

553–554

our callers' StartOffset is not from CodeAddrOffsets .

uint64_t Offset = StartOffset;
while ()
    bool Disassembled = DisAsm->getInstruction...
    if (Disassembled) { 
       // CodeAddrOffsets[Offset] = ..
     }
}

FuncStartOffsetMap.emplace(StartOffset,...);

if the Disassembled for StartOffset is false, it will be an invalid start. Perhaps we can see if the warning is triggered in the future.

llvm/tools/llvm-profgen/ProfileGenerator.h

public:
   }
 
 private:
   void generateLineNumBasedProfile();
   // Lookup or create FunctionSamples for the context
   FunctionSamples &
   getFunctionProfileForContext(const SampleContextFrameVector &Context,
                                bool WasLeafInlined = false);
+  // For profiled only functions, on-demand compute their inline context
+  // function byte size which is used by the pre-inliner.
+  void computeSizeForProfiledFunctions();
   // Post processing for profiles before writing out, such as mermining
   // and trimming cold profiles, running preinliner on profiles.
   void postProcessProfiles();
 
   void computeSummaryAndThreshold();
 
   void populateBodySamplesForFunction(FunctionSamples &FunctionProfile,
                                       const RangeSample &RangeCounters);

llvm/tools/llvm-profgen/ProfileGenerator.cpp

if (I == ProfileMap.end()) {
     FProfile.setContext(FContext);
     return Ret.first->second;
   }
   return I->second;
 }
 
 void CSProfileGenerator::generateProfile() {
   FunctionSamples::ProfileIsCS = true;
+
+  if (Binary->getTrackFuncContextSize())
+    computeSizeForProfiledFunctions();
+
   if (Binary->usePseudoProbes()) {
     // Enable pseudo probe functionalities in SampleProf
     FunctionSamples::ProfileIsProbeBased = true;
     generateProbeBasedProfile();
   } else {
     generateLineNumBasedProfile();
   }
   postProcessProfiles();
 }
 
+void CSProfileGenerator::computeSizeForProfiledFunctions() {
+  // Hash map to deduplicate the function range and the item is a pair of
+  // function start and end offset.
+  std::unordered_map<uint64_t, uint64_t> FuncRanges;
+  // Go through all the ranges in the CS counters, use the start of the range to
+  // look up the function it belongs and record the function range.
+  for (const auto &CI : SampleCounters) {
+    for (auto Item : CI.second.RangeCounter) {
+      // FIXME: Filter the bogus crossing function range.
+      uint64_t RangeStartOffset = Item.first.first;
+      auto FuncRange = Binary->findFuncOffsetRange(RangeStartOffset);
+      if (FuncRange.second != 0)
+        FuncRanges[FuncRange.first] = FuncRange.second;
+    }
+  }
+
+  for (auto I : FuncRanges) {
+    uint64_t StartOffset = I.first;
+    uint64_t EndOffset = I.second;
+    Binary->computeInlinedContextSizeForRange(StartOffset, EndOffset);
+  }
+}
+
 void CSProfileGenerator::generateLineNumBasedProfile() {
   for (const auto &CI : SampleCounters) {
     const StringBasedCtxKey *CtxKey =
         dyn_cast<StringBasedCtxKey>(CI.first.getPtr());
     // Get or create function profile for the range
     FunctionSamples &FunctionProfile =
         getFunctionProfileForContext(CtxKey->Context, CtxKey->WasLeafInlined);
 

llvm/tools/llvm-profgen/ProfiledBinary.h

class ProfiledBinary {
   // Offset to context location map. Used to expand the context.
   std::unordered_map<uint64_t, SampleContextFrameVector> Offset2LocStackMap;
 
   // Offset to instruction size map. Also used for quick offset lookup.
   std::unordered_map<uint64_t, uint64_t> Offset2InstSizeMap;
 
   // An array of offsets of all instructions sorted in increasing order. The
   // sorting is needed to fast advance to the next forward/backward instruction.
-  std::vector<uint64_t> CodeAddrs;
+  std::vector<uint64_t> CodeAddrOffsets;
   // A set of call instruction offsets. Used by virtual unwinding.
   std::unordered_set<uint64_t> CallAddrs;
   // A set of return instruction offsets. Used by virtual unwinding.
   std::unordered_set<uint64_t> RetAddrs;
 
   // Estimate and track function prolog and epilog ranges.
   PrologEpilogTracker ProEpilogTracker;
 
class ProfiledBinary {
 
   /// Helper function to dissassemble the symbol and extract info for unwinding
   bool dissassembleSymbol(std::size_t SI, ArrayRef<uint8_t> Bytes,
                           SectionSymbolsTy &Symbols, const SectionRef &Section);
   /// Symbolize a given instruction pointer and return a full call context.
   SampleContextFrameVector symbolize(const InstructionPointer &IP,
                                      bool UseCanonicalFnName = false,
                                      bool UseProbeDiscriminator = false);
-
   /// Decode the interesting parts of the binary and build internal data
   /// structures. On high level, the parts of interest are:
   ///   1. Text sections, including the main code section and the PLT
   ///   entries that will be used to handle cross-module call transitions.
   ///   2. The .debug_line section, used by Dwarf-based profile generation.
   ///   3. Pseudo probe related sections, used by probe-based profile
   ///   generation.
   void load();
bool addressIsReturn(uint64_t Address) const {
     return RetAddrs.count(Offset);
   }
   bool addressInPrologEpilog(uint64_t Address) const {
     uint64_t Offset = virtualAddrToOffset(Address);
     return ProEpilogTracker.PrologEpilogSet.count(Offset);
   }
 
   uint64_t getAddressforIndex(uint64_t Index) const {
-    return offsetToVirtualAddr(CodeAddrs[Index]);
+    return offsetToVirtualAddr(CodeAddrOffsets[Index]);
   }
 
   bool usePseudoProbes() const { return UsePseudoProbes; }
-  // Get the index in CodeAddrs for the address
+  // 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();
+  }
   uint32_t getIndexForAddr(uint64_t Address) const {
     uint64_t Offset = virtualAddrToOffset(Address);
-    auto Low = llvm::lower_bound(CodeAddrs, Offset);
-    return Low - CodeAddrs.begin();
+    return getIndexForOffset(Offset);
   }
 
   uint64_t getCallAddrFromFrameAddr(uint64_t FrameAddr) const {
     auto I = getIndexForAddr(FrameAddr);
     FrameAddr = I ? getAddressforIndex(I - 1) : 0;
     if (FrameAddr && addressIsCall(FrameAddr))
       return FrameAddr;
     return 0;
public:
   bool inlineContextEqual(uint64_t Add1, uint64_t Add2);
 
   // Get the full context of the current stack with inline context filled in.
   // It will search the disassembling info stored in Offset2LocStackMap. This is
   // used as the key of function sample map
   SampleContextFrameVector
   getExpandedContext(const SmallVectorImpl<uint64_t> &Stack,
                      bool &WasLeafInlined);
+  // Go through instructions among the given range and record its size for the
+  // inline context.
+  void computeInlinedContextSizeForRange(uint64_t StartOffset,
+                                         uint64_t EndOffset);
 
   const MCDecodedPseudoProbe *getCallProbeForAddr(uint64_t Address) const {
     return ProbeDecoder.getCallProbeForAddr(Address);
   }
 
   void getInlineContextForProbe(const MCDecodedPseudoProbe *Probe,
                                 SampleContextFrameVector &InlineContextStack,
                                 bool IncludeLeaf = false) const {
const MCPseudoProbeFuncDesc *getFuncDescForGUID(uint64_t GUID) {
     return ProbeDecoder.getFuncDescForGUID(GUID);
   }
 
   const MCPseudoProbeFuncDesc *
   getInlinerDescForProbe(const MCDecodedPseudoProbe *Probe) {
     return ProbeDecoder.getInlinerDescForProbe(Probe);
   }
 
+  bool getTrackFuncContextSize() { return TrackFuncContextSize; }
+
   bool getIsLoadedByMMap() { return IsLoadedByMMap; }
 
   void setIsLoadedByMMap(bool Value) { IsLoadedByMMap = Value; }
 
   bool getMissingMMapWarned() { return MissingMMapWarned; }
 
   void setMissingMMapWarned(bool Value) { MissingMMapWarned = Value; }
 };
 
 } // end namespace sampleprof
 } // end namespace llvm
 
 #endif

llvm/tools/llvm-profgen/ProfiledBinary.cpp

 
 bool ProfiledBinary::dissassembleSymbol(std::size_t SI, ArrayRef<uint8_t> Bytes,
                                         SectionSymbolsTy &Symbols,
                                         const SectionRef &Section) {
   std::size_t SE = Symbols.size();
   uint64_t SectionOffset = Section.getAddress() - getPreferredBaseAddress();
   uint64_t SectSize = Section.getSize();
   uint64_t StartOffset = Symbols[SI].Addr - getPreferredBaseAddress();
-  uint64_t EndOffset = (SI + 1 < SE)
-                           ? Symbols[SI + 1].Addr - getPreferredBaseAddress()
-                           : SectionOffset + SectSize;
-  if (StartOffset >= EndOffset)
+  uint64_t NextStartOffset =
+      (SI + 1 < SE) ? Symbols[SI + 1].Addr - getPreferredBaseAddress()
+  if (StartOffset >= NextStartOffset)
     return true;
 
   StringRef SymbolName =
       ShowCanonicalFnName
           ? FunctionSamples::getCanonicalFnName(Symbols[SI].Name)
           : Symbols[SI].Name;
   bool ShowDisassembly =
       ShowDisassemblyOnly && (DisassembleFunctionSet.empty() ||
                               DisassembleFunctionSet.count(SymbolName));
   if (ShowDisassembly)
     outs() << '<' << SymbolName << ">:\n";
 
   auto WarnInvalidInsts = [](uint64_t Start, uint64_t End) {
     WithColor::warning() << "Invalid instructions at "
                          << format("%8" PRIx64, Start) << " - "
                          << format("%8" PRIx64, End) << "\n";
   };
 
   uint64_t Offset = StartOffset;
+  uint64_t EndOffset = 0;
   // Size of a consecutive invalid instruction range starting from Offset -1
   // backwards.
   uint64_t InvalidInstLength = 0;
-  while (Offset < EndOffset) {
+  while (Offset < NextStartOffset) {
     MCInst Inst;
     uint64_t Size;
     // Disassemble an instruction.
     bool Disassembled =
         DisAsm->getInstruction(Inst, Size, Bytes.slice(Offset - SectionOffset),
                                Offset + getPreferredBaseAddress(), nulls());
     if (Size == 0)
       Size = 1;
if (ShowDisassembly) {
         outs() << getReversedLocWithContext(
             symbolize(IP, ShowCanonicalFnName, ShowPseudoProbe));
       }
       outs() << "\n";
     }
 
     if (Disassembled) {
       const MCInstrDesc &MCDesc = MII->get(Inst.getOpcode());
-      // Populate a vector of the symbolized callsite at this location
-      // We don't need symbolized info for probe-based profile, just use an
-      // empty stack as an entry to indicate a valid binary offset
-      SampleContextFrameVector SymbolizedCallStack;
-      if (TrackFuncContextSize) {
-        InstructionPointer IP(this, Offset);
-        // TODO: reallocation of Offset2LocStackMap will lead to dangling
-        // strings We need ProfiledBinary to owned these string.
-        Offset2LocStackMap[Offset] = symbolize(IP, true, UsePseudoProbes);
-        SampleContextFrameVector &SymbolizedCallStack =
-            Offset2LocStackMap[Offset];
-        // Record instruction size for the corresponding context
-        if (TrackFuncContextSize && !SymbolizedCallStack.empty())
-          FuncSizeTracker.addInstructionForContext(Offset2LocStackMap[Offset],
-                                                   Size);
-      }
+
       // Record instruction size.
       Offset2InstSizeMap[Offset] = Size;
 
       // Populate address maps.
-      CodeAddrs.push_back(Offset);
+      CodeAddrOffsets.push_back(Offset);
       if (MCDesc.isCall())
         CallAddrs.insert(Offset);
       else if (MCDesc.isReturn())
         RetAddrs.insert(Offset);
 
+      EndOffset = Offset;
+
       if (InvalidInstLength) {
         WarnInvalidInsts(Offset - InvalidInstLength, Offset - 1);
         InvalidInstLength = 0;
       }
     } else {
       InvalidInstLength += Size;
     }
 
for (int32_t I = InlineStack.getNumberOfFrames() - 1; I >= 0; I--) {
     LineLocation Line(LineOffset, Discriminator);
     auto It = NameStrings.insert(FunctionName.str());
     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)
+    WithColor::warning() << "Invalid start instruction at "
+                         << format("%8" PRIx64, StartOffset) << "\n";
+
+  uint64_t Offset = CodeAddrOffsets[Index];
+  while (Offset <= EndOffset) {
+    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];
+  }
+}
+
 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
     this->Address = Binary->getAddressforIndex(Index);