#![recursion_limit = "1024"]
#![allow(non_upper_case_globals)]
extern crate libloading;
extern crate nvml_wrapper_sys as ffi;
pub mod bitmasks;
pub mod device;
pub mod enum_wrappers;
pub mod enums;
pub mod error;
pub mod event;
pub mod high_level;
pub mod nv_link;
pub mod struct_wrappers;
pub mod structs;
#[cfg(test)]
mod test_utils;
pub mod unit;
pub use crate::device::Device;
pub use crate::event::EventSet;
pub use crate::nv_link::NvLink;
pub use crate::unit::Unit;
pub mod sys_exports {
pub mod field_id {
pub use crate::ffi::bindings::field_id::*;
}
}
#[cfg(target_os = "linux")]
use std::convert::TryInto;
#[cfg(target_os = "linux")]
use std::ptr;
use std::{
convert::TryFrom,
ffi::{CStr, CString, OsStr},
mem::{self, ManuallyDrop},
os::raw::{c_int, c_uint},
};
use static_assertions::assert_impl_all;
#[cfg(target_os = "linux")]
use crate::enum_wrappers::device::TopologyLevel;
use crate::error::{nvml_sym, nvml_try, NvmlError};
use crate::ffi::bindings::*;
use crate::struct_wrappers::ExcludedDeviceInfo;
#[cfg(target_os = "linux")]
use crate::struct_wrappers::device::PciInfo;
use crate::struct_wrappers::unit::HwbcEntry;
use crate::bitmasks::InitFlags;
#[cfg(not(target_os = "linux"))]
const LIB_PATH: &str = "nvml.dll";
#[cfg(target_os = "linux")]
const LIB_PATH: &str = "libnvidia-ml.so";
pub fn cuda_driver_version_major(version: i32) -> i32 {
version / 1000
}
pub fn cuda_driver_version_minor(version: i32) -> i32 {
(version % 1000) / 10
}
pub struct Nvml {
lib: ManuallyDrop<NvmlLib>,
}
assert_impl_all!(Nvml: Send, Sync);
impl std::fmt::Debug for Nvml {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str("NVML")
}
}
impl Nvml {
pub fn init() -> Result<Self, NvmlError> {
Self::init_internal(LIB_PATH)
}
fn init_internal(path: impl AsRef<std::ffi::OsStr>) -> Result<Self, NvmlError> {
let lib = unsafe {
let lib = NvmlLib::new(path)?;
let sym = nvml_sym(lib.nvmlInit_v2.as_ref())?;
nvml_try(sym())?;
ManuallyDrop::new(lib)
};
Ok(Self { lib })
}
pub fn init_with_flags(flags: InitFlags) -> Result<Self, NvmlError> {
Self::init_with_flags_internal(LIB_PATH, flags)
}
fn init_with_flags_internal(
path: impl AsRef<std::ffi::OsStr>,
flags: InitFlags,
) -> Result<Self, NvmlError> {
let lib = unsafe {
let lib = NvmlLib::new(path)?;
let sym = nvml_sym(lib.nvmlInitWithFlags.as_ref())?;
nvml_try(sym(flags.bits()))?;
ManuallyDrop::new(lib)
};
Ok(Self { lib })
}
pub fn builder<'a>() -> NvmlBuilder<'a> {
NvmlBuilder::default()
}
pub fn shutdown(mut self) -> Result<(), NvmlError> {
let sym = nvml_sym(self.lib.nvmlShutdown.as_ref())?;
unsafe {
nvml_try(sym())?;
}
let lib = unsafe { ManuallyDrop::take(&mut self.lib) };
mem::forget(self);
Ok(lib.__library.close()?)
}
pub fn device_count(&self) -> Result<u32, NvmlError> {
let sym = nvml_sym(self.lib.nvmlDeviceGetCount_v2.as_ref())?;
unsafe {
let mut count: c_uint = mem::zeroed();
nvml_try(sym(&mut count))?;
Ok(count as u32)
}
}
pub fn sys_driver_version(&self) -> Result<String, NvmlError> {
let sym = nvml_sym(self.lib.nvmlSystemGetDriverVersion.as_ref())?;
unsafe {
let mut version_vec = vec![0; NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE as usize];
nvml_try(sym(
version_vec.as_mut_ptr(),
NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE,
))?;
let version_raw = CStr::from_ptr(version_vec.as_ptr());
Ok(version_raw.to_str()?.into())
}
}
pub fn sys_nvml_version(&self) -> Result<String, NvmlError> {
let sym = nvml_sym(self.lib.nvmlSystemGetNVMLVersion.as_ref())?;
unsafe {
let mut version_vec = vec![0; NVML_SYSTEM_NVML_VERSION_BUFFER_SIZE as usize];
nvml_try(sym(
version_vec.as_mut_ptr(),
NVML_SYSTEM_NVML_VERSION_BUFFER_SIZE,
))?;
let version_raw = CStr::from_ptr(version_vec.as_ptr());
Ok(version_raw.to_str()?.into())
}
}
pub fn sys_cuda_driver_version(&self) -> Result<i32, NvmlError> {
let sym = nvml_sym(self.lib.nvmlSystemGetCudaDriverVersion_v2.as_ref())?;
unsafe {
let mut version: c_int = mem::zeroed();
nvml_try(sym(&mut version))?;
Ok(version)
}
}
pub fn sys_process_name(&self, pid: u32, length: usize) -> Result<String, NvmlError> {
let sym = nvml_sym(self.lib.nvmlSystemGetProcessName.as_ref())?;
unsafe {
let mut name_vec = vec![0; length];
nvml_try(sym(pid, name_vec.as_mut_ptr(), length as c_uint))?;
let name_raw = CStr::from_ptr(name_vec.as_ptr());
Ok(name_raw.to_str()?.into())
}
}
pub fn device_by_index(&self, index: u32) -> Result<Device, NvmlError> {
let sym = nvml_sym(self.lib.nvmlDeviceGetHandleByIndex_v2.as_ref())?;
unsafe {
let mut device: nvmlDevice_t = mem::zeroed();
nvml_try(sym(index, &mut device))?;
Ok(Device::new(device, self))
}
}
pub fn device_by_pci_bus_id<S: AsRef<str>>(&self, pci_bus_id: S) -> Result<Device, NvmlError>
where
Vec<u8>: From<S>,
{
let sym = nvml_sym(self.lib.nvmlDeviceGetHandleByPciBusId_v2.as_ref())?;
unsafe {
let c_string = CString::new(pci_bus_id)?;
let mut device: nvmlDevice_t = mem::zeroed();
nvml_try(sym(c_string.as_ptr(), &mut device))?;
Ok(Device::new(device, self))
}
}
#[deprecated(note = "use `.device_by_uuid()`, this errors on dual GPU boards")]
pub fn device_by_serial<S: AsRef<str>>(&self, board_serial: S) -> Result<Device, NvmlError>
where
Vec<u8>: From<S>,
{
let sym = nvml_sym(self.lib.nvmlDeviceGetHandleBySerial.as_ref())?;
unsafe {
let c_string = CString::new(board_serial)?;
let mut device: nvmlDevice_t = mem::zeroed();
nvml_try(sym(c_string.as_ptr(), &mut device))?;
Ok(Device::new(device, self))
}
}
pub fn device_by_uuid<S: AsRef<str>>(&self, uuid: S) -> Result<Device, NvmlError>
where
Vec<u8>: From<S>,
{
let sym = nvml_sym(self.lib.nvmlDeviceGetHandleByUUID.as_ref())?;
unsafe {
let c_string = CString::new(uuid)?;
let mut device: nvmlDevice_t = mem::zeroed();
nvml_try(sym(c_string.as_ptr(), &mut device))?;
Ok(Device::new(device, self))
}
}
#[cfg(target_os = "linux")]
pub fn topology_common_ancestor(
&self,
device1: &Device,
device2: &Device,
) -> Result<TopologyLevel, NvmlError> {
let sym = nvml_sym(self.lib.nvmlDeviceGetTopologyCommonAncestor.as_ref())?;
unsafe {
let mut level: nvmlGpuTopologyLevel_t = mem::zeroed();
nvml_try(sym(device1.handle(), device2.handle(), &mut level))?;
TopologyLevel::try_from(level)
}
}
pub fn unit_by_index(&self, index: u32) -> Result<Unit, NvmlError> {
let sym = nvml_sym(self.lib.nvmlUnitGetHandleByIndex.as_ref())?;
unsafe {
let mut unit: nvmlUnit_t = mem::zeroed();
nvml_try(sym(index as c_uint, &mut unit))?;
Ok(Unit::new(unit, self))
}
}
pub fn are_devices_on_same_board(
&self,
device1: &Device,
device2: &Device,
) -> Result<bool, NvmlError> {
let sym = nvml_sym(self.lib.nvmlDeviceOnSameBoard.as_ref())?;
unsafe {
let mut bool_int: c_int = mem::zeroed();
nvml_try(sym(device1.handle(), device2.handle(), &mut bool_int))?;
match bool_int {
0 => Ok(false),
_ => Ok(true),
}
}
}
#[cfg(target_os = "linux")]
pub fn topology_gpu_set(&self, cpu_number: u32) -> Result<Vec<Device>, NvmlError> {
let sym = nvml_sym(self.lib.nvmlSystemGetTopologyGpuSet.as_ref())?;
unsafe {
let mut count = match self.topology_gpu_set_count(cpu_number)? {
0 => return Ok(vec![]),
value => value,
};
let mut devices: Vec<nvmlDevice_t> = vec![mem::zeroed(); count as usize];
nvml_try(sym(cpu_number, &mut count, devices.as_mut_ptr()))?;
Ok(devices.into_iter().map(|d| Device::new(d, self)).collect())
}
}
#[cfg(target_os = "linux")]
fn topology_gpu_set_count(&self, cpu_number: u32) -> Result<c_uint, NvmlError> {
let sym = nvml_sym(self.lib.nvmlSystemGetTopologyGpuSet.as_ref())?;
unsafe {
let mut count: c_uint = 0;
nvml_try(sym(cpu_number, &mut count, ptr::null_mut()))?;
Ok(count)
}
}
pub fn hic_versions(&self) -> Result<Vec<HwbcEntry>, NvmlError> {
let sym = nvml_sym(self.lib.nvmlSystemGetHicVersion.as_ref())?;
unsafe {
let mut count: c_uint = match self.hic_count()? {
0 => return Ok(vec![]),
value => value,
};
let mut hics: Vec<nvmlHwbcEntry_t> = vec![mem::zeroed(); count as usize];
nvml_try(sym(&mut count, hics.as_mut_ptr()))?;
hics.into_iter().map(HwbcEntry::try_from).collect()
}
}
pub fn hic_count(&self) -> Result<u32, NvmlError> {
let sym = nvml_sym(self.lib.nvmlSystemGetHicVersion.as_ref())?;
unsafe {
let mut count: c_uint = 1;
let mut hics: [nvmlHwbcEntry_t; 1] = [mem::zeroed()];
match sym(&mut count, hics.as_mut_ptr()) {
nvmlReturn_enum_NVML_SUCCESS | nvmlReturn_enum_NVML_ERROR_INSUFFICIENT_SIZE => {
Ok(count)
}
other => nvml_try(other).map(|_| 0),
}
}
}
pub fn unit_count(&self) -> Result<u32, NvmlError> {
let sym = nvml_sym(self.lib.nvmlUnitGetCount.as_ref())?;
unsafe {
let mut count: c_uint = mem::zeroed();
nvml_try(sym(&mut count))?;
Ok(count)
}
}
pub fn create_event_set(&self) -> Result<EventSet, NvmlError> {
let sym = nvml_sym(self.lib.nvmlEventSetCreate.as_ref())?;
unsafe {
let mut set: nvmlEventSet_t = mem::zeroed();
nvml_try(sym(&mut set))?;
Ok(EventSet::new(set, self))
}
}
#[cfg(target_os = "linux")]
pub fn discover_gpus(&self, pci_info: PciInfo) -> Result<(), NvmlError> {
let sym = nvml_sym(self.lib.nvmlDeviceDiscoverGpus.as_ref())?;
unsafe { nvml_try(sym(&mut pci_info.try_into()?)) }
}
pub fn excluded_device_count(&self) -> Result<u32, NvmlError> {
let sym = nvml_sym(self.lib.nvmlGetExcludedDeviceCount.as_ref())?;
unsafe {
let mut count: c_uint = mem::zeroed();
nvml_try(sym(&mut count))?;
Ok(count)
}
}
pub fn excluded_device_info(&self, index: u32) -> Result<ExcludedDeviceInfo, NvmlError> {
let sym = nvml_sym(self.lib.nvmlGetExcludedDeviceInfoByIndex.as_ref())?;
unsafe {
let mut info: nvmlExcludedDeviceInfo_t = mem::zeroed();
nvml_try(sym(index, &mut info))?;
ExcludedDeviceInfo::try_from(info)
}
}
}
impl Drop for Nvml {
fn drop(&mut self) {
unsafe {
self.lib.nvmlShutdown();
ManuallyDrop::drop(&mut self.lib);
}
}
}
#[derive(Debug, Clone, Eq, PartialEq, Default)]
pub struct NvmlBuilder<'a> {
lib_path: Option<&'a OsStr>,
flags: InitFlags,
}
impl<'a> NvmlBuilder<'a> {
pub fn lib_path(&mut self, path: &'a OsStr) -> &mut Self {
self.lib_path = Some(path);
self
}
pub fn flags(&mut self, flags: InitFlags) -> &mut Self {
self.flags = flags;
self
}
pub fn init(&self) -> Result<Nvml, NvmlError> {
let lib_path = self.lib_path.unwrap_or_else(|| LIB_PATH.as_ref());
if self.flags.is_empty() {
Nvml::init_internal(lib_path)
} else {
Nvml::init_with_flags_internal(lib_path, self.flags)
}
}
}
#[cfg(test)]
mod test {
use super::*;
use crate::bitmasks::InitFlags;
use crate::error::NvmlError;
use crate::test_utils::*;
#[test]
fn init_with_flags() {
Nvml::init_with_flags(InitFlags::NO_GPUS).unwrap();
}
#[test]
fn shutdown() {
test(3, || nvml().shutdown())
}
#[test]
fn device_count() {
test(3, || nvml().device_count())
}
#[test]
fn sys_driver_version() {
test(3, || nvml().sys_driver_version())
}
#[test]
fn sys_nvml_version() {
test(3, || nvml().sys_nvml_version())
}
#[test]
fn sys_cuda_driver_version() {
test(3, || nvml().sys_cuda_driver_version())
}
#[test]
fn sys_cuda_driver_version_major() {
test(3, || {
Ok(cuda_driver_version_major(nvml().sys_cuda_driver_version()?))
})
}
#[test]
fn sys_cuda_driver_version_minor() {
test(3, || {
Ok(cuda_driver_version_minor(nvml().sys_cuda_driver_version()?))
})
}
#[test]
fn sys_process_name() {
let nvml = nvml();
test_with_device(3, &nvml, |device| {
let processes = device.running_graphics_processes()?;
match nvml.sys_process_name(processes[0].pid, 64) {
Err(NvmlError::NoPermission) => Ok("No permission error".into()),
v => v,
}
})
}
#[test]
fn device_by_index() {
let nvml = nvml();
test(3, || nvml.device_by_index(0))
}
#[test]
fn device_by_pci_bus_id() {
let nvml = nvml();
test_with_device(3, &nvml, |device| {
let id = device.pci_info()?.bus_id;
nvml.device_by_pci_bus_id(id)
})
}
#[ignore = "my machine does not support this call"]
#[test]
fn device_by_serial() {
let nvml = nvml();
#[allow(deprecated)]
test_with_device(3, &nvml, |device| {
let serial = device.serial()?;
nvml.device_by_serial(serial)
})
}
#[test]
fn device_by_uuid() {
let nvml = nvml();
test_with_device(3, &nvml, |device| {
let uuid = device.uuid()?;
nvml.device_by_uuid(uuid)
})
}
#[ignore = "my machine does not support this call"]
#[cfg(target_os = "linux")]
#[test]
fn topology_common_ancestor() {
let nvml = nvml();
let device1 = device(&nvml);
let device2 = nvml.device_by_index(1).expect("device");
nvml.topology_common_ancestor(&device1, &device2)
.expect("TopologyLevel");
}
#[test]
#[ignore = "my machine does not support this call"]
fn unit_by_index() {
let nvml = nvml();
test(3, || nvml.unit_by_index(0))
}
#[ignore = "my machine does not support this call"]
#[test]
fn are_devices_on_same_board() {
let nvml = nvml();
let device1 = device(&nvml);
let device2 = nvml.device_by_index(1).expect("device");
nvml.are_devices_on_same_board(&device1, &device2)
.expect("bool");
}
#[cfg(target_os = "linux")]
#[test]
fn topology_gpu_set() {
let nvml = nvml();
test(3, || nvml.topology_gpu_set(0))
}
#[test]
fn hic_version() {
let nvml = nvml();
test(3, || nvml.hic_versions())
}
#[test]
fn unit_count() {
test(3, || nvml().unit_count())
}
#[test]
fn create_event_set() {
let nvml = nvml();
test(3, || nvml.create_event_set())
}
#[cfg(target_os = "linux")]
#[should_panic(expected = "OperatingSystem")]
#[test]
fn discover_gpus() {
let nvml = nvml();
test_with_device(3, &nvml, |device| {
let pci_info = device.pci_info()?;
match nvml.discover_gpus(pci_info) {
Err(NvmlError::NoPermission) => panic!("NoPermission"),
other => other,
}
})
}
#[test]
fn excluded_device_count() {
let nvml = nvml();
test(3, || nvml.excluded_device_count())
}
#[test]
fn excluded_device_info() {
let nvml = nvml();
if nvml.excluded_device_count().unwrap() > 0 {
test(3, || nvml.excluded_device_info(0))
}
}
}