Skip to content

Commit

Permalink
Update the way the seed nodes are picked when generating the layout g…
Browse files Browse the repository at this point in the history
…raph.

PiperOrigin-RevId: 721786286
  • Loading branch information
Google AI Edge authored and copybara-github committed Jan 31, 2025
1 parent 9327916 commit 23e45ea
Showing 1 changed file with 19 additions and 12 deletions.
31 changes: 19 additions & 12 deletions src/ui/src/components/visualizer/worker/graph_processor.ts
Original file line number Diff line number Diff line change
Expand Up @@ -366,30 +366,37 @@ export class GraphProcessor {
generateLayoutGraphConnections(modelGraph: ModelGraph) {
modelGraph.layoutGraphEdges = {};

// Find all op nodes that don't have incoming edges, or have both incoming
// edges and outgoing edges.
const seedOpNodes: OpNode[] = [];
// Find all op nodes that don't have incoming edges.
let seedOpNodes: OpNode[] = [];
const allNonHiddenOpNodes: OpNode[] = [];
for (const node of modelGraph.nodes) {
if (!isOpNode(node) || node.hideInLayout) {
continue;
}
allNonHiddenOpNodes.push(node);
const filteredIncomingEdges = (node.incomingEdges || []).filter(
(edge) =>
!(modelGraph.nodesById[edge.sourceNodeId] as OpNode).hideInLayout,
);
const filteredOutgoingEdges = (node.outgoingEdges || []).filter(
(edge) =>
!(modelGraph.nodesById[edge.targetNodeId] as OpNode).hideInLayout,
);
if (
filteredIncomingEdges.length === 0 ||
(filteredIncomingEdges.length > 0 && filteredOutgoingEdges.length > 0)
) {
if (filteredIncomingEdges.length === 0) {
seedOpNodes.push(node);
}
}

// Do a BFS from opNodesWithoutIncomingEdges.
// If seedOpNodes is empty, it means all the nodes in the graph have
// incoming edges. This indicates that the graph must contain at least one
// full cycle without any "root" nodes. For example, the graph might have
// one circle, or two disjoint circles, etc.
//
// Instead of picking one node from each of these disjointed cycles (which
// might be expensive to calculate), we will just use all the nodes in the
// graph as the seed nodes. The DFS procedure below will handle the dedup
// logic correctly.
if (seedOpNodes.length === 0 && allNonHiddenOpNodes.length > 0) {
seedOpNodes = allNonHiddenOpNodes;
}

// Do a BFS from seedOpNodes.
const queue: OpNode[] = [...seedOpNodes];
const seenNodeIds = new Set<string>();
while (queue.length > 0) {
Expand Down

0 comments on commit 23e45ea

Please sign in to comment.