已合并
[Feature] generate aligned trace in NPUGraph debug dump #43288
Lyric创建于 7月29日
[Feature] generate aligned trace in NPUGraph debug dump #43288
已合并
共 3 个文件变更+495-3
| @@ -1,6 +1,5 @@ | |||
| 1 | -import unittest | 1 | +from unittest import mock |
| 2 | from dataclasses import dataclass | 2 | from dataclasses import dataclass |
| 3 | -from itertools import chain | ||
| 4 | import math | 3 | import math |
| 5 | import os | 4 | import os |
| 6 | 5 | ||
| @@ -329,11 +328,52 @@ class TestIFAAclgraphUpdate(TestCase): | |||
| 329 | static_y_pred = model(static_input) | 328 | static_y_pred = model(static_input) |
| 330 | 329 | ||
| 331 | file_path = os.path.join(os.getcwd(), "jsonPrint.json") | 330 | file_path = os.path.join(os.getcwd(), "jsonPrint.json") |
| 331 | + aligned_file_path = os.path.join(os.getcwd(), "jsonPrint.aligned.json") | ||
| 332 | if os.path.exists(file_path) and os.path.isfile(file_path): | 332 | if os.path.exists(file_path) and os.path.isfile(file_path): |
| 333 | os.remove(file_path) | 333 | os.remove(file_path) |
| 334 | + if os.path.exists(aligned_file_path) and os.path.isfile(aligned_file_path): | ||
| 335 | + os.remove(aligned_file_path) | ||
| 334 | 336 | ||
| 335 | g.debug_dump(file_path) | 337 | g.debug_dump(file_path) |
| 338 | + self.assertTrue(os.path.exists(file_path), "npugraph debug dump file does not exist") | ||
| 336 | self.assertTrue(os.path.getsize(file_path) > 0, "npugraph debug dump assert error") | 339 | self.assertTrue(os.path.getsize(file_path) > 0, "npugraph debug dump assert error") |
| 340 | + self.assertTrue(os.path.exists(aligned_file_path), "npugraph aligned debug dump file does not exist") | ||
| 341 | + self.assertTrue( | ||
| 342 | + os.path.getsize(aligned_file_path) > 0, | ||
| 343 | + "npugraph aligned debug dump assert error", | ||
| 344 | + ) | ||
| 345 | + | ||
| 346 | + os.remove(aligned_file_path) | ||
| 347 | + failure_graph = torch.npu.NPUGraph() | ||
| 348 | + with torch.npu.graph(failure_graph): | ||
| 349 | + static_y_pred = model(static_input) | ||
| 350 | + failure_file_path = os.path.join(os.getcwd(), "jsonPrint.failure.json") | ||
| 351 | + failure_aligned_file_path = os.path.join(os.getcwd(), "jsonPrint.failure.aligned.json") | ||
| 352 | + if os.path.exists(failure_file_path) and os.path.isfile(failure_file_path): | ||
| 353 | + os.remove(failure_file_path) | ||
| 354 | + if os.path.exists(failure_aligned_file_path) and os.path.isfile(failure_aligned_file_path): | ||
| 355 | + os.remove(failure_aligned_file_path) | ||
| 356 | + with mock.patch( | ||
| 357 | + "torch_npu.npu._npugraph_utils.align_trace_json", | ||
| 358 | + side_effect=RuntimeError("alignment failed"), | ||
| 359 | + ) as mock_align: | ||
| 360 | + result = failure_graph.debug_dump(failure_file_path) | ||
| 361 | + self.assertIsNone(result) | ||
| 362 | + mock_align.assert_called_once_with(failure_file_path, None) | ||
| 363 | + self.assertTrue(os.path.exists(failure_file_path), "original graph dump file does not exist") | ||
| 364 | + self.assertTrue(os.path.getsize(failure_file_path) > 0, "original graph dump was not preserved") | ||
| 365 | + self.assertFalse(os.path.exists(failure_aligned_file_path)) | ||
| 366 | + os.remove(failure_file_path) | ||
| 367 | + | ||
| 368 | + with mock.patch( | ||
| 369 | + "torch_npu.npu._npugraph_utils.align_trace_json", | ||
| 370 | + ) as mock_align: | ||
| 371 | + result = torch.npu.NPUGraph().debug_dump(file_path) | ||
| 372 | + self.assertIsNone(result) | ||
| 373 | + mock_align.assert_not_called() | ||
| 374 | + self.assertTrue(os.path.exists(file_path), "existing graph dump file does not exist") | ||
| 375 | + self.assertTrue(os.path.getsize(file_path) > 0, "existing graph dump was not preserved") | ||
| 376 | + | ||
| 337 | os.remove(file_path) | 377 | os.remove(file_path) |
| 338 | 378 | ||
| 339 | 379 | ||
| @@ -0,0 +1,414 @@ | |||
| 1 | +"""Align NPUGraph trace events for clearer Chrome Trace visualization. | ||
| 2 | + | ||
| 3 | +The utility merges virtual streams, rebuilds per-stream timestamps, aligns | ||
| 4 | +cross-stream record/wait pairs, adds flow events, and validates the result. | ||
| 5 | +The generated ``*.aligned.json`` is derived data; the raw dump is unchanged. | ||
| 6 | +""" | ||
| 7 | + | ||
| 8 | +import json | ||
| 9 | +import re | ||
| 10 | +from pathlib import Path as _Path | ||
| 11 | + | ||
| 12 | +import torch | ||
| 13 | + | ||
| 14 | + | ||
| 15 | +log = torch._logging.getArtifactLogger("torch_npu.npugraph", "cudagraphs") | ||
| 16 | + | ||
| 17 | + | ||
| 18 | +# Virtual stream preprocessing. | ||
| 19 | + | ||
| 20 | + | ||
| 21 | +def _stream_id_to_tid(stream_id): | ||
| 22 | + """Convert a numeric stream id to a trace tid.""" | ||
| 23 | + return f"stream{stream_id}" | ||
| 24 | + | ||
| 25 | + | ||
| 26 | +def _tid_to_stream_id(tid): | ||
| 27 | + """Convert a trace tid to a numeric stream id, or return None.""" | ||
| 28 | + if tid.startswith("stream"): | ||
| 29 | + try: | ||
| 30 | + return int(tid[len("stream") :]) | ||
| 31 | + except ValueError: | ||
| 32 | + return None | ||
| 33 | + return None | ||
| 34 | + | ||
| 35 | + | ||
| 36 | +def _merge_stream_active(events, gap=0.0): | ||
| 37 | + """Merge streams activated by STREAM_ACTIVE events in place. | ||
| 38 | + | ||
| 39 | + Returns the modified events and the ``(source_tid, target_tid)`` pairs | ||
| 40 | + that were merged. | ||
| 41 | + """ | ||
| 42 | + tid_events = {} | ||
| 43 | + for ev in events: | ||
| 44 | + tid_events.setdefault(ev.get("tid"), []).append(ev) | ||
| 45 | + for lst in tid_events.values(): | ||
| 46 | + lst.sort(key=lambda e: e.get("ts", 0)) | ||
| 47 | + | ||
| 48 | + stream_active = [ | ||
| 49 | + ev | ||
| 50 | + for ev in events | ||
| 51 | + if ((ev.get("args") or {}).get("Task Type") == "STREAM_ACTIVE" or ev.get("name") == "STREAM_ACTIVE") | ||
| 52 | + and "Active Stream Id" in (ev.get("args") or {}) | ||
| 53 | + ] | ||
| 54 | + stream_active.sort(key=lambda e: e.get("ts", 0)) | ||
| 55 | + | ||
| 56 | + moved = [] | ||
| 57 | + for ev in stream_active: | ||
| 58 | + args = ev.get("args") or {} # pylint: disable=redefined-outer-name | ||
| 59 | + active_id = args.get("Active Stream Id") | ||
| 60 | + if active_id is None: | ||
| 61 | + continue | ||
| 62 | + source_tid = _stream_id_to_tid(active_id) | ||
| 63 | + target_tid = ev.get("tid") | ||
| 64 | + if not target_tid or source_tid == target_tid: | ||
| 65 | + continue | ||
| 66 | + if source_tid not in tid_events: | ||
| 67 | + continue | ||
| 68 | + target_evs = tid_events.get(target_tid, []) | ||
| 69 | + source_evs = tid_events.get(source_tid, []) | ||
| 70 | + if not source_evs: | ||
| 71 | + continue | ||
| 72 | + target_end = max(e.get("ts", 0) + e.get("dur", 0) for e in target_evs) if target_evs else 0 | ||
| 73 | + source_min = min(e.get("ts", 0) for e in source_evs) | ||
| 74 | + offset = (target_end + gap) - source_min | ||
| 75 | + target_stream_id = _tid_to_stream_id(target_tid) | ||
| 76 | + for e in source_evs: | ||
| 77 | + e["ts"] = e.get("ts", 0) + offset | ||
| 78 | + e["tid"] = target_tid | ||
| 79 | + if target_stream_id is not None: | ||
| 80 | + args_e = e.get("args") | ||
| 81 | + if isinstance(args_e, dict) and "Stream Id" in args_e: | ||
| 82 | + args_e["Stream Id"] = target_stream_id | ||
| 83 | + tid_events[target_tid] = target_evs + source_evs | ||
| 84 | + tid_events.pop(source_tid, None) | ||
| 85 | + moved.append((source_tid, target_tid)) | ||
| 86 | + | ||
| 87 | + return events, moved | ||
| 88 | + | ||
| 89 | + | ||
| 90 | +# Control-event prefixes are excluded from name cleanup. | ||
| 91 | +CONTROL_PREFIXES = ( | ||
| 92 | + "EVENT_RECORD_", | ||
| 93 | + "EVENT_WAIT_", | ||
| 94 | + "MEM_WRITE_VALUE_", | ||
| 95 | + "MEM_WAIT_VALUE_", | ||
| 96 | + "EVENT_RESET_", | ||
| 97 | +) | ||
| 98 | + | ||
| 99 | + | ||
| 100 | +def _parse_control(name: str): | ||
| 101 | + """Return the control-event kind and handle, or ``(None, None)``.""" | ||
| 102 | + if name.startswith("EVENT_RECORD_"): | ||
| 103 | + return "EVENT_RECORD", name.split("_")[-1] | ||
| 104 | + if name.startswith("EVENT_WAIT_"): | ||
| 105 | + return "EVENT_WAIT", name.split("_")[-1] | ||
| 106 | + if name.startswith("MEM_WRITE_VALUE_"): | ||
| 107 | + return "MEM_WRITE_VALUE", name.split("_")[-1] | ||
| 108 | + if name.startswith("MEM_WAIT_VALUE_"): | ||
| 109 | + return "MEM_WAIT_VALUE", name.split("_")[-1] | ||
| 110 | + return None, None | ||
| 111 | + | ||
| 112 | + | ||
| 113 | +def _clean_name(name: str) -> str: | ||
| 114 | + """Remove hash-like underscore-separated components from an operator name.""" | ||
| 115 | + parts = name.split("_") | ||
| 116 | + new_parts = [] | ||
| 117 | + for part in parts: | ||
| 118 | + if re.fullmatch(r"[a-z0-9]+", part) and sum(c.isdigit() for c in part) > 5: | ||
| 119 | + continue | ||
| 120 | + new_parts.append(part) | ||
| 121 | + return "_".join(new_parts) if new_parts else name | ||
| 122 | + | ||
| 123 | + | ||
| 124 | +def _align_trace(events): | ||
| 125 | + """Align events and return aligned events plus pair diagnostics.""" | ||
| 126 | + min_gap = 0.2 # Fixed visualization gap within a stream. | ||
| 127 | + eps = 1e-9 # Floating-point tolerance. | ||
| 128 | + | ||
| 129 | + # Normalize reset duration. | ||
| 130 | + for ev in events: | ||
| 131 | + name = ev.get("name", "") | ||
| 132 | + if name.startswith("EVENT_RESET_"): | ||
| 133 | + ev["dur"] = 0.2 | ||
| 134 | + | ||
| 135 | + # Group events by tid and preserve their original order. | ||
| 136 | + tid_indices = {} | ||
| 137 | + for idx, ev in enumerate(events): | ||
| 138 | + tid = ev.get("tid") | ||
| 139 | + tid_indices.setdefault(tid, []).append(idx) | ||
| 140 | + for tid, indices in tid_indices.items(): | ||
| 141 | + indices.sort(key=lambda i: (events[i].get("ts", 0), i)) | ||
| 142 | + | ||
| 143 | + tid_pos = {tid: {idx: pos for pos, idx in enumerate(indices)} for tid, indices in tid_indices.items()} | ||
| 144 | + | ||
| 145 | + # Rebuild timestamps from zero with a fixed gap. | ||
| 146 | + for tid, indices in tid_indices.items(): | ||
| 147 | + prev_end = 0.0 | ||
| 148 | + for pos, idx in enumerate(indices): | ||
| 149 | + ev = events[idx] | ||
| 150 | + gap = 0.0 if pos == 0 else min_gap | ||
| 151 | + ev["ts"] = prev_end + gap | ||
| 152 | + prev_end = ev["ts"] + ev["dur"] | ||
| 153 | + | ||
| 154 | + # Pair control events by handle. | ||
| 155 | + pair_map = {} | ||
| 156 | + for idx, ev in enumerate(events): | ||
| 157 | + kind, cid = _parse_control(ev.get("name", "")) | ||
| 158 | + if not kind: | ||
| 159 | + continue | ||
| 160 | + slot = pair_map.setdefault(cid, {}) | ||
| 161 | + if kind in ("EVENT_RECORD", "MEM_WRITE_VALUE"): | ||
| 162 | + slot["record"] = idx | ||
| 163 | + elif kind in ("EVENT_WAIT", "MEM_WAIT_VALUE"): | ||
| 164 | + slot["wait"] = idx | ||
| 165 | + | ||
| 166 | + # Separate same-stream pairs from cross-stream pairs. | ||
| 167 | + same_tid_pairs = [] | ||
| 168 | + cross_tid_pairs = [] | ||
| 169 | + for cid, slot in pair_map.items(): | ||
| 170 | + if "record" in slot and "wait" in slot: | ||
| 171 | + r_idx = slot["record"] | ||
| 172 | + w_idx = slot["wait"] | ||
| 173 | + if events[r_idx]["tid"] == events[w_idx]["tid"]: | ||
| 174 | + same_tid_pairs.append((events[r_idx]["name"], events[w_idx]["name"])) | ||
| 175 | + else: | ||
| 176 | + cross_tid_pairs.append((cid, r_idx, w_idx)) | ||
| 177 | + | ||
| 178 | + # Iteratively align cross-stream pairs. | ||
| 179 | + max_passes = 200 | ||
| 180 | + for _ in range(max_passes): | ||
| 181 | + changed = False | ||
| 182 | + for cid, r_idx, w_idx in cross_tid_pairs: | ||
| 183 | + record = events[r_idx] | ||
| 184 | + wait = events[w_idx] | ||
| 185 | + record_end = record["ts"] + record["dur"] | ||
| 186 | + if wait["ts"] >= record_end - eps: | ||
| 187 | + continue # The dependency is already satisfied. | ||
| 188 | + needed_dur = record_end - wait["ts"] | ||
| 189 | + delta = needed_dur - wait["dur"] | ||
| 190 | + if abs(delta) > eps: | ||
| 191 | + wait["dur"] = needed_dur | ||
| 192 | + # Shift tasks after the wait by the same delta. | ||
| 193 | + tid = wait["tid"] | ||
| 194 | + start_pos = tid_pos[tid][w_idx] | ||
| 195 | + indices = tid_indices[tid] | ||
| 196 | + for idx in indices[start_pos + 1 :]: | ||
| 197 | + events[idx]["ts"] += delta | ||
| 198 | + changed = True | ||
| 199 | + if not changed: | ||
| 200 | + break | ||
| 201 | + else: | ||
| 202 | + raise RuntimeError("Alignment did not converge within pass limit.") | ||
| 203 | + | ||
| 204 | + # Record pairs that already satisfy the dependency. | ||
| 205 | + skipped_pairs = [] | ||
| 206 | + for cid, r_idx, w_idx in cross_tid_pairs: | ||
| 207 | + record = events[r_idx] | ||
| 208 | + wait = events[w_idx] | ||
| 209 | + if wait["ts"] >= record["ts"] + record["dur"] - eps: | ||
| 210 | + skipped_pairs.append((record["name"], wait["name"])) | ||
| 211 | + | ||
| 212 | + # Add Chrome Trace flow start/end and instant marker events. | ||
| 213 | + flow_events = [] | ||
| 214 | + for cid, slot in pair_map.items(): | ||
| 215 | + if "record" in slot and "wait" in slot: | ||
| 216 | + r = events[slot["record"]] | ||
| 217 | + w = events[slot["wait"]] | ||
| 218 | + try: | ||
| 219 | + flow_id = int(cid) | ||
| 220 | + except ValueError: | ||
| 221 | + flow_id = abs(hash(cid)) % 1000000000000 | ||
| 222 | + flow_events.append( | ||
| 223 | + { | ||
| 224 | + "name": f"PAIR_{cid}", | ||
| 225 | + "cat": "event_record->wait", | ||
| 226 | + "ph": "s", | ||
| 227 | + "pid": r.get("pid"), | ||
| 228 | + "tid": r.get("tid"), | ||
| 229 | + "ts": r["ts"] + r["dur"], | ||
| 230 | + "id": flow_id, | ||
| 231 | + "bp": "e", | ||
| 232 | + } | ||
| 233 | + ) | ||
| 234 | + flow_events.append( | ||
| 235 | + { | ||
| 236 | + "name": f"PAIR_{cid}", | ||
| 237 | + "cat": "event_record->wait", | ||
| 238 | + "ph": "f", | ||
| 239 | + "pid": w.get("pid"), | ||
| 240 | + "tid": w.get("tid"), | ||
| 241 | + "ts": w["ts"], | ||
| 242 | + "id": flow_id, | ||
| 243 | + "bp": "e", | ||
| 244 | + } | ||
| 245 | + ) | ||
| 246 | + flow_events.append( | ||
| 247 | + { | ||
| 248 | + "name": f"PAIR_{cid}_START", | ||
| 249 | + "cat": "event_record->wait", | ||
| 250 | + "ph": "i", | ||
| 251 | + "s": "t", | ||
| 252 | + "pid": r.get("pid"), | ||
| 253 | + "tid": r.get("tid"), | ||
| 254 | + "ts": r["ts"] + r["dur"], | ||
| 255 | + } | ||
| 256 | + ) | ||
| 257 | + flow_events.append( | ||
| 258 | + { | ||
| 259 | + "name": f"PAIR_{cid}_END", | ||
| 260 | + "cat": "event_record->wait", | ||
| 261 | + "ph": "i", | ||
| 262 | + "s": "t", | ||
| 263 | + "pid": w.get("pid"), | ||
| 264 | + "tid": w.get("tid"), | ||
| 265 | + "ts": w["ts"], | ||
| 266 | + } | ||
| 267 | + ) | ||
| 268 | + | ||
| 269 | + # Clean non-control operator names. | ||
| 270 | + for ev in events: | ||
| 271 | + name = ev.get("name") | ||
| 272 | + if isinstance(name, str): | ||
| 273 | + if name.startswith(CONTROL_PREFIXES): | ||
| 274 | + continue | ||
| 275 | + ev["name"] = _clean_name(name) | ||
| 276 | + | ||
| 277 | + events.extend(flow_events) | ||
| 278 | + | ||
| 279 | + return events, same_tid_pairs, skipped_pairs | ||
| 280 | + | ||
| 281 | + | ||
| 282 | +def _validate_alignment(events, skipped_pairs=None): | ||
| 283 | + import math | ||
| 284 | + | ||
| 285 | + if skipped_pairs is None: | ||
| 286 | + skipped_pairs = [] | ||
| 287 | + skipped_set = {tuple(p) for p in skipped_pairs} | ||
| 288 | + | ||
| 289 | + non_finite = [] | ||
| 290 | + for ev in events: | ||
| 291 | + if ev.get("ph") != "X": | ||
| 292 | + continue | ||
| 293 | + ts = ev.get("ts") | ||
| 294 | + dur = ev.get("dur") | ||
| 295 | + if isinstance(ts, float) and not math.isfinite(ts): | ||
| 296 | + non_finite.append(("ts", ev.get("name"))) | ||
| 297 | + if isinstance(dur, float) and not math.isfinite(dur): | ||
| 298 | + non_finite.append(("dur", ev.get("name"))) | ||
| 299 | + if non_finite: | ||
| 300 | + preview = "\n".join(f"{field} {name}" for field, name in non_finite[:10]) | ||
| 301 | + raise RuntimeError(f"Alignment check failed: non-finite values detected ({len(non_finite)}).\n{preview}") | ||
| 302 | + | ||
| 303 | + pairs = {} | ||
| 304 | + for idx, ev in enumerate(events): | ||
| 305 | + if ev.get("ph") != "X": | ||
| 306 | + continue | ||
| 307 | + kind, cid = _parse_control(ev.get("name", "")) | ||
| 308 | + if not kind: | ||
| 309 | + continue | ||
| 310 | + slot = pairs.setdefault(cid, {}) | ||
| 311 | + if kind in ("EVENT_RECORD", "MEM_WRITE_VALUE"): | ||
| 312 | + slot["record"] = idx | ||
| 313 | + elif kind in ("EVENT_WAIT", "MEM_WAIT_VALUE"): | ||
| 314 | + slot["wait"] = idx | ||
| 315 | + | ||
| 316 | + mismatches = [] | ||
| 317 | + for cid, slot in pairs.items(): | ||
| 318 | + if "record" in slot and "wait" in slot: | ||
| 319 | + r = events[slot["record"]] | ||
| 320 | + w = events[slot["wait"]] | ||
| 321 | + if r.get("tid") == w.get("tid"): | ||
| 322 | + continue | ||
| 323 | + if (r.get("name"), w.get("name")) in skipped_set: | ||
| 324 | + continue | ||
| 325 | + r_end = r["ts"] + r["dur"] | ||
| 326 | + w_end = w["ts"] + w["dur"] | ||
| 327 | + if w["ts"] < r_end - 1e-6 and abs(r_end - w_end) > 1e-6: | ||
| 328 | + mismatches.append((cid, r_end, w_end)) | ||
| 329 | + | ||
| 330 | + overlaps = [] | ||
| 331 | + tid_indices = {} | ||
| 332 | + for idx, ev in enumerate(events): | ||
| 333 | + if ev.get("ph") != "X": | ||
| 334 | + continue | ||
| 335 | + tid = ev.get("tid") | ||
| 336 | + tid_indices.setdefault(tid, []).append(idx) | ||
| 337 | + for tid, indices in tid_indices.items(): | ||
| 338 | + indices.sort(key=lambda i: (events[i].get("ts", 0), i)) | ||
| 339 | + prev_end = None | ||
| 340 | + for idx in indices: | ||
| 341 | + ev = events[idx] | ||
| 342 | + ts = ev.get("ts", 0) | ||
| 343 | + dur = ev.get("dur", 0) | ||
| 344 | + if prev_end is None: | ||
| 345 | + if abs(ts - 0.0) > 1e-6: | ||
| 346 | + overlaps.append((tid, None, ev.get("name"))) | ||
| 347 | + prev_end = ts + dur | ||
| 348 | + continue | ||
| 349 | + expected_ts = prev_end + 0.2 | ||
| 350 | + if abs(ts - expected_ts) > 1e-6: | ||
| 351 | + overlaps.append((tid, None, ev.get("name"))) | ||
| 352 | + if len(overlaps) >= 10: | ||
| 353 | + break | ||
| 354 | + prev_end = ts + dur | ||
| 355 | + if len(overlaps) >= 10: | ||
| 356 | + break | ||
| 357 | + | ||
| 358 | + if overlaps: | ||
| 359 | + preview = "\n".join(f"{tid}: {right}" for tid, _, right in overlaps) | ||
| 360 | + raise RuntimeError(f"Alignment check failed: gap mismatch detected ({len(overlaps)}+).\n{preview}") | ||
| 361 | + | ||
| 362 | + return mismatches | ||
| 363 | + | ||
| 364 | + | ||
| 365 | +def align_trace_json(input, output, merge_gap=0.0): # pylint: disable=redefined-builtin | ||
| 366 | + in_path = _Path(input) | ||
| 367 | + out_path = _Path(output) if output else in_path.with_suffix(in_path.suffix.replace(".json", "") + ".aligned.json") | ||
| 368 | + | ||
| 369 | + with in_path.open("r", encoding="utf-8") as f: | ||
| 370 | + events = json.load(f) | ||
| 371 | + | ||
| 372 | + # Merge virtual streams before alignment. | ||
| 373 | + events, moved = _merge_stream_active(events, gap=merge_gap) | ||
| 374 | + if moved: | ||
| 375 | + log.debug("Merged streams: %s", moved) | ||
| 376 | + | ||
| 377 | + aligned, same_tid_pairs, skipped_pairs = _align_trace(events) | ||
| 378 | + mismatches = _validate_alignment(aligned, skipped_pairs=skipped_pairs) | ||
| 379 | + if mismatches: | ||
| 380 | + preview = "\n".join(f"{cid} record_end={r_end} wait_end={w_end}" for cid, r_end, w_end in mismatches[:10]) | ||
| 381 | + raise RuntimeError(f"Alignment check failed: {len(mismatches)} mismatches.\n{preview}") | ||
| 382 | + if same_tid_pairs: | ||
| 383 | + log.debug("Same-tid pairs (no alignment): %s", same_tid_pairs) | ||
| 384 | + if skipped_pairs: | ||
| 385 | + log.debug( | ||
| 386 | + "Cross-tid pairs skipped (wait starts after record end): %s", | ||
| 387 | + skipped_pairs, | ||
| 388 | + ) | ||
| 389 | + | ||
| 390 | + with out_path.open("w", encoding="utf-8") as f: | ||
| 391 | + json.dump(aligned, f, ensure_ascii=False, indent=2) | ||
| 392 | + | ||
| 393 | + log.debug("Aligned trace was written to %s", out_path) | ||
| 394 | + return out_path | ||
| 395 | + | ||
| 396 | + | ||
| 397 | +if __name__ == "__main__": | ||
| 398 | + import argparse | ||
| 399 | + | ||
| 400 | + parser = argparse.ArgumentParser(description="Align trace control tasks and clean task names.") | ||
| 401 | + parser.add_argument("input", help="input trace json path") | ||
| 402 | + parser.add_argument( | ||
| 403 | + "-o", | ||
| 404 | + "--output", | ||
| 405 | + help="output json path (default: input with .aligned suffix)", | ||
| 406 | + ) | ||
| 407 | + parser.add_argument( | ||
| 408 | + "--merge-gap", | ||
| 409 | + type=float, | ||
| 410 | + default=0.0, | ||
| 411 | + help="time gap inserted when merging virtual streams (default: 0.0)", | ||
| 412 | + ) | ||
| 413 | + args = parser.parse_args() | ||
| 414 | + align_trace_json(args.input, args.output, args.merge_gap) | ||
| @@ -667,8 +667,46 @@ class NPUGraph(torch_npu._C._NPUGraph): | |||
| 667 | 667 | ||
| 668 | Arguments: | 668 | Arguments: |
| 669 | debug_path (required): Path to dump the graph to. | 669 | debug_path (required): Path to dump the graph to. |
| 670 | + | ||
| 671 | + An additional ``.aligned.json`` file is generated for visualization | ||
| 672 | + when the graph is successfully dumped. | ||
| 670 | """ | 673 | """ |
| 671 | - return super().debug_dump(debug_path) | 674 | + # C++ debug_dump does not rewrite the file without a successful capture, |
| 675 | + # so avoid aligning a stale dump from an earlier call. | ||
| 676 | + def file_state(path): | ||
| 677 | + try: | ||
| 678 | + stat = os.stat(path) | ||
| 679 | + return ( | ||
| 680 | + stat.st_dev, | ||
| 681 | + stat.st_ino, | ||
| 682 | + stat.st_size, | ||
| 683 | + stat.st_mtime_ns, | ||
| 684 | + stat.st_ctime_ns, | ||
| 685 | + ) | ||
| 686 | + except Exception: # The optional post-processing must not affect debug_dump. | ||
| 687 | + return None | ||
| 688 | + | ||
| 689 | + previous_file_state = file_state(debug_path) | ||
| 690 | + | ||
| 691 | + result = super().debug_dump(debug_path) | ||
| 692 | + | ||
| 693 | + current_file_state = file_state(debug_path) | ||
| 694 | + if current_file_state is None or current_file_state == previous_file_state: | ||
| 695 | + return result | ||
| 696 | + | ||
| 697 | + try: | ||
| 698 | + from ._npugraph_utils import align_trace_json | ||
| 699 | + | ||
| 700 | + align_trace_json(debug_path, None) | ||
| 701 | + except Exception: | ||
| 702 | + log.warning( | ||
| 703 | + "Failed to generate aligned trace for %s; an existing aligned " | ||
| 704 | + "trace, if any, was not updated", | ||
| 705 | + debug_path, | ||
| 706 | + exc_info=True, | ||
| 707 | + ) | ||
| 708 | + | ||
| 709 | + return result | ||
| 672 | 710 | ||
| 673 | def super_kernel_optimize(self, optimize_options=None, debug_options=None): | 711 | def super_kernel_optimize(self, optimize_options=None, debug_options=None): |
| 674 | r"""Calls a function to optimize graph by super kernel. | 712 | r"""Calls a function to optimize graph by super kernel. |