zenyx-engine-telemetry/subcrates/telemetry/src/modules/gpu/vram.rs

use super::super::UNKNOWN;
use ash::vk::{API_VERSION_1_2, ApplicationInfo, InstanceCreateInfo, MemoryHeapFlags};
use humansize::{DECIMAL, make_format};

#[cfg(not(target_os = "macos"))]
pub fn get_metal() -> u64 {
    0
}

#[cfg(target_os = "macos")]
pub fn get_metal() -> u64 {
    use metal::Device as MetalDevice;

    let device = MetalDevice::system_default().expect("No Metal-compatible GPU found");
    device.recommended_max_working_set_size()
}

pub fn get_vulkan(device_id: u32) -> u64 {
    let entry = unsafe { ash::Entry::load().unwrap() };
    let app_info = ApplicationInfo {
        p_application_name: std::ptr::null(),
        application_version: 0,
        p_engine_name: std::ptr::null(),
        engine_version: 0,
        api_version: API_VERSION_1_2,
        ..Default::default()
    };

    let create_info = InstanceCreateInfo {
        p_application_info: &app_info,
        ..Default::default()
    };

    let instance = unsafe { entry.create_instance(&create_info, None).unwrap() };

    let physical_devices = unsafe { instance.enumerate_physical_devices().unwrap() };
    let mut total_vram = 0;

    for device in physical_devices {
        let memory_properties = unsafe { instance.get_physical_device_memory_properties(device) };
        let device_properties = unsafe { instance.get_physical_device_properties(device) };

        if device_id != device_properties.device_id {
            continue;
        }

        for heap in memory_properties.memory_heaps {
            if heap.flags.contains(MemoryHeapFlags::DEVICE_LOCAL) {
                total_vram += heap.size;
            }
        }

        break;
    }

    total_vram
}

pub fn get(vendor: u32, device_id: u32) -> String {
    let formatter = make_format(DECIMAL);

    match vendor {
        0x10DE | 0x1002 | 0x8086 | 0x5143 => formatter(get_vulkan(device_id)),
        0x1010 => formatter(get_metal()),
        _ => UNKNOWN.to_string(),
    }
}