analyze_trace.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
PyTorch Profiler Trace Analyzer

Parse and analyze .pt.trace.json files generated by PyTorch Profiler.
Provides summary statistics and bottleneck identification.

Usage:
    python analyze_trace.py gpu-public033059049197.na175_562.1774457948452218067.pt.trace.json
"""

import json
import sys
from collections import defaultdict
from typing import Dict, List, Tuple


def load_trace(filepath: str) -> Dict:
    """Load Chrome trace JSON file."""
    print(f"Loading trace file: {filepath}")
    with open(filepath, "r") as f:
        trace_data = json.load(f)
    print(f"Loaded {len(trace_data.get('traceEvents', []))} events")
    return trace_data


def analyze_gpu_kernels(trace_data: Dict) -> Dict:
    """Analyze GPU kernel execution patterns."""
    events = trace_data.get("traceEvents", [])

    gpu_kernels = []
    sync_events = []  # Track synchronization events specifically

    # Synchronization function names that cause GPU stalls
    sync_functions = {
        "cudaDeviceSynchronize",
        "cudaStreamSynchronize",
        "cudaEventSynchronize",
        "cudaStreamWaitEvent",
        "cudaMemcpy",
        "cudaMemcpyAsync",  # Can trigger sync in some cases
    }

    for event in events:
        if event.get("ph") == "X":
            name = event.get("name", "")
            duration = event.get("dur", 0)

            # Check if it's a sync-related event
            is_sync = any(sync_fn in name for sync_fn in sync_functions)

            if is_sync:
                sync_events.append(
                    {
                        "name": name,
                        "duration_us": duration,
                        "start_us": event.get("ts", 0),
                        "pid": event.get("pid"),
                        "tid": event.get("tid"),
                        "args": event.get("args", {}),
                    }
                )

            # Also track CUDA-related events
            if "cuda" in name.lower() or "kernel" in name.lower():
                gpu_kernels.append(
                    {
                        "name": name,
                        "duration_us": duration,
                        "start_us": event.get("ts", 0),
                        "pid": event.get("pid"),
                        "tid": event.get("tid"),
                        "is_sync": is_sync,
                    }
                )

    if not gpu_kernels:
        print("No GPU kernel events found")
        return {}

    # Sort by duration
    gpu_kernels.sort(key=lambda x: x["duration_us"], reverse=True)
    sync_events.sort(key=lambda x: x["duration_us"], reverse=True)

    # Statistics
    total_time = sum(k["duration_us"] for k in gpu_kernels)
    sync_time = sum(e["duration_us"] for e in sync_events)
    compute_time = total_time - sync_time

    unique_kernels = defaultdict(lambda: {"count": 0, "total_time": 0})
    sync_breakdown = defaultdict(
        lambda: {"count": 0, "total_time": 0, "max_time": 0, "events": []}
    )

    for kernel in gpu_kernels:
        name = kernel["name"]
        unique_kernels[name]["count"] += 1
        unique_kernels[name]["total_time"] += kernel["duration_us"]

    for event in sync_events:
        name = event["name"].split("(")[0]  # Get base function name
        sync_breakdown[name]["count"] += 1
        sync_breakdown[name]["total_time"] += event["duration_us"]
        sync_breakdown[name]["max_time"] = max(
            sync_breakdown[name]["max_time"], event["duration_us"]
        )
        # Keep top 5 longest events for analysis
        sync_breakdown[name]["events"].append(event)
        sync_breakdown[name]["events"] = sorted(
            sync_breakdown[name]["events"], key=lambda x: x["duration_us"], reverse=True
        )[:5]

    # Top kernels by total time
    top_kernels = sorted(
        unique_kernels.items(), key=lambda x: x[1]["total_time"], reverse=True
    )[:20]

    return {
        "total_kernel_time_us": total_time,
        "sync_time_us": sync_time,
        "compute_time_us": compute_time,
        "sync_percentage": (sync_time / total_time * 100) if total_time > 0 else 0,
        "num_kernels": len(gpu_kernels),
        "num_sync_events": len(sync_events),
        "unique_kernels": len(unique_kernels),
        "top_kernels": top_kernels,
        "sync_breakdown": dict(sync_breakdown),
        "avg_kernel_time_us": total_time / len(gpu_kernels) if gpu_kernels else 0,
        "longest_sync_events": sync_events[:10],  # Top 10 longest sync events
    }


def analyze_cpu_phases(trace_data: Dict) -> Dict:
    """Analyze CPU-side operation phases."""
    events = trace_data.get("traceEvents", [])

    phases = {
        "data_to_device": [],
        "forward_pass": [],
        "backward_pass": [],
        "optimizer_step": [],
        "zero_grad": [],
        "other": [],
    }

    for event in events:
        if event.get("ph") != "X":
            continue

        name = event.get("name", "")
        duration = event.get("dur", 0)

        if "data_to_device" in name:
            phases["data_to_device"].append(duration)
        elif "forward_pass" in name:
            phases["forward_pass"].append(duration)
        elif "backward_pass" in name:
            phases["backward_pass"].append(duration)
        elif "optimizer_step" in name:
            phases["optimizer_step"].append(duration)
        elif "zero_grad" in name:
            phases["zero_grad"].append(duration)
        elif any(kw in name for kw in ["aten::", "cuda"]):
            phases["other"].append(duration)

    # Calculate statistics
    phase_stats = {}
    for phase, durations in phases.items():
        if durations:
            phase_stats[phase] = {
                "count": len(durations),
                "total_us": sum(durations),
                "avg_us": sum(durations) / len(durations),
                "max_us": max(durations),
                "min_us": min(durations),
            }

    return phase_stats


def analyze_call_stacks(trace_data: Dict) -> Dict:
    """Analyze call stacks to find synchronization sources."""
    events = trace_data.get("traceEvents", [])

    # Look for Python frames and correlate with sync events
    python_frames = []
    sync_with_stacks = []

    sync_functions = {
        "cudaDeviceSynchronize",
        "cudaStreamSynchronize",
        "cudaEventSynchronize",
    }

    for event in events:
        name = event.get("name", "")

        # Look for Python stack traces
        if (
            event.get("cat") == "python_function"
            or "python" in str(event.get("args", {})).lower()
        ):
            python_frames.append(event)

        # Look for sync events with stack info
        if any(sync_fn in name for sync_fn in sync_functions):
            args = event.get("args", {})
            if args:
                sync_with_stacks.append(
                    {
                        "name": name,
                        "duration_us": event.get("dur", 0),
                        "stack": args.get("External id", args.get("stack", str(args))),
                        "ts": event.get("ts", 0),
                    }
                )

    # Group by stack trace to find common sources
    stack_groups = defaultdict(lambda: {"count": 0, "total_time": 0})
    for event in sync_with_stacks:
        stack_key = str(event.get("stack", "unknown"))[:200]  # Truncate for grouping
        stack_groups[stack_key]["count"] += 1
        stack_groups[stack_key]["total_time"] += event["duration_us"]

    return {
        "python_frames_count": len(python_frames),
        "sync_with_stacks": sync_with_stacks[:20],
        "stack_groups": dict(stack_groups),
    }


def analyze_timeline_gaps(trace_data: Dict) -> Dict:
    """Analyze gaps in GPU timeline to find idle periods."""
    events = trace_data.get("traceEvents", [])

    # Get GPU kernel events sorted by start time
    gpu_events = []
    for event in events:
        if event.get("ph") == "X":
            name = event.get("name", "")
            # Look for actual GPU compute kernels (not sync operations)
            if (
                "kernel" in name.lower()
                or "gemm" in name.lower()
                or "elementwise" in name.lower()
                or "backward" in name.lower()
            ):
                if "Synchronize" not in name:
                    gpu_events.append(
                        {
                            "name": name,
                            "start": event.get("ts", 0),
                            "end": event.get("ts", 0) + event.get("dur", 0),
                            "duration": event.get("dur", 0),
                        }
                    )

    if len(gpu_events) < 2:
        return {"gaps": [], "total_gap_time": 0}

    # Sort by start time
    gpu_events.sort(key=lambda x: x["start"])

    # Find gaps between consecutive GPU operations
    gaps = []
    for i in range(1, len(gpu_events)):
        gap = gpu_events[i]["start"] - gpu_events[i - 1]["end"]
        if gap > 1000:  # Only track gaps > 1ms
            gaps.append(
                {
                    "gap_us": gap,
                    "after_kernel": gpu_events[i - 1]["name"][:50],
                    "before_kernel": gpu_events[i]["name"][:50],
                    "position": i,
                }
            )

    # Sort by gap size
    gaps.sort(key=lambda x: x["gap_us"], reverse=True)

    return {
        "num_gaps": len(gaps),
        "total_gap_time_us": sum(g["gap_us"] for g in gaps),
        "largest_gaps": gaps[:20],
        "avg_gap_us": sum(g["gap_us"] for g in gaps) / len(gaps) if gaps else 0,
    }


def analyze_memory_events(trace_data: Dict) -> Dict:
    """Analyze memory allocation patterns."""
    events = trace_data.get("traceEvents", [])

    alloc_events = []
    for event in events:
        if "allocator" in event.get("name", "").lower():
            if event.get("ph") == "X":
                alloc_events.append(
                    {
                        "name": event["name"],
                        "duration_us": event.get("dur", 0),
                        "bytes": event.get("args", {}).get("Bytes", 0),
                    }
                )

    if not alloc_events:
        return {"num_alloc_events": 0}

    total_allocated = sum(e["bytes"] for e in alloc_events)

    return {
        "num_alloc_events": len(alloc_events),
        "total_allocated_bytes": total_allocated,
        "total_allocated_gb": total_allocated / (1024**3),
        "avg_alloc_bytes": total_allocated / len(alloc_events),
    }


def print_summary(
    gpu_stats: Dict,
    cpu_stats: Dict,
    mem_stats: Dict,
    stack_stats: Dict = None,
    gap_stats: Dict = None,
):
    """Print comprehensive analysis summary."""
    print("\n" + "=" * 80)
    print("PYTORCH PROFILER TRACE ANALYSIS SUMMARY")
    print("=" * 80)

    # GPU Kernel Analysis
    if gpu_stats:
        print("\n" + "-" * 80)
        print("🔹 GPU KERNEL STATISTICS")
        print("-" * 80)
        print(f"   Total kernels executed: {gpu_stats['num_kernels']}")
        print(f"   Unique kernel types: {gpu_stats['unique_kernels']}")
        print(f"   Total GPU time: {gpu_stats['total_kernel_time_us'] / 1000:.2f} ms")
        print(f"   Average kernel time: {gpu_stats['avg_kernel_time_us']:.2f} μs")

        # Synchronization breakdown - THIS IS THE KEY INSIGHT
        if gpu_stats.get("sync_time_us"):
            sync_time = gpu_stats["sync_time_us"]
            total_time = gpu_stats["total_kernel_time_us"]
            compute_time = gpu_stats.get("compute_time_us", total_time - sync_time)
            sync_pct = gpu_stats.get("sync_percentage", 0)

            print("\n   ⚠️  SYNCHRONIZATION ANALYSIS (Critical for GPU Efficiency):")
            print(f"   ├── Sync overhead: {sync_time / 1000:.2f} ms ({sync_pct:.1f}%)")
            print(
                f"   ├── Actual compute: {compute_time / 1000:.2f} ms ({100-sync_pct:.1f}%)"
            )
            print(f"   └── Sync events: {gpu_stats.get('num_sync_events', 0)}")

            if sync_pct > 30:
                print(f"\n   🔴 CRITICAL: GPU spends {sync_pct:.1f}% of time waiting!")
                print(
                    f"      This indicates CPU-bound bottleneck or excessive .item() calls."
                )

            # Breakdown by sync type
            if gpu_stats.get("sync_breakdown"):
                print("\n   Synchronization breakdown by type:")
                print(
                    f"   {'Type':<30} {'Count':>8} {'Total (ms)':>12} {'Max (ms)':>10}"
                )
                print("   " + "-" * 65)
                for sync_type, stats in sorted(
                    gpu_stats["sync_breakdown"].items(),
                    key=lambda x: x[1]["total_time"],
                    reverse=True,
                ):
                    print(
                        f"   {sync_type:<30} {stats['count']:>8} "
                        f"{stats['total_time'] / 1000:>12.2f} {stats['max_time'] / 1000:>10.2f}"
                    )

        print("\n   Top 10 kernels by total time:")
        print(f"   {'Kernel Name':<50} {'Count':>6} {'Total Time (ms)':>15}")
        print("   " + "-" * 75)
        for name, stats in gpu_stats["top_kernels"][:10]:
            short_name = name[:48] + ".." if len(name) > 50 else name
            print(
                f"   {short_name:<50} {stats['count']:>6} {stats['total_time'] / 1000:>15.2f}"
            )

    # GPU Timeline Gaps
    if gap_stats and gap_stats.get("num_gaps", 0) > 0:
        print("\n" + "-" * 80)
        print("🔹 GPU TIMELINE GAP ANALYSIS")
        print("-" * 80)
        print(f"   Number of significant gaps (>1ms): {gap_stats['num_gaps']}")
        print(f"   Total gap time: {gap_stats['total_gap_time_us'] / 1000:.2f} ms")
        print(f"   Average gap: {gap_stats['avg_gap_us'] / 1000:.2f} ms")

        if gap_stats.get("largest_gaps"):
            print("\n   Largest GPU idle gaps:")
            print(f"   {'Gap (ms)':>10} | After kernel → Before kernel")
            print("   " + "-" * 70)
            for gap in gap_stats["largest_gaps"][:5]:
                print(
                    f"   {gap['gap_us'] / 1000:>10.2f} | {gap['after_kernel'][:30]} → {gap['before_kernel'][:30]}"
                )

    # CPU Phase Analysis
    if cpu_stats:
        print("\n" + "-" * 80)
        print("🔹 CPU PHASE STATISTICS")
        print("-" * 80)
        print(
            f"   {'Phase':<25} {'Count':>6} {'Total (ms)':>12} {'Avg (ms)':>10} {'Max (ms)':>10}"
        )
        print("   " + "-" * 65)
        for phase, stats in sorted(
            cpu_stats.items(), key=lambda x: x[1]["total_us"], reverse=True
        ):
            print(
                f"   {phase:<25} {stats['count']:>6} {stats['total_us'] / 1000:>12.2f} "
                f"{stats['avg_us'] / 1000:>10.2f} {stats['max_us'] / 1000:>10.2f}"
            )

    # Memory Analysis
    if mem_stats.get("num_alloc_events", 0) > 0:
        print("\n" + "-" * 80)
        print("🔹 MEMORY ALLOCATION STATISTICS")
        print("-" * 80)
        print(f"   Total allocation events: {mem_stats['num_alloc_events']}")
        print(f"   Total memory allocated: {mem_stats['total_allocated_gb']:.2f} GB")
        print(
            f"   Average allocation: {mem_stats['avg_alloc_bytes'] / (1024**2):.2f} MB"
        )

    # Bottleneck Analysis
    print("\n" + "-" * 80)
    print("🔍 BOTTLENECK DIAGNOSIS")
    print("-" * 80)

    bottlenecks_found = []

    if gpu_stats:
        sync_pct = gpu_stats.get("sync_percentage", 0)
        if sync_pct > 50:
            bottlenecks_found.append(
                f"🔴 CRITICAL: cudaDeviceSynchronize takes {sync_pct:.1f}% of GPU time!\n"
                f"      Causes: .item() calls, .cpu() transfers, dist.reduce()\n"
                f"      Solution: Batch .item() calls, reduce logging frequency"
            )
        elif sync_pct > 30:
            bottlenecks_found.append(
                f"🟡 WARNING: Sync overhead is {sync_pct:.1f}% of GPU time\n"
                f"      Consider reducing .item() calls and logging frequency"
            )

        avg_kernel = gpu_stats.get("avg_kernel_time_us", 0)
        if avg_kernel < 50:
            bottlenecks_found.append(
                f"🟡 Small average kernel time ({avg_kernel:.1f} μs)\n"
                f"      Many small operations may cause kernel launch overhead\n"
                f"      Consider: torch.compile(), larger batch size, kernel fusion"
            )

    if cpu_stats:
        total_cpu_time = sum(s["total_us"] for s in cpu_stats.values())
        if total_cpu_time > 0:
            for phase, stats in cpu_stats.items():
                pct = (stats["total_us"] / total_cpu_time) * 100
                if phase == "data_to_device" and pct > 20:
                    bottlenecks_found.append(
                        f"🟡 Data transfer takes {pct:.1f}% of CPU time\n"
                        f"      Consider: pin_memory=True, async data loading"
                    )

    if gap_stats and gap_stats.get("total_gap_time_us", 0) > 100000:  # > 100ms
        bottlenecks_found.append(
            f"🟡 GPU idle gaps total {gap_stats['total_gap_time_us'] / 1000:.2f} ms\n"
            f"      GPU is waiting for CPU between kernel launches"
        )

    if bottlenecks_found:
        for bottleneck in bottlenecks_found:
            print(f"   {bottleneck}\n")
    else:
        print("   ✅ No major bottlenecks detected.")

    print("\n" + "=" * 80)
    print("\n💡 RECOMMENDATIONS:")
    print(
        "   1. Reduce cudaDeviceSynchronize: batch .item() calls using get_loss_values()"
    )
    print("   2. Increase logging_steps in config (e.g., 500 instead of 100)")
    print("   3. Use async data loading: pin_memory=True, num_workers > 0")
    print("   4. Consider torch.compile() for kernel fusion")
    print("   5. View detailed timeline: chrome://tracing → Load the .json file")
    print("   6. For TensorBoard: tensorboard --logdir=<profiler_dir>")
    print("=" * 80 + "\n")


def main():
    if len(sys.argv) < 2:
        print("Usage: python analyze_trace.py <trace_file.json>")
        print(
            "Example: python analyze_trace.py gpu-public033059049197.na175_562.1774457948452218067.pt.trace.json"
        )
        sys.exit(1)

    filepath = sys.argv[1]

    try:
        # Load trace
        trace_data = load_trace(filepath)

        # Analyze
        print("\nAnalyzing GPU kernels and synchronization...")
        gpu_stats = analyze_gpu_kernels(trace_data)

        print("Analyzing CPU phases...")
        cpu_stats = analyze_cpu_phases(trace_data)

        print("Analyzing memory events...")
        mem_stats = analyze_memory_events(trace_data)

        print("Analyzing call stacks...")
        stack_stats = analyze_call_stacks(trace_data)

        print("Analyzing GPU timeline gaps...")
        gap_stats = analyze_timeline_gaps(trace_data)

        # Print summary
        print_summary(gpu_stats, cpu_stats, mem_stats, stack_stats, gap_stats)

        # Save analysis to file
        output_file = filepath.replace(".json", "_analysis.json")
        analysis = {
            "gpu_stats": {
                k: v
                for k, v in gpu_stats.items()
                if k not in ["longest_sync_events"]  # Don't save large event lists
            },
            "cpu_stats": cpu_stats,
            "memory_stats": mem_stats,
            "gap_stats": {
                k: v
                for k, v in gap_stats.items()
                if k != "largest_gaps"  # Simplify for JSON
            }
            if gap_stats
            else {},
            "summary": {
                "sync_percentage": gpu_stats.get("sync_percentage", 0),
                "total_sync_time_ms": gpu_stats.get("sync_time_us", 0) / 1000,
                "total_compute_time_ms": gpu_stats.get("compute_time_us", 0) / 1000,
                "total_gap_time_ms": gap_stats.get("total_gap_time_us", 0) / 1000
                if gap_stats
                else 0,
                "bottleneck": "CPU-GPU Synchronization"
                if gpu_stats.get("sync_percentage", 0) > 30
                else "None detected",
            },
        }

        # Clean up sync_breakdown for JSON serialization
        if "sync_breakdown" in analysis["gpu_stats"]:
            for k, v in analysis["gpu_stats"]["sync_breakdown"].items():
                if "events" in v:
                    del v["events"]

        with open(output_file, "w") as f:
            json.dump(analysis, f, indent=2)
        print(f"\n📄 Detailed analysis saved to: {output_file}")

        # Also save a text report
        report_file = filepath.replace(".json", "_report.txt")
        with open(report_file, "w") as f:
            import io
            from contextlib import redirect_stdout

            buffer = io.StringIO()
            with redirect_stdout(buffer):
                print_summary(gpu_stats, cpu_stats, mem_stats, stack_stats, gap_stats)
            f.write(buffer.getvalue())
        print(f"📄 Text report saved to: {report_file}\n")

    except Exception as e:
        print(f"❌ Error analyzing trace: {e}")
        import traceback

        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()