# 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:
```bash
neuronx-cc compile module.mlir --framework XLA --target trn2 --output module.neff
```
```rust
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.
```bash
cargo test
cargo run --release --example run_neff -- module.neff
```