Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / diagnostics / dump / l1.py: 90%
327 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"""Host planning and decoding for task-level Torch MegaKernel Tensor dumps."""
17from __future__ import annotations
19from collections.abc import Sequence
20from dataclasses import dataclass, field
21from math import prod
22import struct
23from typing import Any
25import torch
27from hyper_multicore.diagnostics.common.graph import operator_diagnostic_name
28from hyper_multicore.scheduler.config import (
29 RuntimeConfigC,
30 TaskAiCoreType,
31 TensorDescC,
32)
35L1_DUMP_MAGIC = 0x4C314450
36L1_DUMP_VERSION = 2
37L1_DUMP_PHASE_BEFORE = 1
38L1_DUMP_PHASE_AFTER = 2
39L1_DIAGNOSTIC_MODE = 2
40L1_DUMP_STATUS_PENDING = 0
41L1_DUMP_STATUS_CAPTURED = 1
42L1_DUMP_HEADER = struct.Struct("<8I")
43L1_DUMP_RECORD = struct.Struct("<10I24x")
44L1_DUMP_ALIGNMENT = 64
45L1_DUMP_RECORDS_OFFSET = 64
46L1_DUMP_MAX_RECORDS = 4096
48_CORE_TYPE_NAMES = {
49 TaskAiCoreType.TASK_AICORE_CUBE: "AIC",
50 TaskAiCoreType.TASK_AICORE_VECTOR: "AIV",
51 TaskAiCoreType.TASK_AICORE_MIX: "MIX",
52}
53def _align_up(value: int, alignment: int = L1_DUMP_ALIGNMENT) -> int:
54 return (value + alignment - 1) // alignment * alignment
57def _dtype_name(dtype: Any) -> str:
58 return str(dtype).removeprefix("torch.")
61@dataclass(frozen=True)
62class L1TensorView:
63 """Resolved task operand view copied by the common Device worker.
65 Args:
66 source_offset_bytes: Byte offset from the selected ABI parameter.
67 shape: Logical task-local shape.
68 note: Optional semantic note emitted in the manifest.
69 """
71 source_offset_bytes: int
72 shape: tuple[int, ...]
73 note: str = ""
76@dataclass
77class L1TensorViewContext:
78 """Information available to an optional operator-specific view resolver."""
80 operator: Any
81 task_id: int
82 owner_id: int | None
83 task_desc: Any
84 tensor_desc: Any
85 tensor_spec: Any
86 io: str
87 index: int
88 parameters: Sequence[Any]
89 runtime_config: RuntimeConfigC
90 rank: int
91 stage: str
92 default_view: L1TensorView
93 parameter_cache: dict[int, Any]
94 core_id: int = 0
95 owner_context: Any = None
97 def parameter_cpu(self, position: int) -> Any:
98 """Return one ABI parameter on CPU, caching the debug-only transfer.
100 Args:
101 position: Zero-based position in the native ABI parameter list.
102 """
103 if position not in self.parameter_cache:
104 parameter = self.parameters[position]
105 if not torch.is_tensor(parameter):
106 raise TypeError(
107 f"L1 dump parameter {position} must be torch.Tensor, got {type(parameter).__name__}"
108 )
109 self.parameter_cache[position] = parameter.detach().contiguous().cpu()
110 return self.parameter_cache[position]
112 def make_view(
113 self,
114 shape: Sequence[int],
115 *,
116 source_offset_bytes: int | None = None,
117 note: str = "",
118 ) -> L1TensorView:
119 """Build a validated view relative to this operand's ABI parameter.
121 Args:
122 shape: Local contiguous shape of the selected task operand.
123 source_offset_bytes: Byte offset relative to the bound ABI tensor's data pointer.
124 note: Optional explanation retained in Host-side capture metadata.
125 """
126 resolved_shape = tuple(int(dim) for dim in shape)
127 if any(dim < 0 for dim in resolved_shape):
128 raise ValueError(f"L1 dump view shape must be non-negative, got {resolved_shape}")
129 resolved_offset = self.default_view.source_offset_bytes
130 if source_offset_bytes is not None:
131 resolved_offset = source_offset_bytes
132 if not isinstance(resolved_offset, int) or isinstance(resolved_offset, bool) or resolved_offset < 0:
133 raise ValueError(f"L1 dump source offset must be a non-negative integer, got {resolved_offset!r}")
134 return L1TensorView(resolved_offset, resolved_shape, note)
137@dataclass(frozen=True)
138class _L1TaskLocation:
139 """Device worker assignment for one stable RuntimeConfig task ID."""
141 task_id: int
142 core_type: str
143 core_id: int
146@dataclass(frozen=True)
147class _L1RecordPlan:
148 """Host metadata and fixed Device payload slot for one task operand."""
150 task_id: int
151 task_index: int
152 stage: str
153 owner_id: int | None
154 owner_label: str
155 core_type: str
156 core_id: int
157 io: str
158 index: int
159 phase: str
160 parameter_position: int
161 tensor_name: str
162 dtype: Any
163 source_offset_bytes: int
164 actual_shape: tuple[int, ...]
165 payload_offset: int
166 payload_bytes: int
167 full_bytes: int
168 truncated: bool
169 note: str
170 record_offset: int
171 worker_id: int = 0xFFFFFFFF
174@dataclass
175class _PreparedL1Dump:
176 """Runtime override, Device buffer and immutable record plan for one launch."""
178 runtime_config: Any
179 buffer: Any
180 records: tuple[_L1RecordPlan, ...]
181 total_bytes: int
184@dataclass(frozen=True)
185class _L1Snapshot:
186 """One decoded task-level payload ready for Torch serialization."""
188 record: _L1RecordPlan
189 original_dtype: str
190 stored_dtype: str
191 stored_shape: tuple[int, ...]
192 data: torch.Tensor
195def _operator_stage_name(operator: Any) -> str:
196 return operator_diagnostic_name(operator)
199def _task_locations(
200 runtime_config: RuntimeConfigC, *, vector_worker_stride: int = 2, broadcast_tasks: bool = False,
201) -> dict[Any, _L1TaskLocation]:
202 """Map scheduled task IDs to actual workers, preserving broadcast worker identity."""
203 active_core_count = int(runtime_config.num_workers) if broadcast_tasks else int(runtime_config.num_workers) // 2
204 if active_core_count <= 0:
205 raise ValueError("L1 dump requires RuntimeConfig.num_workers to identify active AIC/AIV workers")
206 schedules = (
207 (
208 TaskAiCoreType.TASK_AICORE_CUBE,
209 runtime_config.cube_task_indices,
210 int(runtime_config.task_index_num[0]),
211 ),
212 (
213 TaskAiCoreType.TASK_AICORE_VECTOR,
214 runtime_config.vector_task_indices,
215 int(runtime_config.task_index_num[1]),
216 ),
217 (
218 TaskAiCoreType.TASK_AICORE_MIX,
219 runtime_config.mix_task_indices,
220 int(runtime_config.task_index_num[2]),
221 ),
222 )
223 locations = {}
224 for core_type, task_indices, task_count in schedules:
225 if task_count < 0 or task_count > len(task_indices):
226 raise ValueError(f"L1 dump task schedule count is outside RuntimeConfig capacity: {task_count}")
227 for schedule_index in range(task_count):
228 task_id = int(task_indices[schedule_index])
229 if broadcast_tasks:
230 for worker_id in range(int(runtime_config.num_workers)):
231 locations[(task_id, worker_id)] = _L1TaskLocation(task_id, "AIV", worker_id)
232 continue
233 worker_count = (int(runtime_config.num_workers) // vector_worker_stride
234 if core_type == TaskAiCoreType.TASK_AICORE_VECTOR else active_core_count)
235 location = _L1TaskLocation(
236 task_id=task_id,
237 core_type=_CORE_TYPE_NAMES[core_type],
238 core_id=schedule_index % worker_count,
239 )
240 previous = locations.setdefault(task_id, location)
241 if previous != location:
242 raise ValueError(f"L1 dump task {task_id} appears in multiple Device schedules")
243 return locations
246def _split_dimension(operator: Any, io: str, index: int) -> int | None:
247 """Return the configured split axis for a selected operator port."""
248 split_spec = operator.split_spec
249 if io == "input":
250 for input_index, dimension in split_spec.split_inputs or ():
251 if input_index == index:
252 return None if dimension == -1 else int(dimension)
253 return None
254 if index >= len(split_spec.split_output_dims):
255 return None
256 dimension = int(split_spec.split_output_dims[index])
257 return None if dimension == -1 else dimension
260def _default_tensor_view(
261 operator: Any,
262 task_desc: Any,
263 tensor_desc: Any,
264 tensor_spec: Any,
265 io: str,
266 index: int,
267 *, allow_strided: bool = False,
268) -> L1TensorView:
269 """Resolve a contiguous static task range and clip its final partial tile."""
270 if tensor_desc.dynamic_shape:
271 raise ValueError(
272 f"L1 dump stage {_operator_stage_name(operator)!r} Tensor {tensor_spec.name!r} has dynamic shape; "
273 "configure OperatorNode.diagnostic_view_resolver"
274 )
275 shape = [int(dim) for dim in tensor_spec.shape]
276 split_dimension = _split_dimension(operator, io, index)
277 if split_dimension is not None and shape:
278 if not 0 <= split_dimension < len(shape):
279 raise ValueError(
280 f"L1 dump split dimension {split_dimension} is invalid for Tensor {tensor_spec.name!r} shape {shape}"
281 )
282 width = shape[split_dimension]
283 count = min(width, int(task_desc.task_split_value))
284 if count < width and prod(shape[:split_dimension]) > 1 and not allow_strided:
285 raise ValueError("L1 default view cannot copy a non-contiguous split; configure diagnostic_view_resolver")
286 row_elements = prod(shape[split_dimension + 1:])
287 if row_elements and width:
288 base_offset = int(tensor_desc.base_ptr_offset)
289 if base_offset % row_elements and not allow_strided:
290 raise ValueError("L1 default split offset is not row-aligned; configure diagnostic_view_resolver")
291 start = base_offset // row_elements % width
292 count = min(count, width - start)
293 shape[split_dimension] = count
294 return L1TensorView(
295 source_offset_bytes=int(tensor_desc.base_ptr_offset) * int(tensor_desc.data_type),
296 shape=tuple(shape),
297 )
300def _selected(config: Any, location: _L1TaskLocation, task_desc: Any, stage: str, owner_id: int | None) -> bool:
301 """Apply task, owner, stage and worker selection predicates."""
302 if config.stages is not None and stage not in config.stages:
303 return False
304 if config.task_ids is not None and location.task_id not in config.task_ids:
305 return False
306 if config.stage_task_indices is not None and int(task_desc.task_index) not in config.stage_task_indices:
307 return False
308 if config.owner_ids is not None and owner_id not in config.owner_ids:
309 return False
310 if config.core_types is not None and location.core_type not in config.core_types:
311 return False
312 return config.core_ids is None or location.core_id in config.core_ids
315def _resolve_record_view(
316 *,
317 operator: Any,
318 task_id: int,
319 owner_id: int | None,
320 task_desc: Any,
321 tensor_desc: Any,
322 tensor_spec: Any,
323 io: str,
324 index: int,
325 parameters: Sequence[Any],
326 runtime_config: RuntimeConfigC,
327 rank: int,
328 stage: str,
329 parameter_cache: dict[int, Any],
330 core_id: int = 0,
331 owner_context: Any = None,
332) -> L1TensorView:
333 """Resolve one selected port's Host view before allocating its payload slot."""
334 resolver = getattr(operator, "diagnostic_view_resolver", None)
335 default_view = _default_tensor_view(
336 operator,
337 task_desc,
338 tensor_desc,
339 tensor_spec,
340 io,
341 index,
342 allow_strided=resolver is not None,
343 ) if not tensor_desc.dynamic_shape else L1TensorView(
344 int(tensor_desc.base_ptr_offset) * int(tensor_desc.data_type),
345 tuple(int(dim) for dim in tensor_spec.shape),
346 )
347 if resolver is None:
348 if tensor_desc.dynamic_shape:
349 raise ValueError(
350 f"L1 dump stage {stage!r} Tensor {tensor_spec.name!r} needs a diagnostic_view_resolver"
351 )
352 return default_view
353 context = L1TensorViewContext(
354 operator=operator,
355 task_id=task_id,
356 owner_id=owner_id,
357 task_desc=task_desc,
358 tensor_desc=tensor_desc,
359 tensor_spec=tensor_spec,
360 io=io,
361 index=index,
362 parameters=parameters,
363 runtime_config=runtime_config,
364 rank=rank,
365 stage=stage,
366 default_view=default_view,
367 parameter_cache=parameter_cache,
368 core_id=core_id,
369 owner_context=owner_context,
370 )
371 if not callable(resolver):
372 raise TypeError(f"L1 dump diagnostic_view_resolver for stage {stage!r} must be callable")
373 view = resolver(context)
374 if not isinstance(view, L1TensorView):
375 raise TypeError(
376 f"L1 dump resolver for stage {stage!r} must return L1TensorView, got {type(view).__name__}"
377 )
378 return view
381def _validate_view_bounds(
382 parameter: Any,
383 view: L1TensorView,
384 tensor_name: str,
385) -> tuple[Any, int]:
386 """Check dtype, contiguity, alignment and the full source range against its ABI tensor."""
387 if not torch.is_tensor(parameter):
388 raise TypeError(f"L1 dump Tensor {tensor_name!r} must map to torch.Tensor, got {type(parameter).__name__}")
389 if not parameter.is_contiguous():
390 raise ValueError(f"L1 dump Tensor {tensor_name!r} ABI parameter must be contiguous")
391 if (isinstance(view.source_offset_bytes, bool) or not isinstance(view.source_offset_bytes, int)
392 or not 0 <= view.source_offset_bytes < 1 << 64):
393 raise ValueError(f"L1 dump Tensor {tensor_name!r} source offset must be a non-negative uint64")
394 if any(isinstance(dim, bool) or not isinstance(dim, int) or dim < 0 for dim in view.shape):
395 raise ValueError(f"L1 dump Tensor {tensor_name!r} view shape must contain non-negative integers")
396 element_size = int(parameter.element_size())
397 if view.source_offset_bytes % element_size:
398 raise ValueError(f"L1 dump Tensor {tensor_name!r} source offset must align to its element byte size")
399 full_bytes = prod(view.shape) * element_size
400 parameter_bytes = int(parameter.numel()) * element_size
401 if view.source_offset_bytes + full_bytes > parameter_bytes:
402 raise ValueError(
403 f"L1 dump view for Tensor {tensor_name!r} exceeds its ABI parameter: "
404 f"offset={view.source_offset_bytes}, bytes={full_bytes}, parameter_bytes={parameter_bytes}"
405 )
406 return parameter.dtype, full_bytes
409def _operator_ports(operator: Any, task: Any) -> tuple[Any, ...]:
410 """Resolve Host ABI ports without modifying the native TaskDesc arrays."""
411 if operator.diagnostic_host_ports:
412 def descriptors(specs: Any) -> list[TensorDescC]:
413 """Build Host-only source descriptors for explicitly bound ABI ports.
415 Args:
416 specs: Native-bound graph operands requiring Host descriptors.
417 """
418 return [TensorDescC(input_position=spec.param_position, data_type=spec.dtype_size,
419 dynamic_shape=1) for spec in specs]
421 return (("input", descriptors(operator.inputs), operator.inputs, len(operator.inputs)),
422 ("output", descriptors(operator.outputs), operator.outputs, len(operator.outputs)))
423 return (("input", task.inputs, operator.inputs, int(task.num_inputs)),
424 ("output", task.outputs, operator.outputs, int(task.num_outputs)))
427@dataclass
428class _L1PlanBuilder:
429 """Own one launch's record list, parameter cache and monotonically assigned slots."""
431 runtime: Any
432 parameters: Sequence[Any]
433 config: Any
434 rank: int
435 broadcast: bool
436 records: list[_L1RecordPlan] = field(default_factory=list)
437 parameter_cache: dict[int, Any] = field(default_factory=dict)
438 next_payload_offset: int = L1_DUMP_RECORDS_OFFSET + L1_DUMP_MAX_RECORDS * L1_DUMP_RECORD.size
440 def add(
441 self, operator: Any, location: _L1TaskLocation, stage: str, owner_id: int | None,
442 io: str, index: int, tensor_desc: Any, tensor_spec: Any,
443 ) -> None:
444 """Validate and append one selected operand, before allocating any payload.
446 Args:
447 operator: Graph operator associated with the selected task.
448 location: Task identity and worker placement in the native schedule.
449 stage: Diagnostic stage name associated with the operator.
450 owner_id: Optional resolved business owner, such as an Expert ID.
451 io: Input or output capture boundary.
452 index: Zero-based operand index within the selected input or output list.
453 tensor_desc: Native task operand descriptor containing source addressing information.
454 tensor_spec: Graph operand contract containing its name and ABI position.
455 """
456 task_id = location.task_id
457 task = self.runtime.host_config.all_tasks[task_id]
458 position = int(tensor_desc.input_position)
459 if position != int(tensor_spec.param_position):
460 raise ValueError(f"L1 task {task_id} {io}{index} parameter position disagrees with graph Tensor")
461 if not 0 <= position < len(self.parameters):
462 raise ValueError(f"L1 task {task_id} requires parameter {position}, got {len(self.parameters)} parameters")
463 parameter = self.parameters[position]
464 if torch.is_tensor(parameter) and parameter.element_size() != int(tensor_desc.data_type):
465 raise ValueError(f"L1 task {task_id} Tensor {tensor_spec.name!r} ABI element byte size mismatch")
466 view = _resolve_record_view(
467 operator=operator, task_id=task_id, owner_id=owner_id, task_desc=task, tensor_desc=tensor_desc,
468 tensor_spec=tensor_spec, io=io, index=index, parameters=self.parameters,
469 runtime_config=self.runtime.host_config, rank=self.rank, stage=stage,
470 parameter_cache=self.parameter_cache, core_id=location.core_id,
471 owner_context=getattr(getattr(self.runtime, "spec", None), "owner_context", None),
472 )
473 dtype, full_bytes = _validate_view_bounds(parameter, view, tensor_spec.name)
474 if self.broadcast and full_bytes == 0:
475 return # No Device worker is launched for this empty row shard.
476 payload_bytes = full_bytes
477 if self.config.max_bytes_per_tensor is not None and payload_bytes > self.config.max_bytes_per_tensor:
478 element_size = int(parameter.element_size())
479 payload_bytes = self.config.max_bytes_per_tensor // element_size * element_size
480 if payload_bytes <= 0:
481 raise ValueError(f"max_bytes_per_tensor={self.config.max_bytes_per_tensor} cannot store one element")
482 self.records.append(_L1RecordPlan(
483 task_id=task_id, task_index=int(task.task_index), stage=stage, owner_id=owner_id,
484 owner_label=self.runtime.owner_label, core_type=location.core_type, core_id=location.core_id,
485 io=io, index=index, phase="before" if io == "input" else "after", parameter_position=position,
486 tensor_name=tensor_spec.name, dtype=dtype, source_offset_bytes=view.source_offset_bytes,
487 actual_shape=view.shape, payload_offset=self.next_payload_offset, payload_bytes=payload_bytes,
488 full_bytes=full_bytes, truncated=payload_bytes < full_bytes, note=view.note,
489 record_offset=L1_DUMP_RECORDS_OFFSET + len(self.records) * L1_DUMP_RECORD.size,
490 worker_id=location.core_id if self.broadcast else 0xFFFFFFFF,
491 ))
492 self.next_payload_offset = _align_up(self.next_payload_offset + payload_bytes)
493 if len(self.records) > L1_DUMP_MAX_RECORDS:
494 raise ValueError(f"L1 dump record count exceeds {L1_DUMP_MAX_RECORDS}")
497def _build_l1_record_plans(
498 runtime: Any, graph: Any, parameters: Sequence[Any], config: Any, rank: int,
499) -> tuple[_L1RecordPlan, ...]:
500 """Select task/port pairs, then compile their independent snapshot slots."""
501 runtime_config = runtime.host_config
502 operators_by_stage = {}
503 for operator in graph.topological_sort():
504 stage = _operator_stage_name(operator)
505 if stage in operators_by_stage:
506 raise ValueError(f"duplicate L1 diagnostic stage name {stage!r}")
507 operators_by_stage[stage] = operator
508 spec = getattr(runtime, "spec", None)
509 broadcast = getattr(spec, "execution_mode", "event_driven") == "aiv_pipeline"
510 locations = _task_locations(runtime_config, vector_worker_stride=getattr(spec, "vector_worker_stride", 2),
511 broadcast_tasks=broadcast)
512 builder = _L1PlanBuilder(runtime, parameters, config, rank, broadcast)
513 for location in sorted(locations.values(), key=lambda item: (item.task_id, item.core_id)):
514 task = runtime_config.all_tasks[location.task_id]
515 stage = runtime.task_stage_names.get(location.task_id)
516 operator = operators_by_stage.get(stage)
517 if operator is None:
518 continue
519 owner_id = runtime.metadata.task_owner_ids.get(location.task_id)
520 if not _selected(config, location, task, stage, owner_id):
521 continue
522 for io, descriptors, specs, count in _operator_ports(operator, task):
523 if io not in config.io:
524 continue
525 if count > len(specs):
526 raise ValueError(f"L1 task {location.task_id} declares too many {io} ports for its graph stage")
527 for index in range(count):
528 if config.tensor_names is None or specs[index].name in config.tensor_names:
529 builder.add(operator, location, stage, owner_id, io, index, descriptors[index], specs[index])
530 if not builder.records:
531 raise ValueError("L1 dump filters selected no task Tensor records")
532 return tuple(builder.records)
535def _pack_l1_buffer(records: tuple[_L1RecordPlan, ...]) -> bytearray:
536 """Encode the versioned header, padded record table and independent payload slots."""
537 payload_offset = min(record.payload_offset for record in records)
538 total_bytes = _align_up(max(record.payload_offset + record.payload_bytes for record in records))
539 raw_buffer = bytearray(total_bytes)
540 L1_DUMP_HEADER.pack_into(
541 raw_buffer,
542 0,
543 L1_DUMP_MAGIC,
544 L1_DUMP_VERSION,
545 1,
546 len(records),
547 L1_DUMP_RECORD.size,
548 L1_DUMP_RECORDS_OFFSET,
549 payload_offset,
550 total_bytes,
551 )
552 phase_values = {"before": L1_DUMP_PHASE_BEFORE, "after": L1_DUMP_PHASE_AFTER}
553 for record in records:
554 source_low = record.source_offset_bytes & 0xFFFFFFFF
555 source_high = record.source_offset_bytes >> 32
556 L1_DUMP_RECORD.pack_into(
557 raw_buffer,
558 record.record_offset,
559 record.task_id,
560 phase_values[record.phase],
561 record.parameter_position,
562 source_low,
563 source_high,
564 record.payload_offset,
565 record.payload_bytes,
566 L1_DUMP_STATUS_PENDING,
567 record.task_index,
568 record.worker_id,
569 )
570 return raw_buffer
573def _prepare_l1_dump(
574 runtime: Any,
575 graph: Any,
576 parameters: Sequence[Any],
577 config: Any,
578 *,
579 rank: int,
580 device: Any,
581) -> _PreparedL1Dump:
582 """Compile selected task views and allocate one Device snapshot buffer."""
583 if runtime is None:
584 raise ValueError("L1 Tensor dump requires the prepared MegaKernel runtime")
585 records = _build_l1_record_plans(runtime, graph, parameters, config, rank)
586 budget_bytes = config.device_memory_budget_mb * 1024 * 1024
587 total_bytes = _align_up(max(record.payload_offset + record.payload_bytes for record in records))
588 if total_bytes >= 1 << 32:
589 raise ValueError("L1 dump buffer exceeds the 32-bit payload-offset ABI")
590 if total_bytes > budget_bytes:
591 raise MemoryError(
592 f"L1 dump requires {total_bytes} Device bytes, exceeding "
593 f"device_memory_budget_mb={config.device_memory_budget_mb} ({budget_bytes} bytes)"
594 )
595 raw_buffer = _pack_l1_buffer(records)
596 device_buffer = torch.frombuffer(raw_buffer, dtype=torch.uint8).clone().to(device=device)
597 runtime_config = runtime.normal_tensor.clone()
598 runtime_config[RuntimeConfigC.diagnostic_mode.offset] = L1_DIAGNOSTIC_MODE
599 return _PreparedL1Dump(
600 runtime_config=runtime_config,
601 buffer=device_buffer,
602 records=records,
603 total_bytes=len(raw_buffer),
604 )
607def _payload_tensor(
608 host_buffer: torch.Tensor,
609 record: _L1RecordPlan,
610) -> tuple[torch.Tensor, tuple[int, ...]]:
611 """Decode one payload slot as an owned CPU Tensor without dtype conversion."""
612 if record.payload_bytes == 0:
613 tensor = torch.empty((0,), dtype=record.dtype)
614 else:
615 payload = host_buffer.narrow(0, record.payload_offset, record.payload_bytes)
616 tensor = payload.view(record.dtype).clone()
617 stored_shape = (tensor.numel(),) if record.truncated else record.actual_shape
618 return tensor.reshape(stored_shape), stored_shape
621def _decode_l1_dump(prepared: _PreparedL1Dump) -> tuple[_L1Snapshot, ...]:
622 """Copy the completed Device buffer to Host and decode all fixed slots."""
623 host_buffer = prepared.buffer.detach().contiguous().cpu()
624 header = tuple(int(value) for value in host_buffer[:L1_DUMP_HEADER.size].view(torch.int32).tolist())
625 if header[:2] != (L1_DUMP_MAGIC, L1_DUMP_VERSION):
626 raise ValueError(f"invalid L1 dump buffer header: magic={header[0]:#x}, version={header[1]}")
627 snapshots = []
628 for record in prepared.records:
629 status_offset = record.record_offset + 7 * 4
630 status = int(host_buffer[status_offset:status_offset + 4].view(torch.int32).item())
631 if status != L1_DUMP_STATUS_CAPTURED:
632 raise RuntimeError(
633 f"L1 dump record was not captured: task_id={record.task_id}, stage={record.stage!r}, "
634 f"io={record.io}, index={record.index}, status={status}, header_enabled={header[2]}"
635 )
636 data, stored_shape = _payload_tensor(host_buffer, record)
637 dtype_name = _dtype_name(data.dtype)
638 snapshots.append(
639 _L1Snapshot(
640 record=record,
641 original_dtype=dtype_name,
642 stored_dtype=dtype_name,
643 stored_shape=stored_shape,
644 data=data,
645 )
646 )
647 return tuple(snapshots)