gpu-trace-perf 1.8.2

use std::{
    collections::{HashMap, HashSet},
    io::Read,
    path::Path,
    sync::atomic::AtomicUsize,
};

use anyhow::{Context, Result, anyhow, bail};
use log::error;
use serde::{Deserialize, de::Error};
use tempfile::NamedTempFile;

// These draw/compute event names from various drivers that appear in utrace
// events for draws
pub const UTRACE_SHADER_DRAW_EVENTS: &[&str] = &[
    "draw",         // tu, intel, radeonsi
    "draw_indexed", // intel
    "draw_indexed_multi",
    "draw_indirect_byte_count",
    "draw_indirect",
    "draw_indexed_indirect",
    "draw_indirect_count",
    "draw_indexed_indirect_count",
    "compute",
    "compute_indirect",
];

// Turnip and intel shader hash names that appear as params in draw events (with
// a bit of extension in case intel starts logging tess/geom).
//
// Note!  We're not considering tu binning shaders here, should probably fix
// that.
pub const UTRACE_SHADER_STAGES: &[&str] = &[
    "vs_sha1", "tcs_sha1", "tes_sha1", "gs_sha1", "fs_sha1", "vs_hash", "tcs_hash", "tes_hash",
    "gs_hash", "fs_hash", "cs_hash",
];

#[derive(Deserialize, Default)]
pub struct Frame {
    batches: Vec<Batch>,
}

impl Frame {
    pub fn event_times(&self, event: &str) -> Result<FrameTimes> {
        let start_strs: HashSet<String> = HashSet::from_iter([
            format!("start_{event}"),       // tu
            format!("intel_begin_{event}"), // intel
            format!("begin_{event}"),       // panvk
            format!("si_begin_{event}"),    // radeonsi
        ]);
        let end_strs: HashSet<String> = HashSet::from_iter([
            format!("end_{event}"),
            format!("intel_end_{event}"),
            format!("si_end_{event}"),
        ]);

        let mut times = Vec::new();

        // Note: radeonsi can split event start/ends across batch IDs.
        let mut start = None;

        for batch in &self.batches {
            for event in &batch.events {
                if start_strs.contains(&event.event) {
                    start = Some(event);
                } else if end_strs.contains(&event.event) {
                    let start_unwrap = start.with_context(|| {
                        format!("missing start event at end event time ({event:?})")
                    })?;
                    times.push((
                        event
                            .time_ns
                            .checked_sub(start_unwrap.time_ns)
                            .with_context(|| {
                                format!(
                                    "end {} before start {}",
                                    event.time_ns, start_unwrap.time_ns
                                )
                            })?,
                        start_unwrap.params.clone(),
                    ));
                    start = None;
                }
            }
        }

        Ok(FrameTimes { times })
    }

    /// Returns an FPS for the recorded frame, if only the events given
    /// contributed to the runtime.
    fn fps(&self, events: &[String]) -> Result<f64> {
        let mut total_ns = 0u64;
        for event in events {
            let frame_times = self.event_times(event)?;

            total_ns += frame_times.times.iter().map(|x| x.0).sum::<u64>();
        }
        if total_ns == 0 {
            bail!("No events captured for the given utrace events ({events:?})");
        }
        Ok(1_000_000_000.0 / (total_ns as f64))
    }
}

#[derive(Deserialize, Default)]
struct Batch {
    events: Vec<Event>,
}

fn u64_parse<'de, D>(deserializer: D) -> Result<u64, D::Error>
where
    D: serde::Deserializer<'de>,
{
    let s: &str = Deserialize::deserialize(deserializer)?;
    // do better hex decoding than this
    str::parse::<u64>(s).map_err(D::Error::custom)
}
#[derive(Deserialize, Debug)]
struct Event {
    event: String,
    #[serde(deserialize_with = "u64_parse")]
    time_ns: u64,
    params: HashMap<String, String>,
}

fn read_utrace_json(data: &[u8]) -> Result<Vec<Frame>> {
    let frames: Vec<Frame> = serde_json::from_slice(data)?;

    Ok(frames)
}

#[derive(Debug, serde::Deserialize)]
struct CsvRecord {
    frame: usize,
    batch: usize,
    event: String,
    time_ns: u64,
    #[serde(rename = "")]
    params: Vec<String>,
}

fn get_or_insert<T: Default>(v: &mut Vec<T>, index: usize) -> &mut T {
    if index >= v.len() {
        v.resize_with(index + 1, T::default);
    }

    &mut v[index]
}

/// Strip the first line that gets repeated from the CSV.
///
/// Radeonsi ends up printing the CSV header multiple times due to multiple
/// utrace contexts being created even for a single GL context, so delete the
/// duplicates.
fn strip_repeated_csv_header(data: &[u8]) -> Vec<u8> {
    let mut result = Vec::new();

    let mut first_line = None;

    for non_repeated_line in data
        .split(|x| *x == b'\n')
        .filter(|line| match &mut first_line {
            None => {
                first_line = Some(line.to_vec());
                true
            }
            Some(first_line) => first_line != *line,
        })
    {
        result.extend_from_slice(non_repeated_line);
        result.push(b'\n');
    }

    result
}

fn read_utrace_csv(data: &[u8]) -> Result<Vec<Frame>> {
    let data = strip_repeated_csv_header(data);

    let mut frames: Vec<Frame> = Vec::new();
    let mut csv_reader = csv::ReaderBuilder::new().flexible(true).from_reader(&*data);

    for result in csv_reader.deserialize() {
        let record: CsvRecord = result?;
        let frame = get_or_insert(&mut frames, record.frame);
        let batch = get_or_insert(&mut frame.batches, record.batch);

        let params = record
            .params
            .iter()
            .map(|s| s.trim())
            .filter(|s| !s.is_empty())
            .map(|s| {
                s.split_once('=')
                    .ok_or_else(|| anyhow!("illegal param: {s}"))
                    .map(|(k, v)| (k.to_owned(), v.to_owned()))
            })
            .collect::<Result<HashMap<_, _>>>()?;

        batch.events.push(Event {
            event: record.event,
            time_ns: record.time_ns,
            params,
        });
    }

    Ok(frames)
}

#[derive(Debug)]
pub struct FrameTimes {
    pub times: Vec<(u64, HashMap<String, String>)>,
}

pub fn u_trace_events(events: &str) -> Vec<String> {
    if events == "drawcalls" {
        return UTRACE_SHADER_DRAW_EVENTS
            .iter()
            .map(|x| (*x).to_owned())
            .collect();
    }

    events.split(',').map(|x| x.to_owned()).collect()
}

/// Parser for Mesa's u_trace events.  If you pass it an event name, it will
/// return .env() necessary to set up utrace, and .last_frame_fps() will return
/// a non-None fps result.
pub struct UTraceParser {
    file: Option<NamedTempFile>,
    active: bool,
    events: Vec<String>,
    use_csv: bool,
}

impl UTraceParser {
    pub fn new(u_trace_event: &str, use_csv: bool) -> UTraceParser {
        let active = !u_trace_event.is_empty();

        UTraceParser {
            file: if !active {
                None
            } else {
                Some(NamedTempFile::new().expect("creating a temp file"))
            },
            active,
            events: u_trace_events(u_trace_event),
            use_csv,
        }
    }

    pub fn env(&self) -> Vec<(String, String)> {
        if let Some(file) = &self.file {
            let trace_fmt = if self.use_csv {
                "print_csv"
            } else {
                "print_json"
            };

            let mut env = vec![
                ("MESA_GPU_TRACES".to_owned(), trace_fmt.to_owned()),
                (
                    "MESA_GPU_TRACEFILE".to_owned(),
                    file.path()
                        .as_os_str()
                        .to_str()
                        .expect("tempfile as String")
                        .to_owned(),
                ),
            ];

            let driver_enables = [
                "TU_GPU_TRACEPOINT",
                "SI_GPU_TRACEPOINT",
                "ANV_GPU_TRACEPOINT",
                "FD_GPU_TRACEPOINT",
            ];
            let mut tracepoint_enable = String::new();
            for event in &self.events {
                tracepoint_enable.push('+');
                tracepoint_enable.push_str(event);
                tracepoint_enable.push(',')
            }

            for driver in driver_enables {
                env.push((driver.to_owned(), tracepoint_enable.clone()));
            }

            env
        } else {
            vec![]
        }
    }

    pub fn active(&self) -> bool {
        self.active
    }

    fn file_data(file: &Path) -> serde_json::Result<Vec<u8>> {
        let mut reader =
            std::io::BufReader::new(std::fs::File::open(file).expect("opening utrace temp file"));

        let mut buf = Vec::new();
        reader
            .read_to_end(&mut buf)
            .expect("reading utrace temp file");

        Ok(buf)
    }

    pub fn results(&mut self) -> Result<Vec<Frame>> {
        if let Some(file) = &mut self.file {
            // Uncomment to save off the file for debugging purposes.
            //
            // file.disable_cleanup(true);
            // println!("saving at {}", file.path().display());

            let data = Self::file_data(file.path()).expect("reading utrace");

            let frames = if self.use_csv {
                read_utrace_csv(&data)
            } else {
                read_utrace_json(&data)
            };

            // Save the first couple of errors for debugging, but don't save
            // forever or you'll end up filling up /tmp.
            static SAVED: AtomicUsize = AtomicUsize::new(0);
            if let Err(err) = &frames {
                if SAVED.fetch_add(1, std::sync::atomic::Ordering::SeqCst) < 3 {
                    file.disable_cleanup(true);
                    error!(
                        "Failed to parse {}: saved source to {}:\n{}",
                        if self.use_csv { "csv" } else { "json" },
                        file.path().display(),
                        err
                    )
                }
            }

            frames
        } else {
            Ok(Vec::new())
        }
    }

    pub fn middle_frame_fps(&mut self) -> Result<f64> {
        let frames = self.results().context("parsing json")?;

        let index = frames.len() / 2;

        if frames.is_empty() {
            bail!("No frames captured from utrace");
        }

        frames[index].fps(&self.events)
    }
}

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

    static TURNIP_JSON: &[u8] = include_bytes!("test_data/u_trace_turnip.json");
    static TURNIP_CSV: &[u8] = include_bytes!("test_data/u_trace_turnip.csv");
    static RADEONSI_CSV: &[u8] = include_bytes!("test_data/u_trace_radeonsi.csv");

    #[test]
    fn test_turnip() {
        let result = read_utrace_json(TURNIP_JSON).expect("parsing");
        assert_eq!(result.len(), 6);

        let renderpass_times = &result[0].event_times("render_pass").unwrap();

        assert_eq!(
            renderpass_times.times.first().unwrap().0,
            535519335156u64 - 535519313212u64
        );

        let cmdbuf_times = &result[5].event_times("cmd_buffer").unwrap();

        assert_eq!(
            cmdbuf_times.times.first().unwrap().0,
            (535523574768 - 535523572220)
        );

        // This was a capture of glxgears, which goes at ~18k fps with swapping,
        // but the renderpass portion is less than that.
        assert_approx_eq!(
            &result[5].fps(&["render_pass".to_owned()]).unwrap(),
            47250.04725004725
        );
        assert_approx_eq!(
            &result[5].fps(&["cmd_buffer".to_owned()]).unwrap(),
            36770.11325194882
        );
        assert_approx_eq!(
            &result[5]
                .fps(&["cmd_buffer".to_owned(), "event_not_present".to_owned()])
                .unwrap(),
            36770.11325194882
        );
    }

    #[test]
    fn test_turnip_csv() {
        let result = read_utrace_csv(TURNIP_CSV).expect("parsing");
        assert_eq!(result.len(), 1);

        let frame = &result[0];
        assert_eq!(frame.batches.len(), 13);

        let renderpass_times = &frame.event_times("render_pass").unwrap();

        assert_eq!(
            renderpass_times.times.first().unwrap().0,
            817658920u64 - 817572600u64
        );

        let cmdbuf_times = &frame.event_times("cmd_buffer").unwrap();

        assert_eq!(
            cmdbuf_times.times.first().unwrap().0,
            803740964u64 - 803737272u64,
        );

        let batch = &frame.batches[4];
        assert_eq!(batch.events.len(), 2);

        let event = &batch.events[1];
        assert_eq!(event.event, "end_cmd_buffer");
        assert_eq!(
            event.params,
            HashMap::from([
                ("renderpasses".to_owned(), "0".to_owned()),
                ("dispatches".to_owned(), "0".to_owned())
            ])
        );
    }

    #[test]
    fn test_radeonsi_csv() {
        let result = read_utrace_csv(RADEONSI_CSV).expect("parsing");
        assert_eq!(result.len(), 1);

        let frame = &result[0];
        assert_eq!(frame.batches.len(), 60);

        let compute_times = &frame.event_times("compute").unwrap();

        let expected = [
            7853639052730u64 - 7853638675570u64,
            7853639110260u64 - 7853639052730u64,
            7853645398620u64 - 7853645394260u64,
            7853645495820u64 - 7853645398620u64,
            7853647200080u64 - 7853647195600u64,
            7853647246820u64 - 7853647200080u64,
        ];
        assert_eq!(compute_times.times.len(), expected.len());

        for (i, time) in compute_times.times.iter().enumerate() {
            assert_eq!(time.0, expected[i]);
        }
    }

    #[test]
    fn test_some_failure() {
        static FAIL_JSON: &[u8] = include_bytes!("test_data/fail.json");
        let result = read_utrace_json(FAIL_JSON);
        assert!(result.is_err());
    }
}