diff --git a/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp b/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp index f8c150a675e64..5c09bb1800cb2 100644 --- a/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp +++ b/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp @@ -1467,14 +1467,26 @@ void CallsiteContextGraph::updateStackNodes() { // of length, and within each length, lexicographically by stack id. The // latter is so that we can specially handle calls that have identical stack // id sequences (either due to cloning or artificially because of the MIB - // context pruning). - std::stable_sort(Calls.begin(), Calls.end(), - [](const CallContextInfo &A, const CallContextInfo &B) { - auto &IdsA = std::get<1>(A); - auto &IdsB = std::get<1>(B); - return IdsA.size() > IdsB.size() || - (IdsA.size() == IdsB.size() && IdsA < IdsB); - }); + // context pruning). Those with the same Ids are then sorted by function to + // facilitate efficiently mapping them to the same context node. + // Because the functions are pointers, to ensure a stable sort first assign + // each function pointer to its first index in the Calls array, and then use + // that to sort by. + DenseMap FuncToIndex; + for (const auto &[Idx, CallCtxInfo] : enumerate(Calls)) + FuncToIndex.insert({std::get<2>(CallCtxInfo), Idx}); + std::stable_sort( + Calls.begin(), Calls.end(), + [&FuncToIndex](const CallContextInfo &A, const CallContextInfo &B) { + auto &IdsA = std::get<1>(A); + auto &IdsB = std::get<1>(B); + auto *FuncA = std::get<2>(A); + auto *FuncB = std::get<2>(B); + return IdsA.size() > IdsB.size() || + (IdsA.size() == IdsB.size() && + (IdsA < IdsB || + (IdsA == IdsB && FuncToIndex[FuncA] < FuncToIndex[FuncB]))); + }); // Find the node for the last stack id, which should be the same // across all calls recorded for this id, and is the id for this @@ -1492,18 +1504,35 @@ void CallsiteContextGraph::updateStackNodes() { DenseSet LastNodeContextIds = LastNode->getContextIds(); assert(!LastNodeContextIds.empty()); - // Map from function to the first call from the below list (with matching - // stack ids) found in that function. Note that calls from different - // functions can have the same stack ids because this is the list of stack - // ids that had (possibly pruned) nodes after building the graph from the - // allocation MIBs. - DenseMap FuncToCallMap; +#ifndef NDEBUG + // Save the set of functions seen for a particular set of the same stack + // ids. This is used to ensure that they have been correctly sorted to be + // adjacent in the Calls list, since we rely on that to efficiently place + // all such matching calls onto the same context node. + DenseSet MatchingIdsFuncSet; +#endif for (unsigned I = 0; I < Calls.size(); I++) { auto &[Call, Ids, Func, SavedContextIds] = Calls[I]; assert(SavedContextIds.empty()); assert(LastId == Ids.back()); +#ifndef NDEBUG + // If this call has a different set of ids than the last one, clear the + // set used to ensure they are sorted properly. + if (I > 0 && Ids != std::get<1>(Calls[I - 1])) + MatchingIdsFuncSet.clear(); + else + // If the prior call had the same stack ids this set would not be empty. + // Check if we already have a call that "matches" because it is located + // in the same function. If the Calls list was sorted properly we should + // not encounter this situation as all such entries should be adjacent + // and processed in bulk further below. + assert(!MatchingIdsFuncSet.contains(Func)); + + MatchingIdsFuncSet.insert(Func); +#endif + // First compute the context ids for this stack id sequence (the // intersection of the context ids of the corresponding nodes). // Start with the remaining saved ids for the last node. @@ -1572,22 +1601,26 @@ void CallsiteContextGraph::updateStackNodes() { continue; } - // If the prior call had the same stack ids this map would not be empty. - // Check if we already have a call that "matches" because it is located - // in the same function. - if (FuncToCallMap.contains(Func)) { - // Record the matching call found for this call, and skip it. We - // will subsequently combine it into the same node. - CallToMatchingCall[Call] = FuncToCallMap[Func]; - continue; - } - // Check if the next set of stack ids is the same (since the Calls vector // of tuples is sorted by the stack ids we can just look at the next one). + // If so, save them in the CallToMatchingCall map so that they get + // assigned to the same context node, and skip them. bool DuplicateContextIds = false; - if (I + 1 < Calls.size()) { - auto NextIds = std::get<1>(Calls[I + 1]); - DuplicateContextIds = Ids == NextIds; + for (unsigned J = I + 1; J < Calls.size(); J++) { + auto &NextIds = std::get<1>(Calls[J]); + if (NextIds != Ids) + break; + auto *NextFunc = std::get<2>(Calls[J]); + if (NextFunc != Func) { + // We have another Call with the same ids but that cannot share this + // node, must duplicate ids for it. + DuplicateContextIds = true; + break; + } + auto &NextCall = std::get<0>(Calls[J]); + CallToMatchingCall[NextCall] = Call; + // Update I so that it gets incremented correctly to skip this call. + I = J; } // If we don't have duplicate context ids, then we can assign all the @@ -1611,14 +1644,7 @@ void CallsiteContextGraph::updateStackNodes() { set_subtract(LastNodeContextIds, StackSequenceContextIds); if (LastNodeContextIds.empty()) break; - // No longer possibly in a sequence of calls with duplicate stack ids, - // clear the map. - FuncToCallMap.clear(); - } else - // Record the call with its function, so we can locate it the next time - // we find a call from this function when processing the calls with the - // same stack ids. - FuncToCallMap[Func] = Call; + } } }