gpu-trace-perf 1.9.0

use anyhow::{Context as _, Result, bail};
use gpu_trace_perf::rdc_file::{
    d3d11_swapchain_size_from_chunk, parse_rdc_driver, parse_rdc_frame_capture_chunks,
};
use log::{error, info};
use regex::Regex;
use std::{
    ffi::OsStr,
    fs::File,
    io::{BufReader, prelude::*},
    path::{Path, PathBuf},
    time::Duration,
};

use crate::{ReplayOutput, TraceTool, replay_command, snapshot::SnapshotResult};

pub struct RenderdocPyTrace {
    file: PathBuf,
    name: String,
    is_directx: bool,
    resolution: Option<(u32, u32)>,
}

fn get_directx_resolution(file: &Path) -> Option<(u32, u32)> {
    let f = match File::open(file) {
        Ok(f) => BufReader::new(f),
        Err(e) => {
            error!(
                "Failed to open {} for resolution parsing: {e}",
                file.display()
            );
            return None;
        }
    };

    match parse_rdc_frame_capture_chunks(f) {
        Ok(chunks) => chunks
            .filter_map(|(id, data)| d3d11_swapchain_size_from_chunk(id, &data))
            .next(),

        Err(e) => {
            error!("Failed to parse {} for swapchain size: {e}", file.display());
            None
        }
    }
}

impl RenderdocPyTrace {
    pub fn new(root: &Path, file: &Path) -> RenderdocPyTrace {
        let is_directx = parse_rdc_driver(file)
            .map(|d| d.is_directx())
            .unwrap_or_else(|_| guess_filename_directx(file));

        let resolution = if is_directx {
            // renderdoccmd replay defaults to 1280x720, and doesn't resize
            // according to the actual swapchain.   Given that we use the
            // screenshot layer on the swapchain to get our images (instead of
            // querying renderdoc through python), we need to tell it the acutal
            // resolution to use on the
            get_directx_resolution(file)
        } else {
            None
        };

        RenderdocPyTrace {
            file: file.to_owned(),
            name: crate::relative_test_name(root, file),
            resolution,
            is_directx,
        }
    }
}

static RENDERDOC_WRAPPER_SCRIPT: &[u8] = include_bytes!("gpu-trace-perf-renderdoc-wrapper.py");
static RENDERDOC_SNAPSHOT_SCRIPT: &[u8] = include_bytes!("gpu-trace-perf-renderdoc-snapshot.py");

/// If we're not using u_trace, then we use a python script that wraps the renderdoc events (draws,
/// blits, etc) in timestamps and takes the diff.  This won't work well for tilers.
impl TraceTool for RenderdocPyTrace {
    fn replay(&self, wrapper: Option<&str>, envs: &[(String, String)]) -> Result<ReplayOutput> {
        if self.is_directx {
            bail!(
                "Wrapper script support untested for DirectX, and renderdoccmd.exe doesn't ship with python support (use utrace, instead)"
            );
        }

        let renderdoc_command: &[_] = &[
            OsStr::new("python3"),
            OsStr::new("-"), // script from stdin
            self.file.as_os_str(),
        ];

        let mut command = replay_command(renderdoc_command, wrapper, envs);
        command.stdin(std::process::Stdio::piped());
        command.stdout(std::process::Stdio::piped());
        command.stderr(std::process::Stdio::piped());

        let mut child = command.spawn().expect("failed to start python3");
        let stdin = child.stdin.as_mut().expect("Failed to open stdin");
        stdin
            .write_all(RENDERDOC_WRAPPER_SCRIPT)
            .expect("failed to write to python's stdin");
        let output = child.wait_with_output().expect("failed to read stdout");

        if !output.status.success() {
            let stderr = std::str::from_utf8(&output.stderr).unwrap();
            if stderr.contains("API is unsupported") {
                bail!("renderdoc reported API (likely window system) unsupported, skipping trace");
            }

            error!("Failed to start renderdoc:");
            error!("{}", stderr);
            error!("command: {:?}", command);

            if stderr.contains("FileNotFound") {
                info!(
                    "TIP: Failure to find a file with a space in its name probably means your wrapper script didn't quote the arguments"
                )
            }
            bail!("Failed to start renderdoc");
        }
        Ok(ReplayOutput::from(output))
    }

    fn fps(&self, output: &ReplayOutput) -> Result<f64> {
        parse_renderdoc_wrapper_output(&output.stdout).map(|x| x as f64)
    }

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

    fn can_snapshot(&self) -> bool {
        true
    }

    fn snapshot(&self, output_dir: &str, loops: u32, timeout: Duration) -> Result<SnapshotResult> {
        let output_dir_path = self.output_dir(output_dir)?.unwrap();

        if self.is_directx {
            return self.snapshot_directx(&output_dir_path, loops, timeout);
        }

        let loops_str = loops.to_string();

        let renderdoc_command: &[_] = &[
            OsStr::new("python3"),
            OsStr::new("-"), // script from stdin
            self.file.as_os_str(),
            output_dir_path.as_os_str(),
            OsStr::new("--loops"),
            OsStr::new(&loops_str),
        ];

        let command = replay_command(renderdoc_command, None, &[]);

        let start_time = std::time::Instant::now();

        let output =
            self.run_replay_command_with_timeout(command, Some(RENDERDOC_SNAPSHOT_SCRIPT), timeout);

        if output.exit_code != 0 {
            if output.stderr.contains("API is unsupported") {
                bail!("renderdoc reported API (likely window system) unsupported, skipping trace");
            }
            // On other failures (including timeout), fall through with empty
            // files so the caller records the failure with the log intact.
            error!("Failed to run renderdoc snapshot: {}", output.stderr);
        }

        // Look for the output snapshots where we expect them to be.  If we just
        // try to parse it out from gpu-trace-perf-renderdoc-snapshot.py output,
        // the stdout getting interleaved with renderdoc's ends up causing
        // frames to be lost.
        let mut files = Vec::new();
        for i in 1..=loops {
            let relative = PathBuf::from(format!("snapshot{i:04}.png"));
            if output_dir_path.join(&relative).exists() {
                files.push(relative);
            }
        }

        Ok(SnapshotResult {
            files,
            output,
            runtime: start_time.elapsed(),
        })
    }
}

impl RenderdocPyTrace {
    fn snapshot_directx(
        &self,
        output_dir_path: &Path,
        loops: u32,
        timeout: Duration,
    ) -> Result<SnapshotResult> {
        let loops_str = loops.to_string();

        // Have to canonicalize, it seems that wine changes cwd to its drive_c.
        let output_dir_path = output_dir_path
            .canonicalize()
            .with_context(|| "canonicalizing {output_dir_path}")?;

        // renderdoccmd defaults to 1280x720 and doesn't look at the actual
        // swapchain created.  Default to something more likely to have been
        // used in captures, if we didn't successfully parse it out.
        let (w, h) = self.resolution.unwrap_or((1920, 1080));
        let w = w.to_string();
        let h = h.to_string();

        let args = &[
            OsStr::new("wine"),
            OsStr::new("renderdoccmd.exe"),
            OsStr::new("replay"),
            self.file.as_os_str(),
            OsStr::new("-l"),
            OsStr::new(&loops_str),
            OsStr::new("-w"),
            OsStr::new(&w),
            OsStr::new("-h"),
            OsStr::new(&h),
        ];

        let mesa_config = format!("output_dir={},frames=all", output_dir_path.display());
        let envs = vec![
            (
                "VK_LOADER_LAYERS_ENABLE".to_string(),
                "VK_LAYER_MESA_screenshot".to_string(),
            ),
            ("VK_LAYER_MESA_SCREENSHOT_CONFIG".to_string(), mesa_config),
        ];

        let command = replay_command(args, None, &envs);

        let start_time = std::time::Instant::now();

        let output = self.run_replay_command_with_timeout(command, None, timeout);

        if output.exit_code != 0 && output.stderr.contains("API is unsupported") {
            bail!("renderdoc reported API (likely window system) unsupported, skipping trace");
        }

        // The Mesa screenshot layer dumps <frame>.png files (0-indexed).
        // Include any that got created.
        let mut files = Vec::new();
        for i in 0..loops {
            let relative = format!("{i}.png");
            if std::fs::exists(output_dir_path.join(&relative))
                .context("checking for snapshot output")?
            {
                files.push(Path::new(&relative).to_owned());
            }
        }

        Ok(SnapshotResult {
            files,
            output,
            runtime: start_time.elapsed(),
        })
    }
}

/// If we're measuring times with u_trace, we don't need timestamps, but we do
/// need to loop the frame because there's a bunch of setup at the start that
/// we're not trying to measure.  renderdoccmd is just the tool for that.
pub struct RenderdocUtraceTrace {
    file: PathBuf,
    name: String,
    is_directx: bool,
}

fn guess_filename_directx(file: &Path) -> bool {
    let file_str = file.to_string_lossy();
    [
        "dx8", "dx9", "dx10", "dx11", "dx12", "d3d8", "d3d9", "d3d10", "d3d11", "d3d12", "dxgi",
    ]
    .iter()
    .any(|x| file_str.contains(x))
}

impl RenderdocUtraceTrace {
    pub fn new(root: &Path, file: &Path) -> RenderdocUtraceTrace {
        let is_directx = parse_rdc_driver(file)
            .map(|d| d.is_directx())
            .unwrap_or_else(|_| guess_filename_directx(file));
        RenderdocUtraceTrace {
            file: file.to_owned(),
            name: crate::relative_test_name(root, file),
            is_directx,
        }
    }
}

/// If we're not using u_trace, then we use a python script that wraps the renderdoc events (draws,
/// blits, etc) in timestamps and takes the diff.  This won't work well for tilers.
impl TraceTool for RenderdocUtraceTrace {
    fn replay(&self, wrapper: Option<&str>, envs: &[(String, String)]) -> Result<ReplayOutput> {
        // Play 3 frames, so we can take one of the middle one as the most
        // representative -- frame 0 has a bunch of setup, but the final frame may
        // not have any draws in it for some reason.
        let args: &[_] = if self.is_directx {
            &[
                OsStr::new("wine"),
                OsStr::new("renderdoccmd.exe"),
                OsStr::new("replay"),
                self.file.as_os_str(),
                OsStr::new("-l"),
                OsStr::new("3"),
            ]
        } else {
            &[
                OsStr::new("renderdoccmd"),
                OsStr::new("replay"),
                self.file.as_os_str(),
                OsStr::new("-l"),
                OsStr::new("3"),
            ]
        };

        let command = replay_command(args, wrapper, envs);
        let output: ReplayOutput = self.run_replay_command(command);

        if output.exit_code != 0 {
            if output.stderr.contains("API is unsupported") {
                bail!("renderdoc reported API (likely window system) unsupported, skipping trace");
            }

            if output.stderr.contains("FileNotFound") {
                println!(
                    "TIP: Failure to find a file with a space in its name probably means your wrapper script didn't quote the arguments"
                )
            }
        }

        Ok(output)
    }

    fn fps(&self, _output: &ReplayOutput) -> Result<f64> {
        unreachable!("shouldn't be called");
    }

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

// Returns the FPS for the frame from gpu-trace-perf-renderdoc-wrapper.py output
fn parse_renderdoc_wrapper_output(output: &str) -> Result<f32> {
    lazy_static! {
        static ref CALL_RE: Regex = Regex::new("EID [0-9]*: (.*)").unwrap();
    }

    let mut total = 0.0;
    for line in output.lines() {
        if let Some(cap) = CALL_RE.captures(line) {
            match cap[1].parse::<f32>() {
                Ok(time) => total += time,
                _ => {
                    bail!("Failed to parse renderdoc time event '{line}'");
                }
            }
        }
    }

    if total == 0.0 || total.is_nan() {
        bail!("Bad total time {total}");
    }

    Ok(1.0 / total)
}

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

    #[test]
    fn test_renderdoc_parsing() {
        // Actual renderdoc output, trimmed down for the testcase.
        let renderdoc_input = "
Counter 1 (GPU Duration):
    Time taken for this event on the GPU, as measured by delta between two GPU timestamps.
    Returns 8 byte CompType.Double, representing CounterUnit.Seconds
Counter 2000000 (N vertices submitted):
    N vertices submitted
    Returns 8 byte CompType.UInt, representing CounterUnit.Absolute
EID 52: 0.000045
EID 370: 0.000004
EID 407: 0.000006
";

        assert_approx_eq!(
            parse_renderdoc_wrapper_output(renderdoc_input).unwrap(),
            1.0 / (0.000_045 + 0.000_004 + 0.000_006),
            0.000_001
        );
    }

    #[test]
    fn test_renderdoc_nan_parsing() {
        // Actual renderdoc output, trimmed down for the testcase.
        let renderdoc_input = "
EID 52: 0.000045
EID 370: nan
EID 407: 0.000006
";

        assert!(parse_renderdoc_wrapper_output(renderdoc_input).is_err());
    }

    #[test]
    fn test_guess_filename_directx() {
        assert!(guess_filename_directx(Path::new(
            "d3d11-renderdoc/witcher3_medium_1.rdc"
        )));
        assert!(!guess_filename_directx(Path::new(
            "/home/anholt/src/traces-db/supertuxkart/supertuxkart-menu.rdc"
        )));
        assert!(!guess_filename_directx(Path::new(
            "/home/anholt/src/traces-db/godot/Material Testers.x86_64_2020.04.08_13.38_frame799.rdc"
        )));
    }
}