gpu-trace-perf 1.4.0

#[macro_use]
extern crate lazy_static;
extern crate assert_approx_eq;
extern crate statrs;
extern crate stderrlog;
extern crate walkdir;

mod angle;
mod apitrace;
mod gfxreconstruct;
mod log_error;
mod renderdoc;
mod shader_analyze;
mod shader_parser;
mod shader_parser_ir3;
mod stats;
mod u_trace;

use anyhow::{Context, Result, bail};
use apitrace::*;
use clap::{Args, Parser, Subcommand};
use gfxreconstruct::*;
use log::error;
use log_error::LogError;
use rand::seq::IndexedRandom;
use renderdoc::*;
use stats::{BootstrappedRelativeAndMaxChange, ResultStats};
use std::fs::{OpenOptions, create_dir_all};
use std::io::{self, prelude::*};
use std::path::{Path, PathBuf};
use std::process::{Command, exit};
use u_trace::UTraceParser;
use walkdir::WalkDir;

use crate::angle::AngleTrace;
use crate::shader_parser::find_shaders;
use crate::u_trace::{UTRACE_SHADER_DRAW_EVENTS, UTRACE_SHADER_STAGES};

#[derive(Debug, Parser)]
#[clap(
    version = "1.4.0",
    author = "Emma Anholt <emma@anholt.net>",
    about = "Plays a collection of GPU traces under different wrappers to evaluate driver changes on performance"
)]
struct Cli {
    /// A level of verbosity, and can be used multiple times
    #[arg(short, long, action = clap::ArgAction::Count)]
    verbose: u8,

    #[command(subcommand)]
    subcmd: SubCommand,
}

#[derive(Default)]
struct ReplayOutput {
    pub stdout: String,
    pub stderr: String,
}

impl From<std::process::Output> for ReplayOutput {
    fn from(value: std::process::Output) -> Self {
        ReplayOutput {
            stdout: String::from_utf8_lossy(&value.stdout).to_string(),
            stderr: String::from_utf8_lossy(&value.stderr).to_string(),
        }
    }
}

fn create_dir(path: PathBuf) -> Result<PathBuf> {
    match create_dir_all(&path) {
        Ok(()) => {}

        Err(err) => {
            if err.kind() != io::ErrorKind::AlreadyExists {
                bail!("failed to create {}: {:?}", path.display(), err);
            }
        }
    }

    Ok(path)
}

trait Replay {
    fn replay(&self, wrapper: Option<&str>, envs: &[(String, String)]) -> Result<ReplayOutput>;

    fn fps(&self, output: &ReplayOutput) -> Result<f64>;

    fn name(&self) -> &str;

    fn capture_draw_times(
        &self,
        output_dir: &str,
        run_name: &str,
        utrace: &mut UTraceParser,
    ) -> Result<()> {
        let path = self
            .output_dir(output_dir)
            .unwrap()
            .unwrap()
            .join(format!("{}-draws.txt", run_name));

        let mut file = std::fs::OpenOptions::new()
            .create(true)
            .append(true)
            .open(&path)
            .with_context(|| format!("appending to {}", path.display()))?;

        for frame in utrace.results()? {
            for events in UTRACE_SHADER_DRAW_EVENTS
                .iter()
                .flat_map(|event| frame.event_times(event))
            {
                for (time, params) in events.times {
                    for param in UTRACE_SHADER_STAGES {
                        if let Some(param) = params.get(*param) {
                            writeln!(file, "{}: {}", param, time)?;
                        }
                    }
                }
            }
        }

        Ok(())
    }

    fn replay_get_fps(
        &self,
        wrapper: Option<&str>,
        output_dir: &str,
        u_trace_event: &str,
        utrace_use_csv: bool,
        capture_shaders_path: &str,
    ) -> Result<f64> {
        let u_trace_event = if u_trace_event.is_empty() && !capture_shaders_path.is_empty() {
            /* Make sure utrace is active if we're capturing draw times. */
            UTRACE_SHADER_DRAW_EVENTS.join(",")
        } else {
            u_trace_event.to_owned()
        };
        let mut utrace = UTraceParser::new(&u_trace_event, utrace_use_csv);
        let envs = utrace.env();

        let output = self.replay(wrapper, &envs)?;
        if utrace.active() {
            if !capture_shaders_path.is_empty() {
                self.capture_draw_times(output_dir, capture_shaders_path, &mut utrace)?;
            }
            utrace.middle_frame_fps()
        } else {
            self.fps(&output)
        }
    }

    /// Runs the trace under the `a` and `b` wrappers, returning the fps.
    fn collect_a_b_fps(&self, run_opts: &Run, a_first: bool) -> Result<(f64, f64)> {
        println!("Running {}", self.name());

        let do_run = |name, wrapper| {
            self.replay_get_fps(
                Some(wrapper),
                &run_opts.output,
                &run_opts.utrace,
                run_opts.utrace_use_csv,
                if run_opts.capture_shaders { name } else { "" },
            )
        };

        let (a, b) = if a_first {
            let a = do_run("a", &run_opts.a);
            let b = do_run("b", &run_opts.b);
            (a, b)
        } else {
            let b = do_run("b", &run_opts.b);
            let a = do_run("a", &run_opts.a);
            (a, b)
        };

        let a = a.context("getting 'a' fps")?;
        let b = b.context("getting 'b' fps")?;

        // Only push the pair now that we got valid results from both.  (our
        // stats are expecting a.len() == b.len()
        if let Some(output_dir) = self
            .output_dir(&run_opts.output)
            .expect("creating output directory")
        {
            append_fps(&output_dir, "a", a).log_error();
            append_fps(&output_dir, "b", b).log_error();
        }

        Ok((a, b))
    }

    fn check_debug_filter(
        &self,
        a: &str,
        b: &str,
        envs: &FilterEnv,
        output_dir: &str,
        capture_shaders: bool,
    ) -> Result<bool> {
        let mut envs: Vec<_> = envs.to_vec();

        if capture_shaders {
            envs.push((
                "IR3_SHADER_DEBUG".to_owned(),
                "vs,fs,gs,tes,tcs,cs".to_owned(),
            ));
            envs.push(("MESA_SHADER_CACHE_DISABLE".to_owned(), "1".to_owned()));
            // Serialize zink so that the dumped shaders don't stomp each other.
            envs.push(("ZINK_DEBUG".to_owned(), "nobgc".to_owned()));
        };

        let runs = [
            ("a", self.replay(Some(a), &envs)?),
            ("b", self.replay(Some(b), &envs)?),
        ];

        if !output_dir.is_empty() {
            for (name, output) in &runs {
                let _ = self.write_output(
                    output_dir,
                    &format!("stdout-{}.txt", name),
                    output.stdout.as_bytes(),
                );
                let _ = self.write_output(
                    output_dir,
                    &format!("stderr-{}.txt", name),
                    output.stderr.as_bytes(),
                );

                if capture_shaders {
                    let shaders_dir = create_dir(Path::new(output_dir).join("shaders"))?;
                    for shader in find_shaders(&output.stderr)? {
                        std::fs::write(shaders_dir.join(&shader.sha1), shader.code.as_bytes())
                            .with_context(|| format!("writing {} shader", shader.sha1))?;
                    }
                }
            }
        }
        Ok(runs[0].1.stderr != runs[1].1.stderr)
    }

    fn output_dir(&self, output_dir: &str) -> Result<Option<PathBuf>> {
        if output_dir.is_empty() {
            return Ok(None);
        }

        create_dir(
            Path::new(output_dir).join(self.name().trim_start_matches('/').replace('/', "-")),
        )
        .map(Some)
    }

    fn write_output(&self, output_dir: &str, filename: &str, data: &[u8]) -> Result<()> {
        if let Some(output_dir) = self
            .output_dir(output_dir)
            .expect("creating output directory")
        {
            let output_path = output_dir.join(filename);

            let mut file = std::fs::File::create(&output_path)
                .with_context(|| format!("creating {}", output_path.display()))?;

            file.write_all(data).context("failed to write")?;
        }

        Ok(())
    }
}

pub fn replay_command<S>(args: &[S], wrapper: Option<&str>, envs: &[(String, String)]) -> Command
where
    S: AsRef<str>,
{
    let mut command = if let Some(wrapper) = wrapper {
        let mut command = Command::new(wrapper);
        command.arg(args[0].as_ref());
        command
    } else {
        Command::new(args[0].as_ref())
    };
    for arg in &args[1..] {
        command.arg(arg.as_ref());
    }

    command.env("NIR_VALIDATE", "0");

    for env in envs {
        command.env(&env.0, &env.1);
    }

    command
}

type FilterEnv<'a> = Vec<(String, String)>;

#[derive(Debug, Subcommand)]
enum SubCommand {
    Run(Run),
    ShaderAnalyze(ShaderAnalyze),
}

#[derive(Debug, Args)]
/// Runs a collection of traces, checking for performance differences between two wrapper scripts
struct Run {
    #[arg(short, long, number_of_values = 1)]
    /// paths to directories containing traces (or traces themselves)
    traces: Vec<String>,

    #[arg(long, number_of_values = 1)]
    /// if present, only traces that differ in stderr when these env vars are
    /// set will be tested
    debug_filter: Vec<String>,

    #[arg(long, default_value = "")]
    /// The name of a Mesa utrace event to parse the times of rather than frame
    /// execution wall time.
    utrace: String,

    #[arg(long, default_value_t = false)]
    /// Use utrace's CSV output instead of the default JSON. May be more reliable for apps that
    /// don't properly close their GL context.
    utrace_use_csv: bool,

    #[arg(long = "output", default_value = "")]
    /// if present, path to a directory to store debug output and logs
    output: String,

    /// If set, capture per-draw, per-shader timings (where a whole draw's time
    /// is attributed to each shader stage) from utrace events and save them in
    /// the output dir.
    #[arg(long)]
    capture_shaders: bool,

    #[arg(index = 1)]
    /// command wrapper to set the runtime environment for the 'a' case
    a: String,

    #[arg(index = 2)]
    /// command wrapper to set the runtime environment for the 'b' case
    b: String,
}

#[derive(Debug, Args)]
/// Runs a collection of traces, checking for performance differences between two wrapper scripts
struct Thread64 {
    #[arg(short, long, number_of_values = 1)]
    /// paths to directories containing traces (or traces themselves)
    traces: Vec<String>,

    #[arg(long = "output")]
    /// if present, path to a directory to store debug output and logs
    output: String,
}

#[derive(Debug, Args)]
/// Runs a collection of traces, checking for performance differences between two wrapper scripts
struct ShaderAnalyze {
    /// path to a directory to store debug output and logs
    output: String,
}

// Walk the directories specified in the command line args and make a vec of the actual traces to run
fn collect_traces(paths: Vec<String>, utrace: bool) -> Vec<Box<dyn Replay>> {
    let mut traces: Vec<Box<dyn Replay>> = Vec::new();

    for file in paths {
        let walk = WalkDir::new(&file).follow_links(true);

        for entry in walk.into_iter().filter_map(|e| e.ok()) {
            if let Some(path) = entry.path().to_str().map(|x| x.to_owned()) {
                if path.ends_with("angle_trace_tests") {
                    match AngleTrace::enumerate(&path) {
                        Ok(angles) => {
                            for trace in angles {
                                traces.push(Box::new(trace));
                            }
                        }
                        Err(e) => eprintln!("{}", e),
                    }
                } else if path.ends_with(".trace") {
                    if utrace {
                        traces.push(Box::new(ApitraceUtraceTrace::new(&path)));
                    } else {
                        traces.push(Box::new(ApitraceTrace::new(&path)));
                    }
                } else if path.ends_with(".rdc") {
                    if utrace {
                        traces.push(Box::new(RenderdocUtraceTrace::new(&path)));
                    } else {
                        traces.push(Box::new(RenderdocPyTrace::new(&path)));
                    }
                } else if path.ends_with(".gfxr") {
                    traces.push(Box::new(GfxreconstructTrace::new(&path)));
                };
            } else {
                error!("Path {:?} not UTF-8, skipping", entry.path());
            }
        }

        // Allow specifying a subtest on angle_trace_tests.
        if file.contains("angle_trace_tests/") {
            if let Some((path, test)) = file.rsplit_once('/') {
                traces.push(Box::new(AngleTrace::new(path, test)));
            }
        }
    }

    traces
}

#[derive(Clone)]
struct TraceResults {
    a: Vec<f64>,
    b: Vec<f64>,
    logged: bool,
}

impl TraceResults {
    pub fn new() -> TraceResults {
        TraceResults {
            a: Vec::new(),
            b: Vec::new(),
            logged: false,
        }
    }

    /// Performs a resampling with replacement from the "from" instance.  We
    /// have matching a,b vector lengths as what we're
    /// resampling.
    pub fn resample(&mut self, from: &TraceResults) {
        assert_eq!(self.a.len(), from.a.len());
        assert_eq!(self.b.len(), from.b.len());

        let mut rng = rand::rng();
        for a in &mut self.a {
            *a = *from.a.choose(&mut rng).unwrap();
        }
        for b in &mut self.b {
            *b = *from.b.choose(&mut rng).unwrap();
        }
    }
}

fn append_fps(path: &Path, wrapper: &str, fps: f64) -> Result<()> {
    let path = path.join(format!("fps-{}.txt", wrapper));
    let mut file = OpenOptions::new()
        .append(true)
        .create(true)
        .open(&path)
        .with_context(|| format!(" opening {}", path.display()))?;

    writeln!(file, "{}", fps).with_context(|| format!("writing fps to {}", path.display()))?;

    Ok(())
}

fn run_each_trace(run_opts: &Run, results: &mut [(Box<dyn Replay>, TraceResults)], a_first: bool) {
    for (trace, results) in results.iter_mut() {
        match trace.collect_a_b_fps(run_opts, a_first) {
            Ok((a, b)) => {
                results.a.push(a);
                results.b.push(b);
            }
            Err(e) => {
                if !results.logged {
                    eprintln!("  Failed: {:?}", e);
                    results.logged = true;
                } else {
                    eprintln!("  Failed");
                }
            }
        }
    }
}

fn print_stats(results: &[(Box<dyn Replay>, TraceResults)]) {
    let results: Vec<_> = results
        .iter()
        .filter(|(_, result)| !result.a.is_empty() && !result.b.is_empty())
        .collect();

    let namelen = results
        .iter()
        .map(|(trace, _)| trace.name().len())
        .max()
        .unwrap_or(0);

    // Apply the Bonferroni correction for multiple hypothesis testing.
    let alpha = 0.05 / results.len() as f64;

    let mut stats: Vec<_> = results
        .iter()
        .map(|(trace, results)| {
            (
                trace.name(),
                ResultStats::new(&results.a, &results.b, alpha),
            )
        })
        .collect();

    stats.sort_by(|(_, a), (_, b)| a.change.total_cmp(&b.change));

    for (trace, stat) in &stats {
        let change = if stat.has_fps() {
            format!("{:>7.2}%", stat.change * 100.0)
        } else {
            "no time detected".to_string()
        };

        let error = if stat.n[0] > 1 && stat.has_fps() {
            format!(" (+/- {:5.1}%)", stat.error * 100.0)
        } else {
            "".to_string()
        };

        let (before, after) = if stat.has_fps() {
            (
                format!("{:5.1}", stat.means[0]),
                format!("{:5.1}", stat.means[1]),
            )
        } else {
            ("".to_string(), "".to_string())
        };

        let count = if stat.n[0] == stat.n[1] {
            format!("{}", stat.n[0])
        } else {
            format!("{}/{}", stat.n[0], stat.n[1])
        };

        println!(
            "{path:namelen$}: {before:6} -> {after:6} fps {change}{error} (n={count})",
            path = trace,
            namelen = namelen,
            before = before,
            after = after,
            change = change,
            error = error,
            count = count
        );
    }

    let samples = stats.iter().map(|x| x.1.n[0]).max().unwrap_or(0);
    if samples > 5 {
        let summary_stats = BootstrappedRelativeAndMaxChange::new(
            &results
                .into_iter()
                .map(|x| x.1.clone())
                .collect::<Vec<TraceResults>>(),
            100,
        );
        println!(
            "average fps {:+.1}% (+/- {:.1}%)",
            summary_stats.relative_mean_change * 100.0 - 100.0,
            summary_stats.relative_mean_error * 100.0
        );
        println!(
            "max fps {:+.1}% (+/- {:.1}%)",
            summary_stats.relative_mean_change * 100.0 - 100.0,
            summary_stats.relative_mean_error * 100.0
        );
    }
}

fn run(run_opts: &Run, traces: Vec<Box<dyn Replay>>) {
    let mut results: Vec<_> = traces
        .into_iter()
        .map(|trace| (trace, TraceResults::new()))
        .collect();

    let mut a_first = true;
    loop {
        run_each_trace(run_opts, &mut results, a_first);
        print_stats(&results);

        a_first = !a_first;
    }
}

// Test the wrapper environment scripts people hand us before trying to start up an actual trace
fn test_run_wrapper(wrapper: &str) -> bool {
    let output = Command::new(wrapper).arg("true").output();
    match output {
        Ok(_) => true,
        Err(err) => {
            println!(
                "Failed to spawn test invocation of '{} true': {:?}",
                wrapper, err
            );
            println!(
                "TIP: Exec format error here probably means you need #!/bin/sh in your shell script."
            );
            false
        }
    }
}

fn debug_filter_traces(traces: Vec<Box<dyn Replay>>, run_opts: &Run) -> Vec<Box<dyn Replay>> {
    let mut filters = Vec::new();
    for f in &run_opts.debug_filter {
        if let Some((k, v)) = f.split_once('=') {
            filters.push((k.to_string(), v.to_string()));
        } else {
            error!(
                "Debug filter '{}' should be a KEY=value environment assignment",
                f
            );
            exit(1);
        }
    }
    if filters.is_empty() && !run_opts.capture_shaders {
        return traces;
    }

    println!(
        "Checking for effects of debug filters{}:",
        if run_opts.capture_shaders {
            " and capturing shaders"
        } else {
            ""
        }
    );

    let mut traces = traces;
    traces.retain(|trace| {
        print!("  {}: ", trace.name());
        match trace.check_debug_filter(
            &run_opts.a,
            &run_opts.b,
            &filters,
            &run_opts.output,
            run_opts.capture_shaders,
        ) {
            Ok(true) => {
                println!("affected");
                true
            }
            Ok(false) => {
                println!("skipped");
                false
            }
            Err(e) => {
                println!("failed ({})", e);
                false
            }
        }
    });

    traces
}

fn main() {
    let cli = Cli::parse();
    stderrlog::new()
        .module(module_path!())
        .verbosity(cli.verbose as usize)
        .init()
        .unwrap();

    match &cli.subcmd {
        SubCommand::Run(run_opts) => {
            if !test_run_wrapper(&run_opts.a) || !test_run_wrapper(&run_opts.b) {
                std::process::exit(1);
            }

            let traces = collect_traces(
                run_opts.traces.clone(),
                !run_opts.utrace.is_empty() || run_opts.capture_shaders,
            );
            if traces.is_empty() {
                println!("No traces found in the given directories:");
                for t in &run_opts.traces {
                    println!("  {}", t);
                }
                std::process::exit(1);
            }

            let traces = debug_filter_traces(traces, run_opts);
            if traces.is_empty() {
                println!("All traces filtered out by the debug filter.");
                std::process::exit(0);
            }

            run(run_opts, traces);
        }

        SubCommand::ShaderAnalyze(analyze) => {
            shader_analyze::shader_analyze(analyze).unwrap();
        }
    }
}

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

    struct UnitTestTrace {
        fps: f64,
        name: String,
    }

    impl Replay for UnitTestTrace {
        fn replay(
            &self,
            wrapper: Option<&str>,
            _envs: &[(String, String)],
        ) -> Result<ReplayOutput> {
            Ok(ReplayOutput {
                stdout: wrapper.unwrap_or("bad").to_string(),
                ..Default::default()
            })
        }

        fn fps(&self, output: &ReplayOutput) -> Result<f64> {
            output
                .stdout
                .parse::<f64>()
                .context("parsing unit test wrapper as f64")
                .map(|x| x * self.fps)
        }

        fn name(&self) -> &str {
            &self.name
        }
    }

    impl UnitTestTrace {
        pub fn new(fps: f64) -> UnitTestTrace {
            UnitTestTrace {
                fps,
                name: format!("unit_test_{}", fps),
            }
        }
    }

    #[test]
    fn test_ordering() {
        let mut results: Vec<(Box<dyn Replay>, _)> =
            vec![(Box::new(UnitTestTrace::new(5.0)), TraceResults::new())];

        let run = Run {
            traces: Vec::new(),
            debug_filter: Vec::new(),
            output: String::new(),
            a: "3.0".to_string(),
            b: "4.0".to_string(),
            utrace: "".to_string(),
            utrace_use_csv: false,
            capture_shaders: false,
        };

        run_each_trace(&run, &mut results, true);
        run_each_trace(&run, &mut results, false);

        for (_, results) in results {
            for a in results.a {
                assert_eq!(a, 5.0 * 3.0);
            }
            for b in results.b {
                assert_eq!(b, 5.0 * 4.0);
            }
        }
    }

    #[test]
    fn verify_cli() {
        use clap::CommandFactory;
        Cli::command().debug_assert();
    }
}