visualize.py

#!/usr/bin/env python3
"""
CPU Cache Benchmark Visualizer
Loads benchmark results from CSV or JSON and creates visualizations
"""

import json
import csv
import sys
import argparse
import matplotlib.pyplot as plt
import numpy as np
from pathlib import Path
import subprocess
import platform
import re

def get_cache_sizes():
    """
    Attempt to detect CPU cache sizes (L1, L2, L3) in bytes.
    Returns tuple (l1, l2, l3). Uses defaults if detection fails.
    """
    # Defaults
    l1 = 32 * 1024
    l2 = 256 * 1024
    l3 = 8 * 1024 * 1024
    
    system = platform.system()
    
    try:
        if system == 'Linux':
            # Try reading from sysfs for cpu0 (assuming symmetric cores)
            base_path = Path("/sys/devices/system/cpu/cpu0/cache")
            if base_path.exists():
                detected_l1 = None
                detected_l2 = None
                detected_l3 = None
                
                for index in base_path.glob("index*"):
                    try:
                        with open(index / "level", "r") as f:
                            level = int(f.read().strip())
                        with open(index / "type", "r") as f:
                            ctype = f.read().strip()
                        with open(index / "size", "r") as f:
                            size_str = f.read().strip()
                            
                        # Parse size (e.g., "32K", "256K", "8192K", "8M")
                        unit_mult = 1
                        if size_str.endswith('K'):
                            unit_mult = 1024
                            size_val = int(size_str[:-1])
                        elif size_str.endswith('M'):
                            unit_mult = 1024 * 1024
                            size_val = int(size_str[:-1])
                        else:
                            size_val = int(size_str)
                        
                        size_bytes = size_val * unit_mult
                        
                        if level == 1 and ctype == "Data":
                            detected_l1 = size_bytes
                        elif level == 2:
                            detected_l2 = size_bytes
                        elif level == 3:
                            detected_l3 = size_bytes
                    except (ValueError, OSError):
                        continue
                
                if detected_l1: l1 = detected_l1
                if detected_l2: l2 = detected_l2
                if detected_l3: l3 = detected_l3
                
        elif system == 'Darwin': # macOS
            try:
                # Use sysctl
                def get_sysctl(name):
                    cmd = ['sysctl', '-n', name]
                    return int(subprocess.check_output(cmd).decode().strip())
                
                l1 = get_sysctl('hw.l1dcachesize')
                l2 = get_sysctl('hw.l2cachesize')
                l3 = get_sysctl('hw.l3cachesize')
            except (subprocess.CalledProcessError, ValueError, FileNotFoundError):
                pass
                
        elif system == 'Windows':
            # Try using PowerShell to get Cache Memory info
            cmd = ["powershell", "-Command", "Get-CimInstance -ClassName Win32_CacheMemory | Select-Object -Property Level, MaxCacheSize"]
            proc = subprocess.run(cmd, capture_output=True, text=True)
            
            if proc.returncode == 0:
                detected_l1 = None
                detected_l2 = None
                detected_l3 = None
                
                # Also get core count to handle aggregated sizes if necessary
                core_cmd = ["powershell", "-Command", "Get-CimInstance -ClassName Win32_Processor | Select-Object -Property NumberOfCores"]
                core_proc = subprocess.run(core_cmd, capture_output=True, text=True)
                num_cores = 1
                if core_proc.returncode == 0:
                    match = re.search(r'\d+', core_proc.stdout)
                    if match:
                        num_cores = int(match.group())
                
                lines = proc.stdout.strip().split('\n')
                for line in lines:
                    parts = line.strip().split()
                    if len(parts) >= 2 and parts[0].isdigit() and parts[-1].isdigit():
                        level = int(parts[0])
                        size_kb = int(parts[-1])
                        size_bytes = size_kb * 1024
                        
                        # Win32_CacheMemory Levels: 3=L1, 4=L2, 5=L3 (usually)
                        # Heuristic: If size is very large relative to typical per-core sizes, 
                        # it might be an aggregated total.
                        
                        if level == 3: # L1
                            # L1 is typically 32-64KB. If > 256KB and we have multiple cores, it's likely total.
                            if size_bytes > 256 * 1024 and num_cores > 1:
                                size_bytes //= num_cores
                            detected_l1 = size_bytes
                        elif level == 4: # L2
                            # L2 is typically 256KB-1MB. If > 4MB and multiple cores, check carefully.
                            # But some CPUs have large shared L2. Assuming per-core for L2 is safer if large.
                            if size_bytes > 4 * 1024 * 1024 and num_cores > 1:
                                size_bytes //= num_cores
                            detected_l2 = size_bytes
                        elif level == 5: # L3
                            # L3 is usually shared, so take full size.
                            detected_l3 = size_bytes

                if detected_l1: l1 = detected_l1
                if detected_l2: l2 = detected_l2
                if detected_l3: l3 = detected_l3

    except Exception as e:
        print(f"Warning: Could not detect cache sizes ({e}), using defaults.")
    
    return l1, l2, l3

def get_cpu_name():
    """
    Attempt to detect CPU name.
    Returns string CPU name or "Unknown CPU".
    """
    cpu_name = "Unknown CPU"
    system = platform.system()
    
    try:
        if system == 'Linux':
            try:
                with open("/proc/cpuinfo", "r") as f:
                    for line in f:
                        if "model name" in line:
                            cpu_name = line.split(":")[1].strip()
                            break
            except (IOError, IndexError):
                pass
        elif system == 'Darwin':
            try:
                cmd = ['sysctl', '-n', 'machdep.cpu.brand_string']
                cpu_name = subprocess.check_output(cmd).decode().strip()
            except (subprocess.CalledProcessError, FileNotFoundError):
                pass
        elif system == 'Windows':
            try:
                cmd = ["powershell", "-Command", "Get-CimInstance -ClassName Win32_Processor | Select-Object -ExpandProperty Name"]
                proc = subprocess.run(cmd, capture_output=True, text=True)
                if proc.returncode == 0:
                    name = proc.stdout.strip()
                    if name:
                        cpu_name = name
            except Exception:
                pass
                
        # Fallback to platform.processor() if specific methods fail and it returns something useful
        if cpu_name == "Unknown CPU":
            proc_info = platform.processor()
            if proc_info:
                cpu_name = proc_info
                
    except Exception:
        pass
        
    return cpu_name

def get_os_name():
    """
    Attempt to detect OS name and version.
    """
    try:
        system = platform.system()
        if system == 'Windows':
            return f"{system} {platform.release()}"
        elif system == 'Darwin':
            return f"macOS {platform.mac_ver()[0]}"
        elif system == 'Linux':
            return f"Linux {platform.release()}"
        return system
    except Exception:
        return "Unknown OS"

# Get cache sizes and CPU/OS info
L1_CACHE_SIZE, L2_CACHE_SIZE, L3_CACHE_SIZE = get_cache_sizes()
CPU_NAME = get_cpu_name()
OS_NAME = get_os_name()
SYSTEM_INFO = f"{CPU_NAME} | {OS_NAME}"

print(f"Detected System: {SYSTEM_INFO}")
print(f"Using cache sizes: L1={L1_CACHE_SIZE//1024}KB, L2={L2_CACHE_SIZE//1024}KB, L3={L3_CACHE_SIZE//1024//1024}MB")


def load_data(filename):
    """Load benchmark data from CSV or JSON file"""
    path = Path(filename)
    if not path.exists():
        print(f"Error: File {filename} not found", file=sys.stderr)
        return None
    
    if path.suffix.lower() == '.json':
        with open(path, 'r') as f:
            return json.load(f)
    else:
        # Assume CSV - try different encodings
        data = []
        encodings = ['utf-8-sig', 'utf-8', 'latin-1', 'cp1252']
        for encoding in encodings:
            try:
                with open(path, 'r', encoding=encoding, newline='') as f:
                    reader = csv.DictReader(f)
                    for row in reader:
                        data.append(row)
                return data
            except (UnicodeDecodeError, UnicodeError):
                continue
        # If all encodings fail, try with errors='replace'
        with open(path, 'r', encoding='utf-8', errors='replace', newline='') as f:
            reader = csv.DictReader(f)
            for row in reader:
                data.append(row)
        return data


def parse_csv_data(data):
    """Parse CSV data into structured format"""
    benchmarks = {
        'sequential-vs-random': {'sequential': [], 'random': []},
        'stride': [],
        'aos-soa': {'aos': [], 'soa': []},
        'pointer-chasing': []
    }
    
    for row in data:
        bench_type = row.get('benchmark', '').strip()
        if bench_type == 'sequential-vs-random':
            pattern = row.get('parameter', '').strip()
            if pattern and row.get('time_ns'):
                benchmarks[bench_type][pattern].append({
                    'size': int(row['size']),
                    'time_ns': float(row['time_ns'])
                })
        elif bench_type == 'stride':
            if row.get('time_ns'):
                benchmarks[bench_type].append({
                    'stride': int(row['parameter']),
                    'size': int(row['size']),
                    'time_ns': float(row['time_ns'])
                })
        elif bench_type == 'aos-soa':
            layout = row.get('parameter', '').strip()
            if layout and row.get('time_ns'):
                benchmarks[bench_type][layout].append({
                    'size': int(row['size']),
                    'time_ns': float(row['time_ns'])
                })
        elif bench_type == 'pointer-chasing':
            if row.get('time_ns'):
                benchmarks[bench_type].append({
                    'size': int(row['size']),
                    'time_ns': float(row['time_ns'])
                })
    
    return benchmarks


def parse_json_data(data):
    """Parse JSON data into structured format"""
    benchmarks = {
        'sequential-vs-random': {'sequential': [], 'random': []},
        'stride': [],
        'aos-soa': {'aos': [], 'soa': []},
        'pointer-chasing': []
    }
    
    for bench in data.get('benchmarks', []):
        bench_type = bench.get('type', '')
        if bench_type == 'sequential-vs-random':
            for result in bench.get('results', []):
                pattern = result['pattern']
                benchmarks[bench_type][pattern].append({
                    'size': bench['array_size'],
                    'time_ns': result['time_ns']
                })
        elif bench_type == 'stride':
            benchmarks[bench_type].append({
                'stride': bench['stride'],
                'size': bench['array_size'],
                'time_ns': bench['time_ns']
            })
        elif bench_type == 'aos-soa':
            for result in bench.get('results', []):
                layout = result['layout']
                benchmarks[bench_type][layout].append({
                    'size': bench['count'] * 12,  # 3 ints = 12 bytes
                    'time_ns': result['time_ns']
                })
        elif bench_type == 'pointer-chasing':
            benchmarks[bench_type].append({
                'size': bench['count'] * 16,  # Node structure size
                'time_ns': bench['time_ns']
            })
    
    return benchmarks


def plot_sequential_random(benchmarks, output_file=None):
    """Plot sequential vs random access comparison"""
    seq_data = benchmarks['sequential-vs-random']['sequential']
    rand_data = benchmarks['sequential-vs-random']['random']
    
    if not seq_data or not rand_data:
        print("Warning: No sequential vs random data to plot")
        return
    
    sizes = [d['size'] for d in seq_data]
    seq_times = [d['time_ns'] for d in seq_data]
    rand_times = [d['time_ns'] for d in rand_data]
    
    fig, ax = plt.subplots(figsize=(10, 6))
    ax.plot(sizes, seq_times, 'o-', label='Sequential', linewidth=2, markersize=6)
    ax.plot(sizes, rand_times, 's-', label='Random', linewidth=2, markersize=6)
    
    ax.set_xlabel('Array Size (bytes)', fontsize=12)
    ax.set_ylabel('Access Time (ns)', fontsize=12)
    ax.set_title(f'Sequential vs Random Access Performance\n{SYSTEM_INFO}', fontsize=14, fontweight='bold')
    ax.legend(fontsize=11)
    ax.grid(True, alpha=0.3)
    ax.set_xscale('log', base=2)
    
    # Set x-axis limits to show all data points, extending slightly beyond max for better visibility
    min_size = min(sizes)
    max_size = max(sizes)
    ax.set_xlim(left=min_size * 0.8, right=max_size * 1.2)
    
    # Add cache size annotations
    # Only show cache lines if they're within the actual data range
    for cache_size, label in [(L1_CACHE_SIZE, 'L1'), (L2_CACHE_SIZE, 'L2'), (L3_CACHE_SIZE, 'L3')]:
        # Show cache line only if it's within the data range (cache_size <= max_size)
        # This ensures L3 doesn't show if data doesn't reach it
        if cache_size <= max_size:
            ax.axvline(x=cache_size, color='red', linestyle='--', alpha=0.5, linewidth=1)
            ax.text(cache_size, ax.get_ylim()[1] * 0.9, label, rotation=0, va='bottom', ha='center',
                   fontsize=11, fontweight='bold', color='black',
                   bbox=dict(boxstyle='round,pad=0.3', facecolor='white', edgecolor='black', alpha=0.8))
    
    # Add statistics text box
    stats_text = (
        f"Sequential:\n"
        f"  Mean: {np.mean(seq_times):.2f} ns\n"
        f"  Median: {np.median(seq_times):.2f} ns\n"
        f"  Std Dev: {np.std(seq_times):.2f} ns\n\n"
        f"Random:\n"
        f"  Mean: {np.mean(rand_times):.2f} ns\n"
        f"  Median: {np.median(rand_times):.2f} ns\n"
        f"  Std Dev: {np.std(rand_times):.2f} ns"
    )
    
    props = dict(boxstyle='round', facecolor='white', alpha=0.8)
    ax.text(0.02, 0.98, stats_text, transform=ax.transAxes, fontsize=10,
            verticalalignment='top', bbox=props)
    
    plt.tight_layout()
    if output_file:
        plt.savefig(output_file, dpi=300, bbox_inches='tight')
        print(f"Saved plot to {output_file}")
    else:
        plt.show()


def plot_stride(benchmarks, output_file=None):
    """Plot stride access performance"""
    stride_data = benchmarks['stride']
    
    if not stride_data:
        print("Warning: No stride data to plot")
        return
    
    strides = sorted(set(d['stride'] for d in stride_data))
    times_by_stride = {s: [] for s in strides}
    
    for d in stride_data:
        times_by_stride[d['stride']].append(d['time_ns'])
    
    avg_times = [np.mean(times_by_stride[s]) for s in strides]
    
    fig, ax = plt.subplots(figsize=(10, 6))
    ax.plot(strides, avg_times, 'o-', linewidth=2, markersize=8, color='purple')
    
    ax.set_xlabel('Stride', fontsize=12)
    ax.set_ylabel('Access Time (ns)', fontsize=12)
    ax.set_title(f'Stride Access Performance\n{SYSTEM_INFO}', fontsize=14, fontweight='bold')
    ax.grid(True, alpha=0.3)
    ax.set_xscale('log', base=2)
    
    # Add statistics text box
    stats_text = (
        f"Statistics:\n"
        f"  Mean: {np.mean(avg_times):.2f} ns\n"
        f"  Median: {np.median(avg_times):.2f} ns\n"
        f"  Std Dev: {np.std(avg_times):.2f} ns"
    )
    
    props = dict(boxstyle='round', facecolor='white', alpha=0.8)
    ax.text(0.02, 0.98, stats_text, transform=ax.transAxes, fontsize=10,
            verticalalignment='top', bbox=props)

    plt.tight_layout()
    if output_file:
        plt.savefig(output_file, dpi=300, bbox_inches='tight')
        print(f"Saved plot to {output_file}")
    else:
        plt.show()


def plot_aos_soa(benchmarks, output_file=None):
    """Plot Array of Structs vs Struct of Arrays comparison"""
    aos_data = benchmarks['aos-soa']['aos']
    soa_data = benchmarks['aos-soa']['soa']
    
    if not aos_data or not soa_data:
        print("Warning: No AoS vs SoA data to plot")
        return
    
    # Use sizes from one of them (should be same)
    sizes = [d['size'] for d in aos_data]
    aos_times = [d['time_ns'] for d in aos_data]
    soa_times = [d['time_ns'] for d in soa_data]
    
    x = np.arange(len(sizes))
    width = 0.35
    
    fig, ax = plt.subplots(figsize=(10, 6))
    bars1 = ax.bar(x - width/2, aos_times, width, label='Array of Structs', alpha=0.9, color='#1f77b4')  # Steel Blue
    bars2 = ax.bar(x + width/2, soa_times, width, label='Struct of Arrays', alpha=0.9, color='#ff7f0e')  # Safety Orange
    
    ax.set_xlabel('Data Size (bytes)', fontsize=12)
    ax.set_ylabel('Access Time (ns)', fontsize=12)
    ax.set_title(f'Array of Structs vs Struct of Arrays Performance\n{SYSTEM_INFO}', fontsize=14, fontweight='bold')
    ax.set_xticks(x)
    ax.set_xticklabels([f'{s//1024}KB' if s < 1024*1024 else f'{s//(1024*1024)}MB' for s in sizes], rotation=45, ha='right')
    ax.legend(fontsize=11)
    ax.grid(True, alpha=0.3, axis='y')
    
    # Add cache size annotations
    min_size = min(sizes)
    max_size = max(sizes)
    for cache_size, label in [(L1_CACHE_SIZE, 'L1'), (L2_CACHE_SIZE, 'L2'), (L3_CACHE_SIZE, 'L3')]:
        if min_size <= cache_size <= max_size:
            # Find the x position by interpolating between bar positions
            # Check if cache_size exactly matches a tested size
            if cache_size in sizes:
                x_pos = sizes.index(cache_size)
            else:
                # Find the two adjacent sizes that bracket the cache size
                # Use log interpolation since sizes typically increase exponentially
                for i in range(len(sizes) - 1):
                    if sizes[i] <= cache_size <= sizes[i + 1]:
                        # Linear interpolation on log scale
                        log_cache = np.log(cache_size)
                        log_small = np.log(sizes[i])
                        log_large = np.log(sizes[i + 1])
                        if log_large != log_small:
                            x_pos = i + (log_cache - log_small) / (log_large - log_small)
                        else:
                            x_pos = i
                        break
                else:
                    continue  # Skip if cache_size is outside all tested sizes
            
            ax.axvline(x=x_pos, color='red', linestyle='--', alpha=0.5, linewidth=1)
            # Position label at top of plot
            y_max = max(max(aos_times), max(soa_times))
            ax.text(x_pos, y_max * 0.95, label, rotation=0, va='bottom', ha='center',
                   fontsize=11, fontweight='bold', color='black',
                   bbox=dict(boxstyle='round,pad=0.3', facecolor='white', edgecolor='black', alpha=0.8))
    
    # Add statistics text box
    stats_text = (
        f"Array of Structs (AoS):\n"
        f"  Mean: {np.mean(aos_times):.2f} ns\n"
        f"  Median: {np.median(aos_times):.2f} ns\n"
        f"  Std Dev: {np.std(aos_times):.2f} ns\n\n"
        f"Struct of Arrays (SoA):\n"
        f"  Mean: {np.mean(soa_times):.2f} ns\n"
        f"  Median: {np.median(soa_times):.2f} ns\n"
        f"  Std Dev: {np.std(soa_times):.2f} ns"
    )
    
    props = dict(boxstyle='round', facecolor='white', alpha=0.8)
    ax.text(0.02, 0.98, stats_text, transform=ax.transAxes, fontsize=10,
            verticalalignment='top', bbox=props)

    plt.tight_layout()
    if output_file:
        plt.savefig(output_file, dpi=300, bbox_inches='tight')
        print(f"Saved plot to {output_file}")
    else:
        plt.show()


def plot_pointer_chasing(benchmarks, output_file=None):
    """Plot pointer chasing performance"""
    pc_data = benchmarks['pointer-chasing']
    
    if not pc_data:
        print("Warning: No pointer chasing data to plot")
        return
    
    sizes = [d['size'] for d in pc_data]
    times = [d['time_ns'] for d in pc_data]
    
    fig, ax = plt.subplots(figsize=(10, 6))
    ax.plot(sizes, times, 'o-', linewidth=2, markersize=8, color='#2ca02c')  # Cooked Asparagus Green
    
    ax.set_xlabel('Data Size (bytes)', fontsize=12)
    ax.set_ylabel('Access Time per Pointer Dereference (ns)', fontsize=12)
    ax.set_title(f'Pointer Chasing Performance\n{SYSTEM_INFO}', fontsize=14, fontweight='bold')
    ax.grid(True, alpha=0.3)
    ax.set_xscale('log', base=2)
    
    # Set x-axis limits to show all data points, extending slightly beyond max for better visibility
    min_size = min(sizes)
    max_size = max(sizes)
    ax.set_xlim(left=min_size * 0.8, right=max_size * 1.2)
    
    # Add cache size annotations
    # Only show cache lines if they're within the actual data range
    for cache_size, label in [(L1_CACHE_SIZE, 'L1'), (L2_CACHE_SIZE, 'L2'), (L3_CACHE_SIZE, 'L3')]:
        # Show cache line only if it's within the data range (cache_size <= max_size)
        # This ensures L3 doesn't show if data doesn't reach it
        if cache_size <= max_size:
            ax.axvline(x=cache_size, color='red', linestyle='--', alpha=0.5, linewidth=1)
            ax.text(cache_size, ax.get_ylim()[1] * 0.9, label, rotation=0, va='bottom', ha='center',
                   fontsize=11, fontweight='bold', color='black',
                   bbox=dict(boxstyle='round,pad=0.3', facecolor='white', edgecolor='black', alpha=0.8))
    
    # Add statistics text box
    stats_text = (
        f"Statistics:\n"
        f"  Mean: {np.mean(times):.2f} ns\n"
        f"  Median: {np.median(times):.2f} ns\n"
        f"  Std Dev: {np.std(times):.2f} ns"
    )
    
    props = dict(boxstyle='round', facecolor='white', alpha=0.8)
    ax.text(0.02, 0.98, stats_text, transform=ax.transAxes, fontsize=10,
            verticalalignment='top', bbox=props)
    
    plt.tight_layout()
    if output_file:
        plt.savefig(output_file, dpi=300, bbox_inches='tight')
        print(f"Saved plot to {output_file}")
    else:
        plt.show()


def main():
    parser = argparse.ArgumentParser(description='Visualize CPU cache benchmark results')
    parser.add_argument('input_file', help='Input CSV or JSON file with benchmark results')
    parser.add_argument('--benchmark', choices=['sequential', 'stride', 'aos-soa', 'pointer-chasing', 'all'],
                       default='all', help='Which benchmark to visualize (default: all)')
    parser.add_argument('--output', help='Output file for plot (default: display interactively)')
    
    args = parser.parse_args()
    
    # Load data
    data = load_data(args.input_file)
    if data is None:
        return 1
    
    # Parse data
    if isinstance(data, dict) and 'benchmarks' in data:
        benchmarks = parse_json_data(data)
    else:
        benchmarks = parse_csv_data(data)
    
    # Create visualizations
    benchmark = args.benchmark
    
    if benchmark == 'sequential' or benchmark == 'all':
        output = f"{args.output}_sequential.png" if args.output and benchmark == 'all' else args.output
        plot_sequential_random(benchmarks, output)
    
    if benchmark == 'stride' or benchmark == 'all':
        output = f"{args.output}_stride.png" if args.output and benchmark == 'all' else args.output
        plot_stride(benchmarks, output)
    
    if benchmark == 'aos-soa' or benchmark == 'all':
        output = f"{args.output}_aos_soa.png" if args.output and benchmark == 'all' else args.output
        plot_aos_soa(benchmarks, output)
    
    if benchmark == 'pointer-chasing' or benchmark == 'all':
        output = f"{args.output}_pointer_chasing.png" if args.output and benchmark == 'all' else args.output
        plot_pointer_chasing(benchmarks, output)
    
    return 0


if __name__ == '__main__':
    sys.exit(main())