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

ironaccelerator_neuron/
runtime.rs
