ironaccelerator_neuron/
runtime.rs

1//! Thin wrappers over `nrt_load` / `nrt_execute`. The AWS Neuron
2//! Runtime takes a NEFF binary (produced by the Neuron compiler) plus
3//! two tensor sets (inputs + outputs) and runs it on one or more
4//! NeuronCores. We don't ship the compiler — we expect the caller to
5//! have a `.neff` already.
6
7use core::ffi::{c_void, CStr};
8
9use crate::drv::{self, Loaded, NrtModelHandle, NrtTensorSetHandle, NRT_SUCCESS};
10
11#[derive(Debug)]
12/// Errors from the Neuron runtime wrappers.
13pub enum NeuronError {
14    /// `libnrt` is not loaded.
15    Unavailable,
16    /// The loaded `libnrt` does not export the named symbol.
17    MissingSymbol(&'static str),
18    /// `libnrt` returned a non-success `NRT_STATUS`.
19    Runtime(u32),
20    /// A name could not be converted to a C string.
21    BadName,
22}
23
24/// Compiled Neuron model, loaded from a NEFF binary onto one or more
25/// NeuronCores.
26pub struct Model {
27    l: &'static Loaded,
28    handle: NrtModelHandle,
29}
30
31impl Model {
32    /// Load `neff` onto NeuronCores `[start_nc, start_nc + nc_count)`.
33    /// `nc_count = 1` is the typical case for a single-core inference
34    /// model; multi-core is used for tensor-parallel Trn1/Trn2 graphs.
35    pub fn load(neff: &[u8], start_nc: i32, nc_count: i32) -> Result<Self, NeuronError> {
36        let l = drv::loaded().ok_or(NeuronError::Unavailable)?;
37        let load = l.nrt_load.ok_or(NeuronError::MissingSymbol("nrt_load"))?;
38        let mut handle: NrtModelHandle = core::ptr::null_mut();
39        let r = unsafe { load(neff.as_ptr(), neff.len(), start_nc, nc_count, &mut handle) };
40        if r != NRT_SUCCESS {
41            return Err(NeuronError::Runtime(r));
42        }
43        Ok(Self { l, handle })
44    }
45
46    /// Execute this model with pre-populated input / output tensor
47    /// sets. Blocks until the NeuronCore signals completion.
48    pub fn execute(&self, inputs: &TensorSet, outputs: &TensorSet) -> Result<(), NeuronError> {
49        let exec = self
50            .l
51            .nrt_execute
52            .ok_or(NeuronError::MissingSymbol("nrt_execute"))?;
53        let r = unsafe { exec(self.handle, inputs.handle, outputs.handle) };
54        if r != NRT_SUCCESS {
55            return Err(NeuronError::Runtime(r));
56        }
57        Ok(())
58    }
59
60    /// Returns the raw `nrt_model_t` handle.
61    pub fn raw(&self) -> NrtModelHandle {
62        self.handle
63    }
64}
65
66impl Drop for Model {
67    fn drop(&mut self) {
68        if let Some(unload) = self.l.nrt_unload {
69            unsafe {
70                let _ = unload(self.handle);
71            }
72        }
73    }
74}
75
76/// A named collection of tensors — the input or output side of an
77/// execute call.
78pub struct TensorSet {
79    l: &'static Loaded,
80    handle: NrtTensorSetHandle,
81}
82
83impl TensorSet {
84    /// Allocates an empty tensor set via `nrt_allocate_tensor_set`.
85    pub fn new() -> Result<Self, NeuronError> {
86        let l = drv::loaded().ok_or(NeuronError::Unavailable)?;
87        let alloc = l
88            .nrt_alloc_tensor_set
89            .ok_or(NeuronError::MissingSymbol("nrt_allocate_tensor_set"))?;
90        let mut handle: NrtTensorSetHandle = core::ptr::null_mut();
91        let r = unsafe { alloc(&mut handle) };
92        if r != NRT_SUCCESS {
93            return Err(NeuronError::Runtime(r));
94        }
95        Ok(Self { l, handle })
96    }
97
98    /// Add a pre-allocated `nrt_tensor_t` to this set under `name`.
99    /// Tensor lifetime is managed by the caller (or by the Neuron
100    /// runtime via `nrt_tensor_allocate`, which we don't wrap yet).
101    /// # Safety
102    /// `tensor` must point to a valid `nrt_tensor_t` that outlives this set.
103    pub unsafe fn add(&mut self, name: &CStr, tensor: *mut c_void) -> Result<(), NeuronError> {
104        let add = self
105            .l
106            .nrt_add_tensor_to_set
107            .ok_or(NeuronError::MissingSymbol("nrt_add_tensor_to_tensor_set"))?;
108        let r = unsafe { add(self.handle, name.as_ptr(), tensor) };
109        if r != NRT_SUCCESS {
110            return Err(NeuronError::Runtime(r));
111        }
112        Ok(())
113    }
114
115    /// Returns the raw `nrt_tensor_set_t` handle.
116    pub fn raw(&self) -> NrtTensorSetHandle {
117        self.handle
118    }
119}
120
121impl Drop for TensorSet {
122    fn drop(&mut self) {
123        if let Some(free) = self.l.nrt_free_tensor_set {
124            unsafe {
125                let _ = free(self.handle);
126            }
127        }
128    }
129}
ironaccelerator_neuron/runtime.rs

ironaccelerator_neuron/
runtime.rs
