From 4526baa6ae6d5a58eb2623a37c818a1d7b77d531 Mon Sep 17 00:00:00 2001 From: Robin Leroy Date: Wed, 1 Jan 2025 23:13:52 +0100 Subject: [PATCH] Less clever and more useful colouring --- osaca/semantics/kernel_dg.py | 123 +++++++++++++---------------------- 1 file changed, 45 insertions(+), 78 deletions(-) diff --git a/osaca/semantics/kernel_dg.py b/osaca/semantics/kernel_dg.py index 318bb8d..416ecb7 100644 --- a/osaca/semantics/kernel_dg.py +++ b/osaca/semantics/kernel_dg.py @@ -2,7 +2,7 @@ import copy import time -from itertools import chain +from itertools import chain, groupby from multiprocessing import Manager, Process, cpu_count import networkx as nx @@ -520,11 +520,6 @@ class KernelDG(nx.DiGraph): lcd_line_numbers = {} for dep in lcd: lcd_line_numbers[dep] = [x.line_number for x, lat in lcd[dep]["dependencies"]] - # add color scheme - graph.graph["node"] = {"colorscheme": "spectral9"} - graph.graph["edge"] = {"colorscheme": "spectral9"} - min_color = 2 - available_colors = 8 # create LCD edges for dep in lcd_line_numbers: @@ -566,79 +561,51 @@ class KernelDG(nx.DiGraph): graph.edges[e]["penwidth"] = 3 # Color the cycles created by loop-carried dependencies, longest first, never recoloring - # any node, so that the longest LCD and most long chains that are involved in the loop are - # legible. - for i, dep in enumerate(sorted(lcd, key=lambda dep: -lcd[dep]["latency"])): - # For cycles that are broken by already-colored (longer) cycles, the color need not be - # the same for each yet-uncolored arc. - # Do not use the same color for such an arc as for the cycles that delimit it. This is - # always possible with 3 colors, as each arc is only adjacent to the preceding and - # following interrupting cycles. - # Since we color edges as well as nodes, there would be room for a more interesting - # graph coloring problem: we could avoid having unrelated arcs with the same color - # meeting at the same vertex, and retain the same color between arcs of the same cycle - # that are interrupted by a single vertex. We mostly ignore this problem. - - # The longest cycle will always have color 1, the second longest cycle will always have - # color 2 except where it overlaps with with the longest cycle, etc.; for arcs that are - # part of short cycles, the colors will be less predictable. - default_color = min_color + i % available_colors - arc = [] - arc_source = lcd_line_numbers[dep][-1] - arcs = [] - for n in lcd_line_numbers[dep]: - if "fillcolor" in graph.nodes[n]: - arcs.append((arc, (arc_source, n))) - arc = [] - arc_source = n - else: - arc.append(n) - if not arcs: # Unconstrained cycle. - arcs.append((arc, tuple())) + # any node or edge, so that the longest LCD and most long chains that are involved in the + # loop are legible. + lcd_by_latencies = sorted( + ( + (latency, list(deps)) + for latency, deps in groupby(lcd, lambda dep: lcd[dep]["latency"]) + ), + reverse=True + ) + node_colors = {} + edge_colors = {} + colors_used = 0 + for i, (latency, deps) in enumerate(lcd_by_latencies): + color = None + for dep in deps: + path = lcd_line_numbers[dep] + for n in path: + if n not in node_colors: + if not color: + color = colors_used + 1 + colors_used += 1 + node_colors[n] = color + for u, v in zip(path, path[1:] + [path[0]]): + if (u, v) not in edge_colors: + # Don’t introduce a color just for an edge. + if not color: + color = colors_used + edge_colors[u, v] = color + max_color = min(11, colors_used) + colorscheme = f"spectral{max(3, max_color)}" + graph.graph["node"] = {"colorscheme" : colorscheme} + graph.graph["edge"] = {"colorscheme" : colorscheme} + for n, color in node_colors.items(): + if "style" not in graph.nodes[n]: + graph.nodes[n]["style"] = "filled" else: - arcs.append((arc, (arc_source, lcd_line_numbers[dep][0]))) - # Try to color the whole cycle with its default color, then with a single color, then - # with different colors by arc, preferring the default. - forbidden_colors = set( - graph.nodes[n]["fillcolor"] for arc, extremities in arcs for n in extremities - if "fillcolor" in graph.nodes[n] - ) - global_color = None - if default_color not in forbidden_colors: - global_color = default_color - elif len(forbidden_colors) < available_colors: - global_color = next( - c for c in range(min_color, min_color + available_colors + 1) - if c not in forbidden_colors - ) - for arc, extremities in arcs: - if global_color: - color = global_color - else: - color = default_color - while color in (graph.nodes[n].get("fillcolor") for n in extremities): - color = min_color + (color + 1) % available_colors - for n in arc: - if "style" not in graph.nodes[n]: - graph.nodes[n]["style"] = "filled" - else: - graph.nodes[n]["style"] += ",filled" - graph.nodes[n]["fillcolor"] = color - if extremities: - (source, sink) = extremities - else: - source = sink = arc[0] - arc = arc[1:] - for u, v in zip([source] + arc, arc + [sink]): - # The backward edge of the cycle is represented as the corresponding forward - # edge with the attribute dir=back. - edge = graph.edges[v, u] if (v, u) in graph.edges else graph.edges[u, v] - if arc: - if "color" in edge: - raise AssertionError( - f"Recoloring {u}->{v} in arc ({source}) {arc} ({sink}) of {dep}" - ) - edge["color"] = color + graph.nodes[n]["style"] += ",filled" + graph.nodes[n]["fillcolor"] = color + if (max_color >= 4 and color == 1) or (max_color >= 10 and color in (1, 2, max_color)): + graph.nodes[n]["fontcolor"] = "white" + for (u, v), color in edge_colors.items(): + # The backward edge of the cycle is represented as the corresponding forward + # edge with the attribute dir=back. + edge = graph.edges[u, v] if (u, v) in graph.edges else graph.edges[v, u] + edge["color"] = color # rename node from [idx] to [idx mnemonic] and add shape mapping = {}