gpu-trace-perf 1.5.0

//! Module for shader snapshot and diffing support.s

use anyhow::{Context, Result};
use clap::Args;
use image::{DynamicImage, GenericImageView as _};
use log::{error, info, warn};
use rayon::prelude::*;
use regex::Regex;
use serde::{Deserialize, Serialize};
use std::{
    cmp::{max, min},
    collections::HashSet,
    io::prelude::*,
    path::{Path, PathBuf},
};

use crate::collect_traces;

pub struct SnapshotResult {
    pub files: Vec<String>,
    pub cmdline: String,
    pub stdout: String,
    pub stderr: String,
}

#[derive(Debug, Args)]
/// Runs a collection of traces, dumping representative frames to .pngs and optionally comparing them.
pub struct Snapshot {
    /// path to a directory to store debug output and logs
    #[arg(long = "output")]
    pub output: String,

    /// path to a directory of baseline state (possibly a previous output dir)
    #[arg(long, default_value = "")]
    baseline: String,

    /// Number of threads to use (trace processes to run in parallel)
    #[arg(short, long, default_value = "1")]
    jobs: usize,

    #[arg(short, long, number_of_values = 1)]
    /// paths to directories containing traces (or traces themselves)
    traces: Vec<String>,
}

fn diff_percentage_filter(input: String) -> String {
    format!(
        "{:.2}%",
        str::parse::<f32>(&input)
            .with_context(|| format!("rendering diff percentage for {}", input))
            .unwrap()
            * 100.0
    )
}

fn format_time(input: String) -> String {
    format!(
        "{:.2}s",
        str::parse::<f32>(&input)
            .with_context(|| format!("rendering time for {}", input))
            .unwrap()
    )
}

#[derive(Serialize, Deserialize, Default)]
struct Run {
    image: String,
    image_preview: String,
    comment: String,
    log: String,
    execution_time: f32,
}

#[derive(Serialize, Deserialize)]
struct Diff {
    trace: String,
    frame: String,
    run_a: Run,
    run_b: Run,
    diff: f32,
    diff_image: String,
    diff_image_preview: String,
}

#[derive(Serialize, Deserialize)]
struct RenderContext {
    comparison: Comparison,
    diffs: Vec<Diff>,
}

#[derive(Serialize, Deserialize)]
struct RunInfo {
    comment: String,
    dxvk_versions: Vec<String>,
    devices: Vec<String>,
    drivers: Vec<String>,
}

/// Information parsed about from a logfile
#[derive(Debug, PartialEq, Eq, Hash, Clone)]
struct LogInfo {
    dxvk: Option<String>,
    driver: Option<String>,
    device: Option<String>,
}

impl From<&str> for LogInfo {
    fn from(value: &str) -> Self {
        let mut device = None;
        let mut driver = None;
        let mut dxvk = None;

        lazy_static! {
            static ref DXVK_RE: Regex = Regex::new("^info:  DXVK: v(.*)$").unwrap();
            static ref DEVICE_RE: Regex = Regex::new("^info:  (.+):$").unwrap();
            static ref DRIVER_RE: Regex = Regex::new("^info:    Driver   : (.+)$").unwrap();
        }

        // Extract DXVK version
        for line in value.lines() {
            if let Some(cap) = DXVK_RE.captures(line) {
                dxvk = Some(cap[1].to_string());
                break;
            }
        }

        // Extract device and driver from the "Creating device:" section
        let mut device_create_lines = value
            .lines()
            .skip_while(|line| *line != "info:  Creating device:")
            .skip(1);
        if let Some(line) = device_create_lines.next() {
            if let Some(cap) = DEVICE_RE.captures(line) {
                device = Some(cap[1].to_string());
            }
        }
        if let Some(line) = device_create_lines.next() {
            if let Some(cap) = DRIVER_RE.captures(line) {
                driver = Some(cap[1].to_string());
            }
        }

        LogInfo {
            driver,
            device,
            dxvk,
        }
    }
}
#[derive(Serialize, Deserialize)]
struct Comparison {
    max_diff: f32,
    run_a: RunInfo,
    run_b: RunInfo,
}

fn render(render_context: &RenderContext) -> Result<String> {
    let mut env = minijinja::Environment::new();
    env.add_template("index.html", include_str!("data/index.html"))
        .context("loading index template")?;
    env.add_template("diff-styles.html", include_str!("data/diff-styles.html"))
        .context("loading diff styles template")?;
    env.add_template("diff-scripts.html", include_str!("data/diff-scripts.html"))
        .context("loading diff scripts template")?;
    env.add_template(
        "codemirror-scripts.html",
        include_str!("data/codemirror-scripts.html"),
    )
    .context("loading diff styles template")?;
    env.add_filter("diff_percentage", diff_percentage_filter);
    env.add_filter("format_time", format_time);

    env.get_template("index.html")
        .context("getting diff HTML template")?
        .render(render_context)
        .context("rendering diff HTML template")
}

fn suffix_file(path: &Path, suffix: &str) -> Result<PathBuf> {
    let filename = path.file_name().context("getting file")?;
    let filename_string = filename.to_string_lossy().to_string();
    let (filename, ext) = filename_string
        .split_once('.')
        .context("getting filename ext")?;
    let filename = format!("{filename}{suffix}.{ext}");
    Ok(path.parent().context("getting parent dir")?.join(filename))
}

fn image_diff_fraction(a: &DynamicImage, b: &DynamicImage) -> f32 {
    let (a_w, a_h) = a.dimensions();
    let (b_w, b_h) = b.dimensions();
    let min_w = min(a_w, b_w);
    let min_h = min(a_h, b_h);
    let max_w = max(a_w, b_w);
    let max_h = max(a_h, b_h);

    let mut diff = 0;

    // Count differing pixels between the two.
    for y in 0..min_h {
        for x in 0..min_w {
            if a.get_pixel(x, y) != b.get_pixel(x, y) {
                diff += 1;
            }
        }
    }

    // Consider any pixels outside of the shared range as different, as well.
    diff += max(a_w, b_w) * max(a_h, b_h) - min_w * min_h;

    diff as f32 / (max_w * max_h) as f32
}

fn generate_preview(image: &DynamicImage, path: &Path) -> Result<PathBuf> {
    let w = 400;
    let h = ((image.height() * w) as f32 / image.width() as f32) as u32;
    let scaled_image = image::imageops::resize(image, w, h, image::imageops::FilterType::Gaussian);

    let path = suffix_file(path, "-preview")?;
    scaled_image
        .save(&path)
        .with_context(|| format!("writing {}", path.display()))?;
    Ok(path.to_path_buf())
}

pub fn html_filename(path: &Path, output_path: &Path) -> Result<String> {
    let relative = path
        .strip_prefix(output_path)
        .context("getting relative path for {}")?;
    Ok(html_escape::encode_quoted_attribute(&relative.to_string_lossy()).to_string())
}

fn run_log(result: &SnapshotResult) -> String {
    format!(
        "cmdline: {}\nstdout:\n{}stderr:\n{}",
        result.cmdline, result.stdout, result.stderr
    )
}

/// Entrypoint for gpu-trace-perf snapshot
pub fn snapshot(opts: &Snapshot) -> Result<()> {
    if opts.jobs > 0 {
        rayon::ThreadPoolBuilder::new()
            .num_threads(opts.jobs)
            .build_global()
            .unwrap();
    }

    let traces = collect_traces(opts.traces.clone(), false);

    info!("Capturing snapshots of traces");

    let results: Vec<_> = traces
        .into_par_iter()
        .filter(|trace| {
            if !trace.can_snapshot() {
                warn!("  No support for capturing snapshot of {}", trace.name());
                false
            } else {
                true
            }
        })
        .map(|trace| {
            info!("  Running {}", trace.name());
            let result = trace.snapshot(&opts.output);
            (trace, result)
        })
        .filter_map(|(trace, result)| match result {
            Ok(result) => {
                if result.files.is_empty() {
                    error!("  No snapshots captured for {}", trace.name());
                }
                Some((trace, result))
            }
            Err(e) => {
                error!("{}", e);
                None
            }
        })
        .collect();

    info!("Snapshot capture completed, generating HTML");

    let output_path = Path::new(&opts.output);
    let baseline_path = if opts.baseline.is_empty() {
        None
    } else {
        Some(Path::new(&opts.baseline))
    };

    let mut max_diff = 0.0;
    let mut diffs = Vec::new();

    // Collect all unique DXVK versions, devices, and drivers for run A and B
    let mut run_a_dxvk_versions = HashSet::new();
    let mut run_a_devices = HashSet::new();
    let mut run_a_drivers = HashSet::new();
    let mut run_b_dxvk_versions = HashSet::new();
    let mut run_b_devices = HashSet::new();
    let mut run_b_drivers = HashSet::new();

    for (trace, result) in &results {
        let subdir = trace.output_subdir();
        // Get the paths local to the current trace's subdir.
        let trace_output_path = output_path.join(&subdir);
        let trace_baseline_path = baseline_path.map(|x| x.join(subdir));

        let log_path = trace_output_path.join("log.txt");
        let log_content = run_log(result);
        trace.write_output(&opts.output, "log.txt", log_content.as_bytes())?;

        // Extract driver info from run B (current run)
        let log_info_b = LogInfo::from(log_content.as_str());
        if let Some(dxvk) = log_info_b.dxvk {
            run_b_dxvk_versions.insert(dxvk);
        }
        if let Some(device) = log_info_b.device {
            run_b_devices.insert(device);
        }
        if let Some(driver) = log_info_b.driver {
            run_b_drivers.insert(driver);
        }

        // Copy the baseline log file into our output directory and extract info from run A
        let baseline_log = if let Some(trace_baseline_path) = trace_baseline_path {
            let from = trace_baseline_path.join("log.txt");
            let to = trace_output_path.join("log-baseline.txt");

            match std::fs::read_to_string(&from) {
                Ok(baseline_content) => {
                    // Extract driver info from run A (baseline)
                    let log_info_a = LogInfo::from(baseline_content.as_str());
                    if let Some(dxvk) = log_info_a.dxvk {
                        run_a_dxvk_versions.insert(dxvk);
                    }
                    if let Some(device) = log_info_a.device {
                        run_a_devices.insert(device);
                    }
                    if let Some(driver) = log_info_a.driver {
                        run_a_drivers.insert(driver);
                    }

                    std::fs::write(&to, &baseline_content)
                        .with_context(|| format!("writing baseline log to {}", to.display()))?;

                    html_filename(&to, output_path)?
                }
                Err(e) => {
                    error!(
                        "Failed to read baseline log from {}: {:?}",
                        from.display(),
                        e
                    );
                    "".to_string()
                }
            }
        } else {
            "".to_string()
        };

        for file in &result.files {
            let snapshot_png = Path::new(&file);
            let snapshot_image = image::open(snapshot_png)
                .with_context(|| format!("reading {}", snapshot_png.display()))?;

            let preview_png =
                generate_preview(&snapshot_image, snapshot_png).context("baseline preview")?;
            let mut diff_image_html = "".to_string();
            let mut diff_image_preview_html = "".to_string();
            let mut diff = 0.0;

            let run_a = if let Some(baseline_path) = baseline_path {
                let snapshot_relative_png = snapshot_png
                    .strip_prefix(output_path)
                    .context("getting relative path to snapshot png")?;
                let orig_baseline_file_path = baseline_path.join(snapshot_relative_png);
                let baseline_file_path = suffix_file(snapshot_png, "-baseline")?;
                match std::fs::copy(&orig_baseline_file_path, &baseline_file_path).with_context(
                    || {
                        format!(
                            "copying baseline png from {}",
                            orig_baseline_file_path.display()
                        )
                    },
                ) {
                    Err(e) => {
                        error!("{:?}", e);
                        Run::default()
                    }
                    Ok(_) => {
                        let baseline_image = image::open(&baseline_file_path)
                            .with_context(|| format!("reading {}", baseline_file_path.display()))?;
                        let baseline_preview_png =
                            generate_preview(&baseline_image, &baseline_file_path)
                                .context("baseline preview")?;

                        diff = image_diff_fraction(&baseline_image, &snapshot_image);

                        let diff_image_path = suffix_file(snapshot_png, "-diff")?;
                        let diff_image = image_diff::diff(&baseline_image, &snapshot_image)
                            .context("creating diff image")?;
                        diff_image
                            .save(&diff_image_path)
                            .context("saving diff png")?;
                        let diff_image_preview_path =
                            generate_preview(&diff_image, &diff_image_path)
                                .context("diff preview")?;

                        diff_image_html = html_filename(&diff_image_path, output_path)?;
                        diff_image_preview_html =
                            html_filename(&diff_image_preview_path, output_path)?;
                        Run {
                            image: html_filename(&baseline_file_path, output_path)?,
                            image_preview: html_filename(&baseline_preview_png, output_path)?,
                            comment: String::new(),
                            execution_time: 0.0,
                            log: baseline_log.clone(),
                        }
                    }
                }
            } else {
                Run::default()
            };

            diffs.push(Diff {
                trace: trace.name().to_string(),
                frame: file.split('/').next_back().unwrap().to_string(),
                run_a,
                run_b: Run {
                    image: html_filename(snapshot_png, output_path)?,
                    image_preview: html_filename(&preview_png, output_path)?,
                    comment: String::new(),
                    execution_time: 0.0,
                    log: html_filename(&log_path, output_path)?,
                },
                diff,
                diff_image: diff_image_html,
                diff_image_preview: diff_image_preview_html,
            });

            if diff > max_diff {
                max_diff = diff;
            }
        }
    }

    // Convert HashSets to sorted Vecs for run A
    let mut run_a_dxvk_versions: Vec<_> = run_a_dxvk_versions.into_iter().collect();
    run_a_dxvk_versions.sort();
    let mut run_a_devices: Vec<_> = run_a_devices.into_iter().collect();
    run_a_devices.sort();
    let mut run_a_drivers: Vec<_> = run_a_drivers.into_iter().collect();
    run_a_drivers.sort();

    // Convert HashSets to sorted Vecs for run B
    let mut run_b_dxvk_versions: Vec<_> = run_b_dxvk_versions.into_iter().collect();
    run_b_dxvk_versions.sort();
    let mut run_b_devices: Vec<_> = run_b_devices.into_iter().collect();
    run_b_devices.sort();
    let mut run_b_drivers: Vec<_> = run_b_drivers.into_iter().collect();
    run_b_drivers.sort();

    let render_context = RenderContext {
        comparison: Comparison {
            max_diff,
            run_a: RunInfo {
                comment: format!("output dir: {}", opts.baseline),
                dxvk_versions: run_a_dxvk_versions,
                devices: run_a_devices,
                drivers: run_a_drivers,
            },
            run_b: RunInfo {
                comment: format!("output dir: {}", opts.output),
                dxvk_versions: run_b_dxvk_versions,
                devices: run_b_devices,
                drivers: run_b_drivers,
            },
        },
        diffs,
    };

    let template = render(&render_context).context("rendering diffs HTML")?;

    let mut file = std::fs::File::create(Path::new(&opts.output).join("index.html"))
        .context("opening diff HTML for output")?;
    file.write_all(template.as_bytes())
        .context("writing diff HTML")?;

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_render() {
        let render_context = RenderContext {
            comparison: Comparison {
                max_diff: 0.35,
                run_a: RunInfo {
                    comment: "run a comment".to_string(),
                    dxvk_versions: vec!["2.7.0".to_string()],
                    devices: vec!["AMD Radeon 890M Graphics (RADV GFX1150)".to_string()],
                    drivers: vec!["radv 25.2.98".to_string()],
                },
                run_b: RunInfo {
                    comment: "run b comment".to_string(),
                    dxvk_versions: vec!["2.7.1".to_string()],
                    devices: vec!["AMD Radeon 890M Graphics (RADV GFX1150)".to_string()],
                    drivers: vec!["radv 25.2.99".to_string()],
                },
            },
            diffs: vec![Diff {
                trace: "trace name.trace".to_string(),
                frame: "snapshot-0001".to_string(),
                run_a: Run {
                    image: "filename".to_string(),
                    image_preview: "preview".to_string(),
                    comment: "run a comment".to_string(),
                    execution_time: 0.1,
                    log: "log a".to_string(),
                },
                run_b: Run {
                    image: "filename".to_string(),
                    image_preview: "preview".to_string(),
                    comment: "run b comment".to_string(),
                    execution_time: 0.5,
                    log: "log b".to_string(),
                },
                diff: 0.1,
                diff_image: "diff.png".to_owned(),
                diff_image_preview: "diff_preview.png".to_owned(),
            }],
        };
        render(&render_context).unwrap();
    }

    #[test]
    fn diff_percentage_filter() {
        assert_eq!(super::diff_percentage_filter(".1".to_string()), "10.00%");
    }

    #[test]
    fn diff_fraction_noop() {
        let mut a = image::RgbaImage::new(2, 2);
        let mut b = image::RgbaImage::new(2, 2);

        for y in 0..2 {
            for x in 0..2 {
                a.put_pixel(x, y, image::Rgba([x as u8, y as u8, 0, 1]));
                b.put_pixel(x, y, image::Rgba([x as u8, y as u8, 0, 1]));
            }
        }

        assert_eq!(
            image_diff_fraction(
                &image::DynamicImage::ImageRgba8(a),
                &image::DynamicImage::ImageRgba8(b)
            ),
            0.0
        );
    }

    #[test]
    fn diff_fraction_mismatch() {
        let mut a = image::RgbaImage::new(2, 2);
        let mut b = image::RgbaImage::new(2, 1);

        for y in 0..2 {
            for x in 0..2 {
                a.put_pixel(x, y, image::Rgba([x as u8, y as u8, 0, 1]));
            }
        }

        // One matching pixel, one mismatching.
        b.put_pixel(0, 0, image::Rgba([0, 0, 0, 1]));
        b.put_pixel(1, 0, image::Rgba([0, 0, 0, 1]));

        assert_eq!(
            image_diff_fraction(
                &image::DynamicImage::ImageRgba8(a),
                &image::DynamicImage::ImageRgba8(b)
            ),
            0.75
        );
    }

    #[test]
    fn test_suffix_file() {
        use std::path::PathBuf;

        let path = PathBuf::from("/path/to/file.png");
        let result = suffix_file(&path, "-preview").unwrap();
        assert_eq!(result, PathBuf::from("/path/to/file-preview.png"));

        let path = PathBuf::from("relative/snapshot.jpg");
        let result = suffix_file(&path, "-diff").unwrap();
        assert_eq!(result, PathBuf::from("relative/snapshot-diff.jpg"));

        // Test with multiple dots
        let path = PathBuf::from("test/my.file.name.png");
        let result = suffix_file(&path, "-test").unwrap();
        assert_eq!(result, PathBuf::from("test/my-test.file.name.png"));
    }

    #[test]
    fn test_html_filename() {
        use std::path::PathBuf;

        let output_path = PathBuf::from("/home/user/output");

        let file = PathBuf::from("/home/user/output/trace/snapshot.png");
        let result = html_filename(&file, &output_path).unwrap();
        assert_eq!(result, "trace/snapshot.png");

        let file = PathBuf::from("/home/user/output/a/b/c/file.png");
        let result = html_filename(&file, &output_path).unwrap();
        assert_eq!(result, "a/b/c/file.png");

        // Test HTML escaping
        let file = PathBuf::from("/home/user/output/file <  \" >.png");
        let result = html_filename(&file, &output_path).unwrap();
        assert_eq!(result, "file &lt;  &quot; &gt;.png");
    }

    #[test]
    fn snapshot_log_info() {
        let log = include_str!("test_data/wine-dxvk-log.txt");
        let info = LogInfo::from(log);

        assert_eq!(info.dxvk, Some("2.7.1".to_string()));
        assert_eq!(
            info.device,
            Some("AMD Radeon 890M Graphics (RADV GFX1150)".to_string())
        );
        assert_eq!(info.driver, Some("radv 25.2.99".to_string()));
    }
}