Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / diagnostics / profiler / cycle_trace.py: 81%
372 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"""Plan and decode TaskDesc-based MegaKernel per-core cycle records."""
17from __future__ import annotations
19__all__ = []
21from collections.abc import Mapping
22from dataclasses import dataclass
23from enum import Enum
24import struct
25from typing import Any
27from hyper_multicore.scheduler.config import (
28 EVENT_INVALID_ID,
29 INVALID_PROFILE_DESC_ID,
30 INVALID_PROFILE_OWNER_ID,
31 MIN_CUBE_WORKER_CAPACITY,
32 RuntimeConfigC,
33 TaskDescC,
34 TaskType,
35)
36from hyper_multicore.diagnostics.common.graph import _scheduled_task_ids
37from hyper_multicore.diagnostics.common.runtime import DiagnosticMode, PreparedDiagnosticRuntime
38from hyper_multicore.scheduler.runtime import runtime_cube_worker_capacity
41# Minimum AIC/AIV profile slots.
42CUBE_SLOT_COUNT = MIN_CUBE_WORKER_CAPACITY
43VECTOR_SLOT_COUNT = 2 * CUBE_SLOT_COUNT
44CORE_SLOT_COUNT = CUBE_SLOT_COUNT + VECTOR_SLOT_COUNT
45RECORD_CAPACITY_ALIGNMENT = 16
46CORE_HEADER = struct.Struct("<QIIIIII32x")
47PROFILE_RECORD = struct.Struct("<QQIIII")
48# Keep each slot representable by the Device ABI; allocate only the graph's actual requirement.
49MAX_RECORDS_PER_CORE = ((0xFFFFFFFF - CORE_HEADER.size) // PROFILE_RECORD.size
50 // RECORD_CAPACITY_ALIGNMENT * RECORD_CAPACITY_ALIGNMENT)
51MAX_CORE_STRIDE_BYTES = CORE_HEADER.size + MAX_RECORDS_PER_CORE * PROFILE_RECORD.size
52MAX_PROFILE_BUFFER_BYTES = CORE_SLOT_COUNT * MAX_CORE_STRIDE_BYTES
53INVALID_OWNER_ID = 0xFFFFFFFF
54TASK_TYPE_DESC_BASE = 0x20000
55GRAPH_STAGE_DESC_BASE = 0x30000
56# GetSystemCycle uses the fixed system counter, not the variable AI Core clock.
57_SYSTEM_COUNTER_FREQUENCY_MHZ = {
58 "ascend910b": 50.0,
59 "ascend910c": 50.0,
60 "ascend910_93": 50.0,
61 "ascend950": 1000.0,
62}
64DEFAULT_STAGE_NAMES = {
65 0x10000: "WaitDependency",
66 0x10006: "TerminateTask",
67 0x10007: "TriggerEvent",
68}
69TASK_TYPE_NAMES = {
70 0: "Terminate",
71 10: "BeginTaskGraph",
72 101: "AddCustom",
73 102: "SwiGLU",
74 103: "MatMul",
75 104: "GroupedMatMul",
76 105: "ShmemPutMemSignal",
77 106: "SwiGLUGrad",
78 107: "MhcPost",
79 108: "MhcNormCast",
80 109: "MhcProjection",
81 110: "MhcMapping",
82 111: "RmsNorm",
83 112: "MhcInputMix",
84 113: "RmsNormGrad",
85 114: "MhcGradPrevA",
86 115: "MhcGradMapping",
87 116: "MhcGradPhiRms",
88 117: "MhcGradPrevXAndPost",
89 118: "MhcPostGrad",
90 119: "MhcGradPrevAAndMapping",
91 141: "ShmemGetMem",
92}
93CORE_TYPE_NAMES = {
94 1: "AIC",
95 2: "AIV",
96}
99class _TaskExecutionMode(Enum):
100 """Describe how Device workers execute scheduled task descriptors.
102 ``EVENT_DRIVEN`` assigns a prebuilt task graph across AIC and AIV worker
103 queues. Device events enforce graph dependencies, allowing independent
104 Cube and Vector tasks to overlap.
106 ``AIV_PIPELINE`` makes every active AIV worker execute the same ordered
107 stage sequence on its own data shard without scheduler events.
108 """
110 EVENT_DRIVEN = "event_driven"
111 AIV_PIPELINE = "aiv_pipeline"
114@dataclass(frozen=True)
115class _ProfileMetadata:
116 """Host-only display metadata selected by a concrete MegaKernel builder."""
118 kernel_name: str
119 owner_label: str
120 stage_names: Mapping[int, str]
121 task_stage_names: Mapping[int, str]
122 execution_mode: _TaskExecutionMode
123 vector_worker_stride: int = 2
126@dataclass(frozen=True)
127class _ProfileLayout:
128 """Host-computed per-worker record requirements and bounded buffer layout."""
130 aic_required_records: int
131 aiv_required_records: int
132 aic_record_capacity: int
133 aiv_record_capacity: int
134 cube_slot_count: int = CUBE_SLOT_COUNT
136 @property
137 def buffer_size(self) -> int:
138 """Return the ordinary Device buffer size for per-core cycle records."""
139 return _profile_buffer_bytes_for_capacities(
140 self.aic_record_capacity,
141 self.aiv_record_capacity,
142 self.cube_slot_count,
143 )
146@dataclass(frozen=True)
147class _CycleTraceConfig:
148 """Resolved hardware and display metadata used while decoding one buffer."""
150 rank: int
151 device_id: int | None
152 cycle_frequency_mhz: float
153 detailed_task_names: bool
154 kernel_name: str
155 owner_label: str
156 stage_names: Mapping[int, str]
157 soc_name: str
158 task_stage_names: Mapping[int, str] | None = None
159 cube_slot_count: int = CUBE_SLOT_COUNT
162def _get_soc_name(device_id: int | None) -> str:
163 """Return and validate the current Ascend SoC name through Torch NPU."""
164 # Metadata construction and CPU/meta modules must not import an optional backend.
165 import torch_npu # pylint: disable=import-outside-toplevel
167 get_device_name = torch_npu.npu.get_device_name
168 try:
169 soc_name = get_device_name(device_id)
170 except TypeError:
171 # Older torch_npu releases only expose the current-device form.
172 soc_name = get_device_name()
173 if not isinstance(soc_name, str) or not soc_name.strip():
174 raise RuntimeError(f"Torch NPU returned an invalid Ascend SoC name: {soc_name!r}")
175 return soc_name.strip()
178def _resolve_cycle_frequency_mhz(soc_name: str) -> float:
179 """Resolve the GetSystemCycle frequency for a backend-reported SoC name."""
180 normalized_name = "".join(character for character in soc_name.lower() if character.isalnum() or character == "_")
181 if not normalized_name.startswith("ascend"):
182 normalized_name = f"ascend{normalized_name}"
183 for soc_prefix, frequency_mhz in _SYSTEM_COUNTER_FREQUENCY_MHZ.items():
184 if normalized_name.startswith(soc_prefix):
185 return frequency_mhz
186 supported_soc_names = ", ".join(sorted(_SYSTEM_COUNTER_FREQUENCY_MHZ))
187 raise RuntimeError(
188 f"Unsupported Ascend SoC for MegaKernel profiling: {soc_name!r}; "
189 f"supported SoC prefixes: {supported_soc_names}"
190 )
193def _validate_runtime_config(runtime_config: RuntimeConfigC) -> None:
194 if not isinstance(runtime_config, RuntimeConfigC):
195 raise TypeError(f"runtime_config must be RuntimeConfigC, got {type(runtime_config).__name__}")
198def _profile_buffer_bytes_for_capacities(
199 aic_record_capacity: int, aiv_record_capacity: int, cube_slot_count: int = CUBE_SLOT_COUNT,
200) -> int:
201 aic_stride = CORE_HEADER.size + aic_record_capacity * PROFILE_RECORD.size
202 aiv_stride = CORE_HEADER.size + aiv_record_capacity * PROFILE_RECORD.size
203 return cube_slot_count * aic_stride + 2 * cube_slot_count * aiv_stride
206def _round_up_record_capacity(required_records: int) -> int:
207 """Round the full per-worker requirement without silently dropping records."""
208 aligned_capacity = (
209 (required_records + RECORD_CAPACITY_ALIGNMENT - 1) // RECORD_CAPACITY_ALIGNMENT * RECORD_CAPACITY_ALIGNMENT
210 )
211 _validate_record_capacity(aligned_capacity, "per-worker")
212 return aligned_capacity
215def _records_for_task(task_desc: TaskDescC) -> int:
216 """Mirror the wait, compute, and trigger records emitted by ExecuteTaskProfiled()."""
217 record_count = 1
218 if task_desc.dependent_event != EVENT_INVALID_ID:
219 record_count += 1
220 if task_desc.task_type != TaskType.TASK_SHMEM_PUT_MEM_SIGNAL:
221 record_count += 1
222 return record_count
225def _max_worker_record_count(
226 runtime_config: RuntimeConfigC,
227 task_indices: Any,
228 scheduled_task_count: int,
229 active_worker_count: int,
230) -> int:
231 """Return the largest record count assigned to one active worker."""
232 if scheduled_task_count < 0 or scheduled_task_count > len(task_indices):
233 raise ValueError(f"scheduled task count is outside RuntimeConfig capacity: {scheduled_task_count}")
234 if scheduled_task_count == 0:
235 return 0
236 if active_worker_count <= 0:
237 raise ValueError("RuntimeConfig.num_workers must identify at least one AIC/AIV worker")
239 worker_record_counts = [0] * active_worker_count
240 for schedule_index in range(scheduled_task_count):
241 task_id = int(task_indices[schedule_index])
242 if not 0 <= task_id < len(runtime_config.all_tasks):
243 raise ValueError(f"scheduled task ID is outside RuntimeConfig capacity: {task_id}")
244 worker_id = schedule_index % active_worker_count
245 worker_record_counts[worker_id] += _records_for_task(runtime_config.all_tasks[task_id])
246 return max(worker_record_counts)
249def _calculate_aiv_pipeline_profile_layout(
250 runtime_config: RuntimeConfigC,
251) -> _ProfileLayout:
252 """Size AIV slots for workers that each execute the complete stage sequence."""
253 worker_capacity = int(runtime_config.num_workers)
254 if worker_capacity <= 0:
255 raise ValueError(f"AIV pipeline worker capacity must be positive, got {worker_capacity}")
256 scheduled_task_count = int(runtime_config.task_index_num[1])
257 if not 0 <= scheduled_task_count <= len(runtime_config.vector_task_indices):
258 raise ValueError(f"scheduled task count is outside RuntimeConfig capacity: {scheduled_task_count}")
259 required_records = 0
260 for schedule_index in range(scheduled_task_count):
261 task_id = int(runtime_config.vector_task_indices[schedule_index])
262 if not 0 <= task_id < len(runtime_config.all_tasks):
263 raise ValueError(f"scheduled task ID is outside RuntimeConfig capacity: {task_id}")
264 task = runtime_config.all_tasks[task_id]
265 if task.dependent_event != EVENT_INVALID_ID or task.trigger_event != EVENT_INVALID_ID:
266 raise ValueError("AIV pipelines cannot contain scheduler events")
267 required_records += 1
268 return _ProfileLayout(
269 aic_required_records=0,
270 aiv_required_records=required_records,
271 aic_record_capacity=0,
272 aiv_record_capacity=_round_up_record_capacity(required_records),
273 cube_slot_count=runtime_cube_worker_capacity(worker_capacity),
274 )
277def _calculate_event_driven_profile_layout(
278 runtime_config: RuntimeConfigC, vector_worker_stride: int = 2,
279) -> _ProfileLayout:
280 """Size AIC/AIV slots for event-driven task queues using the busiest worker."""
281 num_workers = int(runtime_config.num_workers)
282 if num_workers < 0 or num_workers % 2 != 0:
283 raise ValueError(f"RuntimeConfig.num_workers must be a nonnegative even value, got {num_workers}")
284 active_worker_count = num_workers // 2
285 aic_required_records = _max_worker_record_count(
286 runtime_config,
287 runtime_config.cube_task_indices,
288 int(runtime_config.task_index_num[0]),
289 active_worker_count,
290 )
291 aiv_required_records = _max_worker_record_count(
292 runtime_config,
293 runtime_config.vector_task_indices,
294 int(runtime_config.task_index_num[1]),
295 num_workers // vector_worker_stride,
296 )
297 return _ProfileLayout(
298 aic_required_records=aic_required_records,
299 aiv_required_records=aiv_required_records,
300 aic_record_capacity=_round_up_record_capacity(aic_required_records),
301 aiv_record_capacity=_round_up_record_capacity(aiv_required_records),
302 cube_slot_count=runtime_cube_worker_capacity(num_workers),
303 )
306def _calculate_profile_layout(
307 runtime_config: RuntimeConfigC, metadata: _ProfileMetadata | None = None,
308) -> _ProfileLayout:
309 """Select the profile layout algorithm for the configured task execution mode."""
310 _validate_runtime_config(runtime_config)
311 if metadata is None:
312 metadata = _ProfileMetadata("MegaKernel", "Owner", {}, {}, _TaskExecutionMode.EVENT_DRIVEN)
313 if metadata.execution_mode is _TaskExecutionMode.AIV_PIPELINE:
314 return _calculate_aiv_pipeline_profile_layout(runtime_config)
315 if metadata.execution_mode is _TaskExecutionMode.EVENT_DRIVEN:
316 return _calculate_event_driven_profile_layout(runtime_config, metadata.vector_worker_stride)
317 raise ValueError(f"unsupported task execution mode: {metadata.execution_mode!r}")
320def _configure_profile_layout(runtime_config: RuntimeConfigC, metadata: _ProfileMetadata) -> _ProfileLayout:
321 """Calculate and serialize the bounded AIC/AIV record capacities."""
322 layout = _calculate_profile_layout(runtime_config, metadata)
323 runtime_config.aic_profile_record_capacity = layout.aic_record_capacity
324 runtime_config.aiv_profile_record_capacity = layout.aiv_record_capacity
325 return layout
328def _resolve_stage_names(stage_names: Mapping[int, str] | None) -> dict[int, str]:
329 """Merge validated kernel-specific names with the common runtime stages."""
330 resolved_stage_names = dict(DEFAULT_STAGE_NAMES)
331 if stage_names is None:
332 return resolved_stage_names
333 for desc_id, stage_name in stage_names.items():
334 if isinstance(desc_id, bool) or not isinstance(desc_id, int):
335 raise TypeError(f"stage_names keys must be integers, got {desc_id!r}")
336 if not isinstance(stage_name, str) or not stage_name.strip():
337 raise ValueError(f"stage_names values must be non-empty strings, got {stage_name!r}")
338 resolved_stage_names[desc_id] = stage_name
339 return resolved_stage_names
342def _configure_trace_metadata(runtime_config: RuntimeConfigC, metadata: Any) -> _ProfileMetadata:
343 """Encode trace-specific descriptor IDs in an independent profiling configuration."""
344 stages = dict(DEFAULT_STAGE_NAMES)
345 stages.update({GRAPH_STAGE_DESC_BASE + key: name for key, name in metadata.stage_names.items()})
346 task_names = dict(metadata.task_stage_names)
347 for task_id in set(range(int(runtime_config.task_num))) | _scheduled_task_ids(runtime_config):
348 task = runtime_config.all_tasks[task_id]
349 stage_id = metadata.task_stage_ids.get(task_id)
350 if stage_id is not None:
351 task.profile_desc_id = GRAPH_STAGE_DESC_BASE + stage_id
352 owner = metadata.task_owner_ids.get(task_id)
353 task.profile_owner_id = INVALID_PROFILE_OWNER_ID if owner is None else owner
354 elif task.task_type == TaskType.TASK_TERMINATE:
355 task.profile_desc_id = 0x10006
356 task.profile_owner_id = INVALID_PROFILE_OWNER_ID
357 task_names[task_id] = "TerminateTask"
358 else:
359 task.profile_desc_id = INVALID_PROFILE_DESC_ID
360 task.profile_owner_id = INVALID_PROFILE_OWNER_ID
361 return _ProfileMetadata(metadata.kernel_name, metadata.owner_label, stages, task_names,
362 _TaskExecutionMode.EVENT_DRIVEN)
365class _PreparedProfileRuntime:
366 """Own profiling-specific preparation and trace decoding for one runtime."""
368 def __init__(self, runtime: PreparedDiagnosticRuntime) -> None:
369 """Build only profiling-owned metadata, layout and Device configuration."""
370 runtime_config = runtime.copy_host_config()
371 self._metadata = _configure_trace_metadata(runtime_config, runtime.metadata)
372 self._metadata = _ProfileMetadata(
373 self._metadata.kernel_name, self._metadata.owner_label, self._metadata.stage_names,
374 self._metadata.task_stage_names, _TaskExecutionMode(runtime.spec.execution_mode),
375 runtime.spec.vector_worker_stride,
376 )
377 self._layout = _configure_profile_layout(runtime_config, self._metadata)
378 self._rank = runtime.rank
379 self._device_id = runtime.device_id
380 self._host_config = runtime_config
381 self._normal_tensor = runtime.normal_tensor
382 self._device = getattr(runtime.normal_tensor, "device", None)
383 runtime_config.diagnostic_mode = DiagnosticMode.PROFILE
384 self._profile_tensor = runtime.tensor_from_config(runtime_config)
386 @property
387 def normal_tensor(self) -> Any:
388 """Return the original normal image without modifying it."""
389 return self._normal_tensor
391 @property
392 def profile_tensor(self) -> Any:
393 """Return the independent profiling configuration."""
394 return self._profile_tensor
396 @property
397 def buffer_size(self) -> int:
398 """Return the exact ordinary Device buffer size required while profiling."""
399 return self._layout.buffer_size
401 @property
402 def rank(self) -> int:
403 """Return the distributed rank that owns this runtime."""
404 return self._rank
406 @property
407 def device(self) -> Any:
408 """Return the device used to allocate runtime-owned tensors."""
409 return self._device
411 @property
412 def device_id(self) -> int:
413 """Return the local device index that executes this runtime."""
414 return self._device_id
416 @property
417 def kernel_name(self) -> str:
418 """Return the concrete kernel display name."""
419 return self._metadata.kernel_name
421 @property
422 def host_config(self) -> RuntimeConfigC:
423 """Return an immutable Host copy retaining the graph-sized runtime arrays."""
424 return self._host_config
426 @property
427 def task_stage_names(self) -> Mapping[int, str]:
428 """Return task ID to human-facing stage-name metadata."""
429 return self._metadata.task_stage_names
431 @property
432 def owner_label(self) -> str:
433 """Return the concrete kernel's task-owner label."""
434 return self._metadata.owner_label
436 def parse(self, buffer: Any, *, detailed_task_names: bool) -> dict[str, Any]:
437 """Decode one completed invocation buffer.
439 Args:
440 buffer: Host-visible bytes copied from the profiling buffer.
441 detailed_task_names: Whether event names include owner and task details.
442 """
443 soc_name = _get_soc_name(self._device_id)
444 trace = _parse_cycle_buffer(
445 buffer,
446 _CycleTraceConfig(
447 rank=self._rank,
448 device_id=self._device_id,
449 cycle_frequency_mhz=_resolve_cycle_frequency_mhz(soc_name),
450 detailed_task_names=detailed_task_names,
451 kernel_name=self._metadata.kernel_name,
452 owner_label=self._metadata.owner_label,
453 stage_names=self._metadata.stage_names,
454 soc_name=soc_name,
455 task_stage_names=self._metadata.task_stage_names,
456 cube_slot_count=self._layout.cube_slot_count,
457 ),
458 aic_record_capacity=self._layout.aic_record_capacity,
459 aiv_record_capacity=self._layout.aiv_record_capacity,
460 )
461 workers = int(self._host_config.num_workers)
462 pipeline = self._metadata.execution_mode is _TaskExecutionMode.AIV_PIPELINE
463 trace["megaKernelCycleTrace"]["workerCapacities"] = {
464 "cubeSlots": self._layout.cube_slot_count,
465 "vectorSlots": 2 * self._layout.cube_slot_count,
466 "activeCubeWorkers": 0 if pipeline else workers // 2,
467 "activeVectorWorkers": workers // self._metadata.vector_worker_stride,
468 "vectorWorkerStride": self._metadata.vector_worker_stride,
469 }
470 return trace
473def _as_bytes(buffer: Any, cube_slot_count: int = CUBE_SLOT_COUNT) -> bytes:
474 """Convert an array-like profile buffer to validated immutable bytes."""
475 if hasattr(buffer, "tobytes"):
476 raw_buffer = buffer.tobytes()
477 else:
478 raw_buffer = bytes(buffer)
479 minimum_buffer_bytes = 3 * cube_slot_count * CORE_HEADER.size
480 if len(raw_buffer) < minimum_buffer_bytes:
481 raise ValueError(
482 f"MegaKernel profile buffer is too small: expected at least {minimum_buffer_bytes} bytes, "
483 f"got {len(raw_buffer)}"
484 )
485 return raw_buffer
488def _validate_record_capacity(record_capacity: int, core_name: str) -> None:
489 if record_capacity < 0:
490 raise ValueError(f"{core_name} profile record capacity must be nonnegative: {record_capacity}")
491 if record_capacity > MAX_RECORDS_PER_CORE:
492 raise ValueError(
493 f"{core_name} profile record capacity exceeds the {MAX_RECORDS_PER_CORE}-record limit: "
494 f"{record_capacity}"
495 )
496 if record_capacity % RECORD_CAPACITY_ALIGNMENT != 0:
497 raise ValueError(
498 f"{core_name} profile record capacity must be a multiple of {RECORD_CAPACITY_ALIGNMENT}: "
499 f"{record_capacity}"
500 )
503def _profile_layout_from_buffer(
504 raw_buffer: bytes, cube_slot_count: int = CUBE_SLOT_COUNT,
505) -> tuple[int, int, int]:
506 """Read AIC/AIV capacities from the first initialized header of each slot region."""
507 aic_header = CORE_HEADER.unpack_from(raw_buffer, 0)
508 aic_record_capacity = aic_header[5]
509 _validate_record_capacity(aic_record_capacity, "AIC")
510 aic_stride = CORE_HEADER.size + aic_record_capacity * PROFILE_RECORD.size
511 aiv_region_offset = cube_slot_count * aic_stride
512 if len(raw_buffer) < aiv_region_offset + CORE_HEADER.size:
513 raise ValueError(
514 f"MegaKernel profile buffer is too small for the AIV header: "
515 f"expected at least {aiv_region_offset + CORE_HEADER.size} bytes, got {len(raw_buffer)}"
516 )
517 aiv_header = CORE_HEADER.unpack_from(raw_buffer, aiv_region_offset)
518 aiv_record_capacity = aiv_header[5]
519 _validate_record_capacity(aiv_record_capacity, "AIV")
520 required_buffer_bytes = _profile_buffer_bytes_for_capacities(
521 aic_record_capacity,
522 aiv_record_capacity,
523 cube_slot_count,
524 )
525 if len(raw_buffer) < required_buffer_bytes:
526 raise ValueError(
527 f"MegaKernel profile buffer is too small for its declared capacities: "
528 f"expected at least {required_buffer_bytes} bytes, got {len(raw_buffer)}"
529 )
530 return aic_record_capacity, aiv_record_capacity, required_buffer_bytes
533def _profile_slot_layouts(
534 aic_record_capacity: int, aiv_record_capacity: int, cube_slot_count: int = CUBE_SLOT_COUNT,
535):
536 aic_stride = CORE_HEADER.size + aic_record_capacity * PROFILE_RECORD.size
537 aiv_stride = CORE_HEADER.size + aiv_record_capacity * PROFILE_RECORD.size
538 for block_id in range(cube_slot_count):
539 yield 1, block_id, block_id * aic_stride, aic_record_capacity
540 aiv_region_offset = cube_slot_count * aic_stride
541 for block_id in range(2 * cube_slot_count):
542 yield 2, block_id, aiv_region_offset + block_id * aiv_stride, aiv_record_capacity
545def _thread_id(core_type: int, block_id: int) -> int:
546 return core_type * 1000 + block_id
549def _event_name(
550 desc_id: int,
551 task_id: int,
552 stage_task_index: int,
553 owner_id: int,
554 detailed_task_names: bool,
555 owner_label: str,
556 stage_names: Mapping[int, str],
557 task_stage_names: Mapping[int, str],
558) -> str:
559 """Build a stage-only or detailed task name while retaining raw IDs in event args."""
560 if desc_id in stage_names:
561 stage_name = stage_names[desc_id]
562 elif desc_id >= TASK_TYPE_DESC_BASE:
563 task_type = desc_id - TASK_TYPE_DESC_BASE
564 stage_name = f"TaskType_{TASK_TYPE_NAMES.get(task_type, task_type)}"
565 else:
566 stage_name = f"DescId_{desc_id}"
567 task_stage_name = task_stage_names.get(task_id)
568 if desc_id in (0x10000, 0x10007) and task_stage_name is not None:
569 stage_name = f"{task_stage_name}_{stage_name}"
570 if not detailed_task_names:
571 return stage_name
572 task_name = f"{stage_name}_task{stage_task_index + 1}"
573 if owner_id == INVALID_OWNER_ID:
574 return task_name
575 return f"{owner_label}{owner_id}_{task_name}"
578def _resolve_buffer_layout(
579 raw_buffer: bytes,
580 aic_record_capacity: int | None,
581 aiv_record_capacity: int | None,
582 cube_slot_count: int = CUBE_SLOT_COUNT,
583) -> tuple[int, int, int]:
584 """Resolve and validate the two per-core capacities and total buffer size."""
585 if aic_record_capacity is None and aiv_record_capacity is None:
586 return _profile_layout_from_buffer(raw_buffer, cube_slot_count)
587 if aic_record_capacity is None or aiv_record_capacity is None:
588 raise ValueError("AIC and AIV profile capacities must be provided together")
589 _validate_record_capacity(aic_record_capacity, "AIC")
590 _validate_record_capacity(aiv_record_capacity, "AIV")
591 required_buffer_bytes = _profile_buffer_bytes_for_capacities(
592 aic_record_capacity,
593 aiv_record_capacity,
594 cube_slot_count,
595 )
596 if len(raw_buffer) < required_buffer_bytes:
597 raise ValueError(
598 "MegaKernel profile buffer is smaller than its Host-computed layout: "
599 f"expected at least {required_buffer_bytes} bytes, got {len(raw_buffer)}"
600 )
601 return aic_record_capacity, aiv_record_capacity, required_buffer_bytes
604def _validate_profile_slot_header(
605 slot: int,
606 core_type: int,
607 block_id: int,
608 expected_core_type: int,
609 expected_block_id: int,
610 header_capacity: int,
611 record_capacity: int,
612 record_count: int,
613 reserved: int,
614) -> None:
615 """Validate the identity, layout, and count fields of one Device slot."""
616 if core_type != expected_core_type or block_id != expected_block_id:
617 raise ValueError(
618 f"Profile slot identity mismatch: slot={slot}, expected=({expected_core_type}, {expected_block_id}), "
619 f"got=({core_type}, {block_id})"
620 )
621 if header_capacity != record_capacity:
622 raise ValueError(
623 f"Profile slot capacity mismatch: slot={slot}, expected={record_capacity}, got={header_capacity}"
624 )
625 if reserved != 0:
626 raise ValueError(f"Profile core header reserved field must be zero: slot={slot}, got={reserved}")
627 if record_count > record_capacity:
628 raise ValueError(
629 f"Profile record count exceeds capacity: slot={slot}, count={record_count}, "
630 f"capacity={record_capacity}"
631 )
634def _decode_slot_records(
635 raw_buffer: bytes,
636 slot: int,
637 slot_offset: int,
638 record_count: int,
639 core_type: int,
640 block_id: int,
641 entry_cycle: int,
642) -> list[dict[str, int]]:
643 """Decode all cycle records stored after one validated core header."""
644 records = []
645 record_offset = slot_offset + CORE_HEADER.size
646 for record_index in range(record_count):
647 record = PROFILE_RECORD.unpack_from(raw_buffer, record_offset + record_index * PROFILE_RECORD.size)
648 if record[1] < record[0]:
649 raise ValueError(
650 f"Profile cycle interval is negative: slot={slot}, record={record_index}, "
651 f"start={record[0]}, end={record[1]}"
652 )
653 records.append(
654 {
655 "start_cycle": record[0],
656 "end_cycle": record[1],
657 "desc_id": record[2],
658 "task_id": record[3],
659 "stage_task_index": record[4],
660 "owner_id": record[5],
661 "core_type": core_type,
662 "block_id": block_id,
663 "entry_cycle": entry_cycle,
664 }
665 )
666 return records
669def _decode_profile_slot(
670 raw_buffer: bytes,
671 slot: int,
672 layout: tuple[int, int, int, int],
673) -> tuple[list[dict[str, int]], int | None, int]:
674 """Decode one initialized slot and return records, entry cycle, and drops."""
675 expected_core_type, expected_block_id, slot_offset, record_capacity = layout
676 header = CORE_HEADER.unpack_from(raw_buffer, slot_offset)
677 if not any(header):
678 return [], None, 0
679 entry_cycle, record_count, dropped_count, core_type, block_id, header_capacity, reserved = header
680 _validate_profile_slot_header(
681 slot,
682 core_type,
683 block_id,
684 expected_core_type,
685 expected_block_id,
686 header_capacity,
687 record_capacity,
688 record_count,
689 reserved,
690 )
691 if record_count == 0:
692 return [], None, dropped_count
693 return (
694 _decode_slot_records(
695 raw_buffer,
696 slot,
697 slot_offset,
698 record_count,
699 core_type,
700 block_id,
701 entry_cycle,
702 ),
703 entry_cycle or None,
704 dropped_count,
705 )
708def _decode_cycle_records(
709 raw_buffer: bytes,
710 aic_record_capacity: int,
711 aiv_record_capacity: int,
712 cube_slot_count: int = CUBE_SLOT_COUNT,
713) -> tuple[list[dict[str, int]], list[int], int]:
714 """Decode all initialized AIC/AIV slots and aggregate drop counters."""
715 raw_records = []
716 active_entry_cycles = []
717 dropped_records = 0
718 for slot, layout in enumerate(_profile_slot_layouts(aic_record_capacity, aiv_record_capacity, cube_slot_count)):
719 records, entry_cycle, dropped_count = _decode_profile_slot(raw_buffer, slot, layout)
720 dropped_records += dropped_count
721 if entry_cycle is not None:
722 active_entry_cycles.append(entry_cycle)
723 raw_records.extend(records)
724 return raw_records, active_entry_cycles, dropped_records
727def _thread_metadata_events(
728 rank: int,
729 device_id: int,
730 kernel_name: str,
731 raw_records: list[dict[str, int]],
732) -> list[dict[str, Any]]:
733 """Build process and thread metadata for all cores that emitted records."""
734 trace_events = [
735 {
736 "name": "process_name",
737 "ph": "M",
738 "pid": rank,
739 "tid": 0,
740 "args": {"name": f"{kernel_name} rank {rank} device {device_id}"},
741 }
742 ]
743 used_threads = sorted({(record["core_type"], record["block_id"]) for record in raw_records})
744 for sort_index, (core_type, block_id) in enumerate(used_threads):
745 thread_id = _thread_id(core_type, block_id)
746 trace_events.extend(
747 [
748 {
749 "name": "thread_name",
750 "ph": "M",
751 "pid": rank,
752 "tid": thread_id,
753 "args": {"name": f"{CORE_TYPE_NAMES[core_type]}/{block_id}"},
754 },
755 {
756 "name": "thread_sort_index",
757 "ph": "M",
758 "pid": rank,
759 "tid": thread_id,
760 "args": {"sort_index": sort_index},
761 },
762 ]
763 )
764 return trace_events
767def _duration_trace_event(
768 record: Mapping[str, int],
769 *,
770 rank: int,
771 device_id: int,
772 anchor_cycle: int,
773 cycle_frequency_mhz: float,
774 detailed_task_names: bool,
775 owner_label: str,
776 stage_names: Mapping[int, str],
777 task_stage_names: Mapping[int, str],
778) -> dict[str, Any]:
779 """Convert one raw cycle record into a Chrome Trace duration event."""
780 desc_id = record["desc_id"]
781 task_id = record["task_id"]
782 args = {
783 "rank": rank,
784 "device_id": device_id,
785 "core_type": CORE_TYPE_NAMES[record["core_type"]],
786 "block_id": record["block_id"],
787 "task_id": task_id,
788 "desc_id": desc_id,
789 "stage_task_index": record["stage_task_index"],
790 "task_number": record["stage_task_index"] + 1,
791 "start_cycle": record["start_cycle"],
792 "end_cycle": record["end_cycle"],
793 "core_entry_cycle": record["entry_cycle"],
794 }
795 if task_id in task_stage_names:
796 args["task_stage"] = task_stage_names[task_id]
797 if desc_id >= TASK_TYPE_DESC_BASE:
798 args["task_type"] = desc_id - TASK_TYPE_DESC_BASE
799 if record["owner_id"] != INVALID_OWNER_ID:
800 args["owner_id"] = record["owner_id"]
801 return {
802 "name": _event_name(
803 desc_id,
804 task_id,
805 record["stage_task_index"],
806 record["owner_id"],
807 detailed_task_names,
808 owner_label,
809 stage_names,
810 task_stage_names,
811 ),
812 "cat": "MegaKernelInternal",
813 "ph": "X",
814 "pid": rank,
815 "tid": _thread_id(record["core_type"], record["block_id"]),
816 "ts": (record["start_cycle"] - anchor_cycle) / cycle_frequency_mhz,
817 "dur": (record["end_cycle"] - record["start_cycle"]) / cycle_frequency_mhz,
818 "args": args,
819 }
822def _parse_cycle_buffer(
823 buffer: Any,
824 config: _CycleTraceConfig,
825 *,
826 aic_record_capacity: int | None = None,
827 aiv_record_capacity: int | None = None,
828) -> dict[str, Any]:
829 """Decode a device cycle buffer with already resolved hardware and display metadata."""
830 if config.cycle_frequency_mhz <= 0:
831 raise ValueError(f"cycle_frequency_mhz must be positive, got {config.cycle_frequency_mhz}")
832 resolved_stage_names = _resolve_stage_names(config.stage_names)
833 resolved_task_stage_names = dict(config.task_stage_names or {})
834 if config.cube_slot_count < CUBE_SLOT_COUNT:
835 raise ValueError(f"profile Cube slot count must be at least {CUBE_SLOT_COUNT}, got {config.cube_slot_count}")
836 raw_buffer = _as_bytes(buffer, config.cube_slot_count)
837 aic_record_capacity, aiv_record_capacity, required_buffer_bytes = _resolve_buffer_layout(
838 raw_buffer,
839 aic_record_capacity,
840 aiv_record_capacity,
841 config.cube_slot_count,
842 )
843 raw_records, active_entry_cycles, dropped_records = _decode_cycle_records(
844 raw_buffer,
845 aic_record_capacity,
846 aiv_record_capacity,
847 config.cube_slot_count,
848 )
849 if not raw_records:
850 raise ValueError(
851 "MegaKernel cycle profile buffer contains no records; verify that the current profiler schedule is active "
852 "and that the profiled operator completed"
853 )
854 anchor_cycle = (
855 min(active_entry_cycles) if active_entry_cycles else min(record["start_cycle"] for record in raw_records)
856 )
857 resolved_device_id = config.rank if config.device_id is None else config.device_id
858 trace_events = _thread_metadata_events(config.rank, resolved_device_id, config.kernel_name, raw_records)
859 for record in sorted(
860 raw_records,
861 key=lambda item: (item["start_cycle"], item["core_type"], item["block_id"]),
862 ):
863 trace_events.append(_duration_trace_event(
864 record,
865 rank=config.rank,
866 device_id=resolved_device_id,
867 anchor_cycle=anchor_cycle,
868 cycle_frequency_mhz=config.cycle_frequency_mhz,
869 detailed_task_names=config.detailed_task_names,
870 owner_label=config.owner_label,
871 stage_names=resolved_stage_names,
872 task_stage_names=resolved_task_stage_names,
873 ))
874 return {
875 "traceEvents": trace_events,
876 "megaKernelCycleTrace": {
877 "schemaVersion": 1,
878 "rank": config.rank,
879 "deviceId": resolved_device_id,
880 "socName": config.soc_name,
881 "cycleFrequencyMHz": config.cycle_frequency_mhz,
882 "anchorCycle": anchor_cycle,
883 "recordCapacityPerCore": {
884 "AIC": aic_record_capacity,
885 "AIV": aiv_record_capacity,
886 },
887 "maxRecordCapacityPerCore": MAX_RECORDS_PER_CORE,
888 "recordCount": len(raw_records),
889 "droppedRecordCount": dropped_records,
890 "profileBufferBytes": required_buffer_bytes,
891 "detailedTaskNames": config.detailed_task_names,
892 "kernelName": config.kernel_name,
893 "ownerLabel": config.owner_label,
894 "warnings": ([] if dropped_records == 0 else [f"Device dropped {dropped_records} cycle trace records"]),
895 },
896 }