gpu-trace-perf 1.8.2

use std::collections::{BTreeMap, HashMap};

use crate::{
    ShaderAnalyze,
    shader_parser::{Shader, find_shaders},
};
use anyhow::{Context, Result};
use statrs::statistics::Statistics;
use walkdir::WalkDir;

fn parse_time(line: &str) -> Result<(String, u64)> {
    let (sha, time) = line
        .split_once(": ")
        .with_context(|| format!("finding delimiter in {line}"))?;
    let time = str::parse::<u64>(time).with_context(|| format!("parsing u64 in {time}"))?;
    Ok((sha.to_owned(), time))
}

type ShaderMap = HashMap<String, Shader>;
type DrawTimeMap = BTreeMap<(String, String), (u64, u64)>;

struct ShaderImprovement {
    trace: String,
    sha1: String,
    relative: TimeImprovement,
}

struct HeuristicResult {
    per_trace_improvement: Vec<f64>,

    // Shaders with the worst relative delta.
    hurt: Vec<ShaderImprovement>,
}

/// Returns the per-trace performance change of using this heuristic instead of
/// always choosing the baseline
fn try_heuristic<H>(
    shaders: &ShaderMap,
    traces: &Vec<(String, DrawTimeMap)>,
    heuristic: H,
) -> HeuristicResult
where
    H: Fn(&Shader, &Shader) -> bool,
{
    let mut per_trace_improvement = Vec::new();
    let mut hurt = Vec::new();
    for (trace, draw_times) in traces {
        let mut time = 0;

        for ((a_sha, b_sha), times) in draw_times {
            let a = &shaders[a_sha];
            let b = &shaders[b_sha];
            let decision = if heuristic(a, b) { times.1 } else { times.0 };
            time += decision;

            // Also return a list of the shaders that got hit the worst by this heuristic.
            if decision > times.0 {
                hurt.push(ShaderImprovement {
                    trace: trace.clone(),
                    sha1: a_sha.clone(),
                    relative: TimeImprovement::new(times.0, decision),
                });
            }
        }
        per_trace_improvement.push(time as f64);
    }

    hurt.sort_by_key(|x| x.relative);

    HeuristicResult {
        per_trace_improvement,
        hurt,
    }
}

/// Given a heuristic, return a list of the improvements still possible if the
/// heuristic was optimal.
fn opportunities<H>(
    shaders: &ShaderMap,
    traces: &Vec<(String, DrawTimeMap)>,
    heuristic: H,
) -> Vec<ShaderImprovement>
where
    H: Fn(&Shader, &Shader) -> bool,
{
    let mut opportunities = Vec::new();
    for (trace, draw_times) in traces {
        for ((a_sha, b_sha), times) in draw_times {
            let a = &shaders[a_sha];
            let b = &shaders[b_sha];
            let decision = if heuristic(a, b) { times.1 } else { times.0 };

            if decision > times.0 {
                opportunities.push(ShaderImprovement {
                    trace: trace.clone(),
                    sha1: a_sha.clone(),
                    relative: TimeImprovement::new(decision, times.0),
                });
            }
        }
    }

    opportunities.sort_by_key(|x| x.relative);

    opportunities
}

fn optimal_times(traces: &Vec<(String, DrawTimeMap)>) -> Vec<f64> {
    let mut result = Vec::new();
    for (_, draw_times) in traces {
        let time: u64 = draw_times.iter().map(|(_, (a, b))| a.min(b)).sum();
        result.push(time as f64);
    }

    result
}

fn comparison(a: &[f64], b: &[f64]) -> Vec<f64> {
    a.iter()
        .zip(b)
        .map(|(before, after)| {
            // We >1 to mean better, but we started with times so invert it
            before / after
        })
        .collect()
}

fn dump_comparisons(names: &[String], times: &[(&str, HeuristicResult)]) {
    let baseline = &times[0].1;

    let namelen = names.iter().map(|x| x.len()).max().unwrap_or(0);

    for (heuristic, result) in &times[1..] {
        if !result.hurt.is_empty() {
            println!("Hurt shaders in {heuristic}:");
            for hurt in result.hurt.iter().take(5) {
                println!("  {} ({}): {}", hurt.sha1, hurt.trace, hurt.relative);
            }
        }
    }

    let comparisons: Vec<_> = times[1..]
        .iter()
        .map(|(heuristic, result)| {
            (
                heuristic,
                comparison(
                    &baseline.per_trace_improvement,
                    &result.per_trace_improvement,
                ),
            )
        })
        .collect();

    for _ in 0..namelen {
        print!(" ");
    }
    print!("  ");
    for (name, _) in &comparisons {
        print!("{name}  ");
    }
    println!();

    for (i, name) in names.iter().enumerate() {
        print!("{name:namelen$}: ");

        for (name, rel) in &comparisons {
            let field_len: usize = name.len();
            print!("{:field_len$.1}  ", rel[i] * 100.0 - 100.0);
        }

        println!();
    }

    print!("{:namelen$}: ", "AVERAGE");
    for (name, rel) in &comparisons {
        let field_len: usize = name.len() - 1;
        print!("{:field_len$.2}%  ", rel.geometric_mean() * 100.0 - 100.0);
    }
    println!();
}

/// Relative time comparison result, displayed as percentage improvement.
#[derive(Copy, Clone)]
struct TimeImprovement {
    improvement: f64,
}

impl TimeImprovement {
    fn new(a: u64, b: u64) -> TimeImprovement {
        TimeImprovement {
            improvement: (a as f64) / (b as f64),
        }
    }
}

impl std::fmt::Display for TimeImprovement {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{:.1}%", self.improvement * 100.0 - 100.0)
    }
}

impl PartialOrd for TimeImprovement {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for TimeImprovement {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.improvement.total_cmp(&other.improvement)
    }
}

impl PartialEq for TimeImprovement {
    fn eq(&self, other: &Self) -> bool {
        if self.improvement.is_nan() && other.improvement.is_nan() {
            true
        } else {
            self.improvement == other.improvement
        }
    }
}

impl Eq for TimeImprovement {}

/// Creates the map of shader hash to parsed shader info using the dumped
/// shaders from a --capture-shaders run.
fn collect_shaders(path: &str) -> Result<HashMap<String, Shader>> {
    let mut path = path.to_owned();
    path.push_str("/shaders");
    let walk = WalkDir::new(path).follow_links(true);

    let mut shaders = HashMap::new();
    for entry in walk
        .into_iter()
        .filter_map(|e| e.ok())
        .filter(|e| e.file_type().is_file())
    {
        let path = entry.path();
        let shader =
            std::fs::read_to_string(path).with_context(|| format!("reading {}", path.display()))?;
        let parsed = find_shaders(&shader)?;
        assert_eq!(parsed.len(), 1);
        shaders.insert(parsed[0].sha1.clone(), parsed[0].clone());
    }

    Ok(shaders)
}

/// Creates the map of before-and-after sha1s and draw times per trace, from the
/// output of a --capture-shaders run.
///
/// This assumes that the same draws happen between the before and after cases.
/// We throw out a trace from analysis if the number of draws we recorded didn't
/// match.
fn collect_trace_times(path: &str) -> Result<Vec<(String, DrawTimeMap)>> {
    let walk = WalkDir::new(path).follow_links(true);

    let mut trace_times = Vec::new();
    for entry in walk.into_iter().filter_map(|e| e.ok()) {
        let path = entry.path();

        if path.file_name().context("getting file name")? != "a-draws.txt" {
            continue;
        }

        let trace = path.parent().unwrap().file_name().unwrap();
        let name = String::from(trace.to_string_lossy());

        let b_path = path.parent().context("finding parent")?.join("b-draws.txt");
        let a =
            std::fs::read_to_string(path).with_context(|| format!("reading {}", path.display()))?;
        let b = std::fs::read_to_string(&b_path)
            .with_context(|| format!("reading {}", b_path.display()))?;

        if a.lines().count() != b.lines().count() {
            println!("Trace {name} has mismatched draw counts, skipping");
            continue;
        }

        let mut times = BTreeMap::new();
        for (a, b) in a.lines().zip(b.lines()) {
            let (a_sha, a_time) = parse_time(a)?;
            let (b_sha, b_time) = parse_time(b)?;

            if a_sha == b_sha {
                continue;
            }

            if a_sha == "bde117c82f7b16c0cb9431752df65544cfc5a617" {
                /* This is an internal renderdoc shader we don't care about, that would skew results with how many
                 * apps captured with it.
                 */
                continue;
            }

            let e = times.entry((a_sha, b_sha)).or_insert((0u64, 0u64));
            e.0 += a_time;
            e.1 += b_time;
        }

        if times.is_empty() {
            println!("no varying shaders found for {name}");
        } else {
            trace_times.push((name, times));
        }
    }

    // Trim off a common path prefix of all the traces to make the output table
    // more legible.
    let prefix_len = if trace_times.len() >= 2 {
        let mut prefix_len = trace_times[0].0.len();
        for (name, _) in &trace_times[1..] {
            prefix_len = prefix_len.min(name.len());
            while trace_times[0].0[0..prefix_len] != name[0..prefix_len] {
                prefix_len -= 1;
            }
            // Continue walking to the next path separator ('/' had been
            // replaced by '-' in path sanitizing.)
            while prefix_len > 0 && name.chars().nth(prefix_len - 1) != Some('-') {
                prefix_len -= 1;
            }
        }
        prefix_len
    } else {
        0
    };

    if prefix_len != 0 {
        for (name, _) in trace_times.iter_mut() {
            *name = (name[prefix_len..]).to_owned();
        }
    }

    Ok(trace_times)
}

/// Entrypoint for gpu-trace-perf shader-analyze
pub fn shader_analyze(opts: &ShaderAnalyze) -> Result<()> {
    let shaders = collect_shaders(&opts.output)?;

    println!("Found {} shaders captured", shaders.len());

    let mut trace_times = collect_trace_times(&opts.output)?;

    if trace_times.is_empty() {
        println!("No traces to analyze");
        return Ok(());
    }

    trace_times.sort();

    let pathlen = trace_times.iter().map(|x| x.0.len()).max().unwrap_or(0);

    for (trace, times) in &trace_times {
        let mut times: Vec<_> = times
            .iter()
            .map(|(shas, times)| (shas, times, TimeImprovement::new(times.0, times.1)))
            .collect();

        times.sort_by_key(|(_, _, x)| *x);

        let no_stage = "???".to_string();
        for ((a_sha, b_sha), (au, bu), improvement) in times {
            let stage = shaders.get(a_sha).map(|x| &x.stage).unwrap_or(&no_stage);
            println!("{trace:pathlen$}: {au} -> {bu} ({improvement}) {stage} ({a_sha} -> {b_sha})",);
        }
    }
    println!();

    let names: Vec<_> = trace_times.iter().map(|x| x.0.clone()).collect();

    let new_heuristic = |a: &Shader, _: &Shader| {
        ['0', '2', '4', '8', 'a', 'c', 'e'].contains(&a.sha1.chars().next().unwrap())
    };
    let results = vec![
        // Result if we just chose the shader from the run's "before" environment every time.
        (
            "baseline",
            try_heuristic(&shaders, &trace_times, |_, _| false),
        ),
        // Result if we just chose the shader from the run's "after" environment every time.
        (
            "b-always",
            try_heuristic(&shaders, &trace_times, |_, _| true),
        ),
        // Result if we could come up with an optimal choice for every shader.
        (
            "optimal",
            HeuristicResult {
                per_trace_improvement: optimal_times(&trace_times),
                hurt: Vec::new(),
            },
        ),
        // Here's where you can hack in various heuristic choices to evaluate other
        // choices between 'a' and 'b', one dummy example provided.
        (
            "fiftyfifty",
            try_heuristic(&shaders, &trace_times, new_heuristic),
        ),
    ];

    dump_comparisons(&names, &results);

    println!();
    println!("Opportunity left in new heuristic:");
    for improvement in opportunities(&shaders, &trace_times, new_heuristic)
        .iter()
        .rev()
        .take(10)
    {
        println!(
            "  {} ({}) {}",
            improvement.sha1, improvement.trace, improvement.relative
        );
    }

    Ok(())
}

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

    #[test]
    fn test_time_improvement() {
        assert_eq!(TimeImprovement::new(4, 3).to_string(), "33.3%".to_string());
        assert_eq!(TimeImprovement::new(3, 4).to_string(), "-25.0%".to_string());
    }
}