use std::sync::Mutex;
use neuronx::{Error, Placement, Runtime, TensorSet, TensorUsage};
static RUNTIME_LOCK: Mutex<()> = Mutex::new(());
fn with_runtime<F: FnOnce(&Runtime) -> neuronx::Result<()>>(f: F) -> neuronx::Result<()> {
let _guard = RUNTIME_LOCK.lock().unwrap_or_else(|e| e.into_inner());
let rt = Runtime::new()?;
f(&rt)
}
#[test]
fn version_and_instance() -> neuronx::Result<()> {
let v = neuronx::version()?;
eprintln!("libnrt {}.{}.{} ({})", v.major, v.minor, v.patch, v.git_hash);
assert_eq!(v.major, 2);
let info = neuronx::instance_info()?;
eprintln!("{info:?}");
Ok(())
}
#[test]
fn single_runtime() -> neuronx::Result<()> {
with_runtime(|rt| {
assert!(rt.visible_vnc_count()? > 0);
assert!(matches!(Runtime::new(), Err(Error::AlreadyInitialized)));
Ok(())
})
}
#[test]
fn tensor_roundtrip() -> neuronx::Result<()> {
with_runtime(|rt| {
let data: Vec<f32> = (0..16).map(|i| i as f32 * 0.5).collect();
let src = rt.tensor_from_slice(0, &data, "src")?;
assert_eq!(src.size(), 64);
assert_eq!(src.to_vec::<f32>()?, data);
let mut dst = rt.allocate(Placement::Host, 0, 64, "dst")?;
dst.memset(0, 0, 64)?;
dst.copy_from(16, &src, 0, 32)?;
let dst = dst.to_vec::<f32>()?;
assert_eq!(&dst[..4], &[0.; 4]);
assert_eq!(&dst[4..12], &data[..8]);
assert_eq!(&dst[12..], &[0.; 4]);
let mut small = rt.allocate(Placement::Device, 0, 8, "small")?;
assert!(matches!(small.write(0, &[1u32, 2, 3]), Err(Error::SizeMismatch { .. })));
Ok(())
})
}
fn compile_neff(mlir: &str) -> Option<Vec<u8>> {
let cc = std::env::var("NEURONX_CC").unwrap_or_else(|_| "neuronx-cc".to_string());
let target = std::env::var("NEURONX_CC_TARGET").ok().unwrap_or_else(|| {
use neuronx::sys::*;
match neuronx::instance_info().unwrap().family {
NRT_INSTANCE_INF2 | NRT_INSTANCE_INF2E => "inf2",
NRT_INSTANCE_TRN1 | NRT_INSTANCE_TRN1N => "trn1",
NRT_INSTANCE_TRN3
| NRT_INSTANCE_TRN3PDS98
| NRT_INSTANCE_TRN3E
| NRT_INSTANCE_TRN3N
| NRT_INSTANCE_TRN3EN => "trn3",
_ => "trn2",
}
.to_string()
});
static COUNTER: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
let idx = COUNTER.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
let dir = std::env::temp_dir().join(format!("neuronx-test-{}-{idx}", std::process::id()));
std::fs::create_dir_all(&dir).unwrap();
std::fs::write(dir.join("module.mlir"), mlir).unwrap();
let output = std::process::Command::new(&cc)
.current_dir(&dir)
.args(["compile", "module.mlir", "--framework", "XLA", "--target", &target])
.args(["--output", "module.neff"])
.output();
let neff = match output {
Err(err) => {
eprintln!("cannot run {cc}, skipping: {err}");
None
}
Ok(output) if !output.status.success() => {
let stdout = String::from_utf8_lossy(&output.stdout);
panic!("{cc} failed:\n{stdout}\n{}", String::from_utf8_lossy(&output.stderr))
}
Ok(_) => Some(std::fs::read(dir.join("module.neff")).unwrap()),
};
let _ = std::fs::remove_dir_all(&dir);
neff
}
#[test]
fn execute_stablehlo() -> neuronx::Result<()> {
let mlir = r#"
module @add_mul {
func.func @main(%x: tensor<2x4xf32>, %y: tensor<2x4xf32>) -> tensor<2x4xf32> {
%0 = stablehlo.add %x, %y : tensor<2x4xf32>
%1 = stablehlo.multiply %0, %y : tensor<2x4xf32>
return %1 : tensor<2x4xf32>
}
}"#;
let Some(neff) = compile_neff(mlir) else { return Ok(()) };
with_runtime(|rt| {
let model = rt.load(&neff, -1)?;
let infos = model.tensor_info()?;
eprintln!("{infos:#?}");
let inputs_info: Vec<_> = infos.iter().filter(|i| i.usage == TensorUsage::Input).collect();
let outputs_info: Vec<_> =
infos.iter().filter(|i| i.usage == TensorUsage::Output).collect();
assert_eq!(inputs_info.len(), 2);
assert_eq!(outputs_info.len(), 1);
assert_eq!(outputs_info[0].shape, [2, 4]);
let x: Vec<f32> = (0..8).map(|i| i as f32).collect();
let y: Vec<f32> = (0..8).map(|i| 1. + i as f32 * 0.25).collect();
let x_t = rt.tensor_from_slice(0, &x, "x")?;
let y_t = rt.tensor_from_slice(0, &y, "y")?;
let out_t = rt.allocate(Placement::Device, 0, outputs_info[0].size, "out")?;
let mut inputs = TensorSet::new()?;
inputs.add(&inputs_info[0].name, &x_t)?;
inputs.add(&inputs_info[1].name, &y_t)?;
let mut outputs = TensorSet::new()?;
outputs.add(&outputs_info[0].name, &out_t)?;
model.execute(&inputs, &mut outputs)?;
let out = out_t.to_vec::<f32>()?;
let expected: Vec<f32> = x.iter().zip(y.iter()).map(|(x, y)| (x + y) * y).collect();
assert_eq!(out, expected);
Ok(())
})
}
#[test]
fn execute_half() -> neuronx::Result<()> {
use neuronx::half::{bf16, f16};
let mlir = r#"
module @half {
func.func @main(%x: tensor<8xf16>, %y: tensor<8xbf16>) -> (tensor<8xf16>, tensor<8xbf16>) {
%0 = stablehlo.add %x, %x : tensor<8xf16>
%1 = stablehlo.multiply %y, %y : tensor<8xbf16>
return %0, %1 : tensor<8xf16>, tensor<8xbf16>
}
}"#;
let Some(neff) = compile_neff(mlir) else { return Ok(()) };
with_runtime(|rt| {
let model = rt.load(&neff, -1)?;
let infos = model.tensor_info()?;
let dtypes: Vec<_> = infos.iter().map(|i| (i.name.as_str(), i.dtype)).collect();
assert_eq!(
dtypes,
[
("input0", Some(neuronx::DType::F16)),
("input1", Some(neuronx::DType::BF16)),
("output0", Some(neuronx::DType::F16)),
("output1", Some(neuronx::DType::BF16)),
]
);
let x: Vec<f16> = (0..8).map(|i| f16::from_f32(i as f32 * 0.5)).collect();
let y: Vec<bf16> = (0..8).map(|i| bf16::from_f32(i as f32 - 3.)).collect();
let x_t = rt.tensor_from_slice(0, &x, "x")?;
let y_t = rt.tensor_from_slice(0, &y, "y")?;
let out0 = rt.allocate(Placement::Device, 0, 16, "out0")?;
let out1 = rt.allocate(Placement::Device, 0, 16, "out1")?;
let mut inputs = TensorSet::new()?;
inputs.add("input0", &x_t)?;
inputs.add("input1", &y_t)?;
let mut outputs = TensorSet::new()?;
outputs.add("output0", &out0)?;
outputs.add("output1", &out1)?;
model.execute(&inputs, &mut outputs)?;
let expected0: Vec<f16> = x.iter().map(|&v| v + v).collect();
let expected1: Vec<bf16> = y.iter().map(|&v| v * v).collect();
assert_eq!(out0.to_vec::<f16>()?, expected0);
assert_eq!(out1.to_vec::<bf16>()?, expected1);
Ok(())
})
}