wafer-core 0.1.25__py3-none-any.whl → 0.1.27__py3-none-any.whl

wafer_core/lib/trace_compare/loader.py

@@ -0,0 +1,635 @@
+"""Trace loading and parsing logic.
+
+Loads JSON trace files from AMD/NVIDIA profilers and extracts kernel execution data,
+Python call stacks, CPU operator mappings, and layer correlations.
+"""
+
+import bisect
+import sys
+from collections import defaultdict
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+try:
+    import ijson
+except ImportError:
+    ijson = None
+
+import orjson
+import pandas as pd
+
+from .architecture import detect_architecture
+from .classifier import classify
+from .layer_segmentation import segment_layers_by_architecture
+
+ProgressCallback = Callable[[str, float], None]
+
+
+@dataclass
+class SinglePassResult:
+    """Collected data from single-pass event processing."""
+    cpu_op_mapping: dict[int, str] = field(default_factory=dict)
+    python_intervals: list[tuple[int, int, int, int | None, str]] = field(default_factory=list)
+    # Raw python events for lazy python_by_id construction (built on-demand in _get_python_stack_full)
+    python_events_raw: list[dict[str, Any]] = field(default_factory=list)
+    phases: list[dict[str, Any]] = field(default_factory=list)
+    correlation_groups: dict[int, dict[str, Any]] = field(default_factory=lambda: defaultdict(
+        lambda: {"count": 0, "has_attention": False, "has_ffn": False}
+    ))
+    kernel_events: list[dict[str, Any]] = field(default_factory=list)
+    # Lazily built when needed for stack resolution
+    _python_by_id: dict[int, dict[str, Any]] | None = field(default=None)
+    
+    @property
+    def python_by_id(self) -> dict[int, dict[str, Any]]:
+        """Lazily build python_by_id from raw events on first access."""
+        if self._python_by_id is None:
+            self._python_by_id = {}
+            for ev in self.python_events_raw:
+                args = ev.get("args")
+                py_id = args.get("Python id") if args else None
+                if py_id is not None:
+                    self._python_by_id[py_id] = {
+                        "name": ev["name"],
+                        "parent_id": args.get("Python parent id") if args else None,
+                    }
+        return self._python_by_id
+
+
+@dataclass
+class LoadedTrace:
+    """Complete trace data loaded once and reused by all analysis functions."""
+    platform: str
+    gpu_name: str
+    device_props: dict[str, Any]
+    df: pd.DataFrame
+    patterns: dict[tuple[str, str], set[str]]
+    layers: dict[int, int]
+    # For fusion/warnings (kept from raw JSON)
+    kernel_events: list[dict[str, Any]]
+    all_events: list[dict[str, Any]]
+    correlation_groups: dict[int, list[dict[str, Any]]]
+    phases: list[dict[str, Any]]  # Phase annotations for alignment
+
+
+def _process_events_single_pass(
+    events: list[dict[str, Any]],
+    include_stacks: bool = True,
+) -> SinglePassResult:
+    """Process all events in a single iteration.
+    
+    Optimizations applied:
+    - Cache list.append methods for 2-3x speedup on hot paths
+    - Store raw python events, build python_by_id lazily (only ~48 lookups due to caching)
+    - Local variable caching for frequently accessed attributes
+    
+    Args:
+        events: List of trace events
+        include_stacks: Whether to collect Python stack info (expensive operation).
+                       When False, skips python_function processing entirely.
+    """
+    result = SinglePassResult()
+    correlation_groups: dict[int, dict[str, Any]] = defaultdict(
+        lambda: {"count": 0, "has_attention": False, "has_ffn": False}
+    )
+    
+    # Cache list.append methods for faster appending (measured 2-3x speedup)
+    kernel_append = result.kernel_events.append
+    python_interval_append = result.python_intervals.append
+    python_raw_append = result.python_events_raw.append
+    phase_append = result.phases.append
+    cpu_op_mapping = result.cpu_op_mapping
+    
+    for ev in events:
+        cat = ev.get("cat")
+        
+        if cat == "kernel":
+            args = ev.get("args")
+            corr_id = args.get("correlation") if args else None
+            if corr_id is not None:
+                kernel_name = ev.get("name", "").lower()
+                grp = correlation_groups[corr_id]
+                grp["count"] += 1
+                if "attention" in kernel_name or "fmha" in kernel_name:
+                    grp["has_attention"] = True
+                if "cijk_" in kernel_name or "nvjet" in kernel_name or "wvsplitk" in kernel_name or "gemm" in kernel_name:
+                    grp["has_ffn"] = True
+            kernel_append(ev)
+        
+        elif cat == "cpu_op":
+            args = ev.get("args")
+            ext_id = args.get("External id") if args else None
+            if ext_id is not None:
+                cpu_op_mapping[ext_id] = ev.get("name", "")
+        
+        elif cat == "python_function" and include_stacks:
+            # Store raw event for lazy python_by_id construction
+            python_raw_append(ev)
+            # Build interval tuple for binary search
+            args = ev.get("args")
+            py_id = args.get("Python id") if args else None
+            ts = ev["ts"]
+            dur = ev.get("dur", 0)
+            python_interval_append((ts, ts + dur, dur, py_id, ev["name"]))
+                
+        elif cat == "user_annotation":
+            name = ev.get("name", "")
+            if name.startswith("execute_context"):
+                tokens = 0
+                parts = name.split("_")
+                for i, p in enumerate(parts):
+                    if i > 0 and parts[i-1] == "context" and "(" in p and ")" in p:
+                        try:
+                            tokens = int(p.split("(")[1].split(")")[0])
+                            break
+                        except Exception:
+                            pass
+                is_prefill = "generation_0" in name and tokens > 0
+                phase_append({
+                    "type": "prefill" if is_prefill else "decode",
+                    "ts_start": ev["ts"],
+                    "ts_end": ev["ts"] + ev["dur"],
+                })
+    
+    if include_stacks and result.python_intervals:
+        result.python_intervals.sort()
+    
+    result.correlation_groups = dict(correlation_groups)
+    
+    return result
+
+
+def _build_layer_mapping(correlation_groups: dict[int, dict[str, Any]]) -> dict[int, int]:
+    """Build layer mapping from correlation groups."""
+    correlation_to_layer = {}
+    layer_num = 0
+    
+    for corr_id in sorted(correlation_groups.keys()):
+        group = correlation_groups[corr_id]
+        
+        is_layer = (
+            group["count"] >= 70 and group["has_attention"] and group["has_ffn"]
+        )
+        
+        if is_layer:
+            correlation_to_layer[corr_id] = layer_num
+            layer_num += 1
+    
+    return correlation_to_layer
+
+
+def _get_python_stack_full(
+    timestamp: int,
+    python_intervals: list[tuple[int, int, int, int | None, str]],
+    python_by_id: dict[int, dict[str, Any]],
+) -> tuple[str | None, list[str]]:
+    """Get full Python call stack for a kernel launch.
+    
+    Computes the chain on-demand by walking parent pointers.
+    This is fast because we only call this ~48 times due to cpu_op caching.
+    """
+    idx = bisect.bisect_right(
+        python_intervals, (timestamp, float("inf"), float("inf"), None, "")
+    )
+    
+    active_funcs = []
+    for i in range(idx - 1, max(0, idx - 1000), -1):
+        ts_start, ts_end, duration, py_id, name = python_intervals[i]
+        if ts_start <= timestamp <= ts_end:
+            active_funcs.append((duration, py_id, name))
+        if ts_end < timestamp - 1000000:
+            break
+    
+    if not active_funcs:
+        return None, []
+    
+    active_funcs.sort()
+    leaf_duration, leaf_id, leaf_name = active_funcs[0]
+    
+    full_stack = []
+    current_id = leaf_id
+    visited: set[int] = set()
+    
+    while current_id is not None and current_id not in visited and current_id in python_by_id:
+        func = python_by_id[current_id]
+        full_stack.append(func["name"])
+        visited.add(current_id)
+        current_id = func["parent_id"]
+        if len(full_stack) >= 50:
+            break
+    
+    full_stack.reverse()
+    
+    summary = None
+    vllm_funcs = [f for f in full_stack if any(x in f.lower() for x in ["vllm/", "model", "<eval_with_key>"])]
+    
+    if vllm_funcs:
+        summary = vllm_funcs[-1]
+        if any("torch/cuda/graphs" in f for f in full_stack):
+            if len(summary) > 45:
+                summary = "vllm/..." + summary.split("/")[-1]
+            summary = f"{summary} [CUDA graph]"
+        elif len(summary) > 53:
+            summary = "vllm/..." + summary.split("/")[-1]
+    else:
+        summary = leaf_name
+    
+    return summary, full_stack
+
+
+def load_trace(
+    file_path: str | Path,
+    include_stacks: bool = True,
+    on_progress: ProgressCallback | None = None,
+) -> tuple[str, str, dict[str, Any], pd.DataFrame, dict[tuple[str, str], set[str]], dict[int, int]]:
+    """Load trace and return platform info, device properties, kernels, patterns, and layer mapping.
+    
+    Args:
+        file_path: Path to the trace JSON file
+        include_stacks: Whether to resolve Python call stacks (slower but more info)
+        on_progress: Optional callback for progress updates: (stage_name, progress_fraction)
+    """
+    def _progress(stage: str, pct: float) -> None:
+        if on_progress:
+            on_progress(stage, pct)
+    
+    _progress("Reading file", 0.0)
+    with open(file_path, "rb") as f:
+        raw = f.read()
+    
+    _progress("Parsing JSON", 0.1)
+    trace = orjson.loads(raw)
+    
+    props = trace.get("deviceProperties", [{}])[0]
+    is_amd = trace.get("roctracer_version") or props.get("warpSize") == 64
+    platform = "AMD" if is_amd else "NVIDIA"
+    gpu_name = props.get("name", "MI300X" if is_amd else "Unknown GPU")
+    
+    device_props = {
+        "name": gpu_name,
+        "compute_capability": f"{props.get('computeMajor', 0)}.{props.get('computeMinor', 0)}",
+        "total_memory_gb": props.get("totalGlobalMem", 0) / (1024**3),
+        "sm_count": props.get("numSms", 0),
+        "warp_size": props.get("warpSize", 32),
+        "max_threads_per_block": props.get("maxThreadsPerBlock", 0),
+        "shared_mem_per_block_kb": props.get("sharedMemPerBlock", 0) / 1024,
+    }
+    
+    _progress("Processing events", 0.4)
+    events = trace.get("traceEvents", [])
+    pass_result = _process_events_single_pass(events, include_stacks=include_stacks)
+    
+    _progress("Detecting architecture", 0.6)
+    kernel_names = [ev.get("name", "") for ev in pass_result.kernel_events]
+    architecture, _ = detect_architecture(kernel_names)
+    
+    layer_kernels_dict, layer_warnings = segment_layers_by_architecture(
+        pass_result.kernel_events,
+        architecture,
+    )
+    
+    # Convert layer_kernels_dict (layer_num -> kernels) to correlation_id -> layer_num mapping
+    # Fall back to old method if architecture-based segmentation fails
+    if layer_kernels_dict:
+        layer_mapping: dict[int, int] = {}
+        for layer_num, kernels in layer_kernels_dict.items():
+            for kernel in kernels:
+                corr_id = kernel.get("args", {}).get("correlation")
+                if corr_id is not None:
+                    layer_mapping[corr_id] = layer_num
+    else:
+        # Fallback to correlation-based method if architecture segmentation failed
+        layer_mapping = _build_layer_mapping(pass_result.correlation_groups)
+    
+    kernel_data = []
+    kernel_patterns: dict[tuple[str, str], set[str]] = defaultdict(set)
+    sorted_phases = sorted(pass_result.phases, key=lambda p: p["ts_start"])
+    
+    phase_starts = [p["ts_start"] for p in sorted_phases]
+    phase_types = [p["type"] for p in sorted_phases]
+    phase_ends = [p["ts_end"] for p in sorted_phases]
+    
+    def _get_phase_for_timestamp(ts: int) -> str:
+        """Get phase for a timestamp using binary search. O(log n)."""
+        if not phase_starts:
+            return "decode"
+        idx = bisect.bisect_right(phase_starts, ts) - 1
+        if idx >= 0 and phase_starts[idx] <= ts <= phase_ends[idx]:
+            return phase_types[idx]
+        return "decode"
+    
+    cpu_op_cache: dict[str, str | None] = {}
+    
+    _progress("Classifying kernels", 0.7)
+    for ev in pass_result.kernel_events:
+        name_raw = ev["name"]
+        name = sys.intern(name_raw)
+        dur, ts = ev.get("dur", 0), ev["ts"]
+        corr_id = ev.get("args", {}).get("correlation")
+        ext_id = ev.get("args", {}).get("External id")
+        
+        phase = _get_phase_for_timestamp(ts)
+        
+        op, pattern = classify(name, platform)
+        kernel_patterns[(op.value, phase)].add(pattern)
+        
+        layer = layer_mapping.get(corr_id) if corr_id is not None else None
+        cpu_op = pass_result.cpu_op_mapping.get(ext_id) if ext_id is not None else None
+        python_stack: list[str] = []
+        
+        if cpu_op is None and include_stacks:
+            if name in cpu_op_cache:
+                cpu_op = cpu_op_cache[name]
+            else:
+                cpu_op, python_stack = _get_python_stack_full(
+                    ts, pass_result.python_intervals, pass_result.python_by_id
+                )
+                cpu_op_cache[name] = cpu_op
+        
+        kernel_data.append({
+            "name": name,
+            "dur_us": dur,
+            "phase": phase,
+            "op": op.value,
+            "pattern": pattern,
+            "layer": layer,
+            "correlation": corr_id,
+            "cpu_op": cpu_op,
+            "python_stack": python_stack,
+        })
+    
+    _progress("Building DataFrame", 0.95)
+    df = pd.DataFrame(kernel_data)
+    _progress("Complete", 1.0)
+    
+    return platform, gpu_name, device_props, df, dict(kernel_patterns), layer_mapping
+
+
+def load_trace_full(
+    file_path: str | Path,
+    include_stacks: bool = True,
+    on_progress: ProgressCallback | None = None,
+) -> LoadedTrace:
+    """Load trace once with all data needed by downstream analysis functions.
+    
+    Args:
+        file_path: Path to the trace JSON file
+        include_stacks: Whether to resolve Python call stacks
+        on_progress: Optional callback for progress updates
+        
+    Returns:
+        LoadedTrace with all trace data
+    """
+    def _progress(stage: str, pct: float) -> None:
+        if on_progress:
+            on_progress(stage, pct)
+    
+    _progress("Reading file", 0.0)
+    with open(file_path, "rb") as f:
+        raw = f.read()
+    
+    _progress("Parsing JSON", 0.1)
+    trace = orjson.loads(raw)
+    all_events = trace.get("traceEvents", [])
+    
+    props = trace.get("deviceProperties", [{}])[0]
+    is_amd = trace.get("roctracer_version") or props.get("warpSize") == 64
+    platform = "AMD" if is_amd else "NVIDIA"
+    gpu_name = props.get("name", "MI300X" if is_amd else "Unknown GPU")
+    
+    device_props = {
+        "name": gpu_name,
+        "compute_capability": f"{props.get('computeMajor', 0)}.{props.get('computeMinor', 0)}",
+        "total_memory_gb": props.get("totalGlobalMem", 0) / (1024**3),
+        "sm_count": props.get("numSms", 0),
+        "warp_size": props.get("warpSize", 32),
+        "max_threads_per_block": props.get("maxThreadsPerBlock", 0),
+        "shared_mem_per_block_kb": props.get("sharedMemPerBlock", 0) / 1024,
+    }
+    
+    _progress("Processing events", 0.4)
+    pass_result = _process_events_single_pass(all_events, include_stacks=include_stacks)
+    
+    _progress("Detecting architecture", 0.6)
+    kernel_names = [ev.get("name", "") for ev in pass_result.kernel_events]
+    architecture, _ = detect_architecture(kernel_names)
+    
+    layer_kernels_dict, layer_warnings = segment_layers_by_architecture(
+        pass_result.kernel_events,
+        architecture,
+    )
+    
+    # Convert layer_kernels_dict (layer_num -> kernels) to correlation_id -> layer_num mapping
+    if layer_kernels_dict:
+        layer_mapping: dict[int, int] = {}
+        for layer_num, kernels in layer_kernels_dict.items():
+            for kernel in kernels:
+                corr_id = kernel.get("args", {}).get("correlation")
+                if corr_id is not None:
+                    layer_mapping[corr_id] = layer_num
+    else:
+        layer_mapping = _build_layer_mapping(pass_result.correlation_groups)
+    
+    kernel_data = []
+    kernel_patterns: dict[tuple[str, str], set[str]] = defaultdict(set)
+    sorted_phases = sorted(pass_result.phases, key=lambda p: p["ts_start"])
+    
+    phase_starts = [p["ts_start"] for p in sorted_phases]
+    phase_types = [p["type"] for p in sorted_phases]
+    phase_ends = [p["ts_end"] for p in sorted_phases]
+    
+    def _get_phase_for_timestamp(ts: int) -> str:
+        """Get phase for a timestamp using binary search. O(log n)."""
+        if not phase_starts:
+            return "decode"
+        idx = bisect.bisect_right(phase_starts, ts) - 1
+        if idx >= 0 and phase_starts[idx] <= ts <= phase_ends[idx]:
+            return phase_types[idx]
+        return "decode"
+    
+    cpu_op_cache: dict[str, str | None] = {}
+    
+    _progress("Classifying kernels", 0.7)
+    for ev in pass_result.kernel_events:
+        name_raw = ev["name"]
+        name = sys.intern(name_raw)
+        dur, ts = ev.get("dur", 0), ev["ts"]
+        corr_id = ev.get("args", {}).get("correlation")
+        ext_id = ev.get("args", {}).get("External id")
+        
+        phase = _get_phase_for_timestamp(ts)
+        
+        op, pattern = classify(name, platform)
+        kernel_patterns[(op.value, phase)].add(pattern)
+        
+        layer = layer_mapping.get(corr_id) if corr_id is not None else None
+        cpu_op = pass_result.cpu_op_mapping.get(ext_id) if ext_id is not None else None
+        python_stack: list[str] = []
+        
+        if cpu_op is None and include_stacks:
+            if name in cpu_op_cache:
+                cpu_op = cpu_op_cache[name]
+            else:
+                cpu_op, python_stack = _get_python_stack_full(
+                    ts, pass_result.python_intervals, pass_result.python_by_id
+                )
+                cpu_op_cache[name] = cpu_op
+        
+        kernel_data.append({
+            "name": name,
+            "dur_us": dur,
+            "phase": phase,
+            "op": op.value,
+            "pattern": pattern,
+            "layer": layer,
+            "correlation": corr_id,
+            "cpu_op": cpu_op,
+            "python_stack": python_stack,
+        })
+    
+    _progress("Building DataFrame", 0.95)
+    df = pd.DataFrame(kernel_data)
+    
+    kernel_events = pass_result.kernel_events
+    correlation_groups: dict[int, list[dict[str, Any]]] = defaultdict(list)
+    for ev in kernel_events:
+        corr_id = ev.get("args", {}).get("correlation")
+        if corr_id is not None:
+            correlation_groups[corr_id].append(ev)
+    
+    _progress("Complete", 1.0)
+    
+    return LoadedTrace(
+        platform=platform,
+        gpu_name=gpu_name,
+        device_props=device_props,
+        df=df,
+        patterns=dict(kernel_patterns),
+        layers=layer_mapping,
+        kernel_events=kernel_events,
+        all_events=all_events,
+        correlation_groups=dict(correlation_groups),
+        phases=pass_result.phases,
+    )
+
+
+@dataclass
+class StreamingMetadata:
+    """Early metadata available before full trace processing."""
+    platform: str
+    gpu_name: str
+    device_props: dict[str, Any]
+    file_size_mb: float
+
+
+def _extract_metadata_fast(file_path: Path) -> StreamingMetadata:
+    """Extract trace metadata instantly using streaming parser (~2ms).
+    
+    Uses ijson to read only the deviceProperties section without
+    parsing the entire file. Falls back to full parse if ijson unavailable.
+    """
+    file_size_mb = file_path.stat().st_size / (1024 * 1024)
+    platform = "Unknown"
+    gpu_name = "Unknown GPU"
+    device_props: dict[str, Any] = {}
+    
+    if ijson is None:
+        # Fallback: parse full file (slower but works)
+        with open(file_path, "rb") as f:
+            trace = orjson.loads(f.read())
+        props = trace.get("deviceProperties", [{}])[0]
+        is_amd = trace.get("roctracer_version") or props.get("warpSize") == 64
+        platform = "AMD" if is_amd else "NVIDIA"
+        gpu_name = props.get("name", "MI300X" if is_amd else "Unknown GPU")
+        device_props = {
+            "name": gpu_name,
+            "compute_capability": f"{props.get('computeMajor', 0)}.{props.get('computeMinor', 0)}",
+            "total_memory_gb": props.get("totalGlobalMem", 0) / (1024**3),
+            "sm_count": props.get("numSms", 0),
+            "warp_size": props.get("warpSize", 32),
+            "max_threads_per_block": props.get("maxThreadsPerBlock", 0),
+            "shared_mem_per_block_kb": props.get("sharedMemPerBlock", 0) / 1024,
+        }
+        return StreamingMetadata(
+            platform=platform,
+            gpu_name=gpu_name,
+            device_props=device_props,
+            file_size_mb=file_size_mb,
+        )
+    
+    with open(file_path, "rb") as f:
+        parser = ijson.parse(f)
+        for prefix, event, value in parser:
+            if prefix == "deviceProperties.item.name":
+                gpu_name = value
+            elif prefix == "deviceProperties.item.warpSize":
+                platform = "AMD" if value == 64 else "NVIDIA"
+            elif prefix == "deviceProperties.item.computeMajor":
+                device_props["compute_major"] = value
+            elif prefix == "deviceProperties.item.computeMinor":
+                device_props["compute_minor"] = value
+            elif prefix == "deviceProperties.item.totalGlobalMem":
+                device_props["total_memory_gb"] = value / (1024**3)
+            elif prefix == "deviceProperties.item.numSms":
+                device_props["sm_count"] = value
+            elif prefix == "deviceProperties.item.maxThreadsPerBlock":
+                device_props["max_threads_per_block"] = value
+            elif prefix == "deviceProperties.item.sharedMemPerBlock":
+                device_props["shared_mem_per_block_kb"] = value / 1024
+            elif prefix == "traceEvents.item":
+                # Hit first event, stop - we only need metadata
+                break
+    
+    # Fallback platform detection
+    if platform == "Unknown":
+        platform = "AMD" if "MI" in gpu_name or "Instinct" in gpu_name else "NVIDIA"
+    
+    device_props["name"] = gpu_name
+    device_props["compute_capability"] = f"{device_props.get('compute_major', 0)}.{device_props.get('compute_minor', 0)}"
+    device_props["warp_size"] = 64 if platform == "AMD" else 32
+    
+    return StreamingMetadata(
+        platform=platform,
+        gpu_name=gpu_name,
+        device_props=device_props,
+        file_size_mb=file_size_mb,
+    )
+
+
+def load_trace_streaming(
+    file_path: str | Path,
+    include_stacks: bool = True,
+    on_metadata: Callable[[StreamingMetadata], None] | None = None,
+    on_progress: ProgressCallback | None = None,
+) -> tuple[str, str, dict[str, Any], pd.DataFrame, dict[tuple[str, str], set[str]], dict[int, int]]:
+    """Load trace with instant metadata feedback.
+    
+    Hybrid approach:
+    1. Phase 1 (~2ms): Extract metadata with ijson streaming
+    2. Phase 2 (full): Parse with orjson and process
+    
+    The on_metadata callback fires immediately with GPU/platform info,
+    allowing the UI to show feedback while the full load continues.
+    
+    Args:
+        file_path: Path to the trace JSON file
+        include_stacks: Whether to resolve Python call stacks
+        on_metadata: Callback for instant metadata (fires in ~2ms)
+        on_progress: Callback for progress updates during processing
+    """
+    file_path = Path(file_path)
+    
+    # Phase 1: Instant metadata extraction (~2ms)
+    metadata = _extract_metadata_fast(file_path)
+    
+    if on_metadata:
+        on_metadata(metadata)
+    
+    # Phase 2: Full load with orjson (fast) + progress updates
+    return load_trace(
+        file_path,
+        include_stacks=include_stacks,
+        on_progress=on_progress,
+    )