1#![allow(non_camel_case_types)]
22
23use std::ffi::c_void;
24use std::sync::atomic::AtomicU64;
25
26use crate::dlpack::DLDataType;
27use crate::dlpack::DLDevice;
28
29#[repr(i32)]
31#[derive(Debug, Copy, Clone, PartialEq, Eq)]
32pub enum TVMFFITypeIndex {
33 kTVMFFINone = 0,
35 kTVMFFIInt = 1,
37 kTVMFFIBool = 2,
39 kTVMFFIFloat = 3,
41 kTVMFFIOpaquePtr = 4,
43 kTVMFFIDataType = 5,
45 kTVMFFIDevice = 6,
47 kTVMFFIDLTensorPtr = 7,
49 kTVMFFIRawStr = 8,
51 kTVMFFIByteArrayPtr = 9,
53 kTVMFFIObjectRValueRef = 10,
55 kTVMFFISmallStr = 11,
57 kTVMFFISmallBytes = 12,
59 kTVMFFIStaticObjectBegin = 64,
61 kTVMFFIStr = 65,
63 kTVMFFIBytes = 66,
65 kTVMFFIError = 67,
67 kTVMFFIFunction = 68,
69 kTVMFFIShape = 69,
71 kTVMFFITensor = 70,
73 kTVMFFIArray = 71,
75 kTVMFFIMap = 72,
80 kTVMFFIModule = 73,
82 kTVMFFIOpaquePyObject = 74,
84}
85
86#[repr(i32)]
87#[derive(Debug, Copy, Clone, PartialEq, Eq)]
88pub enum TVMFFIObjectDeleterFlagBitMask {
89 kTVMFFIObjectDeleterFlagBitMaskStrong = 1 << 0,
90 kTVMFFIObjectDeleterFlagBitMaskWeak = 1 << 1,
91 kTVMFFIObjectDeleterFlagBitMaskBoth = (1 << 0) | (1 << 1),
92}
93
94pub type TVMFFIObjectHandle = *mut c_void;
96pub type TVMFFIObjectDeleter = unsafe extern "C" fn(self_ptr: *mut c_void, flags: i32);
97
98pub const COMBINED_REF_COUNT_MASK_U32: u64 = (1u64 << 32) - 1;
100pub const COMBINED_REF_COUNT_STRONG_ONE: u64 = 1;
101pub const COMBINED_REF_COUNT_WEAK_ONE: u64 = 1u64 << 32;
102pub const COMBINED_REF_COUNT_BOTH_ONE: u64 =
103 COMBINED_REF_COUNT_STRONG_ONE | COMBINED_REF_COUNT_WEAK_ONE;
104
105#[repr(C)]
106pub struct TVMFFIObject {
107 pub combined_ref_count: AtomicU64,
108 pub type_index: i32,
109 pub __padding: u32,
110 pub deleter: Option<TVMFFIObjectDeleter>,
111 #[cfg(target_pointer_width = "32")]
113 __padding: u32,
114}
115
116impl TVMFFIObject {
117 pub fn new() -> Self {
118 Self {
119 combined_ref_count: AtomicU64::new(0),
120 type_index: 0,
121 __padding: 0,
122 deleter: None,
123 }
124 }
125}
126
127#[repr(C)]
129#[derive(Copy, Clone)]
130pub union TVMFFIAnyDataUnion {
131 pub v_int64: i64,
133 pub v_float64: f64,
135 pub v_ptr: *mut c_void,
137 pub v_c_str: *const i8,
139 pub v_obj: *mut TVMFFIObject,
141 pub v_dtype: DLDataType,
143 pub v_device: DLDevice,
145 pub v_bytes: [u8; 8],
147 pub v_uint64: u64,
149}
150
151#[repr(C)]
153#[derive(Copy, Clone)]
154pub struct TVMFFIAny {
155 pub type_index: i32,
158 pub small_str_len: u32,
160 pub data_union: TVMFFIAnyDataUnion,
162}
163
164impl TVMFFIAny {
165 pub fn new() -> Self {
167 Self {
168 type_index: TVMFFITypeIndex::kTVMFFINone as i32,
169 small_str_len: 0,
170 data_union: TVMFFIAnyDataUnion { v_int64: 0 },
171 }
172 }
173}
174
175#[repr(C)]
177pub struct TVMFFIByteArray {
178 pub data: *const u8,
179 pub size: usize,
180}
181
182impl TVMFFIByteArray {
183 pub fn new(data: *const u8, size: usize) -> Self {
184 Self { data, size }
185 }
186 pub fn as_str(&self) -> &str {
194 unsafe { std::str::from_utf8_unchecked(std::slice::from_raw_parts(self.data, self.size)) }
195 }
196 pub unsafe fn from_str(data: &str) -> Self {
206 Self {
207 data: data.as_ptr(),
208 size: data.len(),
209 }
210 }
211}
212
213pub type TVMFFISafeCallType = unsafe extern "C" fn(
215 handle: *mut c_void,
216 args: *const TVMFFIAny,
217 num_args: i32,
218 result: *mut TVMFFIAny,
219) -> i32;
220
221#[repr(C)]
223pub struct TVMFFIFunctionCell {
224 pub safe_call: TVMFFISafeCallType,
226 pub cxx_call: *mut c_void,
227}
228
229unsafe impl Send for TVMFFIFunctionCell {}
230unsafe impl Sync for TVMFFIFunctionCell {}
231
232#[repr(i32)]
233#[derive(Debug, Copy, Clone, PartialEq, Eq)]
234pub enum TVMFFIBacktraceUpdateMode {
235 kTVMFFIBacktraceUpdateModeReplace = 0,
236 kTVMFFIBacktraceUpdateModeAppend = 1,
237}
238
239#[repr(C)]
241pub struct TVMFFIErrorCell {
242 pub kind: TVMFFIByteArray,
243 pub message: TVMFFIByteArray,
244 pub backtrace: TVMFFIByteArray,
245 pub update_backtrace: unsafe extern "C" fn(
246 self_ptr: *mut c_void,
247 backtrace: *const TVMFFIByteArray,
248 update_mode: i32,
249 ),
250}
251
252#[repr(C)]
254pub struct TVMFFIShapeCell {
255 pub data: *const i64,
256 pub size: usize,
257}
258
259pub type TVMFFIFieldGetter =
261 unsafe extern "C" fn(field: *mut c_void, result: *mut TVMFFIAny) -> i32;
262
263pub type TVMFFIFieldSetter =
265 unsafe extern "C" fn(field: *mut c_void, value: *const TVMFFIAny) -> i32;
266
267#[repr(C)]
269pub struct TVMFFIFieldInfo {
270 pub name: TVMFFIByteArray,
272 pub doc: TVMFFIByteArray,
274 pub metadata: TVMFFIByteArray,
276 pub flags: i64,
278 pub size: i64,
280 pub alignment: i64,
282 pub offset: i64,
284 pub getter: Option<TVMFFIFieldGetter>,
286 pub setter: Option<TVMFFIFieldSetter>,
289 pub default_value: TVMFFIAny,
292 pub field_static_type_index: i32,
294}
295
296pub type TVMFFIObjectCreator = unsafe extern "C" fn(result: *mut TVMFFIObjectHandle) -> i32;
298
299#[repr(C)]
301pub struct TVMFFIMethodInfo {
302 pub name: TVMFFIByteArray,
304 pub doc: TVMFFIByteArray,
306 pub metadata: TVMFFIByteArray,
308 pub flags: i64,
310 pub method: TVMFFIAny,
313}
314
315#[repr(C)]
320pub struct TVMFFITypeMetadata {
321 pub doc: TVMFFIByteArray,
323 pub creator: Option<TVMFFIObjectCreator>,
325 pub total_size: i32,
329 pub structural_eq_hash_kind: i32,
331}
332
333#[repr(C)]
339pub struct TVMFFITypeAttrColumn {
340 pub data: *const TVMFFIAny,
342 pub size: usize,
344}
345
346#[repr(C)]
348pub struct TVMFFITypeInfo {
349 pub type_index: i32,
352 pub type_depth: i32,
354 pub type_key: TVMFFIByteArray,
356 pub type_acenstors: *const *const TVMFFITypeInfo,
360 pub type_key_hash: u64,
362 pub num_fields: i32,
364 pub num_methods: i32,
366 pub fields: *const TVMFFIFieldInfo,
368 pub methods: *const TVMFFIMethodInfo,
370 pub metadata: *const TVMFFITypeMetadata,
372}
373
374unsafe extern "C" {
375 pub fn TVMFFITypeKeyToIndex(type_key: *const TVMFFIByteArray, out_tindex: *mut i32) -> i32;
376 pub fn TVMFFIFunctionGetGlobal(
377 name: *const TVMFFIByteArray,
378 out: *mut TVMFFIObjectHandle,
379 ) -> i32;
380 pub fn TVMFFIFunctionSetGlobal(
381 name: *const TVMFFIByteArray,
382 f: TVMFFIObjectHandle,
383 can_override: i32,
384 ) -> i32;
385 pub fn TVMFFIFunctionCreate(
386 self_ptr: *mut c_void,
387 safe_call: TVMFFISafeCallType,
388 deleter: Option<unsafe extern "C" fn(*mut c_void)>,
389 out: *mut TVMFFIObjectHandle,
390 ) -> i32;
391 pub fn TVMFFIAnyViewToOwnedAny(any_view: *const TVMFFIAny, out: *mut TVMFFIAny) -> i32;
392 pub fn TVMFFIFunctionCall(
393 func: TVMFFIObjectHandle,
394 args: *const TVMFFIAny,
395 num_args: i32,
396 result: *mut TVMFFIAny,
397 ) -> i32;
398 pub fn TVMFFIErrorMoveFromRaised(result: *mut TVMFFIObjectHandle);
399 pub fn TVMFFIErrorSetRaised(error: TVMFFIObjectHandle);
400 pub fn TVMFFIErrorSetRaisedFromCStr(kind: *const i8, message: *const i8);
401 pub fn TVMFFIErrorCreate(
402 kind: *const TVMFFIByteArray,
403 message: *const TVMFFIByteArray,
404 backtrace: *const TVMFFIByteArray,
405 out: *mut TVMFFIObjectHandle,
406 ) -> i32;
407 pub fn TVMFFITensorFromDLPack(
408 from: *mut c_void,
409 require_alignment: i32,
410 require_contiguous: i32,
411 out: *mut TVMFFIObjectHandle,
412 ) -> i32;
413 pub fn TVMFFITensorToDLPack(from: TVMFFIObjectHandle, out: *mut *mut c_void) -> i32;
414 pub fn TVMFFITensorFromDLPackVersioned(
415 from: *mut c_void,
416 require_alignment: i32,
417 require_contiguous: i32,
418 out: *mut TVMFFIObjectHandle,
419 ) -> i32;
420 pub fn TVMFFITensorToDLPackVersioned(from: TVMFFIObjectHandle, out: *mut *mut c_void) -> i32;
421 pub fn TVMFFIStringFromByteArray(input: *const TVMFFIByteArray, out: *mut TVMFFIAny) -> i32;
422 pub fn TVMFFIBytesFromByteArray(input: *const TVMFFIByteArray, out: *mut TVMFFIAny) -> i32;
423 pub fn TVMFFIDataTypeFromString(str: *const TVMFFIByteArray, out: *mut DLDataType) -> i32;
424 pub fn TVMFFIDataTypeToString(dtype: *const DLDataType, out: *mut TVMFFIAny) -> i32;
425 pub fn TVMFFITraceback(
426 filename: *const i8,
427 lineno: i32,
428 func: *const i8,
429 cross_ffi_boundary: i32,
430 ) -> *const TVMFFIByteArray;
431 pub fn TVMFFIGetTypeInfo(type_index: i32) -> *const TVMFFITypeInfo;
432 pub fn TVMFFITestingDummyTarget() -> i32;
433}