Skip to content

Commit

Permalink
[SampleFDO][NFC] Refactoring sample reader to support on-demand read …
Browse files Browse the repository at this point in the history
…profiles for given functions (llvm#104654)

Currently in extended binary format, sample reader only read the
profiles when the function are in the current module at initialization
time, this extends the support to read the arbitrary profiles for given
input functions in later stage. It's used for
llvm#101053.
  • Loading branch information
wlei-llvm authored and 5c4lar committed Aug 29, 2024
1 parent 2e14208 commit 4cc7a9e
Show file tree
Hide file tree
Showing 2 changed files with 183 additions and 86 deletions.
43 changes: 43 additions & 0 deletions llvm/include/llvm/ProfileData/SampleProfReader.h
Original file line number Diff line number Diff line change
Expand Up @@ -380,6 +380,17 @@ class SampleProfileReader {
return sampleprof_error::success;
}

/// Read sample profiles for the given functions.
std::error_code read(const DenseSet<StringRef> &FuncsToUse) {
DenseSet<StringRef> S;
for (StringRef F : FuncsToUse)
if (Profiles.find(FunctionId(F)) == Profiles.end())
S.insert(F);
if (std::error_code EC = read(S, Profiles))
return EC;
return sampleprof_error::success;
}

/// The implementaion to read sample profiles from the associated file.
virtual std::error_code readImpl() = 0;

Expand Down Expand Up @@ -520,8 +531,26 @@ class SampleProfileReader {
/// Compute summary for this profile.
void computeSummary();

/// Read sample profiles for the given functions and write them to the given
/// profile map. Currently it's only used for extended binary format to load
/// the profiles on-demand.
virtual std::error_code read(const DenseSet<StringRef> &FuncsToUse,
SampleProfileMap &Profiles) {
return sampleprof_error::not_implemented;
}

std::unique_ptr<SampleProfileReaderItaniumRemapper> Remapper;

// A map from a function's context hash to its meta data section range, used
// for on-demand read function profile metadata.
std::unordered_map<uint64_t, std::pair<const uint8_t *, const uint8_t *>>
FuncMetadataIndex;

std::pair<const uint8_t *, const uint8_t *> ProfileSecRange;

/// Whether the profile has attribute metadata.
bool ProfileHasAttribute = false;

/// \brief Whether samples are collected based on pseudo probes.
bool ProfileIsProbeBased = false;

Expand Down Expand Up @@ -621,6 +650,8 @@ class SampleProfileReaderBinary : public SampleProfileReader {

/// Read the next function profile instance.
std::error_code readFuncProfile(const uint8_t *Start);
std::error_code readFuncProfile(const uint8_t *Start,
SampleProfileMap &Profiles);

/// Read the contents of the given profile instance.
std::error_code readProfile(FunctionSamples &FProfile);
Expand Down Expand Up @@ -720,11 +751,15 @@ class SampleProfileReaderExtBinaryBase : public SampleProfileReaderBinary {
std::error_code readSecHdrTableEntry(uint64_t Idx);
std::error_code readSecHdrTable();

std::error_code readFuncMetadata(bool ProfileHasAttribute,
SampleProfileMap &Profiles);
std::error_code readFuncMetadata(bool ProfileHasAttribute);
std::error_code readFuncMetadata(bool ProfileHasAttribute,
FunctionSamples *FProfile);
std::error_code readFuncOffsetTable();
std::error_code readFuncProfiles();
std::error_code readFuncProfiles(const DenseSet<StringRef> &FuncsToUse,
SampleProfileMap &Profiles);
std::error_code readNameTableSec(bool IsMD5, bool FixedLengthMD5);
std::error_code readCSNameTableSec();
std::error_code readProfileSymbolList();
Expand Down Expand Up @@ -781,6 +816,14 @@ class SampleProfileReaderExtBinaryBase : public SampleProfileReaderBinary {
};

void setSkipFlatProf(bool Skip) override { SkipFlatProf = Skip; }

private:
/// Read the profiles on-demand for the given functions. This is used after
/// stale call graph matching finds new functions whose profiles aren't loaded
/// at the beginning and we need to loaded the profiles explicitly for
/// potential matching.
std::error_code read(const DenseSet<StringRef> &FuncsToUse,
SampleProfileMap &Profiles) override;
};

class SampleProfileReaderExtBinary : public SampleProfileReaderExtBinaryBase {
Expand Down
226 changes: 140 additions & 86 deletions llvm/lib/ProfileData/SampleProfReader.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -653,7 +653,8 @@ SampleProfileReaderBinary::readProfile(FunctionSamples &FProfile) {
}

std::error_code
SampleProfileReaderBinary::readFuncProfile(const uint8_t *Start) {
SampleProfileReaderBinary::readFuncProfile(const uint8_t *Start,
SampleProfileMap &Profiles) {
Data = Start;
auto NumHeadSamples = readNumber<uint64_t>();
if (std::error_code EC = NumHeadSamples.getError())
Expand All @@ -678,6 +679,11 @@ SampleProfileReaderBinary::readFuncProfile(const uint8_t *Start) {
return sampleprof_error::success;
}

std::error_code
SampleProfileReaderBinary::readFuncProfile(const uint8_t *Start) {
return readFuncProfile(Start, Profiles);
}

std::error_code SampleProfileReaderBinary::readImpl() {
ProfileIsFS = ProfileIsFSDisciminator;
FunctionSamples::ProfileIsFS = ProfileIsFS;
Expand Down Expand Up @@ -725,6 +731,7 @@ std::error_code SampleProfileReaderExtBinaryBase::readOneSection(
break;
}
case SecLBRProfile:
ProfileSecRange = std::make_pair(Data, End);
if (std::error_code EC = readFuncProfiles())
return EC;
break;
Expand All @@ -745,9 +752,9 @@ std::error_code SampleProfileReaderExtBinaryBase::readOneSection(
ProfileIsProbeBased =
hasSecFlag(Entry, SecFuncMetadataFlags::SecFlagIsProbeBased);
FunctionSamples::ProfileIsProbeBased = ProfileIsProbeBased;
bool HasAttribute =
ProfileHasAttribute =
hasSecFlag(Entry, SecFuncMetadataFlags::SecFlagHasAttribute);
if (std::error_code EC = readFuncMetadata(HasAttribute))
if (std::error_code EC = readFuncMetadata(ProfileHasAttribute))
return EC;
break;
}
Expand Down Expand Up @@ -791,6 +798,19 @@ bool SampleProfileReaderExtBinaryBase::useFuncOffsetList() const {
return false;
}

std::error_code
SampleProfileReaderExtBinaryBase::read(const DenseSet<StringRef> &FuncsToUse,
SampleProfileMap &Profiles) {
Data = ProfileSecRange.first;
End = ProfileSecRange.second;
if (std::error_code EC = readFuncProfiles(FuncsToUse, Profiles))
return EC;
End = Data;

if (std::error_code EC = readFuncMetadata(ProfileHasAttribute, Profiles))
return EC;
return sampleprof_error::success;
}

bool SampleProfileReaderExtBinaryBase::collectFuncsFromModule() {
if (!M)
Expand Down Expand Up @@ -838,6 +858,97 @@ std::error_code SampleProfileReaderExtBinaryBase::readFuncOffsetTable() {
return sampleprof_error::success;
}

std::error_code SampleProfileReaderExtBinaryBase::readFuncProfiles(
const DenseSet<StringRef> &FuncsToUse, SampleProfileMap &Profiles) {
const uint8_t *Start = Data;

if (Remapper) {
for (auto Name : FuncsToUse) {
Remapper->insert(Name);
}
}

if (ProfileIsCS) {
assert(useFuncOffsetList());
DenseSet<uint64_t> FuncGuidsToUse;
if (useMD5()) {
for (auto Name : FuncsToUse)
FuncGuidsToUse.insert(Function::getGUID(Name));
}

// For each function in current module, load all context profiles for
// the function as well as their callee contexts which can help profile
// guided importing for ThinLTO. This can be achieved by walking
// through an ordered context container, where contexts are laid out
// as if they were walked in preorder of a context trie. While
// traversing the trie, a link to the highest common ancestor node is
// kept so that all of its decendants will be loaded.
const SampleContext *CommonContext = nullptr;
for (const auto &NameOffset : FuncOffsetList) {
const auto &FContext = NameOffset.first;
FunctionId FName = FContext.getFunction();
StringRef FNameString;
if (!useMD5())
FNameString = FName.stringRef();

// For function in the current module, keep its farthest ancestor
// context. This can be used to load itself and its child and
// sibling contexts.
if ((useMD5() && FuncGuidsToUse.count(FName.getHashCode())) ||
(!useMD5() && (FuncsToUse.count(FNameString) ||
(Remapper && Remapper->exist(FNameString))))) {
if (!CommonContext || !CommonContext->isPrefixOf(FContext))
CommonContext = &FContext;
}

if (CommonContext == &FContext ||
(CommonContext && CommonContext->isPrefixOf(FContext))) {
// Load profile for the current context which originated from
// the common ancestor.
const uint8_t *FuncProfileAddr = Start + NameOffset.second;
if (std::error_code EC = readFuncProfile(FuncProfileAddr))
return EC;
}
}
} else if (useMD5()) {
assert(!useFuncOffsetList());
for (auto Name : FuncsToUse) {
auto GUID = MD5Hash(Name);
auto iter = FuncOffsetTable.find(GUID);
if (iter == FuncOffsetTable.end())
continue;
const uint8_t *FuncProfileAddr = Start + iter->second;
if (std::error_code EC = readFuncProfile(FuncProfileAddr, Profiles))
return EC;
}
} else if (Remapper) {
assert(useFuncOffsetList());
for (auto NameOffset : FuncOffsetList) {
SampleContext FContext(NameOffset.first);
auto FuncName = FContext.getFunction();
StringRef FuncNameStr = FuncName.stringRef();
if (!FuncsToUse.count(FuncNameStr) && !Remapper->exist(FuncNameStr))
continue;
const uint8_t *FuncProfileAddr = Start + NameOffset.second;
if (std::error_code EC = readFuncProfile(FuncProfileAddr, Profiles))
return EC;
}
} else {
assert(!useFuncOffsetList());
for (auto Name : FuncsToUse) {

auto iter = FuncOffsetTable.find(MD5Hash(Name));
if (iter == FuncOffsetTable.end())
continue;
const uint8_t *FuncProfileAddr = Start + iter->second;
if (std::error_code EC = readFuncProfile(FuncProfileAddr, Profiles))
return EC;
}
}

return sampleprof_error::success;
}

std::error_code SampleProfileReaderExtBinaryBase::readFuncProfiles() {
// Collect functions used by current module if the Reader has been
// given a module.
Expand All @@ -849,7 +960,6 @@ std::error_code SampleProfileReaderExtBinaryBase::readFuncProfiles() {

// When LoadFuncsToBeUsed is false, we are using LLVM tool, need to read all
// profiles.
const uint8_t *Start = Data;
if (!LoadFuncsToBeUsed) {
while (Data < End) {
if (std::error_code EC = readFuncProfile(Data))
Expand All @@ -858,88 +968,8 @@ std::error_code SampleProfileReaderExtBinaryBase::readFuncProfiles() {
assert(Data == End && "More data is read than expected");
} else {
// Load function profiles on demand.
if (Remapper) {
for (auto Name : FuncsToUse) {
Remapper->insert(Name);
}
}

if (ProfileIsCS) {
assert(useFuncOffsetList());
DenseSet<uint64_t> FuncGuidsToUse;
if (useMD5()) {
for (auto Name : FuncsToUse)
FuncGuidsToUse.insert(Function::getGUID(Name));
}

// For each function in current module, load all context profiles for
// the function as well as their callee contexts which can help profile
// guided importing for ThinLTO. This can be achieved by walking
// through an ordered context container, where contexts are laid out
// as if they were walked in preorder of a context trie. While
// traversing the trie, a link to the highest common ancestor node is
// kept so that all of its decendants will be loaded.
const SampleContext *CommonContext = nullptr;
for (const auto &NameOffset : FuncOffsetList) {
const auto &FContext = NameOffset.first;
FunctionId FName = FContext.getFunction();
StringRef FNameString;
if (!useMD5())
FNameString = FName.stringRef();

// For function in the current module, keep its farthest ancestor
// context. This can be used to load itself and its child and
// sibling contexts.
if ((useMD5() && FuncGuidsToUse.count(FName.getHashCode())) ||
(!useMD5() && (FuncsToUse.count(FNameString) ||
(Remapper && Remapper->exist(FNameString))))) {
if (!CommonContext || !CommonContext->isPrefixOf(FContext))
CommonContext = &FContext;
}

if (CommonContext == &FContext ||
(CommonContext && CommonContext->isPrefixOf(FContext))) {
// Load profile for the current context which originated from
// the common ancestor.
const uint8_t *FuncProfileAddr = Start + NameOffset.second;
if (std::error_code EC = readFuncProfile(FuncProfileAddr))
return EC;
}
}
} else if (useMD5()) {
assert(!useFuncOffsetList());
for (auto Name : FuncsToUse) {
auto GUID = MD5Hash(Name);
auto iter = FuncOffsetTable.find(GUID);
if (iter == FuncOffsetTable.end())
continue;
const uint8_t *FuncProfileAddr = Start + iter->second;
if (std::error_code EC = readFuncProfile(FuncProfileAddr))
return EC;
}
} else if (Remapper) {
assert(useFuncOffsetList());
for (auto NameOffset : FuncOffsetList) {
SampleContext FContext(NameOffset.first);
auto FuncName = FContext.getFunction();
StringRef FuncNameStr = FuncName.stringRef();
if (!FuncsToUse.count(FuncNameStr) && !Remapper->exist(FuncNameStr))
continue;
const uint8_t *FuncProfileAddr = Start + NameOffset.second;
if (std::error_code EC = readFuncProfile(FuncProfileAddr))
return EC;
}
} else {
assert(!useFuncOffsetList());
for (auto Name : FuncsToUse) {
auto iter = FuncOffsetTable.find(MD5Hash(Name));
if (iter == FuncOffsetTable.end())
continue;
const uint8_t *FuncProfileAddr = Start + iter->second;
if (std::error_code EC = readFuncProfile(FuncProfileAddr))
return EC;
}
}
if (std::error_code EC = readFuncProfiles(FuncsToUse, Profiles))
return EC;
Data = End;
}
assert((CSProfileCount == 0 || CSProfileCount == Profiles.size()) &&
Expand Down Expand Up @@ -1245,6 +1275,27 @@ SampleProfileReaderExtBinaryBase::readFuncMetadata(bool ProfileHasAttribute,
return sampleprof_error::success;
}

std::error_code
SampleProfileReaderExtBinaryBase::readFuncMetadata(bool ProfileHasAttribute,
SampleProfileMap &Profiles) {
if (FuncMetadataIndex.empty())
return sampleprof_error::success;

for (auto &I : Profiles) {
FunctionSamples *FProfile = &I.second;
auto R = FuncMetadataIndex.find(FProfile->getContext().getHashCode());
if (R == FuncMetadataIndex.end())
continue;

Data = R->second.first;
End = R->second.second;
if (std::error_code EC = readFuncMetadata(ProfileHasAttribute, FProfile))
return EC;
assert(Data == End && "More data is read than expected");
}
return sampleprof_error::success;
}

std::error_code
SampleProfileReaderExtBinaryBase::readFuncMetadata(bool ProfileHasAttribute) {
while (Data < End) {
Expand All @@ -1257,8 +1308,11 @@ SampleProfileReaderExtBinaryBase::readFuncMetadata(bool ProfileHasAttribute) {
if (It != Profiles.end())
FProfile = &It->second;

const uint8_t *Start = Data;
if (std::error_code EC = readFuncMetadata(ProfileHasAttribute, FProfile))
return EC;

FuncMetadataIndex[FContext.getHashCode()] = {Start, Data};
}

assert(Data == End && "More data is read than expected");
Expand Down

0 comments on commit 4cc7a9e

Please sign in to comment.