Skip to content

[MemProf] Add ability to export or highlight only a portion of graph #128255

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Feb 22, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
145 changes: 132 additions & 13 deletions llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -100,6 +100,34 @@ static cl::opt<bool> ExportToDot("memprof-export-to-dot", cl::init(false),
cl::Hidden,
cl::desc("Export graph to dot files."));

// How much of the graph to export to dot.
enum DotScope {
All, // The full CCG graph.
Alloc, // Only contexts for the specified allocation.
Context, // Only the specified context.
};

static cl::opt<DotScope> DotGraphScope(
"memprof-dot-scope", cl::desc("Scope of graph to export to dot"),
cl::Hidden, cl::init(DotScope::All),
cl::values(
clEnumValN(DotScope::All, "all", "Export full callsite graph"),
clEnumValN(DotScope::Alloc, "alloc",
"Export only nodes with contexts feeding given "
"-memprof-dot-alloc-id"),
clEnumValN(DotScope::Context, "context",
"Export only nodes with given -memprof-dot-context-id")));

static cl::opt<unsigned>
AllocIdForDot("memprof-dot-alloc-id", cl::init(0), cl::Hidden,
cl::desc("Id of alloc to export if -memprof-dot-scope=alloc "
"or to highlight if -memprof-dot-scope=all"));

static cl::opt<unsigned> ContextIdForDot(
"memprof-dot-context-id", cl::init(0), cl::Hidden,
cl::desc("Id of context to export if -memprof-dot-scope=context or to "
"highlight otherwise"));

static cl::opt<bool>
DumpCCG("memprof-dump-ccg", cl::init(false), cl::Hidden,
cl::desc("Dump CallingContextGraph to stdout after each stage."));
Expand Down Expand Up @@ -548,6 +576,10 @@ class CallsiteContextGraph {
/// Map from callsite node to the enclosing caller function.
std::map<const ContextNode *, const FuncTy *> NodeToCallingFunc;

// When exporting to dot, and an allocation id is specified, contains the
// context ids on that allocation.
DenseSet<uint32_t> DotAllocContextIds;

private:
using EdgeIter = typename std::vector<std::shared_ptr<ContextEdge>>::iterator;

Expand Down Expand Up @@ -1324,6 +1356,8 @@ CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::duplicateContextIds(
assert(ContextIdToAllocationType.count(OldId));
// The new context has the same allocation type as original.
ContextIdToAllocationType[LastContextId] = ContextIdToAllocationType[OldId];
if (DotAllocContextIds.contains(OldId))
DotAllocContextIds.insert(LastContextId);
}
return NewContextIds;
}
Expand Down Expand Up @@ -2084,6 +2118,10 @@ ModuleCallsiteContextGraph::ModuleCallsiteContextGraph(
AllocNode, StackContext, CallsiteContext,
getMIBAllocType(MIBMD), ContextSizeInfo);
}
// If exporting the graph to dot and an allocation id of interest was
// specified, record all the context ids for this allocation node.
if (ExportToDot && AllocNode->OrigStackOrAllocId == AllocIdForDot)
DotAllocContextIds = AllocNode->getContextIds();
assert(AllocNode->AllocTypes != (uint8_t)AllocationType::None);
// Memprof and callsite metadata on memory allocations no longer
// needed.
Expand Down Expand Up @@ -2176,6 +2214,10 @@ IndexCallsiteContextGraph::IndexCallsiteContextGraph(
ContextSizeInfo);
I++;
}
// If exporting the graph to dot and an allocation id of interest was
// specified, record all the context ids for this allocation node.
if (ExportToDot && AllocNode->OrigStackOrAllocId == AllocIdForDot)
DotAllocContextIds = AllocNode->getContextIds();
assert(AllocNode->AllocTypes != (uint8_t)AllocationType::None);
// Initialize version 0 on the summary alloc node to the current alloc
// type, unless it has both types in which case make it default, so
Expand Down Expand Up @@ -3021,7 +3063,16 @@ struct GraphTraits<const CallsiteContextGraph<DerivedCCG, FuncTy, CallTy> *> {
template <typename DerivedCCG, typename FuncTy, typename CallTy>
struct DOTGraphTraits<const CallsiteContextGraph<DerivedCCG, FuncTy, CallTy> *>
: public DefaultDOTGraphTraits {
DOTGraphTraits(bool IsSimple = false) : DefaultDOTGraphTraits(IsSimple) {}
DOTGraphTraits(bool IsSimple = false) : DefaultDOTGraphTraits(IsSimple) {
// If the user requested the full graph to be exported, but provided an
// allocation id, or if the user gave a context id and requested more than
// just a specific context to be exported, note that highlighting is
// enabled.
DoHighlight =
(AllocIdForDot.getNumOccurrences() && DotGraphScope == DotScope::All) ||
(ContextIdForDot.getNumOccurrences() &&
DotGraphScope != DotScope::Context);
}

using GraphType = const CallsiteContextGraph<DerivedCCG, FuncTy, CallTy> *;
using GTraits = GraphTraits<GraphType>;
Expand Down Expand Up @@ -3049,12 +3100,29 @@ struct DOTGraphTraits<const CallsiteContextGraph<DerivedCCG, FuncTy, CallTy> *>
return LabelString;
}

static std::string getNodeAttributes(NodeRef Node, GraphType) {
static std::string getNodeAttributes(NodeRef Node, GraphType G) {
auto ContextIds = Node->getContextIds();
// If highlighting enabled, see if this node contains any of the context ids
// of interest. If so, it will use a different color and a larger fontsize
// (which makes the node larger as well).
bool Highlight = false;
if (DoHighlight) {
assert(ContextIdForDot.getNumOccurrences() ||
AllocIdForDot.getNumOccurrences());
if (ContextIdForDot.getNumOccurrences())
Highlight = ContextIds.contains(ContextIdForDot);
else
Highlight = set_intersects(ContextIds, G->DotAllocContextIds);
}
std::string AttributeString = (Twine("tooltip=\"") + getNodeId(Node) + " " +
getContextIds(Node->getContextIds()) + "\"")
getContextIds(ContextIds) + "\"")
.str();
// Default fontsize is 14
if (Highlight)
AttributeString += ",fontsize=\"30\"";
AttributeString +=
(Twine(",fillcolor=\"") + getColor(Node->AllocTypes) + "\"").str();
(Twine(",fillcolor=\"") + getColor(Node->AllocTypes, Highlight) + "\"")
.str();
if (Node->CloneOf) {
AttributeString += ",color=\"blue\"";
AttributeString += ",style=\"filled,bold,dashed\"";
Expand All @@ -3064,9 +3132,22 @@ struct DOTGraphTraits<const CallsiteContextGraph<DerivedCCG, FuncTy, CallTy> *>
}

static std::string getEdgeAttributes(NodeRef, ChildIteratorType ChildIter,
GraphType) {
GraphType G) {
auto &Edge = *(ChildIter.getCurrent());
auto Color = getColor(Edge->AllocTypes);
// If highlighting enabled, see if this edge contains any of the context ids
// of interest. If so, it will use a different color and a heavier arrow
// size and weight (the larger weight makes the highlighted path
// straighter).
bool Highlight = false;
if (DoHighlight) {
assert(ContextIdForDot.getNumOccurrences() ||
AllocIdForDot.getNumOccurrences());
if (ContextIdForDot.getNumOccurrences())
Highlight = Edge->ContextIds.contains(ContextIdForDot);
else
Highlight = set_intersects(Edge->ContextIds, G->DotAllocContextIds);
}
auto Color = getColor(Edge->AllocTypes, Highlight);
std::string AttributeString =
(Twine("tooltip=\"") + getContextIds(Edge->ContextIds) + "\"" +
// fillcolor is the arrow head and color is the line
Expand All @@ -3075,13 +3156,24 @@ struct DOTGraphTraits<const CallsiteContextGraph<DerivedCCG, FuncTy, CallTy> *>
.str();
if (Edge->IsBackedge)
AttributeString += ",style=\"dotted\"";
// Default penwidth and weight are both 1.
if (Highlight)
AttributeString += ",penwidth=\"2.0\",weight=\"2\"";
return AttributeString;
}

// Since the NodeOwners list includes nodes that are no longer connected to
// the graph, skip them here.
static bool isNodeHidden(NodeRef Node, GraphType) {
return Node->isRemoved();
static bool isNodeHidden(NodeRef Node, GraphType G) {
if (Node->isRemoved())
return true;
// If a scope smaller than the full graph was requested, see if this node
// contains any of the context ids of interest.
if (DotGraphScope == DotScope::Alloc)
return !set_intersects(Node->getContextIds(), G->DotAllocContextIds);
if (DotGraphScope == DotScope::Context)
return !Node->getContextIds().contains(ContextIdForDot);
return false;
}

private:
Expand All @@ -3098,16 +3190,20 @@ struct DOTGraphTraits<const CallsiteContextGraph<DerivedCCG, FuncTy, CallTy> *>
return IdString;
}

static std::string getColor(uint8_t AllocTypes) {
static std::string getColor(uint8_t AllocTypes, bool Highlight) {
// If DoHighlight is not enabled, we want to use the highlight colors for
// NotCold and Cold, and the non-highlight color for NotCold+Cold. This is
// both compatible with the color scheme before highlighting was supported,
// and for the NotCold+Cold color the non-highlight color is a bit more
// readable.
if (AllocTypes == (uint8_t)AllocationType::NotCold)
// Color "brown1" actually looks like a lighter red.
return "brown1";
return !DoHighlight || Highlight ? "brown1" : "lightpink";
if (AllocTypes == (uint8_t)AllocationType::Cold)
return "cyan";
return !DoHighlight || Highlight ? "cyan" : "lightskyblue";
if (AllocTypes ==
((uint8_t)AllocationType::NotCold | (uint8_t)AllocationType::Cold))
// Lighter purple.
return "mediumorchid1";
return Highlight ? "magenta" : "mediumorchid1";
return "gray";
}

Expand All @@ -3117,8 +3213,17 @@ struct DOTGraphTraits<const CallsiteContextGraph<DerivedCCG, FuncTy, CallTy> *>
std::string Result = SStream.str();
return Result;
}

// True if we should highlight a specific context or allocation's contexts in
// the emitted graph.
static bool DoHighlight;
};

template <typename DerivedCCG, typename FuncTy, typename CallTy>
bool DOTGraphTraits<
const CallsiteContextGraph<DerivedCCG, FuncTy, CallTy> *>::DoHighlight =
false;

template <typename DerivedCCG, typename FuncTy, typename CallTy>
void CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::exportToDot(
std::string Label) const {
Expand Down Expand Up @@ -5181,6 +5286,20 @@ bool MemProfContextDisambiguation::processModule(
MemProfContextDisambiguation::MemProfContextDisambiguation(
const ModuleSummaryIndex *Summary, bool isSamplePGO)
: ImportSummary(Summary), isSamplePGO(isSamplePGO) {
// Check the dot graph printing options once here, to make sure we have valid
// and expected combinations.
if (DotGraphScope == DotScope::Alloc && !AllocIdForDot.getNumOccurrences())
llvm::report_fatal_error(
"-memprof-dot-scope=alloc requires -memprof-dot-alloc-id");
if (DotGraphScope == DotScope::Context &&
!ContextIdForDot.getNumOccurrences())
llvm::report_fatal_error(
"-memprof-dot-scope=context requires -memprof-dot-context-id");
if (DotGraphScope == DotScope::All && AllocIdForDot.getNumOccurrences() &&
ContextIdForDot.getNumOccurrences())
llvm::report_fatal_error(
"-memprof-dot-scope=all can't have both -memprof-dot-alloc-id and "
"-memprof-dot-context-id");
if (ImportSummary) {
// The MemProfImportSummary should only be used for testing ThinLTO
// distributed backend handling via opt, in which case we don't have a
Expand Down
Loading