test-smi-rs 0.1.2

#![allow(non_upper_case_globals)]

use crate::core_status::{default_core_status, CoreStatus};
use crate::device_error_info::*;
use crate::device_file::*;
use crate::device_info::*;
use crate::error::*;
use crate::generated::binding::*;
use crate::performance::*;
use crate::types::LinkType;
use std::collections::BTreeMap;
use std::sync::{Arc, RwLock};

pub(crate) fn new_device(
    handle: FuriosaSmiDeviceHandle,
    observer_instance: Arc<RwLock<DeviceObserverInstance>>,
) -> Device {
    Device {
        raw: handle,
        observer_instance,
    }
}

/// Abstraction for a single Furiosa NPU device.
pub struct Device {
    raw: FuriosaSmiDeviceHandle,
    observer_instance: Arc<RwLock<DeviceObserverInstance>>,
}

impl Device {
    /// Returns `DeviceInfo` which contains information about NPU device. (e.g. arch, serial, ...)
    pub fn device_info(&self) -> SmiResult<DeviceInfo> {
        let mut output = default_device_info();

        match unsafe {
            furiosa_smi_get_device_info(self.raw, &mut output as *mut FuriosaSmiDeviceInfo)
        } {
            FuriosaSmiReturnCode_FURIOSA_SMI_RETURN_CODE_OK => Ok(new_device_info(output)),
            err => Err(SmiError::from(err)),
        }
    }

    /// List device files under this device.
    pub fn device_files(&self) -> SmiResult<Vec<DeviceFile>> {
        let mut output = default_device_files();

        match unsafe {
            furiosa_smi_get_device_files(self.raw, &mut output as *mut FuriosaSmiDeviceFiles)
        } {
            FuriosaSmiReturnCode_FURIOSA_SMI_RETURN_CODE_OK => {
                let mut device_files = Vec::with_capacity(output.count as usize);
                for i in 0..output.count as usize {
                    device_files.push(new_device_file(output.device_files[i]));
                }

                Ok(device_files)
            }
            err => Err(SmiError::from(err)),
        }
    }

    /// Examine each core of the device, whether it is available or not.
    pub fn core_status(&self) -> SmiResult<BTreeMap<u32, CoreStatus>> {
        let mut output = default_core_status();

        match unsafe {
            furiosa_smi_get_device_core_status(self.raw, &mut output as *mut FuriosaSmiCoreStatuses)
        } {
            FuriosaSmiReturnCode_FURIOSA_SMI_RETURN_CODE_OK => {
                let mut core_status = BTreeMap::new();
                for i in 0..output.count as usize {
                    core_status.insert(i as u32, CoreStatus::from(output.core_status[i]));
                }

                Ok(core_status)
            }
            err => Err(SmiError::from(err)),
        }
    }

    /// Returns error states of the device.
    pub fn device_error_info(&self) -> SmiResult<DeviceErrorInfo> {
        let mut output = default_device_error_info();

        match unsafe {
            furiosa_smi_get_device_error_info(
                self.raw,
                &mut output as *mut FuriosaSmiDeviceErrorInfo,
            )
        } {
            FuriosaSmiReturnCode_FURIOSA_SMI_RETURN_CODE_OK => Ok(new_device_error_info(output)),
            err => Err(SmiError::from(err)),
        }
    }

    /// Returns a liveness state of the device.
    pub fn liveness(&self) -> SmiResult<bool> {
        let mut output = false;

        match unsafe { furiosa_smi_get_device_liveness(self.raw, &mut output as *mut bool) } {
            FuriosaSmiReturnCode_FURIOSA_SMI_RETURN_CODE_OK => Ok(output),
            err => Err(SmiError::from(err)),
        }
    }

    /// Returns a core utilization of the device.
    pub fn core_utilization(&self) -> SmiResult<CoreUtilization> {
        let mut output = default_core_utilization();
        let observer_instance = match self.observer_instance.read() {
            Ok(instance) => instance,
            Err(_) => return Err(SmiError::InternalError),
        };

        match unsafe {
            furiosa_smi_get_core_utilization(
                observer_instance.raw(),
                self.raw,
                &mut output as *mut FuriosaSmiCoreUtilization,
            )
        } {
            FuriosaSmiReturnCode_FURIOSA_SMI_RETURN_CODE_OK => Ok(new_core_utilization(output)),
            err => Err(SmiError::from(err)),
        }
    }

    /// Returns a memory utilization of the device.
    pub fn memory_utilization(&self) -> SmiResult<MemoryUtilization> {
        let mut output = default_memory_utilization();

        match unsafe {
            furiosa_smi_get_memory_utilization(
                self.raw,
                &mut output as *mut FuriosaSmiMemoryUtilization,
            )
        } {
            FuriosaSmiReturnCode_FURIOSA_SMI_RETURN_CODE_OK => Ok(new_memory_utilization(output)),
            err => Err(SmiError::from(err)),
        }
    }

    /// Returns a performance counter of the device.
    pub fn device_performance_counter(&self) -> SmiResult<DevicePerformanceCounter> {
        let mut output = default_device_performance_counter();

        match unsafe {
            furiosa_smi_get_device_performance_counter(
                self.raw,
                &mut output as *mut FuriosaSmiDevicePerformanceCounter,
            )
        } {
            FuriosaSmiReturnCode_FURIOSA_SMI_RETURN_CODE_OK => {
                Ok(new_device_performance_counter(output))
            }
            err => Err(SmiError::from(err)),
        }
    }

    /// Returns a power consumption of the device.
    pub fn power_consumption(&self) -> SmiResult<f64> {
        let mut output = default_power_consumption();

        match unsafe {
            furiosa_smi_get_device_power_consumption(
                self.raw,
                &mut output as *mut FuriosaSmiDevicePowerConsumption,
            )
        } {
            FuriosaSmiReturnCode_FURIOSA_SMI_RETURN_CODE_OK => Ok(output.rms_total),
            err => Err(SmiError::from(err)),
        }
    }

    /// Returns a temperature of the device.
    pub fn device_temperature(&self) -> SmiResult<DeviceTemperature> {
        let mut output = default_temperature();

        match unsafe {
            furiosa_smi_get_device_temperature(
                self.raw,
                &mut output as *mut FuriosaSmiDeviceTemperature,
            )
        } {
            FuriosaSmiReturnCode_FURIOSA_SMI_RETURN_CODE_OK => Ok(new_device_temperature(output)),
            err => Err(SmiError::from(err)),
        }
    }

    /// Returns a device link type between two devices.
    pub fn device_to_device_link_type(&self, target: &Device) -> SmiResult<LinkType> {
        let mut output =
            FuriosaSmiDeviceToDeviceLinkType_FURIOSA_SMI_DEVICE_TO_DEVICE_LINK_TYPE_UNKNOWN;

        match unsafe {
            furiosa_smi_get_device_to_device_link_type(
                self.raw,
                target.raw,
                &mut output as *mut FuriosaSmiDeviceToDeviceLinkType,
            )
        } {
            FuriosaSmiReturnCode_FURIOSA_SMI_RETURN_CODE_OK => Ok(LinkType::from(output)),
            err => Err(SmiError::from(err)),
        }
    }

    /// Returns a P2P accessibility between two devices.
    pub fn p2p_accessible(&self, target: &Device) -> SmiResult<bool> {
        let mut output = false;

        match unsafe {
            furiosa_smi_get_p2p_accessible(self.raw, target.raw, &mut output as *mut bool)
        } {
            FuriosaSmiReturnCode_FURIOSA_SMI_RETURN_CODE_OK => Ok(output),
            err => Err(SmiError::from(err)),
        }
    }
}