jmb-bar/src/backend/gpu_utilization.rs

use std::collections::HashMap;
use std::hash::{Hash, Hasher};
use std::os::fd::{AsRawFd, FromRawFd, OwnedFd};
use std::{fs, io, mem};

#[repr(C)]
#[derive(Default)]
struct PerfEventAttr {
    type_: u32,
    size: u32,
    config: u64,
    sample_period: u64,
    sample_type: u64,
    read_format: u64,
    flags: u64,
    wakeup_events: u32,
    bp_type: u32,
    config1: u64,
    config2: u64,
    branch_sample_type: u64,
    sample_regs_user: u64,
    sample_stack_user: u32,
    clockid: i32,
    sample_regs_intr: u64,
    aux_watermark: u32,
    sample_max_stack: u16,
    reserved2: u16,
    aux_sample_size: u32,
    reserved3: u32,
    sig_data: u64,
}

const PERF_FORMAT_TOTAL_TIME_ENABLED: u64 = 1;
const PERF_FORMAT_TOTAL_TIME_RUNNING: u64 = 2;

#[repr(C)]
struct ReadFormat {
    value: u64,
    time_enabled: u64,
    time_running: u64,
}

impl ReadFormat {
    fn calculate_utilization(&self, previous: &ReadFormat) -> f64 {
        let busy_delta = self.value.saturating_sub(previous.value);
        let elapsed_delta = self.time_enabled.saturating_sub(previous.time_enabled);

        if elapsed_delta == 0 {
            return 0.0;
        }

        (busy_delta as f64 / elapsed_delta as f64 * 100.0).clamp(0.0, 100.0)
    }
}

fn read_number(path: &str) -> io::Result<u64> {
    fs::read_to_string(path)?
        .trim()
        .parse()
        .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
}

fn read_config(path: &str) -> io::Result<u64> {
    let contents = fs::read_to_string(path)?;
    let value = contents
        .trim()
        .strip_prefix("config=")
        .unwrap_or(contents.trim());

    u64::from_str_radix(value.trim_start_matches("0x"), 16)
        .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
}

struct I915Engine {
    name: String,
    fd: OwnedFd,
}

impl Hash for I915Engine {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.name.hash(state)
    }
}

impl PartialEq<Self> for I915Engine {
    fn eq(&self, other: &Self) -> bool {
        self.name.eq(&other.name)
    }
}

impl Eq for I915Engine {}

impl I915Engine {
    fn open(name: &str) -> io::Result<Self> {
        let base = "/sys/bus/event_source/devices/i915";

        let pmu_type = read_number(&format!("{base}/type"))? as u32;
        let config = read_config(&format!("{base}/events/{name}"))?;

        let attr = PerfEventAttr {
            type_: pmu_type,
            size: size_of::<PerfEventAttr>() as u32,
            config,
            read_format: PERF_FORMAT_TOTAL_TIME_ENABLED | PERF_FORMAT_TOTAL_TIME_RUNNING,
            ..Default::default()
        };

        // pid = -1 and cpu = -1: system-wide event, if permitted.
        let fd = unsafe {
            libc::syscall(
                libc::SYS_perf_event_open,
                &attr as *const PerfEventAttr,
                -1i32, // all processes
                0i32,  // CPU 0
                -1i32, // no event group
                0u64,
            ) as i32
        };

        if fd < 0 {
            return Err(io::Error::last_os_error());
        }

        let fd = unsafe { OwnedFd::from_raw_fd(fd) };

        Ok(Self {
            name: name.to_owned(),
            fd,
        })
    }

    fn read(&self) -> io::Result<ReadFormat> {
        let mut result = ReadFormat {
            value: 0,
            time_enabled: 0,
            time_running: 0,
        };

        let size = mem::size_of::<ReadFormat>();
        let read = unsafe {
            libc::read(
                self.fd.as_raw_fd(),
                &mut result as *mut _ as *mut libc::c_void,
                size,
            )
        };

        if read != size as isize {
            return Err(io::Error::last_os_error());
        }

        Ok(result)
    }
}

pub struct I915UtilizationMonitor {
    engines: HashMap<I915Engine, Option<ReadFormat>>,
}

impl I915UtilizationMonitor {
    pub fn new() -> Option<Self> {
        let mut engines: HashMap<_, Option<ReadFormat>> = HashMap::new();
        let Ok(dir) = fs::read_dir("/sys/bus/event_source/devices/i915/events") else {
            return None;
        };
        for entry in dir {
            let Ok(entry) = entry else { continue };
            let Ok(name) = String::from_utf8(entry.file_name().into_encoded_bytes()) else {
                continue;
            };
            if name.ends_with("-busy") {
                let Ok(engine) = I915Engine::open(&name) else {
                    continue;
                };
                engines.insert(engine, None);
            }
        }
        Some(Self { engines })
    }

    pub fn calc_utilization(&mut self) -> u8 {
        let mut values: Vec<u8> = Vec::new();
        for (engine, sample) in &mut self.engines {
            let new_sample = engine.read().unwrap();
            if let Some(last) = sample {
                values.push(new_sample.calculate_utilization(last) as u8);
            }
            let _ = sample.insert(new_sample);
        }
        values
            .iter()
            .fold(0u8, |arg0: u8, other: &u8| arg0.max(*other))
    }
}