test_smi_rs/
lib.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
//! This is rust bindings for Furiosa System Management Interface Library of FuriosaAI's NPU devices.
//! For additional information of furiosa-smi, please refer to its [Documentation](https://furiosa-ai.github.io/furiosa-smi).
//!
//! To learn more about FuriosaAI's NPU, please visit <https://furiosa.ai>.
//!
//! # Requirements
//!
//! This crate requires FuriosaAI's NPU device and its kernel driver.
//! You can refer to
//! [Driver and Firmware Installation](https://furiosa-ai.github.io/docs/latest/en/software/installation.html)
//! to learn the kernel driver installation.
//!
//! Also, You should install Furiosa System Management Interface Library.
//! You can refer to [Installation using package manager](https://furiosa-ai.github.io/furiosa-smi#autotoc_md5) or [Installation from Source Code](https://furiosa-ai.github.io/furiosa-smi#autotoc_md6).
//!
//! # Usage
//!
//! Add this to your 'Cargo.toml':
//! ```toml
//! [dependencies]
//! furiosa-smi-rs = "0.1"
//! ```
//!
//! ## Listing devices from the system
//!
//! The current implementation mainly offers an API, namely
//! [`list_devices`].
//!
//! [`list_devices`] enumerates all Furiosa NPU devices in the system.
//! One can simply call as below:
//! ```rust,no_run
//! # fn main() -> test_smi_rs::SmiResult<()> {
//! let devices = test_smi_rs::list_devices()?;
//! # Ok(())
//! # }
//! ```
//!
//! [Struct `Device`][`Device`] offers methods for further information of each
//! device.
//!

#![allow(non_upper_case_globals)]

pub use crate::core_status::*;
pub use crate::device::*;
pub use crate::device_error_info::*;
pub use crate::device_file::*;
pub use crate::device_info::*;
pub use crate::error::*;
use crate::generated::binding::*;
pub use crate::performance::*;
pub use crate::types::*;
pub use crate::version::*;
use std::sync::{Arc, RwLock};

mod core_status;
mod device;
mod device_error_info;
mod device_file;
mod device_info;
mod error;
mod generated;
mod performance;
mod types;
mod util;
mod version;

#[allow(clippy::arc_with_non_send_sync)]
/// List all Furiosa NPU devices in the system.
pub fn list_devices() -> SmiResult<Vec<Device>> {
    let mut output = FuriosaSmiDeviceHandles {
        count: 0,
        device_handles: [0; FURIOSA_SMI_MAX_DEVICE_HANDLE_SIZE as usize],
    };

    let observer_instance = Arc::new(RwLock::new(create_observer()?));

    match unsafe { furiosa_smi_get_device_handles(&mut output as *mut FuriosaSmiDeviceHandles) } {
        FuriosaSmiReturnCode_FURIOSA_SMI_RETURN_CODE_OK => {
            let mut devices = Vec::with_capacity(output.count as usize);
            for i in 0..output.count as usize {
                devices.push(new_device(
                    output.device_handles[i],
                    observer_instance.clone(),
                ));
            }
            Ok(devices)
        }
        err => Err(SmiError::from(err)),
    }
}

/// Print a driver information of the device.
pub fn driver_info() -> SmiResult<VersionInfo> {
    let mut output = default_version_info();

    match unsafe { furiosa_smi_get_driver_info(&mut output as *mut FuriosaSmiVersion) } {
        FuriosaSmiReturnCode_FURIOSA_SMI_RETURN_CODE_OK => Ok(new_version_info(output)),
        err => Err(SmiError::from(err)),
    }
}

fn create_observer() -> SmiResult<DeviceObserverInstance> {
    let mut observer_instance: FuriosaSmiObserverInstance = default_observer_instance();

    match unsafe {
        furiosa_smi_create_observer(&mut observer_instance as *mut FuriosaSmiObserverInstance)
    } {
        FuriosaSmiReturnCode_FURIOSA_SMI_RETURN_CODE_OK => {
            Ok(new_observer_instance(observer_instance))
        }
        err => Err(SmiError::from(err)),
    }
}