use rknpu_sys::rknn::{self, rknn_query_cmd};
use std::fs::exists;
use std::path::Path;
use std::ptr::{null_mut, addr_of_mut, from_ref};
use std::ffi::{CStr, CString, c_void};
use std::mem::{MaybeUninit, size_of_val};
use std::os::fd::BorrowedFd;
use std::borrow::Cow;
use std::ops::Deref;
use std::sync::Arc;
use super::{Error, error::Errno};
use super::error::rknn_code_to_result;
mod enums {
use super::rknn::{rknn_tensor_format, rknn_tensor_type};
#[repr(u32)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum TensorFormat {
NCHW = rknn_tensor_format::RKNN_TENSOR_NCHW.0,
NHWC = rknn_tensor_format::RKNN_TENSOR_NHWC.0,
NC1HWC2 = rknn_tensor_format::RKNN_TENSOR_NC1HWC2.0,
Undefined = rknn_tensor_format::RKNN_TENSOR_UNDEFINED.0,
}
impl From<rknn_tensor_format> for TensorFormat {
fn from(value: rknn_tensor_format) -> Self {
match value {
rknn_tensor_format::RKNN_TENSOR_NCHW => Self::NCHW,
rknn_tensor_format::RKNN_TENSOR_NHWC => Self::NHWC,
rknn_tensor_format::RKNN_TENSOR_NC1HWC2 => Self::NC1HWC2,
rknn_tensor_format::RKNN_TENSOR_UNDEFINED => Self::Undefined,
_ => Self::Undefined,
}
}
}
#[repr(u32)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum TensorDataType {
Float32 = rknn_tensor_type::RKNN_TENSOR_FLOAT32.0,
Float16 = rknn_tensor_type::RKNN_TENSOR_FLOAT16.0,
Int8 = rknn_tensor_type::RKNN_TENSOR_INT8.0,
UInt8 = rknn_tensor_type::RKNN_TENSOR_UINT8.0,
Int16 = rknn_tensor_type::RKNN_TENSOR_INT16.0,
UInt16 = rknn_tensor_type::RKNN_TENSOR_UINT16.0,
Int32 = rknn_tensor_type::RKNN_TENSOR_INT32.0,
UInt32 = rknn_tensor_type::RKNN_TENSOR_UINT32.0,
Int64 = rknn_tensor_type::RKNN_TENSOR_INT64.0,
Bool = rknn_tensor_type::RKNN_TENSOR_BOOL.0,
Int4 = rknn_tensor_type::RKNN_TENSOR_INT4.0,
BFloat16 = rknn_tensor_type::RKNN_TENSOR_BFLOAT16.0,
}
impl From<rknn_tensor_type> for TensorDataType {
fn from(value: rknn_tensor_type) -> Self {
match value {
rknn_tensor_type::RKNN_TENSOR_FLOAT32 => Self::Float32,
rknn_tensor_type::RKNN_TENSOR_FLOAT16 => Self::Float16,
rknn_tensor_type::RKNN_TENSOR_INT8 => Self::Int8,
rknn_tensor_type::RKNN_TENSOR_UINT8 => Self::UInt8,
rknn_tensor_type::RKNN_TENSOR_INT16 => Self::Int16,
rknn_tensor_type::RKNN_TENSOR_UINT16 => Self::UInt16,
rknn_tensor_type::RKNN_TENSOR_INT32 => Self::Int32,
rknn_tensor_type::RKNN_TENSOR_UINT32 => Self::UInt32,
rknn_tensor_type::RKNN_TENSOR_INT64 => Self::Int64,
rknn_tensor_type::RKNN_TENSOR_BOOL => Self::Bool,
rknn_tensor_type::RKNN_TENSOR_INT4 => Self::Int4,
rknn_tensor_type::RKNN_TENSOR_BFLOAT16 => Self::BFloat16,
_ => unreachable!(),
}
}
}
impl From<TensorDataType> for rknn_tensor_type {
fn from(value: TensorDataType) -> Self {
rknn_tensor_type(value as u32)
}
}
}
pub use enums::*;
type RawCommand = rknn::rknn_query_cmd;
pub type RawTensorAttribute = rknn::rknn_tensor_attr;
pub type RawTensorMemory = rknn::rknn_tensor_mem;
pub type RawMemorySize = rknn::rknn_mem_size;
#[derive(Debug, Clone, PartialEq, Eq)]
struct Handle(rknn::rknn_context);
impl Handle {
fn new(path_or_model_data: &[u8], model: bool) -> Result<Self, Errno> {
let mut context = MaybeUninit::uninit();
let code = unsafe {
rknn::rknn_init(
context.as_mut_ptr(),
path_or_model_data.as_ptr().cast_mut().cast(),
if model { path_or_model_data.len() as u32 } else { 0 },
0,
null_mut(),
)
};
rknn_code_to_result(code)?;
let context = unsafe { context.assume_init() };
Ok(Self(context))
}
fn query<T: Sized>(&self, command: RawCommand) -> Result<T, Errno> {
let mut result = MaybeUninit::<T>::uninit();
let code = unsafe {
rknn::rknn_query(
self.0,
command,
result.as_mut_ptr() as *mut _,
size_of_val(&result) as _,
)
};
rknn_code_to_result(code)?;
Ok(unsafe { result.assume_init() })
}
fn query_with_request<T, F>(&self, command: RawCommand, request_writer: F) -> Result<T, Errno>
where
F: FnOnce(*mut T),
{
let mut result = MaybeUninit::<T>::uninit();
request_writer(result.as_mut_ptr());
let code = unsafe {
rknn::rknn_query(
self.0,
command,
result.as_mut_ptr() as *mut _,
size_of_val(&result) as _,
)
};
rknn_code_to_result(code)?;
Ok(unsafe { result.assume_init() })
}
fn query_tensor_attribute(&self, index: usize, command: RawCommand) -> Result<RawTensorAttribute, Errno> {
let attr = self.query_with_request(command, |ptr: *mut RawTensorAttribute| {
unsafe { addr_of_mut!((*ptr).index).write(index as _); }
})?;
Ok(attr)
}
fn create_memory(&self, size: usize) -> Result<&'static RawTensorMemory, Error> {
let mem = unsafe {
rknn::rknn_create_mem(self.0, size as u32)
.as_ref::<'static>()
.ok_or(Error::AllocationFailure)?
};
Ok(mem)
}
fn set_io_memory(&self, mem: &RawTensorMemory, attr: &RawTensorAttribute) -> Result<(), Errno> {
let code = unsafe {
rknn::rknn_set_io_mem(
self.0,
std::ptr::from_ref(mem).cast_mut(),
std::ptr::from_ref(attr).cast_mut(),
)
};
rknn_code_to_result(code)?;
Ok(())
}
fn run(&self) -> Result<(), Errno> {
let code = unsafe { rknn::rknn_run(self.0, null_mut()) };
rknn_code_to_result(code)?;
Ok(())
}
fn drop_memory(&self, mem: &RawTensorMemory) -> Result<(), Errno> {
let code = unsafe {
rknn::rknn_destroy_mem(self.0, from_ref(mem).cast_mut())
};
rknn_code_to_result(code)?;
Ok(())
}
}
impl Drop for Handle {
fn drop(&mut self) {
unsafe { rknn::rknn_destroy(self.0); }
}
}
#[derive(Debug, Clone)]
pub enum Quantization {
Dfp { fractional_length: i8 },
Affine { zero_point: i32, scale: f32 },
}
impl Quantization {
pub fn quantize_i8(&self, value: f32) -> i8 {
match self {
Self::Dfp { fractional_length } => {
let factor = 2f32.powi(*fractional_length as i32);
(value * factor).round() as i8
}
Self::Affine { zero_point, scale } => {
let scaled = (value / scale).round() as i32;
(scaled + zero_point) as i8
}
}
}
pub fn dequantize_i8(&self, value: i8) -> f32 {
match self {
Self::Dfp { fractional_length } => {
let factor = 2f32.powi(*fractional_length as i32);
value as f32 / factor
}
Self::Affine { zero_point, scale } => {
(value as i32 - zero_point) as f32 * scale
}
}
}
pub fn quantize_u8(&self, value: f32) -> u8 {
match self {
Self::Dfp { fractional_length } => {
let factor = 2f32.powi(*fractional_length as i32);
(value * factor).round() as u8
}
Self::Affine { zero_point, scale } => {
let scaled = (value / scale).round() as i32;
(scaled + zero_point) as u8
}
}
}
pub fn dequantize_u8(&self, value: u8) -> f32 {
match self {
Self::Dfp { fractional_length } => {
let factor = 2f32.powi(*fractional_length as i32);
value as f32 / factor
}
Self::Affine { zero_point, scale } => {
(value as i32 - zero_point) as f32 * scale
}
}
}
}
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash)]
pub enum TensorMode {
#[default]
Default,
Native,
NativeNC1HWC2,
NativeNHWC,
}
impl TensorMode {
fn to_command(self, direction: TensorDirection) -> rknn_query_cmd {
if direction.is_input() {
match self {
Self::Default => rknn_query_cmd::RKNN_QUERY_INPUT_ATTR,
Self::Native => rknn_query_cmd::RKNN_QUERY_NATIVE_INPUT_ATTR,
Self::NativeNC1HWC2 => rknn_query_cmd::RKNN_QUERY_NATIVE_NC1HWC2_INPUT_ATTR,
Self::NativeNHWC => rknn_query_cmd::RKNN_QUERY_NATIVE_NHWC_INPUT_ATTR,
}
} else {
match self {
Self::Default => rknn_query_cmd::RKNN_QUERY_OUTPUT_ATTR,
Self::Native => rknn_query_cmd::RKNN_QUERY_NATIVE_OUTPUT_ATTR,
Self::NativeNC1HWC2 => rknn_query_cmd::RKNN_QUERY_NATIVE_NC1HWC2_OUTPUT_ATTR,
Self::NativeNHWC => rknn_query_cmd::RKNN_QUERY_NATIVE_NHWC_OUTPUT_ATTR,
}
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum TensorDirection {
Input,
Output,
}
impl TensorDirection {
pub fn is_input(&self) -> bool {
*self == TensorDirection::Input
}
pub fn is_output(&self) -> bool {
*self == TensorDirection::Output
}
}
pub struct Tensor {
handle: ManagedHandle,
direction: TensorDirection,
index: usize,
mode: TensorMode,
attr: TensorAttribute,
memory: Option<Memory>,
}
impl Tensor {
fn new(handle: ManagedHandle, direction: TensorDirection, index: usize) -> Result<Self, Errno> {
let mode = TensorMode::default();
let attr = handle.query_tensor_attribute(index, mode.to_command(direction))?;
Ok(Self {
handle,
direction,
index,
mode,
attr: attr.into(),
memory: None,
})
}
pub fn direction(&self) -> TensorDirection {
self.direction
}
pub fn is_input(&self) -> bool {
self.direction.is_input()
}
pub fn is_output(&self) -> bool {
self.direction.is_output()
}
pub fn index(&self) -> usize {
self.index
}
pub fn name<'t>(&'t self) -> Cow<'t, str> {
self.attr.name.to_string_lossy()
}
pub fn mode(&self) -> TensorMode {
self.mode
}
pub fn use_mode(&mut self, mode: TensorMode) -> Result<&TensorAttribute, Error> {
self.attr = self.handle.query_tensor_attribute(self.index, mode.to_command(self.direction))?.into();
Ok(&self.attr)
}
pub fn use_default(&mut self) -> Result<&TensorAttribute, Error> {
self.use_mode(TensorMode::Default)
}
pub fn use_native(&mut self) -> Result<&TensorAttribute, Error> {
self.use_mode(TensorMode::Native)
}
pub fn use_native_nc1hwc2(&mut self) -> Result<&TensorAttribute, Error> {
self.use_mode(TensorMode::NativeNC1HWC2)
}
pub fn use_native_nhwc(&mut self) -> Result<&TensorAttribute, Error> {
self.use_mode(TensorMode::NativeNHWC)
}
pub fn attr(&self) -> &TensorAttribute {
&self.attr
}
pub fn attr_mut(&mut self) -> &mut TensorAttribute {
&mut self.attr
}
pub fn memory(&self) -> Option<&Memory> {
self.memory.as_ref()
}
fn set_memory(&mut self, memory: Memory) -> Result<&Memory, Error> {
let raw_mem = memory.as_inner();
self.handle.set_io_memory(raw_mem, &self.attr.inner)?;
self.memory = Some(memory);
Ok(self.memory.as_ref().unwrap())
}
pub fn allocate(&mut self) -> Result<&Memory, Error> {
let memory = Memory::new(self.handle.clone(), self.attr.size_with_stride)?;
self.set_memory(memory)
}
}
impl std::fmt::Debug for Tensor {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
#[allow(dead_code)]
#[derive(Debug)]
struct Tensor<'a> {
direction: TensorDirection,
index: usize,
name: &'a CStr,
mode: TensorMode,
attr: &'a TensorAttribute,
}
let display = Tensor {
direction: self.direction,
index: self.index,
name: &self.attr.name,
mode: self.mode,
attr: &self.attr,
};
Tensor::fmt(&display, f)
}
}
#[derive(Clone)]
pub struct TensorAttribute {
inner: RawTensorAttribute,
index: usize,
dims: Vec<usize>,
name: CString,
elements: usize,
size: usize,
format: TensorFormat,
data_type: TensorDataType,
quantization: Option<Quantization>,
w_stride: usize,
size_with_stride: usize,
passthrough: bool,
h_stride: usize,
}
impl TensorAttribute {
pub fn as_inner(&self) -> &RawTensorAttribute {
&self.inner
}
pub fn index(&self) -> usize {
self.index
}
pub fn dims(&self) -> &[usize] {
self.dims.as_ref()
}
pub fn name(&self) -> Cow<'_, str> {
self.name.to_string_lossy()
}
pub fn elements(&self) -> usize {
self.elements
}
pub fn size(&self) -> usize {
self.size
}
pub fn format(&self) -> TensorFormat {
self.format
}
pub fn data_type(&self) -> TensorDataType {
self.data_type
}
pub fn set_data_type(&mut self, data_type: TensorDataType) {
self.data_type = data_type;
self.inner.type_ = data_type.into();
}
pub fn quantization(&self) -> Option<&Quantization> {
self.quantization.as_ref()
}
pub fn w_stride(&self) -> usize {
self.w_stride
}
pub fn size_with_stride(&self) -> usize {
self.size_with_stride
}
pub fn passthrough(&self) -> bool {
self.passthrough
}
pub fn h_stride(&self) -> usize {
self.h_stride
}
}
impl From<rknn::rknn_tensor_attr> for TensorAttribute {
fn from(value: rknn::rknn_tensor_attr) -> Self {
use rknn::rknn_tensor_qnt_type;
Self {
inner: value,
index: value.index as _,
dims: value.dims.into_iter()
.take(value.n_dims as _)
.map(|v| v as _)
.collect(),
name: unsafe { CStr::from_ptr(value.name.as_ptr()).to_owned() },
elements: value.n_elems as _,
size: value.size as _,
format: value.fmt.into(),
data_type: value.type_.into(),
quantization: match value.qnt_type {
rknn_tensor_qnt_type::RKNN_TENSOR_QNT_DFP => {
Some(Quantization::Dfp { fractional_length: value.fl })
}
rknn_tensor_qnt_type::RKNN_TENSOR_QNT_AFFINE_ASYMMETRIC => {
Some(Quantization::Affine { zero_point: value.zp, scale: value.scale })
}
rknn_tensor_qnt_type::RKNN_TENSOR_QNT_NONE => None,
_ => None,
},
w_stride: value.w_stride as _,
size_with_stride: value.size_with_stride as _,
passthrough: value.pass_through != 0,
h_stride: value.h_stride as _,
}
}
}
impl std::fmt::Debug for TensorAttribute {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
#[allow(dead_code)]
#[derive(Debug)]
struct TensorAttribute<'a> {
index: usize,
dims: &'a [usize],
name: &'a CStr,
elements: usize,
size: usize,
format: TensorFormat,
data_type: TensorDataType,
quantization: Option<&'a Quantization>,
w_stride: usize,
size_with_stride: usize,
passthrough: bool,
h_stride: usize,
}
let display = TensorAttribute {
index: self.index,
dims: &self.dims,
name: &self.name,
elements: self.elements,
size: self.size,
format: self.format,
data_type: self.data_type,
quantization: self.quantization.as_ref(),
w_stride: self.w_stride,
size_with_stride: self.size_with_stride,
passthrough: self.passthrough,
h_stride: self.h_stride,
};
TensorAttribute::fmt(&display, f)
}
}
pub struct Memory {
handle: ManagedHandle,
inner: &'static rknn::rknn_tensor_mem,
}
impl Memory {
fn new(handle: ManagedHandle, size: usize) -> Result<Self, Error> {
let raw_mem = handle.create_memory(size)?;
Ok(Self { handle, inner: raw_mem })
}
pub fn as_inner(&self) -> &rknn::rknn_tensor_mem {
self.inner
}
pub fn as_slice(&self) -> &[u8] {
unsafe {
std::slice::from_raw_parts(
self.as_ptr(),
self.size(),
)
}
}
pub fn as_mut_slice(&mut self) -> &mut [u8] {
unsafe {
std::slice::from_raw_parts_mut(
self.as_mut_ptr(),
self.size(),
)
}
}
pub fn as_ptr(&self) -> *const u8 {
self.as_inner().virt_addr as *const _
}
pub fn as_mut_ptr(&self) -> *mut u8 {
self.as_inner().virt_addr as *mut _
}
pub fn fd<'a>(&'a self) -> BorrowedFd<'a> {
unsafe { BorrowedFd::borrow_raw(self.inner.fd) }
}
pub fn physical(&self) -> u64 {
self.as_inner().phys_addr
}
pub fn offset(&self) -> isize {
self.as_inner().offset as _
}
pub fn size(&self) -> usize {
self.as_inner().size as _
}
pub fn private(&self) -> *mut c_void {
self.as_inner().priv_data
}
}
impl Drop for Memory {
fn drop(&mut self) {
let _ = self.handle.drop_memory(self.as_inner());
}
}
#[derive(Debug, Clone)]
pub struct MemoryUsage {
total_weight: usize,
total_internal: usize,
total_dma: usize,
total_sram: usize,
free_sram: usize,
}
impl MemoryUsage {
pub fn total_weight(&self) -> usize {
self.total_weight
}
pub fn total_internal(&self) -> usize {
self.total_internal
}
pub fn total_dma(&self) -> usize {
self.total_dma
}
pub fn total_sram(&self) -> usize {
self.total_sram
}
pub fn free_sram(&self) -> usize {
self.free_sram
}
}
impl From<RawMemorySize> for MemoryUsage {
fn from(value: rknn::rknn_mem_size) -> Self {
Self {
total_weight: value.total_weight_size as _,
total_internal: value.total_internal_size as _,
total_dma: value.total_dma_allocated_size as _,
total_sram: value.total_sram_size as _,
free_sram: value.free_sram_size as _,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct ManagedHandle(Arc<Handle>);
impl Deref for ManagedHandle {
type Target = Handle;
fn deref(&self) -> &Self::Target {
self.0.deref()
}
}
impl From<Handle> for ManagedHandle {
fn from(handle: Handle) -> Self {
Self(Arc::new(handle))
}
}
#[derive(Debug)]
pub struct Context {
handle: ManagedHandle,
inputs: Vec<Tensor>,
outputs: Vec<Tensor>,
}
impl Context {
fn ensure_path(path: &Path) -> Result<(), std::io::Error> {
if !exists(path)? {
Err(std::io::Error::from(std::io::ErrorKind::NotFound))
} else {
Ok(())
}
}
fn init(handle: Handle) -> Result<Self, Error> {
let io_num: rknn::rknn_input_output_num = handle.query(
rknn_query_cmd::RKNN_QUERY_IN_OUT_NUM
)?;
let handle: ManagedHandle = handle.into();
let mut inputs = Vec::new();
let mut outputs = Vec::new();
for i in 0..io_num.n_input as usize {
inputs.push(Tensor::new(handle.clone(), TensorDirection::Input, i)?);
}
for i in 0..io_num.n_output as usize {
outputs.push(Tensor::new(handle.clone(), TensorDirection::Output, i)?);
}
Ok(Self { handle, inputs, outputs })
}
pub fn new(path: impl AsRef<Path>) -> Result<Self, Error> {
Self::ensure_path(path.as_ref())?;
let path_str = path.as_ref().to_string_lossy();
let path_cstr = CString::new(path_str.as_bytes()).unwrap();
let handle = Handle::new(path_cstr.as_bytes(), false)?;
Self::init(handle)
}
pub fn with_model(model: impl AsRef<[u8]>) -> Result<Self, Error> {
let bytes = model.as_ref();
if bytes.is_empty() {
return Err(Error::EmptyModel);
}
let handle = Handle::new(bytes, true)?;
Self::init(handle)
}
pub fn inputs(&self) -> &[Tensor] {
&self.inputs
}
pub fn inputs_mut(&mut self) -> &mut [Tensor] {
&mut self.inputs
}
pub fn input(&self, index: usize) -> Option<&Tensor> {
self.inputs.get(index)
}
pub fn input_mut(&mut self, index: usize) -> Option<&mut Tensor> {
self.inputs.get_mut(index)
}
pub fn outputs(&self) -> &[Tensor] {
&self.outputs
}
pub fn outputs_mut(&mut self) -> &mut [Tensor] {
&mut self.outputs
}
pub fn output(&self, index: usize) -> Option<&Tensor> {
self.outputs.get(index)
}
pub fn output_mut(&mut self, index: usize) -> Option<&mut Tensor> {
self.outputs.get_mut(index)
}
pub fn memory_usage(&self) -> Result<MemoryUsage, Error> {
let mem_size: RawMemorySize = self.handle.query(rknn_query_cmd::RKNN_QUERY_MEM_SIZE)?;
Ok(mem_size.into())
}
pub fn run(&self) -> Result<(), Error> {
self.handle.run()?;
Ok(())
}
}
impl Drop for Context {
fn drop(&mut self) {
}
}