Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / diagnostics / profiler / trace_merger.py: 88%
353 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"""Offline merger for framework and standalone MegaKernel Chrome traces."""
17from __future__ import annotations
19__all__ = ["merge_chrome_traces"]
21import copy
22from dataclasses import dataclass
23import json
24import math
25from pathlib import Path
26import re
27from typing import Any, TypeAlias
30DEFAULT_KERNEL_PATTERN = r"(?i)mega[_]?kernel"
31INTERNAL_EVENT_CATEGORY = "MegaKernelInternal"
32MERGE_SCHEMA_VERSION = 1
33SUPPORTED_CYCLE_TRACE_SCHEMA_VERSION = 1
35TraceObject: TypeAlias = dict[str, Any]
36TraceInput: TypeAlias = str | Path | TraceObject | list[dict[str, Any]]
39@dataclass(frozen=True)
40class _InvocationTrace:
41 """Internal events and display metadata for one MegaKernel invocation."""
43 invocation_id: Any
44 events: list[dict[str, Any]]
45 kernel_name: str
46 rank: Any
47 device_id: Any
48 direction: str | None
50 @property
51 def anchor_us(self) -> float:
52 """Return the first internal event timestamp."""
53 return min(_number(event.get("ts"), "ts") for event in self.events)
55 @property
56 def span_us(self) -> float:
57 """Return the relative end-to-end internal event span."""
58 anchor = self.anchor_us
59 return max(
60 _number(event.get("ts"), "ts") + _number(event.get("dur"), "dur") - anchor
61 for event in self.events
62 )
65def _read_trace(source: TraceInput, trace_name: str) -> TraceObject | list[dict[str, Any]]:
66 """Load a trace from a path or accept an in-memory trace object."""
67 if isinstance(source, (str, Path)):
68 path = Path(source)
69 try:
70 return json.loads(path.read_text(encoding="utf-8"))
71 except OSError as exc:
72 raise ValueError(f"Unable to read {trace_name} from {path}: {exc}") from exc
73 except json.JSONDecodeError as exc:
74 raise ValueError(f"Invalid JSON in {trace_name} {path}: {exc}") from exc
75 if isinstance(source, (dict, list)):
76 return source
77 raise TypeError(
78 f"{trace_name} must be a path, Chrome Trace object, or event list, "
79 f"got {type(source).__name__}"
80 )
83def _normalize_framework_trace(trace: Any) -> TraceObject:
84 """Return a detached framework trace object with a valid event list."""
85 if isinstance(trace, list):
86 if not all(isinstance(event, dict) for event in trace):
87 raise ValueError("framework_trace event list must contain only objects")
88 return {"traceEvents": copy.deepcopy(trace)}
89 if isinstance(trace, dict):
90 events = trace.get("traceEvents")
91 if not isinstance(events, list) or not all(isinstance(event, dict) for event in events):
92 raise ValueError("framework_trace must contain a traceEvents object list")
93 return copy.deepcopy(trace)
94 raise ValueError("framework_trace must be an event list or an object containing traceEvents")
97def _events_from_trace(trace: Any, trace_name: str) -> list[dict[str, Any]]:
98 if isinstance(trace, list):
99 events = trace
100 elif isinstance(trace, dict):
101 events = trace.get("traceEvents")
102 else:
103 events = None
104 if not isinstance(events, list) or not all(isinstance(event, dict) for event in events):
105 raise ValueError(f"{trace_name} must contain a traceEvents object list")
106 return events
109def _number(value: Any, field_name: str) -> float:
110 if isinstance(value, bool) or not isinstance(value, (int, float, str)):
111 raise ValueError(f"Trace event {field_name} must be numeric, got {value!r}")
112 try:
113 resolved_value = float(value)
114 except ValueError as exc:
115 raise ValueError(f"Trace event {field_name} must be numeric, got {value!r}") from exc
116 if not math.isfinite(resolved_value):
117 raise ValueError(f"Trace event {field_name} must be finite, got {value!r}")
118 return resolved_value
121def _is_device_kernel_event(event: dict[str, Any]) -> bool:
122 """Return whether an outer event carries evidence that it is a Device kernel."""
123 category = str(event.get("cat", "")).lower()
124 args = event.get("args", {})
125 task_type = ""
126 if isinstance(args, dict):
127 for key, value in args.items():
128 normalized_key = str(key).strip().lower().replace("_", " ")
129 if normalized_key == "task type":
130 task_type = str(value).lower()
131 break
132 category_is_device = any(token in category for token in ("kernel", "npu", "ascend", "aicore"))
133 args_are_device_task = bool(task_type) and any(
134 token in task_type for token in ("kernel", "ai_core", "aicore", "aic", "aiv")
135 )
136 return category_is_device or args_are_device_task
139def _kernel_score(event: dict[str, Any]) -> int:
140 """Score how likely a complete event is to be a Device kernel."""
141 name = str(event.get("name", "")).lower()
142 category = str(event.get("cat", "")).lower()
143 args_text = json.dumps(event.get("args", {}), ensure_ascii=False).lower()
144 score = 0
145 if _is_device_kernel_event(event):
146 score += 6
147 if "pyboostlaunchaclnn" in name:
148 score += 2
149 elif "pynativelaunchtask" in name:
150 score += 1
151 if "getworkspacesize" in name or "tiling" in name:
152 score -= 3
153 if "_op_" in name or "kernel" in name:
154 score += 5
155 if any(token in category for token in ("kernel", "npu", "ascend", "aicore")):
156 score += 3
157 if any(token in args_text for token in ("kernel", "task type", "ai_core", "aicore")):
158 score += 2
159 if any(token in category for token in ("cpu", "python")):
160 score -= 4
161 return score
164def _complete_events(
165 events: list[dict[str, Any]],
166 *,
167 outer_pid: str | None,
168) -> list[dict[str, Any]]:
169 """Select valid complete events, optionally restricted to one process."""
170 complete_events = []
171 for event in events:
172 if event.get("ph") not in (None, "X"):
173 continue
174 if outer_pid is not None and str(event.get("pid")) != outer_pid:
175 continue
176 try:
177 timestamp = _number(event.get("ts"), "ts")
178 duration = _number(event.get("dur"), "dur")
179 except ValueError:
180 continue
181 if timestamp >= 0 and duration > 0:
182 complete_events.append(event)
183 return complete_events
186def _kernel_candidates(
187 events: list[dict[str, Any]],
188 kernel_pattern: str,
189 outer_pid: str | None,
190) -> list[dict[str, Any]]:
191 """Return the best-scoring complete events matching a kernel pattern."""
192 try:
193 name_pattern = re.compile(kernel_pattern)
194 except re.error as exc:
195 raise ValueError(f"Invalid kernel regex {kernel_pattern!r}: {exc}") from exc
197 candidates = [
198 event
199 for event in _complete_events(events, outer_pid=outer_pid)
200 if name_pattern.search(str(event.get("name", "")))
201 ]
202 if not candidates:
203 raise ValueError(
204 f"No complete MegaKernel event matched pattern {kernel_pattern!r}"
205 + (f" and pid {outer_pid!r}" if outer_pid is not None else "")
206 )
208 best_score = max(_kernel_score(event) for event in candidates)
209 return sorted(
210 (event for event in candidates if _kernel_score(event) == best_score),
211 key=lambda event: _number(event.get("ts"), "ts"),
212 )
215def _fallback_kernel_candidates(
216 events: list[dict[str, Any]],
217 invocations: list[_InvocationTrace],
218 outer_pid: str | None,
219 kernel_index: int,
220) -> list[dict[str, Any]]:
221 """Select a duration-similar contiguous Device-kernel window."""
222 device_events = [
223 event
224 for event in _complete_events(events, outer_pid=outer_pid)
225 if _is_device_kernel_event(event)
226 ]
227 if kernel_index + len(invocations) > len(device_events):
228 raise ValueError(
229 "Not enough complete device-kernel events are available for approximate MegaKernel alignment"
230 + (f" with pid {outer_pid!r}" if outer_pid is not None else "")
231 )
232 device_events.sort(key=lambda event: _number(event.get("ts"), "ts"))
233 device_events = device_events[kernel_index:]
234 invocation_spans = [invocation.span_us for invocation in invocations]
236 def window_distance(start_index: int) -> tuple[float, float]:
237 """Rank a contiguous Device-kernel window by duration similarity.
239 Args:
240 start_index: First Device event in the candidate window.
241 """
242 selected = device_events[start_index:start_index + len(invocations)]
243 duration_distance = 0.0
244 score = 0
245 for event, internal_span in zip(selected, invocation_spans):
246 duration = _number(event.get("dur"), "dur")
247 shorter_penalty = 1.0 if duration < internal_span else 0.0
248 duration_distance += shorter_penalty + abs(duration - internal_span) / max(internal_span, 1.0)
249 score += _kernel_score(event)
250 return duration_distance, -float(score)
252 last_start = len(device_events) - len(invocations)
253 best_start = min(range(last_start + 1), key=window_distance)
254 return device_events[best_start:best_start + len(invocations)]
257def _camel_case_pattern(kernel_name: str) -> str:
258 words = re.findall(r"[A-Z]+(?=[A-Z][a-z]|\d|$)|[A-Z]?[a-z]+|\d+", kernel_name)
259 return r"[_]?".join(re.escape(word) for word in words) if words else re.escape(kernel_name)
262def _default_kernel_pattern(invocations: list[_InvocationTrace]) -> str:
263 patterns = []
264 for invocation in invocations:
265 pattern = _camel_case_pattern(invocation.kernel_name)
266 if pattern and pattern not in patterns:
267 patterns.append(pattern)
268 if not patterns:
269 return DEFAULT_KERNEL_PATTERN
270 return rf"(?i)(?:{'|'.join(patterns)})"
273def _invocation_metadata(metadata: dict[str, Any]) -> dict[Any, dict[str, Any]]:
274 """Index optional invocation metadata by invocation ID."""
275 raw_invocations = metadata.get("invocations", [])
276 if raw_invocations is None:
277 return {}
278 if not isinstance(raw_invocations, list):
279 raise ValueError("megaKernelCycleTrace.invocations must be a list")
280 resolved = {}
281 for item in raw_invocations:
282 if not isinstance(item, dict) or "invocationId" not in item:
283 raise ValueError("Each megaKernelCycleTrace invocation must contain invocationId")
284 invocation_id = item["invocationId"]
285 if invocation_id in resolved:
286 raise ValueError(f"Duplicate MegaKernel invocationId {invocation_id!r}")
287 resolved[invocation_id] = item
288 return resolved
291def _load_cycle_trace(mega_kernel_trace: Any) -> tuple[dict[str, Any], list[dict[str, Any]]]:
292 """Validate cycle-trace metadata and return its complete internal events."""
293 if not isinstance(mega_kernel_trace, dict):
294 raise ValueError("mega_kernel_trace must be a Chrome Trace object")
295 metadata = mega_kernel_trace.get("megaKernelCycleTrace")
296 if not isinstance(metadata, dict):
297 raise ValueError("mega_kernel_trace is missing megaKernelCycleTrace metadata")
298 schema_version = metadata.get("schemaVersion")
299 if schema_version != SUPPORTED_CYCLE_TRACE_SCHEMA_VERSION:
300 raise ValueError(
301 "Unsupported MegaKernel cycle trace schemaVersion: "
302 f"expected {SUPPORTED_CYCLE_TRACE_SCHEMA_VERSION}, got {schema_version!r}"
303 )
305 internal_events = [
306 event
307 for event in _events_from_trace(mega_kernel_trace, "mega_kernel_trace")
308 if event.get("ph") == "X" and event.get("cat") == INTERNAL_EVENT_CATEGORY
309 ]
310 if not internal_events:
311 raise ValueError("mega_kernel_trace contains no complete MegaKernel internal events")
312 return metadata, internal_events
315def _group_invocation_events(internal_events: list[dict[str, Any]]) -> dict[Any, list[dict[str, Any]]]:
316 """Validate event arguments and group records by invocation identifier."""
317 event_args = []
318 for event in internal_events:
319 args = event.get("args", {})
320 if not isinstance(args, dict):
321 raise ValueError("MegaKernel internal event args must be an object")
322 event_args.append(args)
324 has_invocation_id = ["invocation_id" in args for args in event_args]
325 if any(has_invocation_id) and not all(has_invocation_id):
326 raise ValueError("MegaKernel internal events must either all contain invocation_id or all omit it")
328 grouped_events: dict[Any, list[dict[str, Any]]] = {}
329 if all(has_invocation_id):
330 for event, args in zip(internal_events, event_args):
331 invocation_id = args["invocation_id"]
332 grouped_events.setdefault(invocation_id, []).append(event)
333 else:
334 grouped_events[None] = internal_events
335 return grouped_events
338def _build_invocation(
339 invocation_id: Any,
340 events: list[dict[str, Any]],
341 metadata: dict[str, Any],
342 metadata_by_invocation: dict[Any, dict[str, Any]],
343) -> _InvocationTrace:
344 """Resolve display metadata for one grouped invocation."""
345 invocation_item = metadata_by_invocation.get(invocation_id, {})
346 kernel_name = str(invocation_item.get("kernelName") or metadata.get("kernelName") or "MegaKernel")
347 raw_direction = invocation_item.get("direction")
348 return _InvocationTrace(
349 invocation_id=invocation_id,
350 events=events,
351 kernel_name=kernel_name,
352 rank=invocation_item.get("rank", metadata.get("rank")),
353 device_id=invocation_item.get("deviceId", metadata.get("deviceId")),
354 direction=None if raw_direction is None else str(raw_direction),
355 )
358def _load_invocations(mega_kernel_trace: Any) -> list[_InvocationTrace]:
359 """Validate and group standalone internal events into invocations."""
360 metadata, internal_events = _load_cycle_trace(mega_kernel_trace)
361 grouped_events = _group_invocation_events(internal_events)
363 metadata_by_invocation = _invocation_metadata(metadata)
364 return [
365 _build_invocation(invocation_id, events, metadata, metadata_by_invocation)
366 for invocation_id, events in grouped_events.items()
367 ]
370def _core_key(event: dict[str, Any]) -> tuple[str, int]:
371 """Return the logical core type and block identifier for an event."""
372 args = event.get("args", {})
373 if not isinstance(args, dict):
374 raise ValueError("MegaKernel internal event args must be an object")
375 core_type = str(args.get("core_type", "UnknownCore"))
376 block_id = args.get("block_id", -1)
377 if isinstance(block_id, bool):
378 raise ValueError(f"MegaKernel internal block_id must be an integer, got {block_id!r}")
379 try:
380 resolved_block_id = int(block_id)
381 except (TypeError, ValueError) as exc:
382 raise ValueError(f"MegaKernel internal block_id must be an integer, got {block_id!r}") from exc
383 return core_type, resolved_block_id
386def _next_thread_id(events: list[dict[str, Any]], pid: Any) -> int:
387 numeric_thread_ids = []
388 for event in events:
389 thread_id = event.get("tid")
390 if event.get("pid") == pid and isinstance(thread_id, int) and not isinstance(thread_id, bool):
391 numeric_thread_ids.append(thread_id)
392 return max(numeric_thread_ids, default=0) + 1
395def _thread_metadata_events(
396 invocation: _InvocationTrace,
397 outer_pid: Any,
398 thread_ids: dict[tuple[str, int], int],
399 first_sort_index: int,
400) -> list[dict[str, Any]]:
401 """Build framework-thread metadata for the logical Device cores."""
402 metadata_events = []
403 for index, (core_key, thread_id) in enumerate(thread_ids.items()):
404 metadata_events.extend([
405 {
406 "name": "thread_name",
407 "ph": "M",
408 "pid": outer_pid,
409 "tid": thread_id,
410 "args": {"name": f"{invocation.kernel_name}/{core_key[0]}/{core_key[1]}"},
411 },
412 {
413 "name": "thread_sort_index",
414 "ph": "M",
415 "pid": outer_pid,
416 "tid": thread_id,
417 "args": {"sort_index": first_sort_index + index},
418 },
419 ])
420 return metadata_events
423def _align_internal_events(
424 invocation: _InvocationTrace,
425 outer_event: dict[str, Any],
426 outer_pid: Any,
427 outer_timestamp: float,
428 outer_duration: float,
429 thread_ids: dict[tuple[str, int], int],
430) -> list[dict[str, Any]]:
431 """Shift internal events to the selected outer-kernel timeline."""
432 aligned_events = []
433 invocation_anchor = invocation.anchor_us
434 for internal_event in invocation.events:
435 internal_args = copy.deepcopy(internal_event.get("args", {}))
436 aligned_event = copy.deepcopy(internal_event)
437 aligned_event["pid"] = outer_pid
438 aligned_event["tid"] = thread_ids[_core_key(internal_event)]
439 aligned_event["ts"] = outer_timestamp + _number(internal_event.get("ts"), "ts") - invocation_anchor
440 internal_args.update({
441 "parent_kernel": outer_event.get("name"),
442 "parent_kernel_ts": outer_timestamp,
443 "parent_kernel_dur": outer_duration,
444 "cycle_trace_rank": invocation.rank,
445 "cycle_trace_device_id": invocation.device_id,
446 })
447 aligned_event["args"] = internal_args
448 aligned_events.append(aligned_event)
449 return aligned_events
452def _alignment_report(
453 invocation: _InvocationTrace,
454 outer_event: dict[str, Any],
455 outer_timestamp: float,
456 outer_duration: float,
457 aligned_events: list[dict[str, Any]],
458) -> dict[str, Any]:
459 """Describe the selected outer kernel and resulting timeline bounds."""
460 internal_end = max(
461 _number(event.get("ts"), "ts") + _number(event.get("dur"), "dur")
462 for event in aligned_events
463 )
464 return {
465 "invocationId": invocation.invocation_id,
466 "direction": invocation.direction,
467 "kernelName": invocation.kernel_name,
468 "rank": invocation.rank,
469 "deviceId": invocation.device_id,
470 "outerKernelName": outer_event.get("name"),
471 "outerKernelPid": outer_event.get("pid", 0),
472 "outerKernelTid": outer_event.get("tid"),
473 "outerKernelTs": outer_timestamp,
474 "outerKernelDur": outer_duration,
475 "outerKernelScore": _kernel_score(outer_event),
476 "outerKernelIsDevice": _is_device_kernel_event(outer_event),
477 "internalEventCount": len(aligned_events),
478 "internalSpanUs": invocation.span_us,
479 "exceedsOuterKernelUs": max(internal_end - outer_timestamp - outer_duration, 0.0),
480 }
483def _initial_thread_sort_indexes(
484 events: list[dict[str, Any]],
485 selected_candidates: list[dict[str, Any]],
486 invocations: list[_InvocationTrace],
487) -> dict[Any, int]:
488 """Reserve sort-index ranges before existing framework threads."""
489 thread_counts = {}
490 for outer_event, invocation in zip(selected_candidates, invocations):
491 pid = outer_event.get("pid", 0)
492 core_keys = {_core_key(event) for event in invocation.events}
493 thread_counts[pid] = thread_counts.get(pid, 0) + len(core_keys)
495 initial_indexes = {}
496 for pid, thread_count in thread_counts.items():
497 existing_indexes = []
498 for event in events:
499 if event.get("pid") != pid or event.get("name") != "thread_sort_index":
500 continue
501 args = event.get("args", {})
502 if not isinstance(args, dict):
503 continue
504 sort_index = args.get("sort_index")
505 if isinstance(sort_index, int) and not isinstance(sort_index, bool):
506 existing_indexes.append(sort_index)
507 initial_indexes[pid] = min(existing_indexes, default=0) - thread_count
508 return initial_indexes
511def _align_invocation(
512 events: list[dict[str, Any]],
513 invocation: _InvocationTrace,
514 outer_event: dict[str, Any],
515 first_sort_index: int,
516) -> tuple[list[dict[str, Any]], list[dict[str, Any]], dict[str, Any], int]:
517 """Align one internal invocation to its selected outer kernel event."""
518 outer_timestamp = _number(outer_event.get("ts"), "ts")
519 outer_duration = _number(outer_event.get("dur"), "dur")
520 outer_pid = outer_event.get("pid", 0)
521 core_keys = sorted({_core_key(event) for event in invocation.events})
522 first_thread_id = _next_thread_id(events, outer_pid)
523 thread_ids = {
524 core_key: first_thread_id + index
525 for index, core_key in enumerate(core_keys)
526 }
527 metadata_events = _thread_metadata_events(invocation, outer_pid, thread_ids, first_sort_index)
528 aligned_events = _align_internal_events(
529 invocation,
530 outer_event,
531 outer_pid,
532 outer_timestamp,
533 outer_duration,
534 thread_ids,
535 )
536 alignment = _alignment_report(
537 invocation,
538 outer_event,
539 outer_timestamp,
540 outer_duration,
541 aligned_events,
542 )
543 return metadata_events, aligned_events, alignment, first_sort_index + len(core_keys)
546def _validate_merge_options(
547 kernel_pattern: str | None,
548 kernel_index: int,
549 allow_host_wrapper: bool,
550) -> None:
551 """Validate trace merge selection options before reading events."""
552 if not isinstance(kernel_index, int) or isinstance(kernel_index, bool) or kernel_index < 0:
553 raise ValueError(f"kernel_index must be a non-negative integer, got {kernel_index!r}")
554 if kernel_pattern is not None and not isinstance(kernel_pattern, str):
555 raise TypeError(f"kernel_pattern must be a string or None, got {type(kernel_pattern).__name__}")
556 if kernel_pattern == "":
557 raise ValueError("kernel_pattern must not be empty")
558 if not isinstance(allow_host_wrapper, bool):
559 raise TypeError(f"allow_host_wrapper must be bool, got {type(allow_host_wrapper).__name__}")
562def _select_kernel_candidates(
563 outer_events: list[dict[str, Any]],
564 invocations: list[_InvocationTrace],
565 *,
566 kernel_pattern: str | None,
567 resolved_kernel_pattern: str,
568 outer_pid: str | None,
569 kernel_index: int,
570) -> tuple[list[dict[str, Any]], bool, list[str]]:
571 """Select outer kernels by name, falling back to duration similarity."""
572 try:
573 candidates = _kernel_candidates(outer_events, resolved_kernel_pattern, outer_pid)
574 if kernel_index + len(invocations) > len(candidates):
575 raise ValueError(
576 "Not enough outer MegaKernel events: requested indexes "
577 f"{kernel_index}..{kernel_index + len(invocations) - 1}, found {len(candidates)}"
578 )
579 return candidates[kernel_index:kernel_index + len(invocations)], False, []
580 except ValueError as exc:
581 if kernel_pattern is not None or "No complete MegaKernel event matched" not in str(exc):
582 raise
583 selected = _fallback_kernel_candidates(outer_events, invocations, outer_pid, kernel_index)
584 warning = (
585 f"No outer kernel name matched inferred pattern {resolved_kernel_pattern!r}; "
586 "selected a contiguous device-kernel window by duration similarity, so alignment is approximate"
587 )
588 return selected, True, [warning]
591def _validate_selected_candidates(
592 selected_candidates: list[dict[str, Any]],
593 allow_host_wrapper: bool,
594) -> list[str]:
595 """Reject host-only alignment unless the caller explicitly allows it."""
596 host_candidates = [event for event in selected_candidates if not _is_device_kernel_event(event)]
597 if host_candidates and not allow_host_wrapper:
598 host_names = sorted({str(event.get("name", "")) for event in host_candidates})
599 raise ValueError(
600 "Refusing to align MegaKernel internal events to host-only outer event(s) "
601 f"{host_names}. Recollect the framework profile with Device activities, or pass "
602 "allow_host_wrapper=True only for an explicitly approximate diagnostic timeline."
603 )
604 if not host_candidates:
605 return []
606 return ["Selected host-only outer event(s); alignment is approximate because Device kernel events were unavailable"]
609def _append_aligned_invocations(
610 outer_events: list[dict[str, Any]],
611 invocations: list[_InvocationTrace],
612 selected_candidates: list[dict[str, Any]],
613 report: dict[str, Any],
614) -> None:
615 """Append aligned invocation events and diagnostics to the merged trace."""
616 next_sort_indexes = _initial_thread_sort_indexes(outer_events, selected_candidates, invocations)
617 for invocation, outer_event in zip(invocations, selected_candidates):
618 outer_pid = outer_event.get("pid", 0)
619 metadata_events, aligned_events, alignment, next_sort_index = _align_invocation(
620 outer_events,
621 invocation,
622 outer_event,
623 next_sort_indexes[outer_pid],
624 )
625 next_sort_indexes[outer_pid] = next_sort_index
626 outer_events.extend(metadata_events)
627 outer_events.extend(aligned_events)
628 report["mergedEventCount"] += len(aligned_events)
629 report["alignments"].append(alignment)
630 if alignment["exceedsOuterKernelUs"] > 0:
631 report["warnings"].append(
632 f"Invocation {invocation.invocation_id!r} internal events exceed outer kernel "
633 f"by {alignment['exceedsOuterKernelUs']:.3f} us"
634 )
637def _merge_trace_objects(
638 framework_trace: Any,
639 mega_kernel_trace: Any,
640 *,
641 kernel_pattern: str | None,
642 kernel_index: int,
643 outer_pid: str | int | None,
644 allow_host_wrapper: bool,
645) -> TraceObject:
646 """Merge validated in-memory framework and MegaKernel trace objects."""
647 _validate_merge_options(kernel_pattern, kernel_index, allow_host_wrapper)
648 merged_trace = _normalize_framework_trace(framework_trace)
649 outer_events = merged_trace["traceEvents"]
650 invocations = _load_invocations(mega_kernel_trace)
651 resolved_outer_pid = None if outer_pid is None else str(outer_pid)
652 resolved_kernel_pattern = kernel_pattern or _default_kernel_pattern(invocations)
653 selected_candidates, used_fallback_selection, warnings = _select_kernel_candidates(
654 outer_events,
655 invocations,
656 kernel_pattern=kernel_pattern,
657 resolved_kernel_pattern=resolved_kernel_pattern,
658 outer_pid=resolved_outer_pid,
659 kernel_index=kernel_index,
660 )
661 warnings.extend(_validate_selected_candidates(selected_candidates, allow_host_wrapper))
662 report = {
663 "schemaVersion": MERGE_SCHEMA_VERSION,
664 "kernelPattern": resolved_kernel_pattern,
665 "kernelIndex": kernel_index,
666 "usedFallbackKernelSelection": used_fallback_selection,
667 "allowHostWrapperAlignment": allow_host_wrapper,
668 "mergedInvocationCount": len(invocations),
669 "mergedEventCount": 0,
670 "alignments": [],
671 "warnings": warnings,
672 }
673 _append_aligned_invocations(outer_events, invocations, selected_candidates, report)
674 merged_trace["megaKernelTraceMerge"] = report
675 return merged_trace
678def merge_chrome_traces(
679 framework_trace: TraceInput,
680 mega_kernel_trace: TraceInput,
681 output: str | Path,
682 *,
683 kernel_pattern: str | None = None,
684 kernel_index: int = 0,
685 outer_pid: str | int | None = None,
686 allow_host_wrapper: bool = False,
687) -> TraceObject:
688 """Merge standalone MegaKernel events into a framework Chrome Trace.
690 Each standalone invocation is aligned to one outer Device kernel in
691 timestamp order. The merge is pure Host post-processing and does not
692 require framework or NPU runtime imports.
694 Args:
695 framework_trace: Framework Chrome Trace path, object, or event list.
696 mega_kernel_trace: Standalone MegaKernel Chrome Trace path or object.
697 output: Destination path for the merged Chrome Trace object.
698 kernel_pattern: Optional outer-kernel name regular expression. By
699 default it is inferred from standalone invocation metadata.
700 kernel_index: First matching outer-kernel occurrence to consume.
701 outer_pid: Optional outer process ID filter.
702 allow_host_wrapper: Allow approximate alignment to matching Host
703 wrapper events when Device kernel events are unavailable.
705 Returns:
706 Merged Chrome Trace object. Alignment diagnostics are stored in its
707 megaKernelTraceMerge metadata.
709 Raises:
710 TypeError: If an option or in-memory trace has an invalid type.
711 ValueError: If a trace is malformed or cannot be aligned safely.
712 """
713 framework_trace_object = _read_trace(framework_trace, "framework_trace")
714 mega_kernel_trace_object = _read_trace(mega_kernel_trace, "mega_kernel_trace")
715 merged_trace = _merge_trace_objects(
716 framework_trace_object,
717 mega_kernel_trace_object,
718 kernel_pattern=kernel_pattern,
719 kernel_index=kernel_index,
720 outer_pid=outer_pid,
721 allow_host_wrapper=allow_host_wrapper,
722 )
723 output_path = Path(output)
724 output_path.parent.mkdir(parents=True, exist_ok=True)
725 output_path.write_text(json.dumps(merged_trace, ensure_ascii=False), encoding="utf-8")
726 return merged_trace