neuronx 0.1.0

use std::marker::PhantomData;

use crate::error::check;
use crate::{sys, Error, Result, Runtime};

/// Element type of a tensor in a NEFF, mirrors `nrt_dtype_t`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum DType {
    Fp8E3,
    Fp8E4,
    Fp8E5,
    F16,
    BF16,
    F32,
    Fp32r,
    U8,
    U16,
    U32,
    U64,
    I8,
    I16,
    I32,
    I64,
}

impl DType {
    pub fn from_raw(dtype: sys::nrt_dtype_t) -> Option<Self> {
        let dtype = match dtype {
            sys::NRT_DTYPE_FP8_E3 => Self::Fp8E3,
            sys::NRT_DTYPE_FP8_E4 => Self::Fp8E4,
            sys::NRT_DTYPE_FP8_E5 => Self::Fp8E5,
            sys::NRT_DTYPE_FLOAT16 => Self::F16,
            sys::NRT_DTYPE_BFLOAT16 => Self::BF16,
            sys::NRT_DTYPE_FLOAT32 => Self::F32,
            sys::NRT_DTYPE_FP32R => Self::Fp32r,
            sys::NRT_DTYPE_UINT8 => Self::U8,
            sys::NRT_DTYPE_UINT16 => Self::U16,
            sys::NRT_DTYPE_UINT32 => Self::U32,
            sys::NRT_DTYPE_UINT64 => Self::U64,
            sys::NRT_DTYPE_INT8 => Self::I8,
            sys::NRT_DTYPE_INT16 => Self::I16,
            sys::NRT_DTYPE_INT32 => Self::I32,
            sys::NRT_DTYPE_INT64 => Self::I64,
            _ => return None,
        };
        Some(dtype)
    }

    pub fn to_raw(self) -> sys::nrt_dtype_t {
        match self {
            Self::Fp8E3 => sys::NRT_DTYPE_FP8_E3,
            Self::Fp8E4 => sys::NRT_DTYPE_FP8_E4,
            Self::Fp8E5 => sys::NRT_DTYPE_FP8_E5,
            Self::F16 => sys::NRT_DTYPE_FLOAT16,
            Self::BF16 => sys::NRT_DTYPE_BFLOAT16,
            Self::F32 => sys::NRT_DTYPE_FLOAT32,
            Self::Fp32r => sys::NRT_DTYPE_FP32R,
            Self::U8 => sys::NRT_DTYPE_UINT8,
            Self::U16 => sys::NRT_DTYPE_UINT16,
            Self::U32 => sys::NRT_DTYPE_UINT32,
            Self::U64 => sys::NRT_DTYPE_UINT64,
            Self::I8 => sys::NRT_DTYPE_INT8,
            Self::I16 => sys::NRT_DTYPE_INT16,
            Self::I32 => sys::NRT_DTYPE_INT32,
            Self::I64 => sys::NRT_DTYPE_INT64,
        }
    }

    pub fn size_in_bytes(self) -> usize {
        match self {
            Self::Fp8E3 | Self::Fp8E4 | Self::Fp8E5 | Self::U8 | Self::I8 => 1,
            Self::F16 | Self::BF16 | Self::U16 | Self::I16 => 2,
            Self::F32 | Self::Fp32r | Self::U32 | Self::I32 => 4,
            Self::U64 | Self::I64 => 8,
        }
    }
}

/// Host types that can be copied to and from tensors as raw bytes.
///
/// # Safety
/// Implementors must be plain old data: no padding and every bit pattern valid.
pub unsafe trait Element: Copy + Default + 'static {
    const DTYPE: DType;
}

macro_rules! element {
    ($ty:ty, $dtype:ident) => {
        unsafe impl Element for $ty {
            const DTYPE: DType = DType::$dtype;
        }
    };
}
element!(u8, U8);
element!(u16, U16);
element!(u32, U32);
element!(u64, U64);
element!(i8, I8);
element!(i16, I16);
element!(i32, I32);
element!(i64, I64);
element!(f32, F32);
element!(half::f16, F16);
element!(half::bf16, BF16);

fn as_bytes<T: Element>(data: &[T]) -> &[u8] {
    unsafe { std::slice::from_raw_parts(data.as_ptr() as *const u8, std::mem::size_of_val(data)) }
}

fn as_bytes_mut<T: Element>(data: &mut [T]) -> &mut [u8] {
    unsafe {
        std::slice::from_raw_parts_mut(data.as_mut_ptr() as *mut u8, std::mem::size_of_val(data))
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Placement {
    Device,
    Host,
}

/// A buffer allocated by libnrt, either in device memory or in host memory.
#[derive(Debug)]
pub struct Tensor<'rt> {
    ptr: *mut sys::nrt_tensor_t,
    _rt: PhantomData<&'rt Runtime>,
}

// libnrt tensors are not tied to the thread that allocated them.
unsafe impl Send for Tensor<'_> {}

impl Runtime {
    /// Allocates `size` bytes on the NeuronCore `vnc`, `name` is only used for debugging.
    pub fn allocate(
        &self,
        placement: Placement,
        vnc: i32,
        size: usize,
        name: &str,
    ) -> Result<Tensor<'_>> {
        let placement = match placement {
            Placement::Device => sys::NRT_TENSOR_PLACEMENT_DEVICE,
            Placement::Host => sys::NRT_TENSOR_PLACEMENT_HOST,
        };
        let name = std::ffi::CString::new(name)?;
        let mut ptr = std::ptr::null_mut();
        let status =
            unsafe { sys::nrt_tensor_allocate(placement, vnc, size, name.as_ptr(), &mut ptr) };
        check(status, "nrt_tensor_allocate")?;
        Ok(Tensor { ptr, _rt: PhantomData })
    }

    /// Allocates a device tensor on `vnc` and copies `data` to it.
    pub fn tensor_from_slice<T: Element>(
        &self,
        vnc: i32,
        data: &[T],
        name: &str,
    ) -> Result<Tensor<'_>> {
        let mut tensor =
            self.allocate(Placement::Device, vnc, std::mem::size_of_val(data), name)?;
        tensor.write(0, data)?;
        Ok(tensor)
    }
}

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

    /// Size of the tensor in bytes.
    pub fn size(&self) -> usize {
        unsafe { sys::nrt_tensor_get_size(self.ptr) }
    }

    fn check_range(&self, op: &'static str, offset: usize, len: usize) -> Result<()> {
        let size = self.size();
        if offset.checked_add(len).is_none_or(|end| end > size) {
            return Err(Error::SizeMismatch {
                op,
                expected: size.saturating_sub(offset),
                got: len,
            });
        }
        Ok(())
    }

    /// Copies `data` to the tensor, starting at byte `offset`.
    pub fn write_bytes(&mut self, offset: usize, data: &[u8]) -> Result<()> {
        self.check_range("nrt_tensor_write", offset, data.len())?;
        let status = unsafe {
            sys::nrt_tensor_write(self.ptr, data.as_ptr() as *const _, offset, data.len())
        };
        check(status, "nrt_tensor_write")
    }

    /// Copies the tensor content starting at byte `offset` to `data`.
    pub fn read_bytes(&self, offset: usize, data: &mut [u8]) -> Result<()> {
        self.check_range("nrt_tensor_read", offset, data.len())?;
        let status = unsafe {
            sys::nrt_tensor_read(self.ptr, data.as_mut_ptr() as *mut _, offset, data.len())
        };
        check(status, "nrt_tensor_read")
    }

    /// Copies `data` to the tensor, `offset` is in bytes.
    pub fn write<T: Element>(&mut self, offset: usize, data: &[T]) -> Result<()> {
        self.write_bytes(offset, as_bytes(data))
    }

    /// Copies the tensor content starting at byte `offset` to `data`.
    pub fn read<T: Element>(&self, offset: usize, data: &mut [T]) -> Result<()> {
        self.read_bytes(offset, as_bytes_mut(data))
    }

    /// Reads the whole tensor as elements of type `T`.
    pub fn to_vec<T: Element>(&self) -> Result<Vec<T>> {
        let elem_size = std::mem::size_of::<T>();
        let size = self.size();
        if !size.is_multiple_of(elem_size) {
            return Err(Error::SizeMismatch { op: "to_vec", expected: size, got: elem_size });
        }
        let mut data = vec![T::default(); size / elem_size];
        self.read(0, &mut data)?;
        Ok(data)
    }

    /// Copies `size` bytes from `src` at `src_offset` to this tensor at `dst_offset`.
    pub fn copy_from(
        &mut self,
        dst_offset: usize,
        src: &Tensor,
        src_offset: usize,
        size: usize,
    ) -> Result<()> {
        src.check_range("nrt_tensor_copy", src_offset, size)?;
        self.check_range("nrt_tensor_copy", dst_offset, size)?;
        let status =
            unsafe { sys::nrt_tensor_copy(src.ptr, src_offset, self.ptr, dst_offset, size) };
        check(status, "nrt_tensor_copy")
    }

    /// Sets `size` bytes starting at `offset` to `value`.
    pub fn memset(&mut self, offset: usize, value: u8, size: usize) -> Result<()> {
        self.check_range("nrt_tensor_memset", offset, size)?;
        let status = unsafe { sys::nrt_tensor_memset(self.ptr, offset as u64, value as i32, size) };
        check(status, "nrt_tensor_memset")
    }
}

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

/// Named tensors bound to the inputs or outputs of a model execution.
#[derive(Debug)]
pub struct TensorSet<'t> {
    ptr: *mut sys::nrt_tensor_set_t,
    _tensors: PhantomData<&'t ()>,
}

impl<'t> TensorSet<'t> {
    pub fn new() -> Result<Self> {
        let mut ptr = std::ptr::null_mut();
        check(unsafe { sys::nrt_allocate_tensor_set(&mut ptr) }, "nrt_allocate_tensor_set")?;
        Ok(Self { ptr, _tensors: PhantomData })
    }

    /// Binds `tensor` to `name`, the name of an input or output tensor of the NEFF.
    pub fn add(&mut self, name: &str, tensor: &'t Tensor) -> Result<()> {
        let name = std::ffi::CString::new(name)?;
        let status =
            unsafe { sys::nrt_add_tensor_to_tensor_set(self.ptr, name.as_ptr(), tensor.ptr) };
        check(status, "nrt_add_tensor_to_tensor_set")
    }

    pub fn as_ptr(&self) -> *mut sys::nrt_tensor_set_t {
        self.ptr
    }
}

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