Skip to main content

rknn_runtime/
ffi.rs

1//! Raw C FFI types and function signatures (for `librknnmrt.so` obviously).
2//!
3//! This module contains the low-level `#[repr(C)]` structs and function pointer
4//! types that mirror the RKNN C API. You normally don't need to use these
5//! directly - [`RknnModel`](crate::RknnModel) wraps everything for you.
6//!
7//! These are exposed as `pub` in case you want to do something the high-level that API doesn't support yet.
8
9use std::ffi::c_void;
10
11/// Raw RKNN context handle (an opaque u32 ID used by the C API).
12pub type RknnContext = u32;
13
14// Query commands
15pub const RKNN_QUERY_IN_OUT_NUM: u32 = 0;
16pub const RKNN_QUERY_INPUT_ATTR: u32 = 1;
17pub const RKNN_QUERY_OUTPUT_ATTR: u32 = 2;
18pub const RKNN_QUERY_SDK_VERSION: u32 = 5;
19pub const RKNN_QUERY_NATIVE_INPUT_ATTR: u32 = 8;
20pub const RKNN_QUERY_NATIVE_OUTPUT_ATTR: u32 = 9;
21pub const RKNN_QUERY_NATIVE_NHWC_INPUT_ATTR: u32 = 10;
22pub const RKNN_QUERY_NATIVE_NHWC_OUTPUT_ATTR: u32 = 11;
23
24// Tensor format constants
25pub const RKNN_TENSOR_FORMAT_NCHW: u32 = 0;
26pub const RKNN_TENSOR_FORMAT_NHWC: u32 = 1;
27pub const RKNN_TENSOR_FORMAT_NC1HWC2: u32 = 2;
28pub const RKNN_TENSOR_FORMAT_UNDEFINED: u32 = 0xFFFF;
29
30// Tensor type constants
31pub const RKNN_TENSOR_FLOAT32: u32 = 0;
32pub const RKNN_TENSOR_FLOAT16: u32 = 1;
33pub const RKNN_TENSOR_INT8: u32 = 2;
34pub const RKNN_TENSOR_UINT8: u32 = 3;
35pub const RKNN_TENSOR_INT16: u32 = 4;
36pub const RKNN_TENSOR_INT32: u32 = 5;
37
38// Quantization type constants
39pub const RKNN_QUANT_NONE: u32 = 0;
40pub const RKNN_QUANT_DFP: u32 = 1;
41pub const RKNN_QUANT_AFFINE: u32 = 2;
42
43// Memory sync flags
44pub const RKNN_MEM_SYNC_TO_DEVICE: i32 = 0;
45pub const RKNN_MEM_SYNC_FROM_DEVICE: i32 = 1;
46
47pub const RKNN_MAX_DIMS: usize = 16;
48pub const RKNN_MAX_NAME_LEN: usize = 256;
49
50// For rknn_inputs_set format/type
51pub const RKNN_TENSOR_UINT8_INPUT: u8 = 2;
52pub const RKNN_TENSOR_NHWC_INPUT: u8 = 0;
53
54#[repr(C)]
55pub struct RknnInputOutputNum {
56    pub n_input: u32,
57    pub n_output: u32,
58}
59
60#[repr(C)]
61pub struct RknnSdkVersion {
62    pub api_version: [u8; 256],
63    pub drv_version: [u8; 256],
64}
65
66#[repr(C)]
67#[derive(Clone)]
68pub struct RknnTensorAttr {
69    pub index: u32,
70    pub n_dims: u32,
71    pub dims: [u32; RKNN_MAX_DIMS],
72    pub name: [u8; RKNN_MAX_NAME_LEN],
73    pub n_elems: u32,
74    pub size: u32,
75    pub fmt: u32,
76    pub type_: u32,
77    pub qnt_type: u32,
78    pub fl: i8,
79    pub zp: i32,
80    pub scale: f32,
81    pub w_stride: u32,
82    pub size_with_stride: u32,
83    pub pass_through: u8,
84    pub h_stride: u32,
85}
86
87impl RknnTensorAttr {
88    pub fn new() -> Self {
89        unsafe { std::mem::zeroed() }
90    }
91
92    pub fn name_str(&self) -> &str {
93        std::str::from_utf8(&self.name)
94            .unwrap_or("")
95            .trim_end_matches('\0')
96    }
97
98    pub fn shape(&self) -> &[u32] {
99        &self.dims[..self.n_dims as usize]
100    }
101}
102
103impl Default for RknnTensorAttr {
104    fn default() -> Self {
105        Self::new()
106    }
107}
108
109#[repr(C)]
110pub struct RknnTensorMem {
111    pub virt_addr: *mut c_void,
112    pub phys_addr: u64,
113    pub fd: i32,
114    pub offset: i32,
115    pub size: u32,
116    pub flags: u32,
117    pub priv_data: *mut c_void,
118}
119
120#[repr(C)]
121#[derive(Clone)]
122pub struct RknnOutput {
123    pub want_float: u8,
124    pub is_prealloc: u8,
125    pub index: u32,
126    pub buf: *mut c_void,
127    pub size: u32,
128}
129
130impl RknnOutput {
131    pub fn new(index: u32) -> Self {
132        Self {
133            want_float: 0,
134            is_prealloc: 0,
135            index,
136            buf: std::ptr::null_mut(),
137            size: 0,
138        }
139    }
140}
141
142#[repr(C)]
143pub struct RknnInput {
144    pub index: u32,
145    pub buf: *mut c_void,
146    pub size: u32,
147    pub pass_through: u8,
148    pub type_: u8,
149    pub fmt: u8,
150}
151
152// Function type signatures for libloading
153pub type FnRknnInit = unsafe extern "C" fn(
154    *mut RknnContext,
155    *const c_void,
156    u32,
157    u32,
158    *const c_void,
159) -> i32;
160
161pub type FnRknnQuery =
162    unsafe extern "C" fn(RknnContext, u32, *mut c_void, u32) -> i32;
163
164pub type FnRknnRun =
165    unsafe extern "C" fn(RknnContext, *const c_void) -> i32;
166
167pub type FnRknnDestroy = unsafe extern "C" fn(RknnContext) -> i32;
168
169pub type FnRknnCreateMem =
170    unsafe extern "C" fn(RknnContext, u32) -> *mut RknnTensorMem;
171
172pub type FnRknnDestroyMem =
173    unsafe extern "C" fn(RknnContext, *mut RknnTensorMem) -> i32;
174
175pub type FnRknnSetIoMem = unsafe extern "C" fn(
176    RknnContext,
177    *mut RknnTensorMem,
178    *mut RknnTensorAttr,
179) -> i32;
180
181pub type FnRknnMemSync =
182    unsafe extern "C" fn(RknnContext, *mut RknnTensorMem, i32) -> i32;
183
184pub type FnRknnInputsSet =
185    unsafe extern "C" fn(RknnContext, u32, *mut RknnInput) -> i32;
186
187pub type FnRknnOutputsGet = unsafe extern "C" fn(
188    RknnContext,
189    u32,
190    *mut RknnOutput,
191    *const c_void,
192) -> i32;
193
194pub type FnRknnOutputsRelease =
195    unsafe extern "C" fn(RknnContext, u32, *mut RknnOutput) -> i32;