Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / diagnostics / dump / api.py: 90%
311 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 05:03 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""Schedule-driven Tensor dump for Torch MegaKernel invocations."""
17from __future__ import annotations
19from collections.abc import Callable, Collection, Sequence
20from dataclasses import dataclass
21import json
22from pathlib import Path
23from typing import Any
25import torch
27from hyper_multicore.diagnostics.common.session import (
28 _activate_session,
29 _deactivate_session,
30)
31from hyper_multicore.diagnostics.common.call import DiagnosticCall, LaunchArguments, LaunchContext
32from hyper_multicore.diagnostics.common.schedule import DiagnosticAction
33from hyper_multicore.diagnostics.dump.layout import _SAFE_FILE_NAME, _stage_names
34from hyper_multicore.diagnostics.dump.l0 import (
35 _GraphTensor, _TensorSnapshot, _describe_graph_tensors, _snapshot_tensor, _write_l0_stage_files,
36)
37from hyper_multicore.diagnostics.dump.l1 import (
38 _L1Snapshot,
39 _PreparedL1Dump,
40 _decode_l1_dump,
41 _prepare_l1_dump,
42)
45__all__ = [
46 "TensorDumpConfig",
47 "TorchMegaKernelTensorDumper",
48 "tensor_dump",
49]
52_VALID_DIRECTIONS = frozenset(("forward", "backward"))
53_VALID_ROLES = frozenset(("external_input", "stage_output", "final_output"))
54_VALID_IO = frozenset(("input", "output"))
55_VALID_CORE_TYPES = frozenset(("AIC", "AIV"))
58def _normalize_string_filter(
59 value: str | Collection[str] | None,
60 *,
61 name: str,
62 allowed: frozenset[str] | None = None,
63) -> tuple[str, ...] | None:
64 """Normalize a scalar or collection of named selection values."""
65 if value is None:
66 return None
67 values = (value,) if isinstance(value, str) else tuple(value)
68 if not values:
69 raise ValueError(f"{name} must not be empty")
70 for item in values:
71 if not isinstance(item, str) or not item:
72 raise TypeError(f"{name} entries must be non-empty strings, got {item!r}")
73 if allowed is not None and item not in allowed:
74 raise ValueError(f"unsupported {name} entry {item!r}; expected one of {sorted(allowed)}")
75 return values
78def _normalize_non_negative_int_filter(
79 value: int | Collection[int] | None,
80 *,
81 name: str,
82) -> tuple[int, ...] | None:
83 """Normalize integer selections while rejecting booleans and negative IDs."""
84 if value is None:
85 return None
86 if isinstance(value, bool):
87 raise ValueError(f"{name} entries must be non-negative integers, got {value!r}")
88 values = (value,) if isinstance(value, int) else tuple(value)
89 if not values:
90 raise ValueError(f"{name} must not be empty")
91 for item in values:
92 if not isinstance(item, int) or isinstance(item, bool) or item < 0:
93 raise ValueError(f"{name} entries must be non-negative integers, got {item!r}")
94 return values
97def _normalize_ranks(value: int | Collection[int] | None) -> tuple[int, ...] | None:
98 return _normalize_non_negative_int_filter(value, name="ranks")
101@dataclass(frozen=True)
102class TensorDumpConfig:
103 """Configuration for MegaKernel Tensor value collection.
105 Args:
106 output_dir: Root directory for per-rank dump windows.
107 level: Dump level. Supported values are ``"L0"`` and ``"L1"``.
108 kernel_names: Optional exact MegaKernel names whose invocations should
109 be retained.
110 tensor_names: Optional exact logical Tensor names selected from the
111 MegaKernel compute graph.
112 roles: Optional L0 Tensor roles to retain. Supported roles are
113 ``external_input``, ``stage_output`` and ``final_output``.
114 directions: Forward and/or backward invocations to retain.
115 ranks: Optional non-negative distributed ranks to retain.
116 stages: Optional exact L1 stage names.
117 task_ids: Optional stable global RuntimeConfig task IDs.
118 stage_task_indices: Optional task indices local to a stage.
119 owner_ids: Optional business diagnostic owner IDs.
120 core_types: Optional L1 Device worker types, ``AIC`` and/or ``AIV``.
121 core_ids: Optional core-local worker IDs.
122 io: L1 operator ports to capture: ``input``, ``output``, or both.
123 device_memory_budget_mb: Maximum Device memory used by one L1 launch.
124 max_bytes_per_tensor: Optional per-record byte cap. A capped Tensor is
125 exported as a flattened prefix and marked as truncated.
126 """
128 output_dir: str | Path
129 level: str = "L0"
130 tensor_names: str | Collection[str] | None = None
131 roles: str | Collection[str] | None = None
132 directions: str | Collection[str] | None = None
133 ranks: int | Collection[int] | None = None
134 kernel_names: str | Collection[str] | None = None
135 stages: str | Collection[str] | None = None
136 task_ids: int | Collection[int] | None = None
137 stage_task_indices: int | Collection[int] | None = None
138 owner_ids: int | Collection[int] | None = None
139 core_types: str | Collection[str] | None = None
140 core_ids: int | Collection[int] | None = None
141 io: str | Collection[str] | None = None
142 device_memory_budget_mb: int = 512
143 max_bytes_per_tensor: int | None = None
145 def __post_init__(self) -> None:
146 """Validate and normalize user-facing collection filters."""
147 if not isinstance(self.output_dir, (str, Path)):
148 raise TypeError(f"output_dir must be str or Path, got {type(self.output_dir).__name__}")
149 level = self.level.upper() if isinstance(self.level, str) else self.level
150 if level not in ("L0", "L1"):
151 raise ValueError(f"unsupported Tensor dump level {self.level!r}; expected 'L0' or 'L1'")
152 output_dir = Path(self.output_dir).expanduser()
153 kernel_names = _normalize_string_filter(self.kernel_names, name="kernel_names")
154 tensor_names = _normalize_string_filter(self.tensor_names, name="tensor_names")
155 roles = _normalize_string_filter(self.roles, name="roles", allowed=_VALID_ROLES)
156 directions = _normalize_string_filter(
157 self.directions,
158 name="directions",
159 allowed=_VALID_DIRECTIONS,
160 )
161 stages = _normalize_string_filter(self.stages, name="stages")
162 task_ids = _normalize_non_negative_int_filter(self.task_ids, name="task_ids")
163 stage_task_indices = _normalize_non_negative_int_filter(
164 self.stage_task_indices,
165 name="stage_task_indices",
166 )
167 owner_ids = _normalize_non_negative_int_filter(self.owner_ids, name="owner_ids")
168 core_types = _normalize_string_filter(
169 self.core_types,
170 name="core_types",
171 allowed=_VALID_CORE_TYPES,
172 )
173 core_ids = _normalize_non_negative_int_filter(self.core_ids, name="core_ids")
174 io = _normalize_string_filter(self.io, name="io", allowed=_VALID_IO)
175 if not isinstance(self.device_memory_budget_mb, int) or isinstance(
176 self.device_memory_budget_mb,
177 bool,
178 ):
179 raise TypeError("device_memory_budget_mb must be a positive integer")
180 if self.device_memory_budget_mb <= 0:
181 raise ValueError("device_memory_budget_mb must be positive")
182 if self.max_bytes_per_tensor is not None:
183 if not isinstance(self.max_bytes_per_tensor, int) or isinstance(
184 self.max_bytes_per_tensor,
185 bool,
186 ):
187 raise TypeError("max_bytes_per_tensor must be a positive integer or None")
188 if self.max_bytes_per_tensor <= 0:
189 raise ValueError("max_bytes_per_tensor must be positive")
190 if level == "L1" and all(
191 value is None
192 for value in (stages, task_ids, stage_task_indices, owner_ids, tensor_names)
193 ):
194 raise ValueError("L1 Tensor dump requires at least one stage, task, owner, or Tensor filter")
196 object.__setattr__(self, "output_dir", output_dir)
197 object.__setattr__(self, "level", level)
198 object.__setattr__(self, "kernel_names", kernel_names)
199 object.__setattr__(self, "tensor_names", tensor_names)
200 object.__setattr__(self, "roles", roles or tuple(sorted(_VALID_ROLES)))
201 object.__setattr__(self, "directions", directions or tuple(sorted(_VALID_DIRECTIONS)))
202 object.__setattr__(self, "ranks", _normalize_ranks(self.ranks))
203 object.__setattr__(self, "stages", stages)
204 object.__setattr__(self, "task_ids", task_ids)
205 object.__setattr__(self, "stage_task_indices", stage_task_indices)
206 object.__setattr__(self, "owner_ids", owner_ids)
207 object.__setattr__(self, "core_types", core_types)
208 object.__setattr__(self, "core_ids", core_ids)
209 object.__setattr__(self, "io", io or tuple(sorted(_VALID_IO)))
211 def selects_invocation(self, direction: str, rank: int, kernel_name: str | None = None) -> bool:
212 """Return whether one invocation passes kernel, direction and rank filters.
214 Args:
215 direction: Forward or backward invocation direction.
216 rank: Job-global distributed rank used for selection and export.
217 kernel_name: Stable diagnostic name of the invoked kernel.
218 """
219 if self.kernel_names is not None and kernel_name not in self.kernel_names:
220 return False
221 if direction not in self.directions:
222 return False
223 return self.ranks is None or rank in self.ranks
225 def selects_tensor(self, tensor_name: str, role: str) -> bool:
226 """Return whether one logical Tensor passes L0 name and role filters.
228 Args:
229 tensor_name: Logical tensor name recorded by the graph.
230 role: Graph-derived external input, stage output or final output role.
231 """
232 if role not in self.roles:
233 return False
234 return self.tensor_names is None or tensor_name in self.tensor_names
237@dataclass
238class _DumpInvocation:
239 """A completed invocation retained until its dump window is finalized."""
241 invocation_id: int
242 step: int
243 kernel_name: str
244 direction: str
245 rank: int
246 device: str
247 level: str
248 snapshots: list[_TensorSnapshot]
249 l1_snapshots: list[_L1Snapshot]
250 device_buffer_bytes: int = 0
253class _TensorDumpCall(DiagnosticCall):
254 """Capture handle surrounding one native MegaKernel launch."""
256 def __init__(
257 self,
258 on_complete: Callable[[_DumpInvocation], None] | None = None,
259 invocation: _DumpInvocation | None = None,
260 graph_tensors: tuple[_GraphTensor, ...] = (),
261 parameters: Sequence[Any] = (),
262 prepared_l1: _PreparedL1Dump | None = None,
263 *,
264 arguments: LaunchArguments,
265 ) -> None:
266 """Store complete launch arguments and L0 or L1 capture state."""
267 super().__init__(arguments)
268 self._on_complete = on_complete
269 self._invocation = invocation
270 self._graph_tensors = graph_tensors
271 self._parameters = parameters
272 self._prepared_l1 = prepared_l1
273 self.finished = False
275 def complete(self) -> None:
276 """Capture post-launch values after a successful native launch."""
277 if self.finished:
278 return
279 if self._on_complete is not None and self._invocation is not None:
280 if self._prepared_l1 is not None:
281 self._invocation.l1_snapshots.extend(_decode_l1_dump(self._prepared_l1))
282 else:
283 for graph_tensor in self._graph_tensors:
284 if graph_tensor.role != "external_input":
285 self._invocation.snapshots.append(
286 _snapshot_tensor(self._parameters[graph_tensor.param_position], graph_tensor, "after")
287 )
288 self._on_complete(self._invocation)
289 self.finished = True
291 def cancel(self) -> None:
292 """Discard retained snapshots when the native launch fails."""
293 if self.finished:
294 return
295 if self._invocation is not None:
296 self._invocation.snapshots.clear()
297 self._invocation.l1_snapshots.clear()
298 self._prepared_l1 = None
299 self.finished = True
302class TorchMegaKernelTensorDumper:
303 """Standalone schedule-driven L0/L1 Tensor dumper for Torch MegaKernels."""
305 def __init__(
306 self,
307 config: TensorDumpConfig,
308 *,
309 schedule: Callable[[int], DiagnosticAction] | None,
310 on_dump_ready: Callable[[Any], None] | None,
311 ) -> None:
312 # pylint: disable=redefined-outer-name
313 """Validate options and initialize dump window state."""
314 if not isinstance(config, TensorDumpConfig):
315 raise TypeError(f"config must be TensorDumpConfig, got {type(config).__name__}")
316 if schedule is not None and not callable(schedule):
317 raise TypeError("schedule must be callable or None")
318 if on_dump_ready is not None and not callable(on_dump_ready):
319 raise TypeError("on_dump_ready must be callable or None")
320 self.config = config
321 self._schedule = schedule
322 self._on_dump_ready = on_dump_ready
323 self._step = 0
324 self._action = self._action_for_step(0)
325 self._entered = False
326 self._next_invocation_id = 0
327 self._window_index = 0
328 self._invocations: list[_DumpInvocation] = []
329 self._last_manifest_paths: tuple[Path, ...] = ()
331 @property
332 def step_num(self) -> int:
333 """Return the zero-based user step currently selected."""
334 return self._step
336 @property
337 def current_action(self) -> DiagnosticAction:
338 """Return the current shared diagnostic schedule action."""
339 return self._action
341 @property
342 def last_manifest_paths(self) -> tuple[Path, ...]:
343 """Return manifest paths from the most recently finalized window."""
344 return self._last_manifest_paths
346 def __enter__(self) -> TorchMegaKernelTensorDumper:
347 """Activate this dumper as the process-local diagnostic session."""
348 _activate_session("tensor_dump", self)
349 self._entered = True
350 return self
352 def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
353 """Finalize successful snapshots or discard them after an exception."""
354 try:
355 if exc_type is None and self._invocations:
356 self._finalize_window()
357 elif exc_type is not None:
358 self._invocations.clear()
359 finally:
360 _deactivate_session(self)
361 self._entered = False
363 def step(self) -> None:
364 """Finish the current user step and advance the dump schedule."""
365 if not self._entered:
366 raise RuntimeError("dumper.step() must be called inside its context")
367 next_action = self._action_for_step(self._step + 1)
368 if self._action is DiagnosticAction.RECORD_AND_SAVE:
369 self._finalize_window()
370 elif self._action is DiagnosticAction.RECORD and next_action not in (
371 DiagnosticAction.RECORD,
372 DiagnosticAction.RECORD_AND_SAVE,
373 ):
374 self._finalize_window()
375 self._step += 1
376 self._action = next_action
378 def _action_for_step(self, step: int) -> DiagnosticAction:
379 if self._schedule is None:
380 return DiagnosticAction.RECORD
381 action = self._schedule(step)
382 if not isinstance(action, DiagnosticAction):
383 raise TypeError(f"MegaKernel schedule must return DiagnosticAction, got {type(action).__name__}")
384 return action
386 def is_capture_enabled(self) -> bool:
387 """Return whether this step collects Tensor values."""
388 return self._action in (DiagnosticAction.RECORD, DiagnosticAction.RECORD_AND_SAVE)
390 def prepare_call(self, context: LaunchContext) -> DiagnosticCall:
391 """Prepare only this dumper's selected invocation, with complete arguments.
393 Args:
394 context: Selected invocation or operand context supplied by the diagnostic framework.
395 """
396 runtime = context.runtime
397 if not self.is_capture_enabled() or not self.config.selects_invocation(
398 context.direction, runtime.rank, runtime.kernel_name,
399 ):
400 return DiagnosticCall(context.ordinary_arguments())
401 # Graph/owner work is lazy and never part of ordinary MegaKernel execution.
402 _ = runtime.metadata
403 _stage_names(runtime.graph.topological_sort())
404 return self._acquire_call(
405 runtime.graph, context.parameters, direction=context.direction, rank=runtime.rank,
406 kernel_name=runtime.kernel_name, device=runtime.normal_tensor.device,
407 runtime=runtime, event_counters=context.fallback_event_counters,
408 )
410 def _acquire_call(
411 self,
412 graph: Any,
413 parameters: Sequence[Any],
414 *,
415 direction: str,
416 rank: int,
417 kernel_name: str,
418 device: Any,
419 runtime: Any,
420 event_counters: Any,
421 ) -> DiagnosticCall:
422 """Plan a selected invocation and retain independent input snapshots before launch."""
423 arguments = LaunchArguments(
424 runtime.normal_tensor, event_counters,
425 runtime.normal_tensor if event_counters is None else event_counters, event_counters,
426 )
427 if self.config.level == "L1":
428 prepared_l1 = _prepare_l1_dump(
429 runtime,
430 graph,
431 parameters,
432 self.config,
433 rank=rank,
434 device=device,
435 )
436 invocation = _DumpInvocation(
437 invocation_id=self._next_invocation_id,
438 step=self._step,
439 direction=direction,
440 kernel_name=kernel_name,
441 rank=rank,
442 device=str(device),
443 level=self.config.level,
444 snapshots=[],
445 l1_snapshots=[],
446 device_buffer_bytes=prepared_l1.total_bytes,
447 )
448 self._next_invocation_id += 1
449 arguments.runtime_config = prepared_l1.runtime_config
450 arguments.diagnostic_buffer = prepared_l1.buffer
451 return _TensorDumpCall(self._invocations.append, invocation, prepared_l1=prepared_l1, arguments=arguments)
453 graph_tensors = tuple(
454 tensor
455 for tensor in _describe_graph_tensors(graph)
456 if self.config.selects_tensor(tensor.name, tensor.role)
457 )
458 if not graph_tensors:
459 return DiagnosticCall(arguments)
460 max_position = max(tensor.param_position for tensor in graph_tensors)
461 if max_position >= len(parameters):
462 raise ValueError(
463 f"MegaKernel dump requires param_position {max_position}, but only {len(parameters)} parameters exist"
464 )
465 invocation = _DumpInvocation(
466 invocation_id=self._next_invocation_id,
467 step=self._step,
468 direction=direction,
469 kernel_name=kernel_name,
470 rank=rank,
471 device=str(device),
472 level=self.config.level,
473 snapshots=[],
474 l1_snapshots=[],
475 )
476 self._next_invocation_id += 1
477 for graph_tensor in graph_tensors:
478 if graph_tensor.role == "external_input":
479 invocation.snapshots.append(
480 _snapshot_tensor(parameters[graph_tensor.param_position], graph_tensor, "before")
481 )
482 return _TensorDumpCall(self._invocations.append, invocation, graph_tensors, parameters, arguments=arguments)
484 def _finalize_window(self) -> None:
485 """Export completed invocations by rank and release the window's CPU snapshots."""
486 if not self._invocations:
487 self._last_manifest_paths = ()
488 self._window_index += 1
489 return
490 manifest_paths = []
491 ranks = sorted({invocation.rank for invocation in self._invocations})
492 for rank in ranks:
493 rank_invocations = [invocation for invocation in self._invocations if invocation.rank == rank]
494 manifest_paths.append(self._write_rank_window(rank, rank_invocations))
495 self._last_manifest_paths = tuple(manifest_paths)
496 self._invocations.clear()
497 self._window_index += 1
498 if self._on_dump_ready is not None:
499 self._on_dump_ready(self)
501 def _write_l0_invocation(
502 self,
503 invocation: _DumpInvocation,
504 invocation_dir: Path,
505 window_dir: Path,
506 ) -> dict[str, Any]:
507 """Write logical tensor snapshots and their stage-port metadata."""
508 tensor_metadata = []
509 for snapshot in invocation.snapshots:
510 safe_name = _SAFE_FILE_NAME.sub("_", snapshot.graph_tensor.name)
511 output_path, port_metadata = _write_l0_stage_files(
512 snapshot,
513 invocation_dir,
514 window_dir,
515 safe_name,
516 )
517 tensor_metadata.append(
518 {
519 "name": snapshot.graph_tensor.name,
520 "paramPosition": snapshot.graph_tensor.param_position,
521 "role": snapshot.graph_tensor.role,
522 "capturePhase": snapshot.phase,
523 "producers": list(snapshot.graph_tensor.producers),
524 "consumers": list(snapshot.graph_tensor.consumers),
525 "ports": port_metadata,
526 "declaredShape": list(snapshot.graph_tensor.declared_shape),
527 "actualShape": list(snapshot.shape),
528 "originalDtype": snapshot.original_dtype,
529 "storedDtype": snapshot.stored_dtype,
530 "file": output_path.relative_to(window_dir).as_posix(),
531 }
532 )
533 return {
534 "invocationId": invocation.invocation_id,
535 "step": invocation.step,
536 "direction": invocation.direction,
537 "kernelName": invocation.kernel_name,
538 "device": invocation.device,
539 "tensorCount": len(tensor_metadata),
540 "tensors": tensor_metadata,
541 }
543 def _write_l1_invocation(
544 self,
545 invocation: _DumpInvocation,
546 invocation_dir: Path,
547 window_dir: Path,
548 ) -> dict[str, Any]:
549 """Write task-local tensor snapshots with their native copy metadata."""
550 snapshots_by_task: dict[tuple[int, str, int], list[_L1Snapshot]] = {}
551 for snapshot in invocation.l1_snapshots:
552 record = snapshot.record
553 snapshots_by_task.setdefault((record.task_id, record.core_type, record.core_id), []).append(snapshot)
555 task_metadata = []
556 for key in sorted(snapshots_by_task):
557 task_id = key[0]
558 snapshots = snapshots_by_task[key]
559 first = snapshots[0].record
560 safe_stage = _SAFE_FILE_NAME.sub("_", first.stage)
561 task_dir = invocation_dir / "stages" / safe_stage / f"task_{task_id:06d}"
562 if first.worker_id != 0xFFFFFFFF:
563 task_dir = task_dir.with_name(f"{task_dir.name}_worker_{first.core_id:03d}")
564 task_dir.mkdir(parents=True, exist_ok=True)
565 tensor_metadata = []
566 for snapshot in snapshots:
567 record = snapshot.record
568 safe_name = _SAFE_FILE_NAME.sub("_", record.tensor_name)
569 file_name = f"{record.io}_{record.index:02d}_{safe_name}.pt"
570 output_path = task_dir / file_name
571 torch.save(snapshot.data, output_path)
572 tensor_metadata.append(
573 {
574 "name": record.tensor_name,
575 "paramPosition": record.parameter_position,
576 "io": record.io,
577 "index": record.index,
578 "capturePhase": record.phase,
579 "sourceOffsetBytes": record.source_offset_bytes,
580 "actualShape": list(record.actual_shape),
581 "storedShape": list(snapshot.stored_shape),
582 "originalDtype": snapshot.original_dtype,
583 "storedDtype": snapshot.stored_dtype,
584 "fullBytes": record.full_bytes,
585 "storedBytes": record.payload_bytes,
586 "truncated": record.truncated,
587 "note": record.note,
588 "file": output_path.relative_to(window_dir).as_posix(),
589 }
590 )
591 metadata_path = task_dir / "task.json"
592 metadata = {
593 "taskId": task_id,
594 "stageTaskIndex": first.task_index,
595 "stage": first.stage,
596 "owner": {"label": first.owner_label, "id": first.owner_id},
597 "coreType": first.core_type,
598 "coreId": first.core_id,
599 "tensorCount": len(tensor_metadata),
600 "tensors": tensor_metadata,
601 }
602 metadata_path.write_text(
603 json.dumps(metadata, ensure_ascii=False, indent=2),
604 encoding="utf-8",
605 )
606 metadata["metadataFile"] = metadata_path.relative_to(window_dir).as_posix()
607 task_metadata.append(metadata)
608 return {
609 "invocationId": invocation.invocation_id,
610 "step": invocation.step,
611 "direction": invocation.direction,
612 "kernelName": invocation.kernel_name,
613 "device": invocation.device,
614 "deviceBufferBytes": invocation.device_buffer_bytes,
615 "taskCount": len(task_metadata),
616 "tasks": task_metadata,
617 }
619 def _write_rank_window(self, rank: int, invocations: list[_DumpInvocation]) -> Path:
620 """Export one rank's window without overwriting an existing directory."""
621 window_dir = self.config.output_dir / f"rank_{rank}" / f"window_{self._window_index:04d}"
622 if window_dir.exists():
623 raise FileExistsError(f"Tensor dump window already exists: {window_dir}")
624 window_dir.mkdir(parents=True)
625 invocation_metadata = []
626 for invocation in invocations:
627 invocation_dir = window_dir / f"invocation_{invocation.invocation_id:06d}_{invocation.direction}"
628 if invocation.level == "L1":
629 metadata = self._write_l1_invocation(invocation, invocation_dir, window_dir)
630 else:
631 metadata = self._write_l0_invocation(invocation, invocation_dir, window_dir)
632 invocation_metadata.append(metadata)
633 is_l1 = self.config.level == "L1"
634 manifest = {
635 "schemaVersion": 4,
636 "level": self.config.level,
637 "tensorSerialization": "torch.save",
638 "tensorFileLayout": "stage_only",
639 "windowIndex": self._window_index,
640 "capturePhaseSemantics": "task_boundary" if is_l1 else "mega_kernel_boundary",
641 "stageViewStorage": "per_task_files" if is_l1 else "primary_file_with_hard_links",
642 "rank": rank,
643 "invocationCount": len(invocation_metadata),
644 "invocations": invocation_metadata,
645 }
646 manifest_path = window_dir / "manifest.json"
647 manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
648 return manifest_path
651def tensor_dump(
652 config: TensorDumpConfig,
653 *,
654 schedule: Callable[[int], DiagnosticAction] | None = None,
655 on_dump_ready: Callable[[Any], None] | None = None,
656) -> TorchMegaKernelTensorDumper:
657 # pylint: disable=redefined-outer-name
658 """Create the standalone tensor MegaKernel Tensor dumper.
660 Args:
661 config: Output and logical Tensor selection configuration.
662 schedule: Optional shared diagnostic step schedule. Warmup actions do
663 not copy Tensor values because L0 collection needs no warmup.
664 on_dump_ready: Optional callback invoked after a window is exported.
666 Returns:
667 tensor MegaKernel Tensor dump context manager.
669 Note:
670 Tensor dumping synchronizes Device-to-Host copies and changes execution
671 timing. It is mutually exclusive with MegaKernel cycle profiling.
672 """
673 return TorchMegaKernelTensorDumper(
674 config,
675 schedule=schedule,
676 on_dump_ready=on_dump_ready,
677 )