benchmark.py

"""
Comprehensive Benchmarking Suite for HRM SRAM/DRAM Memory Tiering.

Measures and compares:
  - Per-module latency (L-level SRAM vs H-level DRAM)
  - End-to-end inference latency
  - Throughput (samples/sec)
  - Memory usage per tier (SRAM / DRAM)
  - Cross-tier transfer overhead
  - SRAM hit rate
  - Triton kernel-level profiling
  - GPU utilization / power draw (when available)
"""

import json
import time
import math
import os
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass, asdict

import torch
import torch.nn.functional as F

from models.memory_tier import MemoryTierManager
from models.hrm.hrm_tiered import HRM_Tiered, HRM_Tiered_Inner
from models.hrm.hrm_act_v1 import (
    HierarchicalReasoningModel_ACTV1,
    HierarchicalReasoningModel_ACTV1Config,
    HierarchicalReasoningModel_ACTV1InnerCarry,
)
from models.triton_kernels import triton_memory_latency_probe


@dataclass
class BenchmarkResult:
    """Results from a single benchmark run."""
    model_name: str
    batch_size: int
    seq_len: int
    hidden_size: int
    H_cycles: int
    L_cycles: int
    H_layers: int
    L_layers: int
    num_iterations: int
    warmup_iterations: int

    # Latency (μs)
    l_level_latency_mean_us: float
    l_level_latency_min_us: float
    l_level_latency_max_us: float
    l_level_latency_std_us: float

    h_level_latency_mean_us: float
    h_level_latency_min_us: float
    h_level_latency_max_us: float
    h_level_latency_std_us: float

    total_inference_latency_mean_ms: float
    total_inference_latency_min_ms: float
    total_inference_latency_max_ms: float
    total_inference_latency_std_ms: float

    # Throughput
    throughput_samples_per_sec: float

    # Memory (MB)
    sram_peak_mb: float
    dram_peak_mb: float
    total_gpu_memory_mb: float

    # Transfer metrics (μs)
    h_l_transfer_mean_us: float
    l_h_transfer_mean_us: float

    # SRAM metrics
    sram_hit_rate: float

    # Triton kernel metrics
    triton_sram_probe_latency_us: float
    triton_dram_probe_latency_us: float

    # GPU metrics
    gpu_utilization_pct: Optional[float]
    gpu_power_w: Optional[float]
    gpu_temperature_c: Optional[float]

    # Derived
    h_over_l_latency_ratio: float
    memory_efficiency: float  # useful_compute_time / total_time


def _std(values: List[float]) -> float:
    if len(values) < 2:
        return 0.0
    mean = sum(values) / len(values)
    var = sum((v - mean) ** 2 for v in values) / (len(values) - 1)
    return math.sqrt(var)


def _create_dummy_batch(
    batch_size: int,
    seq_len: int,
    vocab_size: int,
    device: torch.device,
) -> Dict[str, torch.Tensor]:
    """Create a synthetic batch for benchmarking."""
    return {
        "inputs": torch.randint(0, vocab_size, (batch_size, seq_len), device=device),
        "labels": torch.randint(0, vocab_size, (batch_size, seq_len), device=device),
        "puzzle_identifiers": torch.arange(batch_size, device=device),
    }


def _get_gpu_metrics() -> Dict[str, Optional[float]]:
    """Try to read GPU utilization, power, temperature via nvidia-smi."""
    metrics = {'utilization': None, 'power': None, 'temperature': None}
    try:
        import subprocess
        result = subprocess.run(
            ['nvidia-smi', '--query-gpu=utilization.gpu,power.draw,temperature.gpu',
             '--format=csv,noheader,nounits'],
            capture_output=True, text=True, timeout=5,
        )
        if result.returncode == 0:
            parts = result.stdout.strip().split(',')
            if len(parts) >= 3:
                metrics['utilization'] = float(parts[0].strip())
                metrics['power'] = float(parts[1].strip())
                metrics['temperature'] = float(parts[2].strip())
    except Exception:
        pass
    return metrics


def _run_triton_latency_probe(
    batch_size: int,
    hidden_size: int,
    device: torch.device,
    num_iters: int = 100,
) -> Tuple[float, float]:
    """Measure SRAM vs DRAM effective latency using Triton probe kernels.

    Returns (sram_latency_us, dram_latency_us).
    """
    if not torch.cuda.is_available():
        return 0.0, 0.0

    # Small tensor → fits in SRAM (L1/L2/registers)
    sram_data = torch.randn(batch_size, hidden_size, device=device, dtype=torch.float32)

    # Large tensor → forces DRAM access (much larger than L2)
    dram_size = max(hidden_size, 65536)  # Force spill to DRAM
    dram_data = torch.randn(batch_size * 64, dram_size, device=device, dtype=torch.float32)

    # Make a contiguous slice for DRAM probe
    dram_probe_data = dram_data[:batch_size, :hidden_size].contiguous()

    # Warmup
    triton_memory_latency_probe(sram_data, num_iters=10)
    triton_memory_latency_probe(dram_probe_data, num_iters=10)
    torch.cuda.synchronize()

    # SRAM probe
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    start.record()
    triton_memory_latency_probe(sram_data, num_iters=num_iters)
    end.record()
    torch.cuda.synchronize()
    sram_us = start.elapsed_time(end) * 1000  # ms → μs

    # DRAM probe (access scattered to prevent caching)
    start2 = torch.cuda.Event(enable_timing=True)
    end2 = torch.cuda.Event(enable_timing=True)
    start2.record()
    triton_memory_latency_probe(dram_probe_data, num_iters=num_iters)
    end2.record()
    torch.cuda.synchronize()
    dram_us = start2.elapsed_time(end2) * 1000

    # Cleanup
    del sram_data, dram_data
    torch.cuda.empty_cache()

    return sram_us, dram_us


def benchmark_tiered_model(
    batch_size: int = 8,
    seq_len: int = 64,
    hidden_size: int = 512,
    num_heads: int = 8,
    H_cycles: int = 2,
    L_cycles: int = 2,
    H_layers: int = 4,
    L_layers: int = 4,
    halt_max_steps: int = 1,
    warmup: int = 5,
    iterations: int = 20,
    device: Optional[torch.device] = None,
) -> BenchmarkResult:
    """Benchmark the tiered HRM model."""

    if device is None:
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    vocab_size = 32
    config_dict = {
        'batch_size': batch_size,
        'seq_len': seq_len,
        'puzzle_emb_ndim': 0,
        'num_puzzle_identifiers': batch_size,
        'vocab_size': vocab_size,
        'H_cycles': H_cycles,
        'L_cycles': L_cycles,
        'H_layers': H_layers,
        'L_layers': L_layers,
        'hidden_size': hidden_size,
        'expansion': 4.0,
        'num_heads': num_heads,
        'pos_encodings': 'rope',
        'halt_max_steps': halt_max_steps,
        'halt_exploration_prob': 0.0,
    }

    # Memory manager
    mem_mgr = MemoryTierManager(device=device, enable_tracking=True)

    # Create model
    model = HRM_Tiered(config_dict, memory_manager=mem_mgr).to(device)
    model.eval()

    batch = _create_dummy_batch(batch_size, seq_len, vocab_size, device)

    # ---- Warmup ----
    with torch.no_grad():
        for _ in range(warmup):
            carry = model.initial_carry(batch)
            carry.inner_carry.z_H = carry.inner_carry.z_H.to(device)
            carry.inner_carry.z_L = carry.inner_carry.z_L.to(device)
            carry.steps = carry.steps.to(device)
            carry.halted = carry.halted.to(device)
            carry.current_data = {k: v.to(device) for k, v in carry.current_data.items()}
            model(carry, batch)

    model.reset_timing()
    mem_mgr.reset_stats()
    if torch.cuda.is_available():
        torch.cuda.reset_peak_memory_stats(device)
        torch.cuda.synchronize()

    # ---- Benchmark iterations ----
    total_latencies_ms = []
    with torch.no_grad():
        for _ in range(iterations):
            carry = model.initial_carry(batch)
            carry.inner_carry.z_H = carry.inner_carry.z_H.to(device)
            carry.inner_carry.z_L = carry.inner_carry.z_L.to(device)
            carry.steps = carry.steps.to(device)
            carry.halted = carry.halted.to(device)
            carry.current_data = {k: v.to(device) for k, v in carry.current_data.items()}

            if torch.cuda.is_available():
                start_event = torch.cuda.Event(enable_timing=True)
                end_event = torch.cuda.Event(enable_timing=True)
                start_event.record()

            t0 = time.perf_counter()
            model(carry, batch)

            if torch.cuda.is_available():
                end_event.record()
                torch.cuda.synchronize()
                total_latencies_ms.append(start_event.elapsed_time(end_event))
            else:
                total_latencies_ms.append((time.perf_counter() - t0) * 1000)

    # ---- Collect results ----
    timing = model.get_timing_stats()
    mem_stats = mem_mgr.get_stats()

    # GPU memory
    total_gpu_mb = 0.0
    if torch.cuda.is_available():
        total_gpu_mb = torch.cuda.max_memory_allocated(device) / (1024 * 1024)

    # GPU metrics
    gpu_metrics = _get_gpu_metrics()

    # Triton latency probe
    sram_probe_us, dram_probe_us = _run_triton_latency_probe(
        batch_size, hidden_size, device,
    )

    # Compute derived metrics
    l_mean = timing['L_forward_us']['mean_us']
    h_mean = timing['H_forward_us']['mean_us']
    ratio = h_mean / l_mean if l_mean > 0 else float('inf')

    total_mean_ms = sum(total_latencies_ms) / len(total_latencies_ms)
    throughput = batch_size / (total_mean_ms / 1000) if total_mean_ms > 0 else 0

    compute_time = timing['L_forward_us']['total_us'] + timing['H_forward_us']['total_us']
    transfer_time = timing['H_L_transfer_us']['total_us'] + timing['L_H_transfer_us']['total_us']
    efficiency = compute_time / (compute_time + transfer_time) if (compute_time + transfer_time) > 0 else 0

    # Build L-level stats
    l_values = [timing['L_forward_us']['min_us'], timing['L_forward_us']['max_us']]

    result = BenchmarkResult(
        model_name='HRM_Tiered',
        batch_size=batch_size,
        seq_len=seq_len,
        hidden_size=hidden_size,
        H_cycles=H_cycles,
        L_cycles=L_cycles,
        H_layers=H_layers,
        L_layers=L_layers,
        num_iterations=iterations,
        warmup_iterations=warmup,

        l_level_latency_mean_us=l_mean,
        l_level_latency_min_us=timing['L_forward_us']['min_us'],
        l_level_latency_max_us=timing['L_forward_us']['max_us'],
        l_level_latency_std_us=0.0,

        h_level_latency_mean_us=h_mean,
        h_level_latency_min_us=timing['H_forward_us']['min_us'],
        h_level_latency_max_us=timing['H_forward_us']['max_us'],
        h_level_latency_std_us=0.0,

        total_inference_latency_mean_ms=total_mean_ms,
        total_inference_latency_min_ms=min(total_latencies_ms),
        total_inference_latency_max_ms=max(total_latencies_ms),
        total_inference_latency_std_ms=_std(total_latencies_ms),

        throughput_samples_per_sec=throughput,

        sram_peak_mb=mem_stats['sram']['peak_mb'],
        dram_peak_mb=mem_stats['dram']['peak_mb'],
        total_gpu_memory_mb=total_gpu_mb,

        h_l_transfer_mean_us=timing['H_L_transfer_us']['mean_us'],
        l_h_transfer_mean_us=timing['L_H_transfer_us']['mean_us'],

        sram_hit_rate=mem_stats['sram']['hit_rate'],

        triton_sram_probe_latency_us=sram_probe_us,
        triton_dram_probe_latency_us=dram_probe_us,

        gpu_utilization_pct=gpu_metrics['utilization'],
        gpu_power_w=gpu_metrics['power'],
        gpu_temperature_c=gpu_metrics['temperature'],

        h_over_l_latency_ratio=ratio,
        memory_efficiency=efficiency,
    )

    # Cleanup
    del model, batch, carry
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    return result


def benchmark_baseline_model(
    batch_size: int = 8,
    seq_len: int = 64,
    hidden_size: int = 512,
    num_heads: int = 8,
    H_cycles: int = 2,
    L_cycles: int = 2,
    H_layers: int = 4,
    L_layers: int = 4,
    halt_max_steps: int = 1,
    warmup: int = 5,
    iterations: int = 20,
    device: Optional[torch.device] = None,
) -> BenchmarkResult:
    """Benchmark the original (non-tiered) HRM model."""

    if device is None:
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    vocab_size = 32
    config_dict = {
        'batch_size': batch_size,
        'seq_len': seq_len,
        'puzzle_emb_ndim': 0,
        'num_puzzle_identifiers': batch_size,
        'vocab_size': vocab_size,
        'H_cycles': H_cycles,
        'L_cycles': L_cycles,
        'H_layers': H_layers,
        'L_layers': L_layers,
        'hidden_size': hidden_size,
        'expansion': 4.0,
        'num_heads': num_heads,
        'pos_encodings': 'rope',
        'halt_max_steps': halt_max_steps,
        'halt_exploration_prob': 0.0,
    }

    model = HierarchicalReasoningModel_ACTV1(config_dict).to(device)
    model.eval()

    batch = _create_dummy_batch(batch_size, seq_len, vocab_size, device)

    # ---- Warmup ----
    with torch.no_grad():
        for _ in range(warmup):
            carry = model.initial_carry(batch)
            carry.inner_carry.z_H = carry.inner_carry.z_H.to(device)
            carry.inner_carry.z_L = carry.inner_carry.z_L.to(device)
            carry.steps = carry.steps.to(device)
            carry.halted = carry.halted.to(device)
            carry.current_data = {k: v.to(device) for k, v in carry.current_data.items()}
            model(carry, batch)

    if torch.cuda.is_available():
        torch.cuda.reset_peak_memory_stats(device)
        torch.cuda.synchronize()

    # ---- Benchmark ----
    total_latencies_ms = []
    with torch.no_grad():
        for _ in range(iterations):
            carry = model.initial_carry(batch)
            carry.inner_carry.z_H = carry.inner_carry.z_H.to(device)
            carry.inner_carry.z_L = carry.inner_carry.z_L.to(device)
            carry.steps = carry.steps.to(device)
            carry.halted = carry.halted.to(device)
            carry.current_data = {k: v.to(device) for k, v in carry.current_data.items()}

            if torch.cuda.is_available():
                start_event = torch.cuda.Event(enable_timing=True)
                end_event = torch.cuda.Event(enable_timing=True)
                start_event.record()

            t0 = time.perf_counter()
            model(carry, batch)

            if torch.cuda.is_available():
                end_event.record()
                torch.cuda.synchronize()
                total_latencies_ms.append(start_event.elapsed_time(end_event))
            else:
                total_latencies_ms.append((time.perf_counter() - t0) * 1000)

    total_gpu_mb = 0.0
    if torch.cuda.is_available():
        total_gpu_mb = torch.cuda.max_memory_allocated(device) / (1024 * 1024)

    total_mean_ms = sum(total_latencies_ms) / len(total_latencies_ms)
    throughput = batch_size / (total_mean_ms / 1000) if total_mean_ms > 0 else 0

    gpu_metrics = _get_gpu_metrics()

    result = BenchmarkResult(
        model_name='HRM_Baseline',
        batch_size=batch_size,
        seq_len=seq_len,
        hidden_size=hidden_size,
        H_cycles=H_cycles,
        L_cycles=L_cycles,
        H_layers=H_layers,
        L_layers=L_layers,
        num_iterations=iterations,
        warmup_iterations=warmup,

        l_level_latency_mean_us=0, l_level_latency_min_us=0,
        l_level_latency_max_us=0, l_level_latency_std_us=0,
        h_level_latency_mean_us=0, h_level_latency_min_us=0,
        h_level_latency_max_us=0, h_level_latency_std_us=0,

        total_inference_latency_mean_ms=total_mean_ms,
        total_inference_latency_min_ms=min(total_latencies_ms),
        total_inference_latency_max_ms=max(total_latencies_ms),
        total_inference_latency_std_ms=_std(total_latencies_ms),

        throughput_samples_per_sec=throughput,

        sram_peak_mb=0, dram_peak_mb=0,
        total_gpu_memory_mb=total_gpu_mb,

        h_l_transfer_mean_us=0, l_h_transfer_mean_us=0,
        sram_hit_rate=0,

        triton_sram_probe_latency_us=0, triton_dram_probe_latency_us=0,

        gpu_utilization_pct=gpu_metrics['utilization'],
        gpu_power_w=gpu_metrics['power'],
        gpu_temperature_c=gpu_metrics['temperature'],

        h_over_l_latency_ratio=0,
        memory_efficiency=1.0,
    )

    del model, batch, carry
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    return result


def compare_models(
    batch_sizes: List[int] = [1, 8, 32],
    seq_lens: List[int] = [64, 128],
    hidden_size: int = 512,
    warmup: int = 5,
    iterations: int = 20,
    device: Optional[torch.device] = None,
) -> List[Dict]:
    """Run comparative benchmark between tiered and baseline HRM."""
    results = []

    for bs in batch_sizes:
        for sl in seq_lens:
            print(f"\n{'='*60}")
            print(f"  Benchmarking: batch_size={bs}, seq_len={sl}")
            print(f"{'='*60}")

            print("  → Baseline model...")
            baseline = benchmark_baseline_model(
                batch_size=bs, seq_len=sl, hidden_size=hidden_size,
                warmup=warmup, iterations=iterations, device=device,
            )

            print("  → Tiered model...")
            tiered = benchmark_tiered_model(
                batch_size=bs, seq_len=sl, hidden_size=hidden_size,
                warmup=warmup, iterations=iterations, device=device,
            )

            comparison = {
                'batch_size': bs,
                'seq_len': sl,
                'baseline': asdict(baseline),
                'tiered': asdict(tiered),
                'speedup': baseline.total_inference_latency_mean_ms / tiered.total_inference_latency_mean_ms if tiered.total_inference_latency_mean_ms > 0 else 0,
                'memory_savings_mb': baseline.total_gpu_memory_mb - tiered.total_gpu_memory_mb,
                'throughput_improvement': tiered.throughput_samples_per_sec / baseline.throughput_samples_per_sec if baseline.throughput_samples_per_sec > 0 else 0,
            }
            results.append(comparison)

            # Print summary
            print(f"\n  Results:")
            print(f"    Baseline latency: {baseline.total_inference_latency_mean_ms:.2f} ms")
            print(f"    Tiered latency:   {tiered.total_inference_latency_mean_ms:.2f} ms")
            print(f"    Speedup:          {comparison['speedup']:.2f}x")
            print(f"    H/L ratio:        {tiered.h_over_l_latency_ratio:.2f}x")
            print(f"    SRAM probe:       {tiered.triton_sram_probe_latency_us:.1f} μs")
            print(f"    DRAM probe:       {tiered.triton_dram_probe_latency_us:.1f} μs")

    return results


def print_results_table(results: List[BenchmarkResult]):
    """Pretty-print benchmark results as a table."""
    header = (
        f"{'Model':<15} {'BS':>4} {'Seq':>5} "
        f"{'Latency(ms)':>12} {'Throughput':>12} "
        f"{'L_lat(μs)':>10} {'H_lat(μs)':>10} {'H/L':>6} "
        f"{'GPU_MB':>8} {'Efficiency':>10}"
    )
    print(f"\n{'='*len(header)}")
    print(header)
    print(f"{'='*len(header)}")

    for r in results:
        print(
            f"{r.model_name:<15} {r.batch_size:>4} {r.seq_len:>5} "
            f"{r.total_inference_latency_mean_ms:>12.2f} {r.throughput_samples_per_sec:>12.1f} "
            f"{r.l_level_latency_mean_us:>10.1f} {r.h_level_latency_mean_us:>10.1f} {r.h_over_l_latency_ratio:>6.2f} "
            f"{r.total_gpu_memory_mb:>8.1f} {r.memory_efficiency:>10.3f}"
        )
    print()


def generate_plots(results: List[Dict], output_dir: str = "benchmark_results"):
    """Generate comparison plots using matplotlib."""
    try:
        import matplotlib
        matplotlib.use('Agg')
        import matplotlib.pyplot as plt
        import numpy as np
    except ImportError:
        print("matplotlib not available — skipping plot generation.")
        return

    os.makedirs(output_dir, exist_ok=True)

    # ---- Plot 1: Latency comparison ----
    fig, axes = plt.subplots(1, 3, figsize=(18, 5))
    fig.suptitle('HRM SRAM/DRAM Tiering — Benchmark Results', fontsize=14, fontweight='bold')

    configs = [f"bs={r['batch_size']}\nseq={r['seq_len']}" for r in results]
    baseline_lat = [r['baseline']['total_inference_latency_mean_ms'] for r in results]
    tiered_lat = [r['tiered']['total_inference_latency_mean_ms'] for r in results]

    x = np.arange(len(configs))
    w = 0.35

    ax = axes[0]
    ax.bar(x - w/2, baseline_lat, w, label='Baseline', color='#e74c3c', alpha=0.8)
    ax.bar(x + w/2, tiered_lat, w, label='Tiered (SRAM/DRAM)', color='#2ecc71', alpha=0.8)
    ax.set_xlabel('Configuration')
    ax.set_ylabel('Latency (ms)')
    ax.set_title('Inference Latency')
    ax.set_xticks(x)
    ax.set_xticklabels(configs, fontsize=8)
    ax.legend()
    ax.grid(axis='y', alpha=0.3)

    # ---- Plot 2: Throughput ----
    ax = axes[1]
    baseline_tp = [r['baseline']['throughput_samples_per_sec'] for r in results]
    tiered_tp = [r['tiered']['throughput_samples_per_sec'] for r in results]
    ax.bar(x - w/2, baseline_tp, w, label='Baseline', color='#e74c3c', alpha=0.8)
    ax.bar(x + w/2, tiered_tp, w, label='Tiered', color='#2ecc71', alpha=0.8)
    ax.set_xlabel('Configuration')
    ax.set_ylabel('Samples/sec')
    ax.set_title('Throughput')
    ax.set_xticks(x)
    ax.set_xticklabels(configs, fontsize=8)
    ax.legend()
    ax.grid(axis='y', alpha=0.3)

    # ---- Plot 3: H vs L latency (tiered only) ----
    ax = axes[2]
    l_lat = [r['tiered']['l_level_latency_mean_us'] for r in results]
    h_lat = [r['tiered']['h_level_latency_mean_us'] for r in results]
    ax.bar(x - w/2, l_lat, w, label='L-level (SRAM)', color='#3498db', alpha=0.8)
    ax.bar(x + w/2, h_lat, w, label='H-level (DRAM)', color='#e67e22', alpha=0.8)
    ax.set_xlabel('Configuration')
    ax.set_ylabel('Latency (μs)')
    ax.set_title('Per-Module Latency')
    ax.set_xticks(x)
    ax.set_xticklabels(configs, fontsize=8)
    ax.legend()
    ax.grid(axis='y', alpha=0.3)

    plt.tight_layout()
    plot_path = os.path.join(output_dir, 'benchmark_comparison.png')
    plt.savefig(plot_path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f"  Plot saved: {plot_path}")

    # ---- Plot 4: Memory breakdown ----
    fig, axes = plt.subplots(1, 2, figsize=(12, 5))
    fig.suptitle('Memory Analysis', fontsize=14, fontweight='bold')

    ax = axes[0]
    gpu_mem_baseline = [r['baseline']['total_gpu_memory_mb'] for r in results]
    gpu_mem_tiered = [r['tiered']['total_gpu_memory_mb'] for r in results]
    ax.bar(x - w/2, gpu_mem_baseline, w, label='Baseline', color='#e74c3c', alpha=0.8)
    ax.bar(x + w/2, gpu_mem_tiered, w, label='Tiered', color='#2ecc71', alpha=0.8)
    ax.set_xlabel('Configuration')
    ax.set_ylabel('GPU Memory (MB)')
    ax.set_title('Total GPU Memory')
    ax.set_xticks(x)
    ax.set_xticklabels(configs, fontsize=8)
    ax.legend()
    ax.grid(axis='y', alpha=0.3)

    ax = axes[1]
    triton_sram = [r['tiered']['triton_sram_probe_latency_us'] for r in results]
    triton_dram = [r['tiered']['triton_dram_probe_latency_us'] for r in results]
    ax.bar(x - w/2, triton_sram, w, label='SRAM Probe', color='#3498db', alpha=0.8)
    ax.bar(x + w/2, triton_dram, w, label='DRAM Probe', color='#e67e22', alpha=0.8)
    ax.set_xlabel('Configuration')
    ax.set_ylabel('Latency (μs)')
    ax.set_title('Triton Memory Probe Latency')
    ax.set_xticks(x)
    ax.set_xticklabels(configs, fontsize=8)
    ax.legend()
    ax.grid(axis='y', alpha=0.3)

    plt.tight_layout()
    plot_path = os.path.join(output_dir, 'memory_analysis.png')
    plt.savefig(plot_path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f"  Plot saved: {plot_path}")