diff --git a/bolt/lib/Rewrite/PseudoProbeRewriter.cpp b/bolt/lib/Rewrite/PseudoProbeRewriter.cpp index 7516918b2389fc..4925b4b385d9b1 100644 --- a/bolt/lib/Rewrite/PseudoProbeRewriter.cpp +++ b/bolt/lib/Rewrite/PseudoProbeRewriter.cpp @@ -155,7 +155,8 @@ void PseudoProbeRewriter::parsePseudoProbe() { ProbeDecoder.printProbesForAllAddresses(outs()); } - for (const auto &[GUID, FuncDesc] : ProbeDecoder.getGUID2FuncDescMap()) { + for (const auto &FuncDesc : ProbeDecoder.getGUID2FuncDescMap()) { + uint64_t GUID = FuncDesc.FuncGUID; if (!FuncStartAddrs.contains(GUID)) continue; BinaryFunction *BF = BC.getBinaryFunctionAtAddress(FuncStartAddrs[GUID]); diff --git a/llvm/include/llvm/MC/MCPseudoProbe.h b/llvm/include/llvm/MC/MCPseudoProbe.h index 854f1209c39346..32905c1e9a424a 100644 --- a/llvm/include/llvm/MC/MCPseudoProbe.h +++ b/llvm/include/llvm/MC/MCPseudoProbe.h @@ -61,6 +61,7 @@ #include "llvm/ADT/StringRef.h" #include "llvm/ADT/iterator.h" #include "llvm/IR/PseudoProbe.h" +#include "llvm/Support/Allocator.h" #include "llvm/Support/ErrorOr.h" #include #include @@ -86,7 +87,7 @@ enum class MCPseudoProbeFlag { struct MCPseudoProbeFuncDesc { uint64_t FuncGUID = 0; uint64_t FuncHash = 0; - std::string FuncName; + StringRef FuncName; MCPseudoProbeFuncDesc(uint64_t GUID, uint64_t Hash, StringRef Name) : FuncGUID(GUID), FuncHash(Hash), FuncName(Name){}; @@ -100,8 +101,18 @@ class MCDecodedPseudoProbe; using InlineSite = std::tuple; using MCPseudoProbeInlineStack = SmallVector; // GUID to PseudoProbeFuncDesc map -using GUIDProbeFunctionMap = - std::unordered_map; +class GUIDProbeFunctionMap : public std::vector { +public: + auto find(uint64_t GUID) const { + auto CompareDesc = [](const MCPseudoProbeFuncDesc &Desc, uint64_t GUID) { + return Desc.FuncGUID < GUID; + }; + auto It = llvm::lower_bound(*this, GUID, CompareDesc); + if (It->FuncGUID != GUID) + return end(); + return It; + } +}; class MCDecodedPseudoProbeInlineTree; @@ -389,6 +400,8 @@ class MCPseudoProbeDecoder { // GUID to PseudoProbeFuncDesc map. GUIDProbeFunctionMap GUID2FuncDescMap; + BumpPtrAllocator FuncNameAllocator; + // Address to probes map. AddressProbesMap Address2ProbesMap; diff --git a/llvm/lib/MC/MCPseudoProbe.cpp b/llvm/lib/MC/MCPseudoProbe.cpp index 5951499c0cb280..90d7588407068a 100644 --- a/llvm/lib/MC/MCPseudoProbe.cpp +++ b/llvm/lib/MC/MCPseudoProbe.cpp @@ -274,7 +274,7 @@ static StringRef getProbeFNameForGUID(const GUIDProbeFunctionMap &GUID2FuncMAP, auto It = GUID2FuncMAP.find(GUID); assert(It != GUID2FuncMAP.end() && "Probe function must exist for a valid GUID"); - return It->second.FuncName; + return It->FuncName; } void MCPseudoProbeFuncDesc::print(raw_ostream &OS) { @@ -390,32 +390,46 @@ bool MCPseudoProbeDecoder::buildGUID2FuncDescMap(const uint8_t *Start, Data = Start; End = Data + Size; + uint32_t FuncDescCount = 0; while (Data < End) { - auto ErrorOrGUID = readUnencodedNumber(); - if (!ErrorOrGUID) + // GUID + if (!readUnencodedNumber()) return false; - - auto ErrorOrHash = readUnencodedNumber(); - if (!ErrorOrHash) + // Hash + if (!readUnencodedNumber()) return false; auto ErrorOrNameSize = readUnsignedNumber(); if (!ErrorOrNameSize) return false; - uint32_t NameSize = std::move(*ErrorOrNameSize); - - auto ErrorOrName = readString(NameSize); - if (!ErrorOrName) + // Function name + if (!readString(*ErrorOrNameSize)) return false; + ++FuncDescCount; + } + assert(Data == End && "Have unprocessed data in pseudo_probe_desc section"); + GUID2FuncDescMap.reserve(FuncDescCount); - uint64_t GUID = std::move(*ErrorOrGUID); - uint64_t Hash = std::move(*ErrorOrHash); - StringRef Name = std::move(*ErrorOrName); + Data = Start; + End = Data + Size; + while (Data < End) { + uint64_t GUID = + cantFail(errorOrToExpected(readUnencodedNumber())); + uint64_t Hash = + cantFail(errorOrToExpected(readUnencodedNumber())); + uint32_t NameSize = + cantFail(errorOrToExpected(readUnsignedNumber())); + StringRef Name = cantFail(errorOrToExpected(readString(NameSize))); // Initialize PseudoProbeFuncDesc and populate it into GUID2FuncDescMap - GUID2FuncDescMap.emplace(GUID, MCPseudoProbeFuncDesc(GUID, Hash, Name)); + GUID2FuncDescMap.emplace_back(GUID, Hash, Name.copy(FuncNameAllocator)); } assert(Data == End && "Have unprocessed data in pseudo_probe_desc section"); + assert(GUID2FuncDescMap.size() == FuncDescCount && + "Mismatching function description count pre- and post-parsing"); + llvm::sort(GUID2FuncDescMap, [](const auto &LHS, const auto &RHS) { + return LHS.FuncGUID < RHS.FuncGUID; + }); return true; } @@ -648,12 +662,8 @@ bool MCPseudoProbeDecoder::buildAddress2ProbeMap( void MCPseudoProbeDecoder::printGUID2FuncDescMap(raw_ostream &OS) { OS << "Pseudo Probe Desc:\n"; - // Make the output deterministic - std::map OrderedMap(GUID2FuncDescMap.begin(), - GUID2FuncDescMap.end()); - for (auto &I : OrderedMap) { - I.second.print(OS); - } + for (auto &I : GUID2FuncDescMap) + I.print(OS); } void MCPseudoProbeDecoder::printProbeForAddress(raw_ostream &OS, @@ -705,7 +715,7 @@ const MCPseudoProbeFuncDesc * MCPseudoProbeDecoder::getFuncDescForGUID(uint64_t GUID) const { auto It = GUID2FuncDescMap.find(GUID); assert(It != GUID2FuncDescMap.end() && "Function descriptor doesn't exist"); - return &It->second; + return &*It; } void MCPseudoProbeDecoder::getInlineContextForProbe(