neuronx 0.1.0

//! Loads a NEFF, binds zero-filled tensors to all its inputs and outputs, and
//! times a few executions.
//!
//! cargo run --release --example run_neff -- model.neff [iterations]
use neuronx::{Placement, Runtime, TensorSet, TensorUsage};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args: Vec<String> = std::env::args().collect();
    let path = args.get(1).ok_or("usage: run_neff <model.neff> [iterations]")?;
    let iterations: usize = args.get(2).map_or(Ok(10), |s| s.parse())?;

    let v = neuronx::version()?;
    println!("libnrt {}.{}.{}", v.major, v.minor, v.patch);
    let rt = Runtime::new()?;
    let neff = std::fs::read(path)?;
    let model = rt.load(&neff, -1)?;
    let infos = model.tensor_info()?;

    let tensors = infos
        .iter()
        .map(|info| {
            println!("{:?} {} {:?} {:?}", info.usage, info.name, info.dtype, info.shape);
            let mut t = rt.allocate(Placement::Device, 0, info.size, &info.name)?;
            t.memset(0, 0, info.size)?;
            Ok(t)
        })
        .collect::<neuronx::Result<Vec<_>>>()?;
    let mut inputs = TensorSet::new()?;
    let mut outputs = TensorSet::new()?;
    for (info, t) in infos.iter().zip(tensors.iter()) {
        match info.usage {
            TensorUsage::Input => inputs.add(&info.name, t)?,
            TensorUsage::Output => outputs.add(&info.name, t)?,
        }
    }

    model.execute(&inputs, &mut outputs)?;
    let start = std::time::Instant::now();
    for _ in 0..iterations {
        model.execute(&inputs, &mut outputs)?;
    }
    let elapsed = start.elapsed();
    println!(
        "{iterations} executions, {:.3}ms each",
        elapsed.as_secs_f64() * 1e3 / iterations as f64
    );
    Ok(())
}