neuronx-rs

Rust bindings for libnrt, the AWS Neuron runtime used on Inferentia2 and Trainium instances.

The neuronx crate is a thin layer on top of the C API: raw bindgen bindings in neuronx::sys, plus RAII wrappers for the runtime, models (NEFF files), tensors and tensor sets. It does not compile anything, use neuronx-cc to produce a NEFF, e.g. from a StableHLO module:

neuronx-cc compile module.mlir --framework XLA --target trn2 --output module.neff

use neuronx::{Placement, Runtime, TensorSet};

let rt = Runtime::new()?;
let model = rt.load(&std::fs::read("module.neff")?, -1)?;
let x = rt.tensor_from_slice(0, &[1f32, 2., 3., 4.], "x")?;
let out = rt.allocate(Placement::Device, 0, 16, "out")?;
let mut inputs = TensorSet::new()?;
inputs.add("input0", &x)?;
let mut outputs = TensorSet::new()?;
outputs.add("output0", &out)?;
model.execute(&inputs, &mut outputs)?;
let out: Vec<f32> = out.to_vec()?;

Model::tensor_info returns the names, shapes and dtypes of the NEFF inputs and outputs.

Setup

The build needs the aws-neuronx-runtime-lib package, which installs the headers and libnrt.so in /opt/aws/neuron. Set NEURON_DIR to use another location. bindgen also needs libclang.

Running needs a Neuron device and the aws-neuronx-dkms driver. The tests compile a small StableHLO module when neuronx-cc is on the PATH (or set NEURONX_CC), and skip that test otherwise.

cargo test
cargo run --release --example run_neff -- module.neff