diff --git a/simpler_setup/tools/README.md b/simpler_setup/tools/README.md index e8b63ec381..8580484c0c 100644 --- a/simpler_setup/tools/README.md +++ b/simpler_setup/tools/README.md @@ -334,17 +334,34 @@ python -m simpler_setup.tools.deps_viewer outputs/_/deps.json \ # Redundant-only: select the transitively-implied edges reduced would drop python -m simpler_setup.tools.deps_viewer outputs/_/deps.json \ --edge-mode omitted + +# Conservative lifecycle-aware view: preserve OUTPUT_EXISTING reuse boundaries +# and require direct TensorMap dataflow around every byte of an omitted INOUT +python -m simpler_setup.tools.deps_viewer outputs/_/deps.json \ + --edge-mode omitted_v2 ``` `--edge-mode` selects which structural `(pred, succ)` edges are visible: - `full` (default) — every dependency edge. -- `reduced` — the minimal (transitively-reduced) edge set: every edge already - implied by a longer path is dropped, e.g. `A->C` when `A->B->C` exists. +- `reduced` — the transitively-reduced scheduling edge set: an `explicit` or + `tensormap` edge already implied by a longer path is dropped, e.g. `A->C` + when `A->B->C` exists. A `creator` edge is always retained because it keeps + the task that owns a tensor referenced by the consumer alive; execution order + alone cannot replace that lifetime relationship. - `omitted` — only the redundant edges `reduced` would drop (its complement), for auditing exactly which dependencies are transitively covered. - -`reduced` and `omitted` print the redundant edges to stdout as a +- `reduced_v2` — first computes structural redundancy, then preserves every + `OUTPUT_EXISTING` creator edge as a possible reuse-generation boundary. An + `INOUT` creator edge is omitted only when direct `tensormap` annotations prove + that every occupied byte flows from an earlier Output and continues to a + later `INOUT` owned by the same creator. Regions are derived from the + underlying `buffer_addr`, dtype, shape, start offset, and strides. Missing, + ambiguous, or excessively complex metadata is preserved conservatively. +- `omitted_v2` — only structurally redundant edges that remain safe to omit + after the v2 Output-lifetime check; the complement of `reduced_v2`. + +All reduction modes print the redundant edges to stdout as a ` -> ` list, where each task uses the same label as the rendered graph — the bare `local` counter when every task is in ring 0, or the explicit `(ring, local)` tuple once any task lives in ring >= 1. Text output emits only @@ -355,18 +372,19 @@ selected edge set. Selected edges are drawn above background-colored edges so they stay visible where routes overlap. When `-o` is omitted the graph is written to a mode-specific stem (`deps_viewer_reduced.*` / `deps_viewer_omitted.*`) rather than `deps_viewer.*` so it never clobbers a -full-graph render in the same directory. Reduction is purely structural (it -ignores the per-edge tensor/arg identity) and is skipped with a warning if the -graph contains a cycle. +full-graph render in the same directory. Reachability is computed structurally, +but if any annotation for a `(pred, succ)` pair has `source=creator`, that edge +is protected from reduction. Reduction is skipped with a warning if the graph +contains a cycle. ### Command-Line Options | Option | Short | Description | | ------ | ----- | ----------- | | `input` | | Path to `deps.json` (default: newest under `./outputs/`) | -| `--output` | `-o` | Output path; default stem is `deps_viewer`, or `deps_viewer_{mode}` for `reduced` / `omitted` | +| `--output` | `-o` | Output path; default stem is `deps_viewer`, or `deps_viewer_{mode}` for any reduction mode | | `--format` | | Output format: `text` (default) or `html` | -| `--edge-mode` | | Select visible edges: `full`, `reduced`, or `omitted`; HTML preserves full layout. | +| `--edge-mode` | | Select visible edges: `full`, `reduced`, `omitted`, `reduced_v2`, or `omitted_v2`; HTML preserves full layout. | | `--engine` | | HTML-only Graphviz layout engine: `dot` (default), `sfdp`, `neato`, `fdp`, `circo`, `twopi` | | `--direction` | | HTML-only flow direction for hierarchical layouts: `LR` (default) / `TB` / `BT` / `RL` | | `--show-tensor-info` | | HTML-only: render per-task tensor rows and route edges to specific arg ports | diff --git a/simpler_setup/tools/deps_viewer.py b/simpler_setup/tools/deps_viewer.py index 07d2022b39..865e82903c 100644 --- a/simpler_setup/tools/deps_viewer.py +++ b/simpler_setup/tools/deps_viewer.py @@ -33,18 +33,25 @@ python -m simpler_setup.tools.deps_viewer DEPS_JSON --format html --engine sfdp python -m simpler_setup.tools.deps_viewer DEPS_JSON --edge-mode reduced python -m simpler_setup.tools.deps_viewer DEPS_JSON --edge-mode omitted + python -m simpler_setup.tools.deps_viewer DEPS_JSON --edge-mode reduced_v2 + python -m simpler_setup.tools.deps_viewer DEPS_JSON --edge-mode omitted_v2 -``--edge-mode`` selects which edges are visible (structural transitive reduction, -purely on ``(pred, succ)`` — per-edge tensor identity is ignored, and it is -skipped with a warning if the graph contains a cycle): +``--edge-mode`` selects which edges are visible (transitive reduction on +``(pred, succ)``, while preserving ``source=creator`` tensor-lifetime edges; +reduction is skipped with a warning if the graph contains a cycle): - ``full`` (default) — every edge. -- ``reduced`` — the minimal edge set: drops every edge whose ordering is already - implied by a longer path (e.g. ``A->C`` when ``A->B->C`` exists). +- ``reduced`` — the reduced scheduling-edge set: drops an ``explicit`` or + ``tensormap`` edge whose ordering is already implied by a longer path (e.g. + ``A->C`` when ``A->B->C`` exists), while retaining every ``creator`` edge. - ``omitted`` — only the redundant edges ``reduced`` would drop (its complement), useful for auditing exactly which dependencies are transitively covered. +- ``reduced_v2`` / ``omitted_v2`` — start from structural reduction, then + preserve ``OUTPUT_EXISTING`` reuse boundaries and restore creator edges unless + every byte of an ``INOUT`` has direct TensorMap dataflow from an earlier + Output and to a later ``INOUT`` owned by the same creator. -``reduced`` and ``omitted`` print the redundant edges to stdout. Text output +All reduction modes print the redundant edges to stdout. Text output emits only the selected edge set. HTML output keeps every edge in the Graphviz layout and colors the unselected edges like the page background, preserving the full-graph layout while drawing the selected edge set above unselected edges. @@ -84,6 +91,27 @@ def _normalize_task_id(v): _normalize_tensor_id = _normalize_task_id +_DTYPE_BYTES = { + "FLOAT32": 4, + "FLOAT16": 2, + "INT32": 4, + "INT16": 2, + "INT8": 1, + "UINT8": 1, + "BFLOAT16": 2, + "INT64": 8, + "UINT64": 8, + "UINT16": 2, + "UINT32": 4, + "BOOL": 1, +} +_OUTPUT_LIFETIME_ARG_TYPES = frozenset({"INOUT", "OUTPUT_EXISTING"}) +# Exact expansion of an arbitrary strided view can grow as the product of its +# non-contiguous dimensions. Refuse an expensive proof and preserve the edge +# instead; v2 must never call an unproven lifetime boundary redundant. +_MAX_STRIDED_REGION_INTERVALS = 10_000 + + def _normalize_small_int(v): try: return int(v) @@ -215,15 +243,18 @@ def _load_deps_edges(deps_path): return sorted(edges), sorted(nodes), annotations, tensor_table, task_table -def _transitive_reduction(edges, nodes): - """DAG transitive reduction on structural ``(pred, succ)`` edges. +def _transitive_reduction(edges, nodes, annotations=None): + """DAG transitive reduction that preserves tensor-lifetime references. An edge ``(u, v)`` is redundant when ``v`` is still reachable from ``u`` through some other path that does not use the direct ``(u, v)`` edge — i.e. - the dependency it expresses is already implied by a longer chain. Reduction - is purely structural: it ignores the per-edge tensor / arg annotations, so a - ``(u, v)`` carrying its own tensor is dropped whenever the ordering it - encodes is transitively covered. + the scheduling dependency it expresses is already implied by a longer + chain. A ``source=creator`` annotation is different: it retains the task + that owns a tensor still referenced by the consumer, so it must survive even + when the producer-to-consumer execution order is transitively covered. If + any annotation row for a structural ``(u, v)`` edge has creator source, the + whole rendered edge is protected. Explicit and tensormap dependencies remain + structural-reduction candidates. Runs in ``O(V·E)``: nodes are visited in reverse topological order while a descendant-reachability set is accumulated per node, so each edge is tested @@ -237,6 +268,9 @@ def _transitive_reduction(edges, nodes): ``is_dag=False`` (transitive reduction is only well-defined on a DAG); the caller warns and emits the full graph. """ + annotations = annotations or {} + lifetime_edges = {edge for edge, rows in annotations.items() if any(row.get("source") == "creator" for row in rows)} + succ: dict[int, set[int]] = {n: set() for n in nodes} indeg: dict[int, int] = {n: 0 for n in nodes} for u, v in edges: @@ -271,7 +305,7 @@ def _transitive_reduction(edges, nodes): for v in succ[u]: indirect |= reach[v] for v in succ[u]: - if v in indirect: + if v in indirect and (u, v) not in lifetime_edges: redundant.add((u, v)) node_reach = set(succ[u]) node_reach |= indirect @@ -284,6 +318,263 @@ def _transitive_reduction(edges, nodes): return kept, removed, True +def _merge_intervals(intervals): + """Return sorted, non-overlapping half-open byte intervals.""" + merged = [] + for begin, end in sorted(intervals): + if begin >= end: + continue + if merged and begin <= merged[-1][1]: + merged[-1] = (merged[-1][0], max(merged[-1][1], end)) + else: + merged.append((begin, end)) + return tuple(merged) + + +def _intervals_cover(target, covers): + """Whether the union of ``covers`` completely contains ``target``.""" + cover = _merge_intervals(covers) + cover_idx = 0 + for target_begin, target_end in target: + cursor = target_begin + while cover_idx < len(cover) and cover[cover_idx][1] <= cursor: + cover_idx += 1 + idx = cover_idx + while cursor < target_end and idx < len(cover) and cover[idx][0] <= cursor: + cursor = max(cursor, cover[idx][1]) + idx += 1 + if cursor < target_end: + return False + return True + + +def _strided_byte_region(row, arg, tensor_table, role="consumer"): + """Resolve one annotated tensor view into exact occupied byte intervals. + + Returns ``(buffer_addr, intervals)`` or ``None`` when the metadata is + incomplete or exact expansion would exceed the safety cap. Callers treat + ``None`` conservatively as a lifetime boundary. + """ + tensor_id = _normalize_tensor_id(row.get("tensor_id", arg.get("tensor_id"))) + tensor = tensor_table.get(tensor_id) if tensor_id is not None else None + if not isinstance(tensor, dict): + return None + + try: + buffer_addr = int(tensor["buffer_addr"]) + start_offset = int(row.get(f"{role}_start_offset", arg.get("start_offset"))) + buffer_numel = int(tensor["buffer_numel"]) + except (KeyError, TypeError, ValueError): + return None + if buffer_addr < 0 or start_offset < 0 or buffer_numel < 0: + return None + + shape_raw = row.get(f"{role}_shape", arg.get("shape")) + strides_raw = row.get(f"{role}_strides", arg.get("strides")) + if not isinstance(shape_raw, list) or not isinstance(strides_raw, list) or len(shape_raw) != len(strides_raw): + return None + try: + shape = [int(v) for v in shape_raw] + strides = [int(v) for v in strides_raw] + except (TypeError, ValueError): + return None + if not shape or any(v <= 0 for v in shape) or any(v <= 0 for v in strides): + return None + + dtype = str(row.get(f"{role}_dtype", arg.get("dtype", tensor.get("dtype")))).upper() + buffer_dtype = str(tensor.get("dtype", dtype)).upper() + dtype_bytes = _DTYPE_BYTES.get(dtype) + buffer_dtype_bytes = _DTYPE_BYTES.get(buffer_dtype) + # A reinterpretation would make start_offset/strides and buffer_numel use + # different element units. The runtime normally keeps these equal; malformed + # or future reinterpret metadata must veto the proof rather than guess. + if dtype_bytes is None or buffer_dtype_bytes is None or dtype != buffer_dtype: + return None + + # Collapse the largest canonically-contiguous suffix into one interval; + # enumerate only the outer strided coordinates. If even the last dimension + # is strided, the contiguous block is one element and every coordinate is + # enumerated exactly. + suffix_start = len(shape) + contiguous_elems = 1 + for dim in range(len(shape) - 1, -1, -1): + if strides[dim] != contiguous_elems: + break + contiguous_elems *= shape[dim] + suffix_start = dim + + interval_count = 1 + for dim in range(suffix_start): + interval_count *= shape[dim] + if interval_count > _MAX_STRIDED_REGION_INTERVALS: + return None + + offsets = [start_offset] + for dim in range(suffix_start): + offsets = [base + idx * strides[dim] for base in offsets for idx in range(shape[dim])] + + intervals = _merge_intervals( + ( + buffer_addr + offset * dtype_bytes, + buffer_addr + (offset + contiguous_elems) * dtype_bytes, + ) + for offset in offsets + ) + buffer_end = buffer_addr + buffer_numel * buffer_dtype_bytes + if not intervals or any(begin < buffer_addr or end > buffer_end for begin, end in intervals): + return None + return buffer_addr, intervals + + +def _task_args_by_index(task_table, task_id): + task = task_table.get(task_id) + if not isinstance(task, dict): + return {} + args_by_idx = {} + for arg in task.get("args", []): + if not isinstance(arg, dict): + continue + arg_idx = _normalize_small_int(arg.get("idx")) + if arg_idx is not None: + args_by_idx[arg_idx] = arg + return args_by_idx + + +def _intersect_intervals(left, right): + """Return the exact intersection of two sorted interval collections.""" + intersections = [] + left_idx = 0 + right_idx = 0 + while left_idx < len(left) and right_idx < len(right): + begin = max(left[left_idx][0], right[right_idx][0]) + end = min(left[left_idx][1], right[right_idx][1]) + if begin < end: + intersections.append((begin, end)) + if left[left_idx][1] <= right[right_idx][1]: + left_idx += 1 + else: + right_idx += 1 + return tuple(intersections) + + +def _creator_output_records(annotations, tensor_table, task_table): + """Collect exact existing-output regions carried by creator annotations.""" + by_edge = {} + by_owner_task_buffer = {} + for edge, rows in annotations.items(): + args_by_idx = _task_args_by_index(task_table, edge[1]) + for row in rows: + if row.get("source") != "creator": + continue + arg_idx = _normalize_small_int(row.get("arg")) + arg = args_by_idx.get(arg_idx) + if not isinstance(arg, dict) or arg.get("type") not in _OUTPUT_LIFETIME_ARG_TYPES: + continue + region = _strided_byte_region(row, arg, tensor_table) + if region is None: + continue + buffer_addr, intervals = region + record = { + "edge": edge, + "owner": edge[0], + "task": edge[1], + "arg": arg_idx, + "arg_type": arg.get("type"), + "buffer_addr": buffer_addr, + "intervals": intervals, + } + by_edge.setdefault(edge, []).append(record) + by_owner_task_buffer.setdefault((edge[0], edge[1], buffer_addr), []).append(record) + return by_edge, by_owner_task_buffer + + +def _modifier_records(annotations, tensor_table, task_table): + """Collect exact direct producer-to-INOUT dataflow evidence.""" + incoming = {} + outgoing = {} + for edge, rows in annotations.items(): + args_by_idx = _task_args_by_index(task_table, edge[1]) + for row in rows: + if row.get("source") != "tensormap": + continue + arg_idx = _normalize_small_int(row.get("arg")) + arg = args_by_idx.get(arg_idx) + if not isinstance(arg, dict) or arg.get("type") != "INOUT": + continue + consumer_region = _strided_byte_region(row, arg, tensor_table) + producer_region = _strided_byte_region(row, {}, tensor_table, role="producer") + if consumer_region is None or producer_region is None or consumer_region[0] != producer_region[0]: + continue + overlap = _intersect_intervals(consumer_region[1], producer_region[1]) + if not overlap: + continue + record = { + "edge": edge, + "pred": edge[0], + "succ": edge[1], + "arg": arg_idx, + "buffer_addr": consumer_region[0], + "overlap": overlap, + } + incoming.setdefault((edge[1], arg_idx, consumer_region[0]), []).append(record) + outgoing.setdefault((edge[0], consumer_region[0]), []).append(record) + return incoming, outgoing + + +def _transitive_reduction_v2(edges, nodes, annotations, tensor_table, task_table): + """Structural reduction followed by conservative lifetime restoration. + + A structurally redundant creator edge is removable only when every creator + annotation on that structural pair is an exactly-known INOUT region, and + every byte has direct TensorMap dataflow from an earlier Output and to a + later INOUT owned by the same creator. OUTPUT_EXISTING can begin a reuse + generation, so it is always preserved. Missing or ambiguous evidence also + preserves the edge; graph reachability alone never joins reuse rounds. + """ + kept, structural_removed, is_dag = _transitive_reduction(edges, nodes) + if not is_dag: + return kept, structural_removed, False + + records_by_edge, records_by_owner_task_buffer = _creator_output_records(annotations, tensor_table, task_table) + incoming_modifiers, outgoing_modifiers = _modifier_records(annotations, tensor_table, task_table) + final_removed = set(structural_removed) + for edge in structural_removed: + creator_rows = [row for row in annotations.get(edge, []) if row.get("source") == "creator"] + if not creator_rows: + continue + records = records_by_edge.get(edge, []) + # Only INOUT has direct modifier dataflow that can prove it remains + # within one reuse generation. OUTPUT_EXISTING, INPUT, missing task arg + # metadata, and overly-complex strides all veto removal. + if len(records) != len(creator_rows) or any(record["arg_type"] != "INOUT" for record in records): + final_removed.discard(edge) + continue + for record in records: + owner = record["owner"] + buffer_addr = record["buffer_addr"] + earlier_intervals = [] + for modifier in incoming_modifiers.get((record["task"], record["arg"], buffer_addr), []): + peers = records_by_owner_task_buffer.get((owner, modifier["pred"], buffer_addr), []) + for peer in peers: + earlier_intervals.extend(_intersect_intervals(modifier["overlap"], peer["intervals"])) + + later_intervals = [] + for modifier in outgoing_modifiers.get((record["task"], buffer_addr), []): + peers = records_by_owner_task_buffer.get((owner, modifier["succ"], buffer_addr), []) + for peer in peers: + if peer["arg"] == modifier["arg"] and peer["arg_type"] == "INOUT": + later_intervals.extend(_intersect_intervals(modifier["overlap"], peer["intervals"])) + if not _intervals_cover(record["intervals"], earlier_intervals) or not _intervals_cover( + record["intervals"], later_intervals + ): + final_removed.discard(edge) + break + + kept = [edge for edge in edges if edge not in final_removed] + removed = [edge for edge in edges if edge in final_removed] + return kept, removed, True + + def _backfill_output_tensor_ids(task_table, annotations): """Recover ``tensor_id`` for OUTPUT slots that the runtime hadn't materialized at submit time. @@ -1236,6 +1527,8 @@ def _build_parser(): %(prog)s deps.json --format html --show-tensor-info %(prog)s deps.json --edge-mode reduced # select non-redundant edges, print what was removed %(prog)s deps.json --edge-mode omitted # select transitively-implied (redundant) edges + %(prog)s deps.json --edge-mode reduced_v2 # require direct INOUT dataflow before omitting creator edges + %(prog)s deps.json --edge-mode omitted_v2 # show only lifetime-safe redundant edges """, ) p.add_argument("input", nargs="?", help="Path to deps.json (default: newest under ./outputs/).") @@ -1248,14 +1541,17 @@ def _build_parser(): ) p.add_argument( "--edge-mode", - choices=["full", "reduced", "omitted"], + choices=["full", "reduced", "omitted", "reduced_v2", "omitted_v2"], default="full", help=( "full (default) selects every dependency edge; reduced applies transitive reduction, selecting the " - "minimal edge set; omitted selects only the redundant edges reduced would drop (the complement of " - "reduced). reduced/omitted print the redundant edges to stdout, are structural (pred,succ) level, " - "apply to both text and html, and are skipped with a warning on a cyclic graph. In html, all edges " - "still participate in layout; unselected edges are colored as background and drawn below selected edges." + "minimal scheduling-edge set while preserving creator tensor-lifetime edges; omitted selects only " + "the redundant edges reduced would drop (the complement of reduced). reduced/omitted print the " + "redundant edges to stdout. reduced_v2/omitted_v2 first find structural redundancy, then restore " + "OUTPUT_EXISTING reuse boundaries and require direct TensorMap dataflow into and out of every byte " + "of an INOUT before omitting its creator edge. All reduction modes work for text and html and are " + "skipped with a warning on a cyclic graph. In html, all edges still participate in layout; " + "unselected edges are colored as background and drawn below selected edges." ), ) p.add_argument( @@ -1300,6 +1596,54 @@ def _validate_args(args, argv): return 0 +def _apply_edge_mode(edge_mode, output_format, edges, nodes, annotations, tensor_table, task_table): + """Select rendered edges and HTML visibility metadata for one mode.""" + hidden_html_edges = set() + visible_html_edge_count = len(edges) + if edge_mode == "full": + return edges, annotations, "deps_viewer", hidden_html_edges, visible_html_edge_count + + is_v2 = edge_mode.endswith("_v2") + if is_v2: + kept, removed, is_dag = _transitive_reduction_v2(edges, nodes, annotations, tensor_table, task_table) + else: + kept, removed, is_dag = _transitive_reduction(edges, nodes, annotations) + if not is_dag: + print( + f"warning: dependency graph has a cycle; transitive reduction skipped, " + f"emitting full graph (--edge-mode {edge_mode} ignored)", + file=sys.stderr, + ) + return edges, annotations, "deps_viewer", hidden_html_edges, visible_html_edge_count + + fmt_task = _make_task_formatter(nodes) + is_reduced = edge_mode.startswith("reduced") + shown = kept if is_reduced else removed + if is_reduced: + prefix = "Lifecycle-aware transitive reduction" if is_v2 else "Transitive reduction" + print(f"{prefix}: removed {len(removed)} redundant edge(s) of {len(edges)} ({len(kept)} kept)") + else: + prefix = "Lifecycle-safe redundant edges only" if is_v2 else "Redundant edges only" + print(f"{prefix}: showing {len(removed)} redundant edge(s) of {len(edges)}") + for pred, succ in removed: + print(f" - {fmt_task(pred)} -> {fmt_task(succ)}") + + if output_format == "html": + hidden_html_edges = set(removed if is_reduced else kept) + visible_html_edge_count = len(shown) + print( + f"HTML layout preserves all {len(edges)} edge(s); " + f"{len(hidden_html_edges)} unselected edge(s) are colored as background" + ) + return edges, annotations, f"deps_viewer_{edge_mode}", hidden_html_edges, visible_html_edge_count + + # Keep only annotations for the selected text edges so summary counts stay + # consistent with the rendered edge set instead of counting hidden rows. + shown_set = set(shown) + annotations = {key: rows for key, rows in annotations.items() if key in shown_set} + return shown, annotations, f"deps_viewer_{edge_mode}", hidden_html_edges, visible_html_edge_count + + def main(argv=None): argv = list(sys.argv[1:] if argv is None else argv) parser = _build_parser() @@ -1324,45 +1668,15 @@ def main(argv=None): if meta: nodes = sorted(set(nodes) | set(meta.keys()), key=_sort_task_id_key) - mode_stem = "deps_viewer" - hidden_html_edges = set() - visible_html_edge_count = len(edges) - if args.edge_mode in ("reduced", "omitted"): - kept, removed, is_dag = _transitive_reduction(edges, nodes) - if not is_dag: - print( - f"warning: dependency graph has a cycle; transitive reduction skipped, " - f"emitting full graph (--edge-mode {args.edge_mode} ignored)", - file=sys.stderr, - ) - else: - fmt_task = _make_task_formatter(nodes) - # reduced keeps the minimal edge set; omitted keeps exactly the - # redundant edges that reduced would drop (the two are complements). - shown = kept if args.edge_mode == "reduced" else removed - if args.edge_mode == "reduced": - print( - f"Transitive reduction: removed {len(removed)} redundant edge(s) of {len(edges)} ({len(kept)} kept)" - ) - else: - print(f"Redundant edges only: showing {len(removed)} redundant edge(s) of {len(edges)}") - for u, v in removed: - print(f" - {fmt_task(u)} -> {fmt_task(v)}") - if args.format == "html": - hidden_html_edges = set(removed if args.edge_mode == "reduced" else kept) - visible_html_edge_count = len(shown) - print( - f"HTML layout preserves all {len(edges)} edge(s); " - f"{len(hidden_html_edges)} unselected edge(s) are colored as background" - ) - else: - edges = shown - # Keep only the annotations of the shown edges so annotated-edge - # counts (emit_text summary, the final print) stay consistent with - # the rendered edge set instead of counting the hidden ones. - shown_set = set(shown) - annotations = {k: v for k, v in annotations.items() if k in shown_set} - mode_stem = f"deps_viewer_{args.edge_mode}" + edges, annotations, mode_stem, hidden_html_edges, visible_html_edge_count = _apply_edge_mode( + args.edge_mode, + args.format, + edges, + nodes, + annotations, + tensor_table, + task_table, + ) out = ( Path(args.output) diff --git a/tests/ut/py/test_deps_viewer.py b/tests/ut/py/test_deps_viewer.py index 2f5becc7c9..6a49a526db 100644 --- a/tests/ut/py/test_deps_viewer.py +++ b/tests/ut/py/test_deps_viewer.py @@ -380,6 +380,17 @@ def test_transitive_reduction_drops_multi_hop_shortcut(): assert kept == [(1, 2), (2, 3), (3, 4)] +def test_transitive_reduction_can_drop_tensormap_shortcut(): + edges = [(1, 2), (2, 3), (1, 3)] + annotations = {(1, 3): [{"source": "tensormap", "tensor_id": 7}]} + + kept, removed, is_dag = deps_viewer._transitive_reduction(edges, [1, 2, 3], annotations) + + assert is_dag + assert removed == [(1, 3)] + assert kept == [(1, 2), (2, 3)] + + def test_transitive_reduction_detects_cycle_and_falls_back(): edges = [(1, 2), (2, 1)] kept, removed, is_dag = deps_viewer._transitive_reduction(edges, [1, 2]) @@ -388,11 +399,229 @@ def test_transitive_reduction_detects_cycle_and_falls_back(): assert removed == [] -def _write_deps_edges(tmp_path, edges): +def _output_lifetime_graph(regions): + """alloc->A->B->C plus direct alloc edges to B/C. + + ``regions`` maps A/B/C task ids to ``(start_offset, shape, strides)`` for + one OUTPUT_EXISTING followed by two INOUT views of alloc-owned storage. + """ + alloc, a, b, c = 1, 2, 3, 4 + edges = [(alloc, a), (alloc, b), (alloc, c), (a, b), (b, c)] + annotations = {} + task_table = {} + for task_id in (a, b, c): + start_offset, shape, strides = regions[task_id] + task_table[task_id] = { + "task_id": task_id, + "args": [ + { + "idx": 0, + "type": "OUTPUT_EXISTING" if task_id == a else "INOUT", + "tensor_id": 7, + "dtype": "FLOAT32", + "shape": shape, + "start_offset": str(start_offset), + "strides": strides, + } + ], + } + annotations[(alloc, task_id)] = [ + { + "pred": alloc, + "succ": task_id, + "arg": 0, + "source": "creator", + "tensor_id": 7, + "consumer_dtype": "FLOAT32", + "consumer_shape": shape, + "consumer_start_offset": str(start_offset), + "consumer_strides": strides, + } + ] + for pred, succ in ((a, b), (b, c)): + pred_offset, pred_shape, pred_strides = regions[pred] + succ_offset, succ_shape, succ_strides = regions[succ] + annotations[(pred, succ)] = [ + { + "pred": pred, + "succ": succ, + "arg": 0, + "source": "tensormap", + "overlap": "covered", + "tensor_id": 7, + "consumer_dtype": "FLOAT32", + "consumer_shape": succ_shape, + "consumer_start_offset": str(succ_offset), + "consumer_strides": succ_strides, + "producer_shape": pred_shape, + "producer_start_offset": str(pred_offset), + "producer_strides": pred_strides, + } + ] + tensor_table = { + 7: { + "tensor_id": 7, + "buffer_addr": "4096", + "buffer_numel": "128", + "dtype": "FLOAT32", + } + } + return edges, [alloc, a, b, c], annotations, tensor_table, task_table + + +def test_transitive_reduction_v2_drops_only_full_range_middle_output(): + graph = _output_lifetime_graph( + { + 2: (0, [2, 4], [4, 1]), + 3: (0, [8], [1]), + 4: (0, [4, 2], [2, 1]), + } + ) + + kept, removed, is_dag = deps_viewer._transitive_reduction_v2(*graph) + + assert is_dag + assert removed == [(1, 3)] + assert kept == [(1, 2), (1, 4), (2, 3), (3, 4)] + + +def test_transitive_reduction_v2_does_not_cross_output_existing_reuse_boundary(): + alloc, a1, b1, a2, b2 = 1, 2, 3, 4, 5 + edges = [ + (alloc, a1), + (alloc, b1), + (alloc, a2), + (alloc, b2), + (a1, b1), + (b1, a2), + (a2, b2), + ] + annotations = {} + task_table = {} + for task_id, arg_type in ( + (a1, "OUTPUT_EXISTING"), + (b1, "INOUT"), + (a2, "OUTPUT_EXISTING"), + (b2, "INOUT"), + ): + task_table[task_id] = { + "task_id": task_id, + "args": [ + { + "idx": 0, + "type": arg_type, + "tensor_id": 7, + "dtype": "FLOAT32", + "shape": [8], + "start_offset": "0", + "strides": [1], + } + ], + } + annotations[(alloc, task_id)] = [ + { + "pred": alloc, + "succ": task_id, + "arg": 0, + "source": "creator", + "tensor_id": 7, + "consumer_dtype": "FLOAT32", + "consumer_shape": [8], + "consumer_start_offset": "0", + "consumer_strides": [1], + } + ] + for pred, succ in ((a1, b1), (a2, b2)): + annotations[(pred, succ)] = [ + { + "pred": pred, + "succ": succ, + "arg": 0, + "source": "tensormap", + "overlap": "covered", + "tensor_id": 7, + "consumer_dtype": "FLOAT32", + "consumer_shape": [8], + "consumer_start_offset": "0", + "consumer_strides": [1], + "producer_shape": [8], + "producer_start_offset": "0", + "producer_strides": [1], + } + ] + annotations[(b1, a2)] = [{"pred": b1, "succ": a2, "arg": -1, "source": "explicit"}] + tensor_table = { + 7: { + "tensor_id": 7, + "buffer_addr": "4096", + "buffer_numel": "128", + "dtype": "FLOAT32", + } + } + + kept, removed, is_dag = deps_viewer._transitive_reduction_v2( + edges, [alloc, a1, b1, a2, b2], annotations, tensor_table, task_table + ) + + assert is_dag + assert removed == [] + assert kept == edges + + +def _write_output_lifetime_deps(tmp_path): + edges, nodes, annotations, tensor_table, task_table = _output_lifetime_graph( + { + 2: (0, [8], [1]), + 3: (0, [8], [1]), + 4: (0, [8], [1]), + } + ) + annotated_edges = [] + for pred, succ in edges: + rows = annotations.get((pred, succ)) + if rows: + annotated_edges.extend(rows) + else: + annotated_edges.append({"pred": str(pred), "succ": str(succ), "arg": -1, "source": "explicit"}) + deps = { + "tasks": list(task_table.values()), + "tensors": list(tensor_table.values()), + "edges": annotated_edges, + } + path = tmp_path / "deps.json" + path.write_text(json.dumps(deps)) + return path, nodes + + +def test_main_v2_edge_modes_keep_lifetime_boundaries_and_omit_middle(tmp_path): + deps_path, _nodes = _write_output_lifetime_deps(tmp_path) + reduced_out = tmp_path / "reduced_v2.txt" + omitted_out = tmp_path / "omitted_v2.txt" + + reduced_rc = deps_viewer.main([str(deps_path), "--edge-mode", "reduced_v2", "-o", str(reduced_out)]) + omitted_rc = deps_viewer.main([str(deps_path), "--edge-mode", "omitted_v2", "-o", str(omitted_out)]) + + assert reduced_rc == 0 + assert omitted_rc == 0 + assert "unique_task_edges: 4" in reduced_out.read_text() + assert "unique_task_edges: 1" in omitted_out.read_text() + assert "=== TASK 1" in omitted_out.read_text() + assert " -> 3" in omitted_out.read_text() + + +def _write_deps_edges(tmp_path, edges, sources=None): + sources = sources or {} + annotated_edges = [] + for pred, succ in edges: + source = sources.get((pred, succ), "explicit") + edge = {"pred": str(pred), "succ": str(succ), "arg": -1 if source == "explicit" else 0, "source": source} + if source != "explicit": + edge["tensor_id"] = "7" + annotated_edges.append(edge) deps = { "tasks": [], "tensors": [], - "edges": [{"pred": str(p), "succ": str(s), "arg": 0, "source": "creator"} for p, s in edges], + "edges": annotated_edges, } path = tmp_path / "deps.json" path.write_text(json.dumps(deps)) @@ -420,6 +649,22 @@ def test_main_reduced_mode_drops_edge_and_prints_removed(tmp_path, capsys): assert "(1, 1) -> (1, 3)" in stdout +def test_main_reduced_mode_keeps_transitive_creator_reference(tmp_path, capsys): + # Mirrors vector_example's t0 -> t1 -> t3 -> t4 path plus t0 -> t4: + # t4 directly reads tensor c created by t0, so that creator-lifetime edge + # must survive even though the task ordering is transitively covered. + t0, t1, t3, t4 = 1, 2, 3, 4 + edges = [(t0, t1), (t1, t3), (t3, t4), (t0, t4)] + deps_path = _write_deps_edges(tmp_path, edges, sources={(t0, t4): "creator"}) + out = tmp_path / "graph.txt" + + rc = deps_viewer.main([str(deps_path), "--edge-mode", "reduced", "-o", str(out)]) + + assert rc == 0 + assert "unique_task_edges: 4" in out.read_text() + assert "removed 0 redundant edge(s)" in capsys.readouterr().out + + def test_main_full_mode_keeps_all_edges(tmp_path, capsys): ring = 1 << 32 a, b, c = ring + 1, ring + 2, ring + 3 @@ -456,6 +701,8 @@ def fake_emit_html(edges, _nodes, _meta, **kwargs): cases = [ ("reduced", {(a, c)}, 2), ("omitted", {(a, b), (b, c)}, 1), + ("reduced_v2", {(a, c)}, 2), + ("omitted_v2", {(a, b), (b, c)}, 1), ] for mode, expected_hidden, expected_visible_count in cases: rc = deps_viewer.main(