neuronx 0.1.0

use std::marker::PhantomData;

use crate::error::check;
use crate::{c_chars_to_string, sys, DType, Result, Runtime, TensorSet};

/// A NEFF loaded on one or more NeuronCores.
#[derive(Debug)]
pub struct Model<'rt> {
    ptr: *mut sys::nrt_model_t,
    _rt: PhantomData<&'rt Runtime>,
}

// libnrt supports calling nrt_execute on the same model from several threads.
unsafe impl Send for Model<'_> {}
unsafe impl Sync for Model<'_> {}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TensorUsage {
    Input,
    Output,
}

/// Description of an input or output tensor of a NEFF.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TensorInfo {
    pub name: String,
    pub usage: TensorUsage,
    /// Size in bytes.
    pub size: usize,
    /// `None` when libnrt reports a dtype that these bindings do not know about.
    pub dtype: Option<DType>,
    pub shape: Vec<usize>,
}

impl Runtime {
    /// Loads a NEFF on the NeuronCore `vnc`, -1 lets the runtime pick the first free one.
    pub fn load(&self, neff: &[u8], vnc: i32) -> Result<Model<'_>> {
        let mut ptr = std::ptr::null_mut();
        let status =
            unsafe { sys::nrt_load(neff.as_ptr() as *const _, neff.len(), vnc, -1, &mut ptr) };
        check(status, "nrt_load")?;
        Ok(Model { ptr, _rt: PhantomData })
    }

    /// Loads a NEFF that uses collectives, `device_id` is the rank of this process among
    /// the `device_count` participants.
    pub fn load_collectives(
        &self,
        neff: &[u8],
        vnc: i32,
        device_id: u32,
        device_count: u32,
    ) -> Result<Model<'_>> {
        let mut ptr = std::ptr::null_mut();
        let status = unsafe {
            sys::nrt_load_collectives(
                neff.as_ptr() as *const _,
                neff.len(),
                vnc,
                -1,
                device_id,
                device_count,
                &mut ptr,
            )
        };
        check(status, "nrt_load_collectives")?;
        Ok(Model { ptr, _rt: PhantomData })
    }
}

impl Model<'_> {
    pub fn as_ptr(&self) -> *mut sys::nrt_model_t {
        self.ptr
    }

    /// Runs the model synchronously, all inputs and outputs of the NEFF must be bound.
    pub fn execute(&self, inputs: &TensorSet, outputs: &mut TensorSet) -> Result<()> {
        let status = unsafe { sys::nrt_execute(self.ptr, inputs.as_ptr(), outputs.as_ptr()) };
        check(status, "nrt_execute")
    }

    /// Runs the model `repeat_count` times on the same inputs and outputs.
    pub fn execute_repeat(
        &self,
        inputs: &TensorSet,
        outputs: &mut TensorSet,
        repeat_count: i32,
    ) -> Result<()> {
        let status = unsafe {
            sys::nrt_execute_repeat(self.ptr, inputs.as_ptr(), outputs.as_ptr(), repeat_count)
        };
        check(status, "nrt_execute_repeat")
    }

    /// Number of NeuronCores (logical, i.e. VNCs) used by the model.
    pub fn vnc_count(&self) -> Result<u32> {
        let mut count = 0;
        check(
            unsafe { sys::nrt_get_model_vnc_count(self.ptr, &mut count) },
            "nrt_get_model_vnc_count",
        )?;
        Ok(count)
    }

    /// Names, shapes and dtypes of the inputs and outputs of the NEFF.
    pub fn tensor_info(&self) -> Result<Vec<TensorInfo>> {
        let mut array = std::ptr::null_mut();
        check(
            unsafe { sys::nrt_get_model_tensor_info(self.ptr, &mut array) },
            "nrt_get_model_tensor_info",
        )?;
        let infos = unsafe {
            let count = (*array).tensor_count as usize;
            let infos = (*array).tensor_array.as_slice(count);
            infos.iter().map(|info| tensor_info(info)).collect()
        };
        check(unsafe { sys::nrt_free_model_tensor_info(array) }, "nrt_free_model_tensor_info")?;
        Ok(infos)
    }
}

unsafe fn tensor_info(info: &sys::nrt_tensor_info_t) -> TensorInfo {
    let shape = if info.shape.is_null() {
        vec![]
    } else {
        std::slice::from_raw_parts(info.shape, info.ndim as usize)
            .iter()
            .map(|&d| d as usize)
            .collect()
    };
    let usage = if info.usage == sys::NRT_TENSOR_USAGE_OUTPUT {
        TensorUsage::Output
    } else {
        TensorUsage::Input
    };
    TensorInfo {
        name: c_chars_to_string(&info.name),
        usage,
        size: info.size,
        dtype: DType::from_raw(info.dtype),
        shape,
    }
}

impl Drop for Model<'_> {
    fn drop(&mut self) {
        unsafe { sys::nrt_unload(self.ptr) };
    }
}