benchmark_gpu_vs_cpu.py

import time
import os
import subprocess
import torch
import numpy as np
from decoder import SSTVDecoder

def run_benchmark():
    print("Preparing Benchmark...")
    # 1. Generate heavy payload (1080p)
    # We use 1080p to make the diff noticeable.
    input_image = "test_pattern.png"
    if not os.path.exists(input_image):
        print("Error: test_pattern.png missing")
        return
        
    audio_file = "benchmark_1080p.wav"
    # Generate audio if not exists or if we want to be sure
    # Using existing sstv_filter logic
    # We call sstv_filter.py externally
    if not os.path.exists(audio_file):
        print("Generating 1080p test audio (this may take a moment)...")
        cmd = [
            "python3", "sstv_filter.py",
            input_image, "dummy_bench.png",
            "--preset", "1080p",
            "--save-audio", audio_file
        ]
        subprocess.run(cmd, stdout=subprocess.DEVNULL, check=True)
        
    # 1080p dims
    w, h = 1920, 1080
    
    # 2. Benchmark CPU
    print("\n--- Benchmarking CPU ---")
    start_cpu = time.time()
    
    # We perform decode multiple times? No, 1080p is slow enough. One run.
    decoder_cpu = SSTVDecoder(audio_file, w, h, force_device='cpu')
    decoder_cpu.decode("benchmark_out_cpu.png")
    
    end_cpu = time.time()
    dur_cpu = end_cpu - start_cpu
    print(f"CPU Time: {dur_cpu:.2f}s")
    
    # 3. Benchmark MPS
    print("\n--- Benchmarking MPS ---")
    if not torch.backends.mps.is_available():
        print("MPS not available! Skipping.")
        dur_mps = 0
    else:
        # Warmup (optional, but good for GPU)
        # We assume init cost is part of the experience
        start_mps = time.time()
        
        decoder_mps = SSTVDecoder(audio_file, w, h, force_device='mps')
        decoder_mps.decode("benchmark_out_mps.png")
        
        end_mps = time.time()
        dur_mps = end_mps - start_mps
        print(f"MPS Time: {dur_mps:.2f}s")

    # 4. Report
    with open("benchmark_report.md", "w") as f:
        f.write(f"# CPU vs MPS Benchmark Report\n")
        f.write(f"Device: Apple Silicon (M-Series)\n")
        f.write(f"Resolution: 1080p (1920x1080)\n\n")
        f.write("| Backend | Time (seconds) | Speedup |\n")
        f.write("|---|---|---|\n")
        f.write(f"| **CPU** (SciPy) | {dur_cpu:.2f}s | 1.0x |\n")
        if dur_mps > 0:
            speedup = dur_cpu / dur_mps
            f.write(f"| **MPS** (PyTorch) | {dur_mps:.2f}s | **{speedup:.2f}x** |\n")
        else:
            f.write(f"| MPS | N/A | - |\n")
            
    print(f"\nBenchmark Complete. Report saved to benchmark_report.md")

if __name__ == "__main__":
    run_benchmark()