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

wafer_core/lib/trace_compare/analyzer.py

@@ -0,0 +1,729 @@
+"""Main trace comparison analysis logic.
+
+Compares GPU traces from AMD and NVIDIA platforms, identifying performance differences
+at the operation level and layer level.
+"""
+
+import sys
+from collections import defaultdict
+from concurrent.futures import ProcessPoolExecutor
+from pathlib import Path
+from typing import Any
+
+import pandas as pd
+
+from .aligner import align_traces, TraceAlignment
+from .fusion_analyzer import analyze_fusion_from_alignment
+from .same_kernel_analyzer import analyze_same_kernels_from_alignment
+from .loader import load_trace_full, LoadedTrace
+
+
+def analyze_traces_from_loaded(
+    trace1: LoadedTrace,
+    trace2: LoadedTrace,
+    phase_filter: str = "all",
+    max_stacks: int = 3,
+) -> dict[str, Any]:
+    """Analyze two loaded traces and return comparison data.
+    
+    Args:
+        trace1: First loaded trace
+        trace2: Second loaded trace
+        phase_filter: Filter by phase ('all', 'prefill', or 'decode')
+        max_stacks: Maximum number of Python stack traces to collect per operation (0 for unlimited)
+        
+    Returns:
+        Dictionary containing:
+            - metadata: trace info (GPUs, kernel counts, total times, etc.)
+            - operations: per-operation comparison data
+            - layers: per-layer comparison data (if layers detected)
+    """
+    df1 = trace1.df
+    df2 = trace2.df
+    
+    # Apply phase filter
+    if phase_filter != "all":
+        df1_filtered = df1[df1["phase"] == phase_filter]
+        df2_filtered = df2[df2["phase"] == phase_filter]
+        
+        if len(df1_filtered) == 0 and len(df2_filtered) == 0:
+            trace1_phases = {k: int(v) for k, v in df1["phase"].value_counts().items()}
+            trace2_phases = {k: int(v) for k, v in df2["phase"].value_counts().items()}
+            raise ValueError(
+                f"No {phase_filter} phase found. "
+                f"Trace1 phases: {trace1_phases}, Trace2 phases: {trace2_phases}"
+            )
+        
+        df1, df2 = df1_filtered, df2_filtered
+    
+    # Pre-compute aggregations for both operations and layers in single pass
+    trace1_by_op = df1.groupby("op").agg({
+        "dur_us": ["sum", "mean", "count"],
+        "phase": lambda x: set(x.dropna().unique()),
+        "cpu_op": lambda x: x.dropna().mode()[0] if len(x.dropna()) > 0 else None,
+    })
+    trace1_by_op.columns = ["total_us", "avg_us", "count", "phases", "cpu_op"]
+    
+    trace2_by_op = df2.groupby("op").agg({
+        "dur_us": ["sum", "mean", "count"],
+        "phase": lambda x: set(x.dropna().unique()),
+        "cpu_op": lambda x: x.dropna().mode()[0] if len(x.dropna()) > 0 else None,
+    })
+    trace2_by_op.columns = ["total_us", "avg_us", "count", "phases", "cpu_op"]
+    
+    # Group by layer for layer-level analysis
+    df1_layered = df1[df1["layer"].notna()]
+    df2_layered = df2[df2["layer"].notna()]
+    
+    trace1_by_layer = df1_layered.groupby("layer").agg({
+        "dur_us": ["sum", "count"],
+    }) if len(df1_layered) > 0 else pd.DataFrame()
+    if len(trace1_by_layer) > 0:
+        trace1_by_layer.columns = ["total_us", "count"]
+    
+    trace2_by_layer = df2_layered.groupby("layer").agg({
+        "dur_us": ["sum", "count"],
+    }) if len(df2_layered) > 0 else pd.DataFrame()
+    if len(trace2_by_layer) > 0:
+        trace2_by_layer.columns = ["total_us", "count"]
+    
+    results: dict[str, Any] = {
+        "metadata": {
+            "trace1_name": f"{trace1.platform} {trace1.gpu_name}",
+            "trace2_name": f"{trace2.platform} {trace2.gpu_name}",
+            "trace1_platform": trace1.platform,
+            "trace1_gpu": trace1.gpu_name,
+            "trace1_device": trace1.device_props,
+            "trace2_platform": trace2.platform,
+            "trace2_gpu": trace2.gpu_name,
+            "trace2_device": trace2.device_props,
+            "trace1_kernels": len(df1),
+            "trace2_kernels": len(df2),
+            "trace1_total_ms": df1["dur_us"].sum() / 1000,
+            "trace2_total_ms": df2["dur_us"].sum() / 1000,
+            "phase": phase_filter,
+            "trace1_layers": len(trace1.layers),
+            "trace2_layers": len(trace2.layers),
+        },
+        "operations": [],
+        "layers": [],
+    }
+    
+    # Per-operation comparison
+    all_ops = set(trace1_by_op.index) | set(trace2_by_op.index)
+    rmsnorm_compared = False
+    
+    for op in sorted(all_ops):
+        has_trace1 = op in trace1_by_op.index
+        has_trace2 = op in trace2_by_op.index
+        
+        trace1_op_for_pattern = op
+        trace2_op_for_pattern = op
+        skip_comparison = False
+        
+        if op == "RMSNorm+GEMM" and not has_trace2:
+            has_trace2 = "RMSNorm" in trace2_by_op.index
+            trace2_op_for_pattern = "RMSNorm"
+            rmsnorm_compared = True
+        elif op == "RMSNorm" and not has_trace1:
+            if rmsnorm_compared:
+                skip_comparison = True
+            else:
+                has_trace1 = "RMSNorm+GEMM" in trace1_by_op.index
+                trace1_op_for_pattern = "RMSNorm+GEMM"
+                rmsnorm_compared = True
+        
+        if skip_comparison or not (has_trace1 and has_trace2):
+            continue
+        
+        trace1_agg = trace1_by_op.loc[trace1_op_for_pattern]
+        trace2_agg = trace2_by_op.loc[trace2_op_for_pattern]
+        
+        trace1_avg = trace1_agg["avg_us"]
+        trace2_avg = trace2_agg["avg_us"]
+        trace1_total = trace1_agg["total_us"] / 1000
+        trace2_total = trace2_agg["total_us"] / 1000
+        trace1_count = int(trace1_agg["count"])
+        trace2_count = int(trace2_agg["count"])
+        ratio = trace1_avg / trace2_avg if trace2_avg > 0 else 1
+        gap_ms = trace1_total - trace2_total
+        
+        trace1_pattern = list(
+            trace1.patterns.get(
+                (trace1_op_for_pattern, "decode"),
+                trace1.patterns.get((trace1_op_for_pattern, "prefill"), {"unknown"}),
+            )
+        )[0]
+        trace2_pattern = list(
+            trace2.patterns.get(
+                (trace2_op_for_pattern, "decode"),
+                trace2.patterns.get((trace2_op_for_pattern, "prefill"), {"unknown"}),
+            )
+        )[0]
+        
+        trace1_cpu_op = trace1_agg["cpu_op"]
+        trace2_cpu_op = trace2_agg["cpu_op"]
+        
+        # Get detailed kernel data and stacks only when needed
+        trace1_data = df1[df1["op"] == trace1_op_for_pattern]
+        trace2_data = df2[df2["op"] == trace2_op_for_pattern]
+        
+        # Collect Python stacks if available
+        trace1_python_stacks = []
+        trace2_python_stacks = []
+        
+        if max_stacks != 0:
+            stack_limit = None if max_stacks == 0 else max_stacks
+            for stack_list in trace1_data["python_stack"].head(stack_limit):
+                if stack_list and len(stack_list) > 0:
+                    trace1_python_stacks.append(stack_list)
+            
+            for stack_list in trace2_data["python_stack"].head(stack_limit):
+                if stack_list and len(stack_list) > 0:
+                    trace2_python_stacks.append(stack_list)
+        
+        # Aggregate individual kernels
+        trace1_kernels = trace1_data.groupby("name").agg({"dur_us": ["sum", "count", "mean"]}).reset_index()
+        trace1_kernels.columns = ["name", "total_us", "count", "avg_us"]
+        trace1_kernels = trace1_kernels.sort_values("total_us", ascending=False)
+        trace1_kernels_list = trace1_kernels.to_dict("records")
+        
+        trace2_kernels = trace2_data.groupby("name").agg({"dur_us": ["sum", "count", "mean"]}).reset_index()
+        trace2_kernels.columns = ["name", "total_us", "count", "avg_us"]
+        trace2_kernels = trace2_kernels.sort_values("total_us", ascending=False)
+        trace2_kernels_list = trace2_kernels.to_dict("records")
+        
+        if gap_ms > 5.0:
+            status = "slower"
+        elif gap_ms < -5.0:
+            status = "faster"
+        else:
+            status = "similar"
+        
+        phases = trace1_agg["phases"] | trace2_agg["phases"]
+        
+        results["operations"].append({
+            "operation": op,
+            "trace1_count": trace1_count,
+            "trace2_count": trace2_count,
+            "trace1_avg_us": trace1_avg,
+            "trace2_avg_us": trace2_avg,
+            "trace1_total_ms": trace1_total,
+            "trace2_total_ms": trace2_total,
+            "ratio": ratio,
+            "gap_ms": gap_ms,
+            "status": status,
+            "trace1_kernel": trace1_pattern,
+            "trace2_kernel": trace2_pattern,
+            "trace1_cpu_op": trace1_cpu_op,
+            "trace2_cpu_op": trace2_cpu_op,
+            "trace1_python_stacks": trace1_python_stacks,
+            "trace2_python_stacks": trace2_python_stacks,
+            "trace1_kernels": trace1_kernels_list,
+            "trace2_kernels": trace2_kernels_list,
+            "phases": sorted(list(phases)) if phases else ["all"],
+        })
+    
+    results["operations"].sort(key=lambda x: abs(x["gap_ms"]), reverse=True)
+    
+    # Layer-wise analysis
+    if len(trace1_by_layer) > 0 or len(trace2_by_layer) > 0:
+        all_layers = sorted(set(trace1_by_layer.index) | set(trace2_by_layer.index))
+        
+        for layer_num in all_layers:
+            has_trace1 = layer_num in trace1_by_layer.index
+            has_trace2 = layer_num in trace2_by_layer.index
+            
+            if has_trace1 and has_trace2:
+                trace1_agg = trace1_by_layer.loc[layer_num]
+                trace2_agg = trace2_by_layer.loc[layer_num]
+                
+                trace1_total = trace1_agg["total_us"] / 1000
+                trace2_total = trace2_agg["total_us"] / 1000
+                trace1_count = int(trace1_agg["count"])
+                trace2_count = int(trace2_agg["count"])
+                ratio = trace1_total / trace2_total if trace2_total > 0 else 1
+                gap_ms = trace1_total - trace2_total
+                
+                threshold_ms = 0.1
+                threshold_ratio = 1.2
+                if gap_ms > threshold_ms and ratio > threshold_ratio:
+                    status = "slower"
+                elif gap_ms < -threshold_ms and ratio < (1.0 / threshold_ratio):
+                    status = "faster"
+                else:
+                    status = "similar"
+                
+                results["layers"].append({
+                    "layer": int(layer_num),
+                    "trace1_kernels": trace1_count,
+                    "trace2_kernels": trace2_count,
+                    "trace1_total_ms": trace1_total,
+                    "trace2_total_ms": trace2_total,
+                    "ratio": ratio,
+                    "gap_ms": gap_ms,
+                    "status": status,
+                    "in_both": True,
+                })
+            elif has_trace1:
+                trace1_agg = trace1_by_layer.loc[layer_num]
+                trace1_total = trace1_agg["total_us"] / 1000
+                trace1_count = int(trace1_agg["count"])
+                
+                results["layers"].append({
+                    "layer": int(layer_num),
+                    "trace1_kernels": trace1_count,
+                    "trace2_kernels": 0,
+                    "trace1_total_ms": trace1_total,
+                    "trace2_total_ms": 0.0,
+                    "ratio": 0.0,
+                    "gap_ms": trace1_total,
+                    "status": "trace1_only",
+                    "in_both": False,
+                })
+            elif has_trace2:
+                trace2_agg = trace2_by_layer.loc[layer_num]
+                trace2_total = trace2_agg["total_us"] / 1000
+                trace2_count = int(trace2_agg["count"])
+                
+                results["layers"].append({
+                    "layer": int(layer_num),
+                    "trace1_kernels": 0,
+                    "trace2_kernels": trace2_count,
+                    "trace1_total_ms": 0.0,
+                    "trace2_total_ms": trace2_total,
+                    "ratio": 0.0,
+                    "gap_ms": -trace2_total,
+                    "status": "trace2_only",
+                    "in_both": False,
+                })
+        
+        results["layers"].sort(key=lambda x: (not x["in_both"], abs(x["gap_ms"])), reverse=True)
+    
+    return results
+
+
+def analyze_traces(
+    trace1_path: str | Path,
+    trace2_path: str | Path,
+    phase_filter: str = "all",
+    max_stacks: int = 3,
+    include_stacks: bool = True,
+) -> dict[str, Any]:
+    """Analyze two traces and return comparison data.
+    
+    Args:
+        trace1_path: Path to first trace file
+        trace2_path: Path to second trace file
+        phase_filter: Filter by phase ('all', 'prefill', or 'decode')
+        max_stacks: Maximum number of Python stack traces to collect per operation (0 for unlimited)
+        include_stacks: Whether to include Python stack traces (disable for faster analysis)
+        
+    Returns:
+        Dictionary containing:
+            - metadata: trace info (GPUs, kernel counts, total times, etc.)
+            - operations: per-operation comparison data
+            - layers: per-layer comparison data (if layers detected)
+    """
+    # Load both traces in parallel using separate processes
+    # This provides ~1.7x speedup over sequential loading
+    print("Loading traces in parallel...", file=sys.stderr)
+    
+    with ProcessPoolExecutor(max_workers=2) as executor:
+        future1 = executor.submit(load_trace_full, str(trace1_path), include_stacks)
+        future2 = executor.submit(load_trace_full, str(trace2_path), include_stacks)
+        trace1 = future1.result()
+        trace2 = future2.result()
+    
+    print("Analyzing operations...", file=sys.stderr)
+    
+    result = analyze_traces_from_loaded(trace1, trace2, phase_filter, max_stacks)
+    
+    # Update metadata with file paths for backward compatibility
+    result["metadata"]["trace1_name"] = str(trace1_path)
+    result["metadata"]["trace2_name"] = str(trace2_path)
+    
+    return result
+
+
+def analyze_traces_aligned(
+    trace1: LoadedTrace,
+    trace2: LoadedTrace,
+    phase_filter: str = "all",
+) -> dict[str, Any]:
+    """Analyze traces using kernel-to-kernel alignment.
+    
+    Args:
+        trace1: First loaded trace
+        trace2: Second loaded trace
+        phase_filter: Filter by phase ('all', 'prefill', or 'decode')
+        
+    Returns:
+        Dictionary with alignment-based comparison data
+    """
+    amd_phases = trace1.phases
+    nvidia_phases = trace2.phases
+    
+    if phase_filter != "all":
+        amd_phases = [p for p in amd_phases if p.get("type") == phase_filter]
+        nvidia_phases = [p for p in nvidia_phases if p.get("type") == phase_filter]
+    
+    amd_kernels = trace1.kernel_events
+    nvidia_kernels = trace2.kernel_events
+    
+    if phase_filter != "all" and amd_phases:
+        phase_starts = [p["ts_start"] for p in amd_phases]
+        phase_ends = [p["ts_end"] for p in amd_phases]
+        amd_kernels = [
+            k for k in amd_kernels
+            if any(phase_starts[i] <= k.get("ts", 0) <= phase_ends[i] 
+                   for i in range(len(phase_starts)))
+        ]
+    
+    if phase_filter != "all" and nvidia_phases:
+        phase_starts = [p["ts_start"] for p in nvidia_phases]
+        phase_ends = [p["ts_end"] for p in nvidia_phases]
+        nvidia_kernels = [
+            k for k in nvidia_kernels
+            if any(phase_starts[i] <= k.get("ts", 0) <= phase_ends[i] 
+                   for i in range(len(phase_starts)))
+        ]
+    
+    alignment = align_traces(
+        amd_kernels,
+        nvidia_kernels,
+        amd_phases,
+        nvidia_phases,
+        trace1.platform,
+        trace2.platform,
+    )
+    
+    layer_alignments = []
+    for layer_align in alignment.layer_alignments:
+        kernel_pairs = []
+        for pair in layer_align.kernel_pairs:
+            kernel_pairs.append({
+                "position": pair.position,
+                "operation": pair.operation,
+                "operation_detail": pair.operation_detail,
+                "amd_kernel": pair.amd_kernel,
+                "amd_avg_us": pair.amd_avg_us,
+                "amd_count": pair.amd_count,
+                "amd_total_us": pair.amd_total_us,
+                "nvidia_kernel": pair.nvidia_kernel,
+                "nvidia_avg_us": pair.nvidia_avg_us,
+                "nvidia_count": pair.nvidia_count,
+                "nvidia_total_us": pair.nvidia_total_us,
+                "ratio": pair.ratio,
+                "gap_us": pair.gap_us,
+                "fusion_note": pair.fusion_note,
+                "is_same_kernel": pair.is_same_kernel,
+            })
+        
+        layer_alignments.append({
+            "layer": layer_align.layer,
+            "amd_total_us": layer_align.amd_total_us,
+            "nvidia_total_us": layer_align.nvidia_total_us,
+            "ratio": layer_align.ratio,
+            "gap_us": layer_align.gap_us,
+            "kernel_pairs": kernel_pairs,
+        })
+    
+    fusion_result = analyze_fusion_from_alignment(alignment.layer_alignments)
+    same_kernel_result = analyze_same_kernels_from_alignment(alignment.layer_alignments)
+    
+    if trace1.platform == "AMD":
+        amd_trace, nvidia_trace = trace1, trace2
+    else:
+        amd_trace, nvidia_trace = trace2, trace1
+    
+    return {
+        "metadata": {
+            "amd_gpu": amd_trace.gpu_name,
+            "nvidia_gpu": nvidia_trace.gpu_name,
+            "amd_platform": amd_trace.platform,
+            "nvidia_platform": nvidia_trace.platform,
+            "model_layers": alignment.num_layers,
+            "forward_passes": alignment.num_forward_passes,
+            "phase_breakdown": alignment.phase_breakdown,
+            "phase_filter": phase_filter,
+            "trace1_platform": trace1.platform,
+            "trace1_gpu": trace1.gpu_name,
+            "trace1_device": trace1.device_props,
+            "trace2_platform": trace2.platform,
+            "trace2_gpu": trace2.gpu_name,
+            "trace2_device": trace2.device_props,
+            "trace1_kernels": len(amd_trace.kernel_events),
+            "trace2_kernels": len(nvidia_trace.kernel_events),
+            "trace1_total_ms": sum(k.get("dur", 0) for k in amd_trace.kernel_events) / 1000,
+            "trace2_total_ms": sum(k.get("dur", 0) for k in nvidia_trace.kernel_events) / 1000,
+            "phase": phase_filter,
+            "trace1_layers": alignment.num_layers,
+            "trace2_layers": alignment.num_layers,
+        },
+        "layer_alignments": layer_alignments,
+        "fusion_analysis": fusion_result,
+        "same_kernel_analysis": same_kernel_result,
+    }
+    
+    # Apply phase filter
+    if phase_filter != "all":
+        df1_filtered = df1[df1["phase"] == phase_filter]
+        df2_filtered = df2[df2["phase"] == phase_filter]
+        
+        if len(df1_filtered) == 0 and len(df2_filtered) == 0:
+            trace1_phases = {k: int(v) for k, v in df1["phase"].value_counts().items()}
+            trace2_phases = {k: int(v) for k, v in df2["phase"].value_counts().items()}
+            raise ValueError(
+                f"No {phase_filter} phase found. "
+                f"Trace1 phases: {trace1_phases}, Trace2 phases: {trace2_phases}"
+            )
+        
+        df1, df2 = df1_filtered, df2_filtered
+    
+    # Pre-compute aggregations for both operations and layers in single pass
+    trace1_by_op = df1.groupby("op").agg({
+        "dur_us": ["sum", "mean", "count"],
+        "phase": lambda x: set(x.dropna().unique()),
+        "cpu_op": lambda x: x.dropna().mode()[0] if len(x.dropna()) > 0 else None,
+    })
+    trace1_by_op.columns = ["total_us", "avg_us", "count", "phases", "cpu_op"]
+    
+    trace2_by_op = df2.groupby("op").agg({
+        "dur_us": ["sum", "mean", "count"],
+        "phase": lambda x: set(x.dropna().unique()),
+        "cpu_op": lambda x: x.dropna().mode()[0] if len(x.dropna()) > 0 else None,
+    })
+    trace2_by_op.columns = ["total_us", "avg_us", "count", "phases", "cpu_op"]
+    
+    # Group by layer for layer-level analysis
+    df1_layered = df1[df1["layer"].notna()]
+    df2_layered = df2[df2["layer"].notna()]
+    
+    trace1_by_layer = df1_layered.groupby("layer").agg({
+        "dur_us": ["sum", "count"],
+    }) if len(df1_layered) > 0 else pd.DataFrame()
+    if len(trace1_by_layer) > 0:
+        trace1_by_layer.columns = ["total_us", "count"]
+    
+    trace2_by_layer = df2_layered.groupby("layer").agg({
+        "dur_us": ["sum", "count"],
+    }) if len(df2_layered) > 0 else pd.DataFrame()
+    if len(trace2_by_layer) > 0:
+        trace2_by_layer.columns = ["total_us", "count"]
+    
+    results: dict[str, Any] = {
+        "metadata": {
+            "trace1_name": str(trace1_path),
+            "trace2_name": str(trace2_path),
+            "trace1_platform": p1,
+            "trace1_gpu": gpu1,
+            "trace1_device": dev1,
+            "trace2_platform": p2,
+            "trace2_gpu": gpu2,
+            "trace2_device": dev2,
+            "trace1_kernels": len(df1),
+            "trace2_kernels": len(df2),
+            "trace1_total_ms": df1["dur_us"].sum() / 1000,
+            "trace2_total_ms": df2["dur_us"].sum() / 1000,
+            "phase": phase_filter,
+            "trace1_layers": len(layers1),
+            "trace2_layers": len(layers2),
+        },
+        "operations": [],
+        "layers": [],
+    }
+    
+    # Per-operation comparison
+    all_ops = set(trace1_by_op.index) | set(trace2_by_op.index)
+    rmsnorm_compared = False
+    
+    for op in sorted(all_ops):
+        has_trace1 = op in trace1_by_op.index
+        has_trace2 = op in trace2_by_op.index
+        
+        trace1_op_for_pattern = op
+        trace2_op_for_pattern = op
+        skip_comparison = False
+        
+        if op == "RMSNorm+GEMM" and not has_trace2:
+            has_trace2 = "RMSNorm" in trace2_by_op.index
+            trace2_op_for_pattern = "RMSNorm"
+            rmsnorm_compared = True
+        elif op == "RMSNorm" and not has_trace1:
+            if rmsnorm_compared:
+                skip_comparison = True
+            else:
+                has_trace1 = "RMSNorm+GEMM" in trace1_by_op.index
+                trace1_op_for_pattern = "RMSNorm+GEMM"
+                rmsnorm_compared = True
+        
+        if skip_comparison or not (has_trace1 and has_trace2):
+            continue
+        
+        trace1_agg = trace1_by_op.loc[trace1_op_for_pattern]
+        trace2_agg = trace2_by_op.loc[trace2_op_for_pattern]
+        
+        trace1_avg = trace1_agg["avg_us"]
+        trace2_avg = trace2_agg["avg_us"]
+        trace1_total = trace1_agg["total_us"] / 1000
+        trace2_total = trace2_agg["total_us"] / 1000
+        trace1_count = int(trace1_agg["count"])
+        trace2_count = int(trace2_agg["count"])
+        ratio = trace1_avg / trace2_avg if trace2_avg > 0 else 1
+        gap_ms = trace1_total - trace2_total
+        
+        trace1_pattern = list(
+            patterns1.get(
+                (trace1_op_for_pattern, "decode"),
+                patterns1.get((trace1_op_for_pattern, "prefill"), {"unknown"}),
+            )
+        )[0]
+        trace2_pattern = list(
+            patterns2.get(
+                (trace2_op_for_pattern, "decode"),
+                patterns2.get((trace2_op_for_pattern, "prefill"), {"unknown"}),
+            )
+        )[0]
+        
+        trace1_cpu_op = trace1_agg["cpu_op"]
+        trace2_cpu_op = trace2_agg["cpu_op"]
+        
+        # Get detailed kernel data and stacks only when needed
+        trace1_data = df1[df1["op"] == trace1_op_for_pattern]
+        trace2_data = df2[df2["op"] == trace2_op_for_pattern]
+        
+        # Collect Python stacks if available
+        trace1_python_stacks = []
+        trace2_python_stacks = []
+        
+        if include_stacks:
+            stack_limit = None if max_stacks == 0 else max_stacks
+            for stack_list in trace1_data["python_stack"].head(stack_limit):
+                if stack_list and len(stack_list) > 0:
+                    trace1_python_stacks.append(stack_list)
+            
+            for stack_list in trace2_data["python_stack"].head(stack_limit):
+                if stack_list and len(stack_list) > 0:
+                    trace2_python_stacks.append(stack_list)
+        
+        # Aggregate individual kernels
+        trace1_kernels = trace1_data.groupby("name").agg({"dur_us": ["sum", "count", "mean"]}).reset_index()
+        trace1_kernels.columns = ["name", "total_us", "count", "avg_us"]
+        trace1_kernels = trace1_kernels.sort_values("total_us", ascending=False)
+        trace1_kernels_list = trace1_kernels.to_dict("records")
+        
+        trace2_kernels = trace2_data.groupby("name").agg({"dur_us": ["sum", "count", "mean"]}).reset_index()
+        trace2_kernels.columns = ["name", "total_us", "count", "avg_us"]
+        trace2_kernels = trace2_kernels.sort_values("total_us", ascending=False)
+        trace2_kernels_list = trace2_kernels.to_dict("records")
+        
+        if gap_ms > 5.0:
+            status = "slower"
+        elif gap_ms < -5.0:
+            status = "faster"
+        else:
+            status = "similar"
+        
+        phases = trace1_agg["phases"] | trace2_agg["phases"]
+        
+        results["operations"].append({
+            "operation": op,
+            "trace1_count": trace1_count,
+            "trace2_count": trace2_count,
+            "trace1_avg_us": trace1_avg,
+            "trace2_avg_us": trace2_avg,
+            "trace1_total_ms": trace1_total,
+            "trace2_total_ms": trace2_total,
+            "ratio": ratio,
+            "gap_ms": gap_ms,
+            "status": status,
+            "trace1_kernel": trace1_pattern,
+            "trace2_kernel": trace2_pattern,
+            "trace1_cpu_op": trace1_cpu_op,
+            "trace2_cpu_op": trace2_cpu_op,
+            "trace1_python_stacks": trace1_python_stacks,
+            "trace2_python_stacks": trace2_python_stacks,
+            "trace1_kernels": trace1_kernels_list,
+            "trace2_kernels": trace2_kernels_list,
+            "phases": sorted(list(phases)) if phases else ["all"],
+        })
+    
+    results["operations"].sort(key=lambda x: abs(x["gap_ms"]), reverse=True)
+    
+    # Layer-wise analysis
+    if len(trace1_by_layer) > 0 or len(trace2_by_layer) > 0:
+        all_layers = sorted(set(trace1_by_layer.index) | set(trace2_by_layer.index))
+        
+        for layer_num in all_layers:
+            has_trace1 = layer_num in trace1_by_layer.index
+            has_trace2 = layer_num in trace2_by_layer.index
+            
+            if has_trace1 and has_trace2:
+                trace1_agg = trace1_by_layer.loc[layer_num]
+                trace2_agg = trace2_by_layer.loc[layer_num]
+                
+                trace1_total = trace1_agg["total_us"] / 1000
+                trace2_total = trace2_agg["total_us"] / 1000
+                trace1_count = int(trace1_agg["count"])
+                trace2_count = int(trace2_agg["count"])
+                ratio = trace1_total / trace2_total if trace2_total > 0 else 1
+                gap_ms = trace1_total - trace2_total
+                
+                threshold_ms = 0.1
+                threshold_ratio = 1.2
+                if gap_ms > threshold_ms and ratio > threshold_ratio:
+                    status = "slower"
+                elif gap_ms < -threshold_ms and ratio < (1.0 / threshold_ratio):
+                    status = "faster"
+                else:
+                    status = "similar"
+                
+                results["layers"].append({
+                    "layer": int(layer_num),
+                    "trace1_kernels": trace1_count,
+                    "trace2_kernels": trace2_count,
+                    "trace1_total_ms": trace1_total,
+                    "trace2_total_ms": trace2_total,
+                    "ratio": ratio,
+                    "gap_ms": gap_ms,
+                    "status": status,
+                    "in_both": True,
+                })
+            elif has_trace1:
+                trace1_agg = trace1_by_layer.loc[layer_num]
+                trace1_total = trace1_agg["total_us"] / 1000
+                trace1_count = int(trace1_agg["count"])
+                
+                results["layers"].append({
+                    "layer": int(layer_num),
+                    "trace1_kernels": trace1_count,
+                    "trace2_kernels": 0,
+                    "trace1_total_ms": trace1_total,
+                    "trace2_total_ms": 0.0,
+                    "ratio": 0.0,
+                    "gap_ms": trace1_total,
+                    "status": "trace1_only",
+                    "in_both": False,
+                })
+            elif has_trace2:
+                trace2_agg = trace2_by_layer.loc[layer_num]
+                trace2_total = trace2_agg["total_us"] / 1000
+                trace2_count = int(trace2_agg["count"])
+                
+                results["layers"].append({
+                    "layer": int(layer_num),
+                    "trace1_kernels": 0,
+                    "trace2_kernels": trace2_count,
+                    "trace1_total_ms": 0.0,
+                    "trace2_total_ms": trace2_total,
+                    "ratio": 0.0,
+                    "gap_ms": -trace2_total,
+                    "status": "trace2_only",
+                    "in_both": False,
+                })
+        
+        results["layers"].sort(key=lambda x: (not x["in_both"], abs(x["gap_ms"])), reverse=True)
+    
+    print("Analysis complete.", file=sys.stderr)
+    return results