gpu-trace-perf 1.6.0

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

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

use crate::collect_traces;

#[derive(Debug, Serialize, Deserialize)]
pub struct SnapshotResult {
    /// Paths relative to the output directory of the frames captured.
    pub files: Vec<PathBuf>,
    pub cmdline: String,
    pub stdout: String,
    pub stderr: String,
    pub runtime: Duration,
}

impl SnapshotResult {
    pub fn log(&self) -> String {
        format!(
            "cmdline: {}\nstdout:\n{}stderr:\n{}",
            self.cmdline, self.stdout, self.stderr
        )
    }

    pub fn write_to(&self, filename: &Path) -> Result<()> {
        std::fs::write(
            filename,
            serde_json::to_string_pretty(self)
                .context("serializing snapshot result")?
                .as_bytes(),
        )
        .with_context(|| format!("Writing snapshot result JSON to {})", filename.display()))
    }

    pub fn import_baseline(baseline_path: &Path, output_path: &Path) -> Result<SnapshotResult> {
        let mut result: SnapshotResult = {
            let result_path = baseline_path.join("result.json");
            serde_json::from_reader(BufReader::new(
                std::fs::File::open(&result_path).with_context(|| {
                    format!("opening baseline result path {}", result_path.display())
                })?,
            ))
            .context("parsing baseline snapshot result JSON")?
        };

        // Copy baseline pngs into output directory and update their names in the result we save.
        for file in &mut result.files {
            let from = baseline_path.join(&file);
            let to = output_path.join(suffix_file(file, "-baseline")?);

            std::fs::copy(&from, &to).with_context(|| {
                format!(
                    "copying baseline png from {} to {}",
                    from.display(),
                    to.display()
                )
            })?;
            *file = to;
        }

        let result_path = output_path.join("result-baseline.json");
        result.write_to(&result_path).with_context(|| {
            format!(
                "writing baseline snapshot result to {}",
                result_path.display()
            )
        })?;

        Ok(result)
    }
}

#[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,

    /// Number of times to loop over the last frame (default 5 - normal swap plus 4 loops)
    #[arg(short, long = "loop", default_value = "5", value_parser = clap::value_parser!(u32).range(1..))]
    loops: u32,

    #[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,
}

impl Run {
    fn new(
        output_path: &Path,
        image: &Path,
        preview: &Path,
        log: &Path,
        result: &SnapshotResult,
    ) -> Self {
        Run {
            image: html_filename(image, output_path).unwrap(),
            image_preview: html_filename(preview, output_path).unwrap(),
            comment: String::new(),
            execution_time: result.runtime.as_secs_f32(),
            log: html_filename(log, output_path).unwrap(),
        }
    }
}

#[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>,
}

impl RunInfo {
    fn new(comment: String, versions: Versions) -> Self {
        // Convert HashSets to sorted Vecs
        let mut dxvk_versions: Vec<_> = versions.dxvk.into_iter().collect();
        dxvk_versions.sort();
        let mut devices: Vec<_> = versions.devices.into_iter().collect();
        devices.sort();
        let mut drivers: Vec<_> = versions.drivers.into_iter().collect();
        drivers.sort();

        RunInfo {
            comment,
            dxvk_versions,
            devices,
            drivers,
        }
    }
}

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

impl<T: AsRef<str>> From<T> for LogInfo {
    fn from(value: T) -> 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.as_ref().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
            .as_ref()
            .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(Default)]
struct Versions {
    dxvk: HashSet<String>,
    devices: HashSet<String>,
    drivers: HashSet<String>,
}

impl Versions {
    fn new() -> Self {
        Versions::default()
    }

    fn add(&mut self, info: LogInfo) {
        if let Some(dxvk) = info.dxvk {
            self.dxvk.insert(dxvk);
        }
        if let Some(device) = info.device {
            self.devices.insert(device);
        }
        if let Some(driver) = info.driver {
            self.drivers.insert(driver);
        }
    }
}

#[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::Triangle);

    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 summarize_errors(diffs: &[Diff]) -> bool {
    let mut printed_errors = 0;

    for diff in diffs {
        if diff.diff == 0.0 {
            break;
        }

        if printed_errors == 0 {
            warn!("Rendering differences in:");
        }

        if printed_errors == 5 {
            warn!("  ... and more");
            break;
        }

        warn!("  {} - {}", diff.trace, diff.frame);

        printed_errors += 1;
    }

    if printed_errors == 0 {
        info!("No rendering differences found.");
        false
    } else {
        true
    }
}

/// 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, opts.loops);

            if let Ok(result) = &result {
                let json_filename = trace
                    .output_dir(&opts.output)
                    .unwrap()
                    .unwrap()
                    .join("result.json");
                result
                    .write_to(&json_filename)
                    .expect("writing snapshot result JSON");
            };

            (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 images");

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

    let mut diffs = Vec::new();

    // Collect all unique DXVK versions, devices, and drivers for run A and B
    let mut versions_a = Versions::new();
    let mut versions_b = Versions::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");
        std::fs::write(&log_path, result.log()).context("writing log file")?;

        // Extract driver info from run B (current run)
        versions_b.add(LogInfo::from(&result.stdout));

        // Try loading the baseline from the previous run (reporting an error
        // but continuing if we fail).
        let baseline_result =
            trace_baseline_path.and_then(
                |trace_baseline_path| match SnapshotResult::import_baseline(
                    &trace_baseline_path,
                    &trace_output_path,
                ) {
                    Ok(baseline_result) => Some(baseline_result),
                    Err(e) => {
                        error!("Failed to load results.json for {}: {}", trace.name(), e);
                        None
                    }
                },
            );

        // Copy the baseline result file into our output directory and extract info from run A
        if let Some(baseline_result) = baseline_result {
            // Extract driver info from run A (baseline)
            versions_a.add(LogInfo::from(baseline_result.stdout.as_str()));

            let baseline_log = trace_output_path.join("log-baseline.txt");
            std::fs::write(&baseline_log, baseline_result.log())?;

            let mut save_passing_frames = true;
            for (file, baseline_file) in result.files.iter().zip(baseline_result.files.iter()) {
                debug!("Preparing diff for {}/{}", trace.name(), file.display());

                let snapshot_png = trace_output_path.join(file);
                let baseline_png = Path::new(baseline_file);

                let snapshot_image = image::open(&snapshot_png)
                    .with_context(|| format!("reading {}", snapshot_png.display()))?;

                let baseline_image = image::open(baseline_png)
                    .with_context(|| format!("reading {}", baseline_png.display()))?;

                let diff = image_diff_fraction(&baseline_image, &snapshot_image);

                if diff == 0.0 {
                    if !save_passing_frames {
                        std::fs::remove_file(&snapshot_png).unwrap_or_else(|e| {
                            error!(
                                "Failed to remove {} for passing frame: {e}",
                                snapshot_png.display()
                            )
                        });
                        std::fs::remove_file(baseline_png).unwrap_or_else(|e| {
                            error!(
                                "Failed to remove {} for passing frame: {e}",
                                baseline_png.display()
                            )
                        });
                        continue;
                    } else {
                        // Once we've recorded a passing frame for the trace,
                        // skip the other loops if they also pass.
                        save_passing_frames = false;
                    }
                }

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

                let baseline_preview_png =
                    generate_preview(&baseline_image, baseline_png).context("baseline preview")?;

                let preview_png =
                    generate_preview(&snapshot_image, &snapshot_png).context("baseline preview")?;

                let diff_preview_png =
                    generate_preview(&diff_image, &diff_png).context("diff preview")?;

                diffs.push(Diff {
                    trace: trace.name().to_string(),
                    frame: file.to_string_lossy().to_string(),
                    run_a: Run::new(
                        output_path,
                        baseline_png,
                        &baseline_preview_png,
                        &baseline_log,
                        result,
                    ),
                    run_b: Run::new(output_path, &snapshot_png, &preview_png, &log_path, result),
                    diff,
                    diff_image: html_filename(&diff_png, output_path)?,
                    diff_image_preview: html_filename(&diff_preview_png, output_path)?,
                });
            }
        } else {
            for file in &result.files {
                debug!("Preparing snapshot of {}/{}", trace.name(), file.display());

                let snapshot_png = trace_output_path.join(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")?;

                diffs.push(Diff {
                    trace: trace.name().to_string(),
                    frame: file.to_string_lossy().to_string(),
                    run_a: Run::default(),
                    run_b: Run::new(output_path, &snapshot_png, &preview_png, &log_path, result),
                    diff: 0.0,
                    diff_image: "".to_string(),
                    diff_image_preview: "".to_string(),
                });
            }
        };
    }

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

    diffs.sort_by(|a, b| b.diff.total_cmp(&a.diff).then(a.trace.cmp(&b.trace)));

    let render_context = RenderContext {
        comparison: Comparison {
            max_diff: diffs.iter().map(|d| d.diff).fold(0.0, |a, b| a.max(b)),
            run_a: RunInfo::new(format!("output dir: {}", opts.baseline), versions_a),
            run_b: RunInfo::new(format!("output dir: {}", opts.output), versions_b),
        },
        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")?;

    if summarize_errors(&render_context.diffs) {
        anyhow::bail!("Rendering differences found");
    }

    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()));
    }
}