nuvec_gpu/metal/
profile.rs

1//! Per-kernel GPU time, for finding where time goes (`NUVEC_PROFILE=1`).
2//!
3//! While enabled, array evaluation commits every dispatch in a command buffer of its own
4//! and waits for it, and each completed command buffer's GPU time (Metal's GPU start and
5//! end timestamps) is added to its kernel's total. Commands then never overlap (nor wait
6//! for one another on the GPU), so totals show each kernel's cost alone, not the
7//! overlapped wall time. `NUVEC_PROFILE=shapes` keeps each
8//! grid size (threads per dimension) apart. The GPU idles between command buffers, so its
9//! clocks drop: absolute times are inflated (most for short kernels); compare shares and
10//! orders of magnitude, and measure throughput with the benchmarks.
11
12use std::collections::HashMap;
13use std::sync::{Arc, LazyLock, Mutex, PoisonError};
14
15static ENABLED: LazyLock<bool> =
16    LazyLock::new(|| std::env::var_os("NUVEC_PROFILE").is_some_and(|v| v != "0"));
17
18/// Per kernel: dispatches and GPU seconds.
19type Table = HashMap<Arc<str>, (u64, f64)>;
20
21static TABLE: Mutex<Option<Table>> = Mutex::new(None);
22
23/// Whether profiling is on (`NUVEC_PROFILE` set to anything but `0`, read once).
24pub fn enabled() -> bool {
25    *ENABLED
26}
27
28/// Whether grid sizes are kept apart (`NUVEC_PROFILE=shapes`).
29pub(crate) fn shapes() -> bool {
30    static SHAPES: LazyLock<bool> =
31        LazyLock::new(|| std::env::var_os("NUVEC_PROFILE").is_some_and(|v| v == "shapes"));
32    *SHAPES
33}
34
35/// Adds `seconds` of GPU time, shared equally by `kernels` (one per dispatch).
36pub(crate) fn record(kernels: &[Arc<str>], seconds: f64) {
37    if kernels.is_empty() {
38        return;
39    }
40    let share = seconds / kernels.len() as f64;
41    let mut table = TABLE.lock().unwrap_or_else(PoisonError::into_inner);
42    let table = table.get_or_insert_with(HashMap::new);
43    for k in kernels {
44        let entry = table.entry(k.clone()).or_default();
45        entry.0 += 1;
46        entry.1 += share;
47    }
48}
49
50/// `(kernel, dispatches, GPU seconds)` recorded so far, most expensive first.
51pub fn report() -> Vec<(String, u64, f64)> {
52    let table = TABLE.lock().unwrap_or_else(PoisonError::into_inner);
53    let mut rows: Vec<_> = table
54        .iter()
55        .flatten()
56        .map(|(k, &(n, s))| (k.to_string(), n, s))
57        .collect();
58    rows.sort_by(|a, b| b.2.total_cmp(&a.2));
59    rows
60}
61
62/// Forgets everything recorded so far.
63pub fn reset() {
64    *TABLE.lock().unwrap_or_else(PoisonError::into_inner) = None;
65}
nuvec_gpu/metal/profile.rs

nuvec_gpu/metal/
profile.rs
