Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / diagnostics / profiler / api.py: 90%
245 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"""Standalone Torch MegaKernel profiling API and capture lifecycle."""
17from __future__ import annotations
19__all__ = ["TorchMegaKernelProfiler", "mega_kernel_profile"]
21from collections.abc import Callable
22from dataclasses import dataclass, field
23import json
24from pathlib import Path
25from typing import Any
27import torch
29from hyper_multicore.diagnostics.common.session import (
30 _activate_session,
31 _deactivate_session,
32)
33from hyper_multicore.diagnostics.profiler.cycle_trace import _PreparedProfileRuntime
34from hyper_multicore.diagnostics.common.call import DiagnosticCall, LaunchArguments, LaunchContext
35from hyper_multicore.diagnostics.common.schedule import DiagnosticAction
38@dataclass
39class _RawInvocation:
40 """One completed invocation copied from Device to Host."""
42 runtime: _PreparedProfileRuntime
43 data: bytes
44 direction: str
45 step: int
46 invocation_id: int
49@dataclass
50class _TraceAccumulator:
51 """Aggregate parsed invocations into one capture window."""
53 global_anchor: int
54 trace_events: list[dict[str, Any]] = field(default_factory=list)
55 metadata_keys: set[tuple[Any, Any, Any]] = field(default_factory=set)
56 dropped_count: int = 0
57 record_count: int = 0
58 invocation_metadata: list[dict[str, Any]] = field(default_factory=list)
60 def append(self, invocation: _RawInvocation, trace: dict[str, Any]) -> None:
61 """Append one parsed invocation and shift it to the window anchor.
63 Args:
64 invocation: Invocation identity and captured bytes.
65 trace: Parsed standalone trace for the invocation.
66 """
67 metadata = trace["megaKernelCycleTrace"]
68 shift_us = (metadata["anchorCycle"] - self.global_anchor) / metadata["cycleFrequencyMHz"]
69 self.dropped_count += metadata["droppedRecordCount"]
70 self.record_count += metadata["recordCount"]
71 invocation_metadata = {
72 "invocationId": invocation.invocation_id,
73 "step": invocation.step,
74 "direction": invocation.direction,
75 "kernelName": invocation.runtime.kernel_name,
76 "rank": metadata.get("rank", invocation.runtime.rank),
77 "deviceId": metadata.get("deviceId", invocation.runtime.device_id),
78 "anchorCycle": metadata["anchorCycle"],
79 "recordCount": metadata["recordCount"],
80 "droppedRecordCount": metadata["droppedRecordCount"],
81 }
82 for key in ('workerCapacities', 'profileBufferBytes'):
83 if key in metadata:
84 invocation_metadata[key] = metadata[key]
85 self.invocation_metadata.append(invocation_metadata)
86 self._append_events(trace["traceEvents"], invocation, shift_us)
88 def _append_events(
89 self,
90 events: list[dict[str, Any]],
91 invocation: _RawInvocation,
92 shift_us: float,
93 ) -> None:
94 """Deduplicate metadata events and shift duration events."""
95 for event in events:
96 if event["ph"] == "M":
97 key = (event["name"], event["pid"], event["tid"])
98 if key not in self.metadata_keys:
99 self.metadata_keys.add(key)
100 self.trace_events.append(event)
101 continue
102 copied_event = dict(event)
103 copied_event["ts"] += shift_us
104 copied_event["args"] = {
105 **event.get("args", {}),
106 "invocation_id": invocation.invocation_id,
107 "step": invocation.step,
108 "direction": invocation.direction,
109 }
110 self.trace_events.append(copied_event)
113class _ProfileCall(DiagnosticCall):
114 """Internal launch arguments and profiling-buffer ownership for one kernel call."""
116 def __init__(
117 self,
118 *,
119 runtime_config: Any,
120 event_counters: Any,
121 diagnostic_buffer: Any,
122 clear_event_counters: Any | None = None,
123 profiler: TorchMegaKernelProfiler | None = None,
124 runtime: _PreparedProfileRuntime | None = None,
125 profile_buffer_key: tuple[int, int] | None = None,
126 direction: str = "",
127 step: int = 0,
128 invocation_id: int = 0,
129 ) -> None:
130 """Store launch arguments and optional active-session bookkeeping."""
131 super().__init__(LaunchArguments(
132 runtime_config, event_counters, diagnostic_buffer,
133 event_counters if clear_event_counters is None else clear_event_counters,
134 ))
135 self._profiler = profiler
136 self.runtime = runtime
137 self.profile_buffer_key = profile_buffer_key
138 self.direction = direction
139 self.step = step
140 self.invocation_id = invocation_id
141 self.finished = False
143 def complete(self) -> None:
144 """Mark the direct-write profile buffer as complete after a successful launch."""
145 if self.finished:
146 return
147 if self._profiler is not None:
148 self._profiler.complete_call(self)
149 self.finished = True
151 def cancel(self) -> None:
152 """Discard an unlaunched call and release its private profile buffer."""
153 if self.finished:
154 return
155 if self._profiler is not None:
156 self._profiler.cancel_call(self)
157 self.finished = True
160class TorchMegaKernelProfiler:
161 """Standalone schedule-driven MegaKernel internal profiler for Torch."""
163 def __init__(
164 self,
165 *,
166 schedule: Callable[[int], DiagnosticAction] | None,
167 on_trace_ready: Callable[[Any], None] | None,
168 detailed_task_names: bool,
169 max_pending_calls: int,
170 ) -> None:
171 # pylint: disable=redefined-outer-name
172 """Validate options and initialize the schedule-driven capture state."""
173 if schedule is not None and not callable(schedule):
174 raise TypeError("schedule must be callable or None")
175 if on_trace_ready is not None and not callable(on_trace_ready):
176 raise TypeError("on_trace_ready must be callable or None")
177 if not isinstance(detailed_task_names, bool):
178 raise TypeError("detailed_task_names must be bool")
179 if not isinstance(max_pending_calls, int) or isinstance(max_pending_calls, bool) or max_pending_calls <= 0:
180 raise ValueError("max_pending_calls must be a positive integer, " f"got {max_pending_calls!r}")
181 self._schedule = schedule
182 self._on_trace_ready = on_trace_ready
183 self._detailed_task_names = detailed_task_names
184 self._max_pending_calls = max_pending_calls
185 self._step = 0
186 self._action = self._action_for_step(0)
187 self._entered = False
188 self._next_invocation_id = 0
189 self._pending: list[_ProfileCall] = []
190 self._profile_buffer_pool: dict[tuple[int, int], list[Any]] = {}
191 self._profile_runtimes: dict[Any, _PreparedProfileRuntime] = {}
192 self._raw_invocations: list[_RawInvocation] = []
193 self._last_trace: dict[str, Any] | None = None
194 self._window_index = 0
196 @property
197 def step_num(self) -> int:
198 """Return the zero-based step currently being recorded."""
199 return self._step
201 @property
202 def current_action(self) -> DiagnosticAction:
203 """Return the current schedule action."""
204 return self._action
206 def __enter__(self) -> TorchMegaKernelProfiler:
207 """Activate this profiler as the process-local MegaKernel session."""
208 _activate_session("profiler", self)
209 self._entered = True
210 return self
212 def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
213 """Finalize successful captures or discard pending data on failure."""
214 try:
215 if exc_type is None:
216 self._finish_context()
217 else:
218 self._drain(keep=False)
219 finally:
220 self._raw_invocations.clear()
221 for call in self._pending:
222 self._release_profile_buffer(call, cache=False)
223 self._pending.clear()
224 self._release_profile_buffer_pool()
225 _deactivate_session(self)
226 self._profile_runtimes.clear()
227 self._entered = False
229 def step(self) -> None:
230 """Finish the current schedule step and advance to the next one."""
231 if not self._entered:
232 raise RuntimeError("profiler.step() must be called inside its context")
233 next_action = self._action_for_step(self._step + 1)
234 if self._action is DiagnosticAction.WARMUP:
235 self._drain(keep=False)
236 elif self._action is DiagnosticAction.RECORD_AND_SAVE:
237 self._drain(keep=True)
238 self._finalize_window()
239 elif self._action is DiagnosticAction.RECORD and next_action not in (
240 DiagnosticAction.RECORD,
241 DiagnosticAction.RECORD_AND_SAVE,
242 ):
243 self._drain(keep=True)
244 self._finalize_window()
245 self._step += 1
246 self._action = next_action
248 def export_chrome_trace(self, path: str | Path) -> dict[str, Any]:
249 """Write the most recently completed internal capture window.
251 Args:
252 path: Destination path for the Chrome Trace JSON object.
253 """
254 if self._last_trace is None:
255 raise RuntimeError(
256 "no completed MegaKernel capture window is available; "
257 "call step() at the window boundary or leave the context first"
258 )
259 output_path = Path(path)
260 output_path.parent.mkdir(parents=True, exist_ok=True)
261 output_path.write_text(
262 json.dumps(self._last_trace, ensure_ascii=False),
263 encoding="utf-8",
264 )
265 return self._last_trace
267 def _action_for_step(self, step: int) -> DiagnosticAction:
268 if self._schedule is None:
269 return DiagnosticAction.RECORD
270 action = self._schedule(step)
271 if not isinstance(action, DiagnosticAction):
272 raise TypeError("MegaKernel schedule must return DiagnosticAction, " f"got {type(action).__name__}")
273 return action
275 def is_capture_enabled(self) -> bool:
276 """Return whether the current schedule action records Device cycles."""
277 return self._action is not DiagnosticAction.NONE
279 def prepare_call(self, context: LaunchContext) -> DiagnosticCall:
280 """Select profiling independently, without constructing any dump state.
282 Args:
283 context: Selected invocation or operand context supplied by the diagnostic framework.
284 """
285 if not self.is_capture_enabled():
286 return DiagnosticCall(context.ordinary_arguments())
287 runtime = self._profile_runtimes.get(context.runtime)
288 if runtime is None:
289 runtime = _PreparedProfileRuntime(context.runtime)
290 self._profile_runtimes[context.runtime] = runtime
291 return self.acquire_call(runtime, context.fallback_event_counters, context.direction)
293 def acquire_call(
294 self,
295 runtime: _PreparedProfileRuntime,
296 event_counters: Any,
297 direction: str,
298 ) -> _ProfileCall:
299 """Reserve a profiling buffer and register one pending kernel call.
301 Args:
302 runtime: Prepared runtime configuration and profiling metadata.
303 event_counters: Device event-counter storage.
304 direction: Kernel direction label included in trace metadata.
305 """
306 if len(self._pending) >= self._max_pending_calls:
307 self._drain(keep=self._action is not DiagnosticAction.WARMUP)
308 invocation_id = self._next_invocation_id
309 profile_buffer, profile_buffer_key = self._take_profile_buffer(runtime)
310 self._next_invocation_id += 1
311 call = _ProfileCall(
312 runtime_config=runtime.profile_tensor,
313 event_counters=event_counters,
314 clear_event_counters=event_counters,
315 profiler=self,
316 diagnostic_buffer=profile_buffer,
317 runtime=runtime,
318 direction=direction,
319 step=self._step,
320 profile_buffer_key=profile_buffer_key,
321 invocation_id=invocation_id,
322 )
323 self._pending.append(call)
324 return call
326 def complete_call(self, call: _ProfileCall) -> None:
327 """Validate that a successfully launched call belongs to this session.
329 Args:
330 call: Pending kernel call that completed launch.
331 """
332 if call not in self._pending:
333 raise RuntimeError("completed MegaKernel profiling call is not pending")
335 def _take_profile_buffer(
336 self,
337 runtime: _PreparedProfileRuntime,
338 ) -> tuple[Any, tuple[int, int]]:
339 """Take one ordinary NPU buffer that is private to the next invocation."""
340 profile_buffer_key = (runtime.device_id, runtime.buffer_size)
341 available_buffers = self._profile_buffer_pool.get(profile_buffer_key)
342 if available_buffers:
343 profile_buffer = available_buffers.pop()
344 else:
345 profile_buffer = torch.empty(
346 (runtime.buffer_size,),
347 dtype=torch.uint8,
348 device=runtime.device,
349 )
350 profile_buffer.zero_()
351 return profile_buffer, profile_buffer_key
353 def cancel_call(self, call: _ProfileCall) -> None:
354 """Cancel a pending call and release its profiling buffer.
356 Args:
357 call: Pending kernel call to cancel.
358 """
359 if call in self._pending:
360 self._pending.remove(call)
361 self._release_profile_buffer(call, cache=False)
363 def _release_profile_buffer(
364 self,
365 call: _ProfileCall,
366 *,
367 cache: bool,
368 ) -> None:
369 """Return a completed buffer to the session pool or release its storage."""
370 profile_buffer = call.diagnostic_buffer
371 profile_buffer_key = call.profile_buffer_key
372 if profile_buffer is None or profile_buffer_key is None:
373 return
374 if cache:
375 self._profile_buffer_pool.setdefault(profile_buffer_key, []).append(profile_buffer)
376 else:
377 profile_buffer.untyped_storage().resize_(0)
378 call.arguments.diagnostic_buffer = None
379 call.profile_buffer_key = None
381 def _release_profile_buffer_pool(self) -> None:
382 """Release ordinary NPU buffers retained for reuse by this session."""
383 for available_buffers in self._profile_buffer_pool.values():
384 for profile_buffer in available_buffers:
385 profile_buffer.untyped_storage().resize_(0)
386 self._profile_buffer_pool.clear()
388 def _drain(self, *, keep: bool) -> None:
389 """Synchronize pending calls, optionally copy their records to Host."""
390 if not self._pending:
391 if not keep:
392 self._raw_invocations.clear()
393 return
394 unfinished = [call for call in self._pending if not call.finished]
395 if unfinished:
396 raise RuntimeError("cannot copy MegaKernel profiling buffers before launches complete")
398 device_ids = {call.runtime.device_id for call in self._pending if call.runtime is not None}
399 for device_id in device_ids:
400 torch.npu.synchronize(device_id)
402 for call in self._pending:
403 runtime = call.runtime
404 profile_buffer = call.diagnostic_buffer
405 if keep and runtime is not None and profile_buffer is not None:
406 host_data = profile_buffer.cpu().numpy().tobytes()
407 self._raw_invocations.append(
408 _RawInvocation(
409 runtime=runtime,
410 data=host_data,
411 direction=call.direction,
412 step=call.step,
413 invocation_id=call.invocation_id,
414 )
415 )
416 self._release_profile_buffer(call, cache=True)
417 self._pending.clear()
418 if not keep:
419 self._raw_invocations.clear()
421 def _finish_context(self) -> None:
422 if self._action is DiagnosticAction.WARMUP:
423 self._drain(keep=False)
424 return
425 if self._pending or self._raw_invocations:
426 self._drain(keep=True)
427 if self._raw_invocations:
428 self._finalize_window()
430 def _finalize_window(self) -> None:
431 try:
432 self._last_trace = self._build_trace(self._raw_invocations)
433 finally:
434 self._raw_invocations.clear()
435 self._window_index += 1
436 if self._on_trace_ready is not None:
437 self._on_trace_ready(self)
439 def _build_trace(self, invocations: list[_RawInvocation]) -> dict[str, Any]:
440 """Combine retained invocation data into one standalone Chrome Trace."""
441 if not invocations:
442 return {
443 "traceEvents": [],
444 "megaKernelCycleTrace": {
445 "schemaVersion": 1,
446 "windowIndex": self._window_index,
447 "invocationCount": 0,
448 "recordCount": 0,
449 "droppedRecordCount": 0,
450 "detailedTaskNames": self._detailed_task_names,
451 "invocations": [],
452 "warnings": [],
453 },
454 }
456 parsed_invocations = [
457 (
458 invocation,
459 invocation.runtime.parse(
460 invocation.data,
461 detailed_task_names=self._detailed_task_names,
462 ),
463 )
464 for invocation in invocations
465 ]
466 global_anchor = min(trace["megaKernelCycleTrace"]["anchorCycle"] for _, trace in parsed_invocations)
467 accumulator = _TraceAccumulator(global_anchor)
468 for invocation, trace in parsed_invocations:
469 accumulator.append(invocation, trace)
471 metadata_events = [event for event in accumulator.trace_events if event["ph"] == "M"]
472 duration_events = sorted(
473 (event for event in accumulator.trace_events if event["ph"] != "M"),
474 key=lambda event: (event["ts"], event["pid"], event["tid"]),
475 )
476 return {
477 "traceEvents": metadata_events + duration_events,
478 "megaKernelCycleTrace": {
479 "schemaVersion": 1,
480 "windowIndex": self._window_index,
481 "anchorCycle": global_anchor,
482 "invocationCount": len(invocations),
483 "recordCount": accumulator.record_count,
484 "droppedRecordCount": accumulator.dropped_count,
485 "detailedTaskNames": self._detailed_task_names,
486 "invocations": accumulator.invocation_metadata,
487 "warnings": (
488 []
489 if accumulator.dropped_count == 0
490 else [f"Device dropped {accumulator.dropped_count} cycle trace records"]
491 ),
492 },
493 }
496# The argument name intentionally mirrors torch.profiler.profile.
497def mega_kernel_profile(
498 *,
499 schedule: Callable[[int], DiagnosticAction] | None = None,
500 on_trace_ready: Callable[[Any], None] | None = None,
501 detailed_task_names: bool = False,
502 max_pending_calls: int = 16,
503) -> TorchMegaKernelProfiler:
504 # pylint: disable=redefined-outer-name
505 """Create the standalone MegaKernel profiler.
507 The returned context manager records only MegaKernel-internal events. It
508 does not start an external framework profiler, and its export operation
509 does not merge with an external trace.
511 Args:
512 schedule: Optional step schedule. Without one, all calls are retained.
513 on_trace_ready: Optional callback invoked after each completed window.
514 detailed_task_names: Include owner and task numbers in event names.
515 max_pending_calls: Maximum in-flight invocation buffers before an
516 intermediate D2H drain and reuse.
518 Returns:
519 Standalone MegaKernel profiler context manager.
520 """
521 return TorchMegaKernelProfiler(
522 schedule=schedule,
523 on_trace_ready=on_trace_ready,
524 detailed_task_names=detailed_task_names,
525 max_pending_calls=max_pending_calls,
526 )