ironaccelerator_neuron/
drv.rs

1//! `libnrt` loader. We need three symbols — `nrt_init`,
2//! `nrt_get_total_nc_count`, and `nrt_get_version` — and a best-effort
3//! read of the Neuron-specific env vars that AWS containers set.
4//!
5//! `drop(sym)` on a `libloading::Symbol` ends the borrow on its Library;
6//! Symbol doesn't impl Drop but the lifetime parameter does the work.
7
8#![allow(clippy::drop_non_drop)]
9
10use core::ffi::c_void;
11
12use libloading::{Library, Symbol};
13use once_cell::sync::OnceCell;
14
15/// `NRT_STATUS` return code from the Neuron runtime.
16pub type NrtStatus = u32;
17/// Success status (`NRT_SUCCESS`).
18pub const NRT_SUCCESS: NrtStatus = 0;
19
20/// Opaque `nrt_model_t` handle.
21pub type NrtModelHandle = *mut c_void;
22/// Opaque `nrt_tensor_set_t` handle.
23pub type NrtTensorSetHandle = *mut c_void;
24/// Opaque `nrt_tensor_t` handle.
25pub type NrtTensorHandle = *mut c_void;
26
27/// Signature of `nrt_load`.
28pub type NrtLoadFn = unsafe extern "C" fn(
29    neff: *const u8,
30    neff_size: usize,
31    start_nc: i32,
32    nc_count: i32,
33    model: *mut NrtModelHandle,
34) -> NrtStatus;
35/// Signature of `nrt_unload`.
36pub type NrtUnloadFn = unsafe extern "C" fn(model: NrtModelHandle) -> NrtStatus;
37/// Signature of `nrt_execute`.
38pub type NrtExecuteFn = unsafe extern "C" fn(
39    model: NrtModelHandle,
40    input_set: NrtTensorSetHandle,
41    output_set: NrtTensorSetHandle,
42) -> NrtStatus;
43/// Signature of `nrt_allocate_tensor_set`.
44pub type NrtAllocTensorSetFn = unsafe extern "C" fn(out: *mut NrtTensorSetHandle) -> NrtStatus;
45/// Signature of `nrt_free_tensor_set`.
46pub type NrtFreeTensorSetFn = unsafe extern "C" fn(set: NrtTensorSetHandle) -> NrtStatus;
47/// Signature of `nrt_add_tensor_to_tensor_set`.
48pub type NrtAddTensorToSetFn = unsafe extern "C" fn(
49    set: NrtTensorSetHandle,
50    name: *const core::ffi::c_char,
51    tensor: NrtTensorHandle,
52) -> NrtStatus;
53
54type NrtInitFn =
55    unsafe extern "C" fn(framework: u32, framework_version: *const core::ffi::c_char) -> NrtStatus;
56type NrtGetTotalNcCountFn = unsafe extern "C" fn(count: *mut u32) -> NrtStatus;
57type NrtGetVersionFn = unsafe extern "C" fn(out: *mut core::ffi::c_char, out_len: u32) -> NrtStatus;
58
59static LIB: OnceCell<Option<Loaded>> = OnceCell::new();
60
61/// Resolved `libnrt` state: core count, version, and optional entry points.
62pub struct Loaded {
63    _lib: Library,
64    total_cores: u32,
65    version: String,
66    /// `nrt_load` entry point, if the library exports it.
67    pub nrt_load: Option<NrtLoadFn>,
68    /// `nrt_unload` entry point, if the library exports it.
69    pub nrt_unload: Option<NrtUnloadFn>,
70    /// `nrt_execute` entry point, if the library exports it.
71    pub nrt_execute: Option<NrtExecuteFn>,
72    /// `nrt_allocate_tensor_set` entry point, if the library exports it.
73    pub nrt_alloc_tensor_set: Option<NrtAllocTensorSetFn>,
74    /// `nrt_free_tensor_set` entry point, if the library exports it.
75    pub nrt_free_tensor_set: Option<NrtFreeTensorSetFn>,
76    /// `nrt_add_tensor_to_tensor_set` entry point, if the library exports it.
77    pub nrt_add_tensor_to_set: Option<NrtAddTensorToSetFn>,
78}
79
80unsafe impl Send for Loaded {}
81unsafe impl Sync for Loaded {}
82
83#[cfg(any(target_os = "linux", target_os = "android"))]
84const LIB_CANDIDATES: &[&str] = &[
85    "libnrt.so.1",
86    "libnrt.so",
87    "/opt/aws/neuron/lib/libnrt.so.1",
88];
89#[cfg(not(any(target_os = "linux", target_os = "android")))]
90const LIB_CANDIDATES: &[&str] = &[];
91
92pub(crate) fn loaded() -> Option<&'static Loaded> {
93    LIB.get_or_init(load).as_ref()
94}
95
96fn load() -> Option<Loaded> {
97    for name in LIB_CANDIDATES {
98        let lib = match unsafe { Library::new(*name) } {
99            Ok(l) => l,
100            Err(_) => continue,
101        };
102        unsafe {
103            let init: Symbol<NrtInitFn> = match lib.get(b"nrt_init\0") {
104                Ok(s) => s,
105                Err(_) => continue,
106            };
107            // framework=0 (NRT_FRAMEWORK_TYPE_NONE) is accepted for
108            // non-framework embedders like us.
109            let tag = b"ironaccelerator\0" as *const _ as *const core::ffi::c_char;
110            if init(0, tag) != NRT_SUCCESS {
111                continue;
112            }
113            let count_fn: Symbol<NrtGetTotalNcCountFn> = match lib.get(b"nrt_get_total_nc_count\0")
114            {
115                Ok(s) => s,
116                Err(_) => continue,
117            };
118            let mut count: u32 = 0;
119            if count_fn(&mut count) != NRT_SUCCESS {
120                continue;
121            }
122            let version = lib
123                .get::<NrtGetVersionFn>(b"nrt_get_version\0")
124                .ok()
125                .and_then(|f| read_version(*f))
126                .unwrap_or_default();
127            drop(init);
128            drop(count_fn);
129
130            // Optional symbols — present on Neuron SDK 2.x but we don't
131            // want to fail the whole backend if one is missing on an
132            // older runtime.
133            let nrt_load = lib.get::<NrtLoadFn>(b"nrt_load\0").ok().map(|s| *s);
134            let nrt_unload = lib.get::<NrtUnloadFn>(b"nrt_unload\0").ok().map(|s| *s);
135            let nrt_execute = lib.get::<NrtExecuteFn>(b"nrt_execute\0").ok().map(|s| *s);
136            let nrt_alloc_tensor_set = lib
137                .get::<NrtAllocTensorSetFn>(b"nrt_allocate_tensor_set\0")
138                .ok()
139                .map(|s| *s);
140            let nrt_free_tensor_set = lib
141                .get::<NrtFreeTensorSetFn>(b"nrt_free_tensor_set\0")
142                .ok()
143                .map(|s| *s);
144            let nrt_add_tensor_to_set = lib
145                .get::<NrtAddTensorToSetFn>(b"nrt_add_tensor_to_tensor_set\0")
146                .ok()
147                .map(|s| *s);
148
149            return Some(Loaded {
150                _lib: lib,
151                total_cores: count,
152                version,
153                nrt_load,
154                nrt_unload,
155                nrt_execute,
156                nrt_alloc_tensor_set,
157                nrt_free_tensor_set,
158                nrt_add_tensor_to_set,
159            });
160        }
161    }
162    None
163}
164
165unsafe fn read_version(f: NrtGetVersionFn) -> Option<String> {
166    let mut buf = vec![0i8; 128];
167    let status = f(buf.as_mut_ptr() as *mut _, buf.len() as u32);
168    if status != NRT_SUCCESS {
169        return None;
170    }
171    let bytes: Vec<u8> = buf
172        .iter()
173        .take_while(|&&c| c != 0)
174        .map(|&c| c as u8)
175        .collect();
176    String::from_utf8(bytes).ok()
177}
178
179#[derive(Debug, Clone, Copy, PartialEq, Eq)]
180/// Neuron hardware generation, inferred from instance type or runtime version.
181pub enum NeuronGen {
182    /// inf1 — first-gen Inferentia.
183    Inf1,
184    /// trn1 / inf2 — second-gen Trainium / Inferentia.
185    Trn1,
186    /// trn2 — third-gen Trainium.
187    Trn2,
188    /// Generation could not be determined.
189    Unknown,
190}
191
192/// Returns `true` if `libnrt` loaded and initialised.
193pub fn is_available() -> bool {
194    loaded().is_some()
195}
196
197/// Total NeuronCores from `nrt_get_total_nc_count`, or `0` if `libnrt` is unavailable.
198pub fn total_cores() -> u32 {
199    loaded().map(|l| l.total_cores).unwrap_or(0)
200}
201
202/// Runtime version string from `nrt_get_version`, if available.
203pub fn version() -> Option<&'static str> {
204    loaded().map(|l| l.version.as_str())
205}
206
207/// Best-effort generation detection. Prefers `NEURON_INSTANCE_TYPE` env
208/// hint (commonly set by AWS Deep Learning AMIs + containers), falls
209/// back to the Neuron runtime major version.
210pub fn detect_generation() -> NeuronGen {
211    if let Ok(t) = std::env::var("NEURON_INSTANCE_TYPE") {
212        let t = t.to_ascii_lowercase();
213        if t.starts_with("trn2") {
214            return NeuronGen::Trn2;
215        }
216        if t.starts_with("trn1") || t.starts_with("inf2") {
217            return NeuronGen::Trn1;
218        }
219        if t.starts_with("inf1") {
220            return NeuronGen::Inf1;
221        }
222    }
223    match version()
224        .and_then(|v| v.split('.').next())
225        .and_then(|m| m.parse::<u32>().ok())
226    {
227        Some(n) if n >= 3 => NeuronGen::Trn2,
228        Some(2) => NeuronGen::Trn1,
229        Some(1) => NeuronGen::Inf1,
230        _ => NeuronGen::Unknown,
231    }
232}
ironaccelerator_neuron/drv.rs

ironaccelerator_neuron/
drv.rs
