tvm_ffi_sys/
c_api.rs

1/*
2 * Licensed to the Apache Software Foundation (ASF) under one
3 * or more contributor license agreements.  See the NOTICE file
4 * distributed with this work for additional information
5 * regarding copyright ownership.  The ASF licenses this file
6 * to you under the Apache License, Version 2.0 (the
7 * "License"); you may not use this file except in compliance
8 * with the License.  You may obtain a copy of the License at
9 *
10 *   http://www.apache.org/licenses/LICENSE-2.0
11 *
12 * Unless required by applicable law or agreed to in writing,
13 * software distributed under the License is distributed on an
14 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15 * KIND, either express or implied.  See the License for the
16 * specific language governing permissions and limitations
17 * under the License.
18 */
19// NOTE: we manually write the C ABI as they are reasonably minimal
20// and we need to ensure clear control of the atomic access etc.
21#![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///  The index type of the FFI objects
30#[repr(i32)]
31#[derive(Debug, Copy, Clone, PartialEq, Eq)]
32pub enum TVMFFITypeIndex {
33    /// None/nullptr value
34    kTVMFFINone = 0,
35    /// POD int value
36    kTVMFFIInt = 1,
37    /// POD bool value
38    kTVMFFIBool = 2,
39    /// POD float value
40    kTVMFFIFloat = 3,
41    /// Opaque pointer object
42    kTVMFFIOpaquePtr = 4,
43    /// DLDataType
44    kTVMFFIDataType = 5,
45    /// DLDevice
46    kTVMFFIDevice = 6,
47    /// DLTensor*
48    kTVMFFIDLTensorPtr = 7,
49    /// const char*
50    kTVMFFIRawStr = 8,
51    /// TVMFFIByteArray*
52    kTVMFFIByteArrayPtr = 9,
53    /// R-value reference to ObjectRef
54    kTVMFFIObjectRValueRef = 10,
55    /// Small string on stack
56    kTVMFFISmallStr = 11,
57    /// Small bytes on stack
58    kTVMFFISmallBytes = 12,
59    /// Structural-mutation marker indicating that no new value was produced
60    kTVMFFIUnchanged = 13,
61    /// Start of statically defined objects.
62    kTVMFFIStaticObjectBegin = 64,
63    /// String object, layout = { TVMFFIObject, TVMFFIByteArray, ... }
64    kTVMFFIStr = 65,
65    /// Bytes object, layout = { TVMFFIObject, TVMFFIByteArray, ... }
66    kTVMFFIBytes = 66,
67    /// Error object.
68    kTVMFFIError = 67,
69    /// Function object.
70    kTVMFFIFunction = 68,
71    /// Shape object, layout = { TVMFFIObject, { const int64_t*, size_t }, ... }
72    kTVMFFIShape = 69,
73    /// Tensor object, layout = { TVMFFIObject, DLTensor, ... }
74    kTVMFFITensor = 70,
75    /// Array object.
76    kTVMFFIArray = 71,
77    //----------------------------------------------------------------
78    // more complex objects
79    //----------------------------------------------------------------
80    /// Map object.
81    kTVMFFIMap = 72,
82    /// Runtime dynamic loaded module object.
83    kTVMFFIModule = 73,
84    /// Opaque python object.
85    kTVMFFIOpaquePyObject = 74,
86    /// Mutable list object.
87    kTVMFFIList = 75,
88    /// Mutable dict object.
89    kTVMFFIDict = 76,
90    /// Structural visit interrupt object.
91    kTVMFFIVisitInterrupt = 77,
92    /// Arbitrary-precision integer object.
93    kTVMFFIBigInt = 78,
94    /// End of the statically allocated object type-index range.
95    kTVMFFIStaticObjectEnd = 79,
96    /// Start of dynamically allocated object type indices.
97    kTVMFFIDynObjectBegin = 128,
98}
99
100#[repr(i32)]
101#[derive(Debug, Copy, Clone, PartialEq, Eq)]
102pub enum TVMFFIObjectDeleterFlagBitMask {
103    kTVMFFIObjectDeleterFlagBitMaskStrong = 1 << 0,
104    kTVMFFIObjectDeleterFlagBitMaskWeak = 1 << 1,
105    kTVMFFIObjectDeleterFlagBitMaskBoth = (1 << 0) | (1 << 1),
106}
107
108/// Bit flags attached to reflected fields.
109#[repr(i32)]
110#[derive(Debug, Copy, Clone, PartialEq, Eq)]
111pub enum TVMFFIFieldFlagBitMask {
112    kTVMFFIFieldFlagBitMaskWritable = 1 << 0,
113    kTVMFFIFieldFlagBitMaskHasDefault = 1 << 1,
114    kTVMFFIFieldFlagBitMaskIsStaticMethod = 1 << 2,
115    kTVMFFIFieldFlagBitMaskSEqHashIgnore = 1 << 3,
116    kTVMFFIFieldFlagBitMaskSEqHashDefPattern = 1 << 4,
117    kTVMFFIFieldFlagBitMaskDefaultFromFactory = 1 << 5,
118    kTVMFFIFieldFlagBitMaskReprOff = 1 << 6,
119    kTVMFFIFieldFlagBitMaskCompareOff = 1 << 7,
120    kTVMFFIFieldFlagBitMaskHashOff = 1 << 8,
121    kTVMFFIFieldFlagBitMaskInitOff = 1 << 9,
122    kTVMFFIFieldFlagBitMaskKwOnly = 1 << 10,
123    kTVMFFIFieldFlagBitSetterIsFunctionObj = 1 << 11,
124    kTVMFFIFieldFlagBitMaskSEqHashDefSimple = 1 << 12,
125}
126
127/// Definition-region mode used by structural traversal.
128#[repr(i32)]
129#[derive(Debug, Copy, Clone, PartialEq, Eq)]
130pub enum TVMFFIDefRegionKind {
131    kTVMFFIDefRegionKindNone = 0,
132    kTVMFFIDefRegionKindPattern = 1,
133    kTVMFFIDefRegionKindSimple = 2,
134}
135
136/// Structural equality/hash participation kind stored in type metadata.
137#[repr(i32)]
138#[derive(Debug, Copy, Clone, PartialEq, Eq)]
139pub enum TVMFFISEqHashKind {
140    kTVMFFISEqHashKindUnsupported = 0,
141    kTVMFFISEqHashKindTreeNode = 1,
142    kTVMFFISEqHashKindFreeVar = 2,
143    kTVMFFISEqHashKindDAGNode = 3,
144    kTVMFFISEqHashKindConstTreeNode = 4,
145    kTVMFFISEqHashKindUniqueInstance = 5,
146}
147
148/// Handle to Object from C API's pov
149pub type TVMFFIObjectHandle = *mut c_void;
150pub type TVMFFIObjectDeleter = unsafe extern "C" fn(self_ptr: *mut c_void, flags: i32);
151
152// constants for working with combined reference count
153pub const COMBINED_REF_COUNT_MASK_U32: u64 = (1u64 << 32) - 1;
154pub const COMBINED_REF_COUNT_STRONG_ONE: u64 = 1;
155pub const COMBINED_REF_COUNT_WEAK_ONE: u64 = 1u64 << 32;
156pub const COMBINED_REF_COUNT_BOTH_ONE: u64 =
157    COMBINED_REF_COUNT_STRONG_ONE | COMBINED_REF_COUNT_WEAK_ONE;
158
159#[repr(C)]
160pub struct TVMFFIObject {
161    pub combined_ref_count: AtomicU64,
162    pub type_index: i32,
163    pub __padding: u32,
164    pub deleter: Option<TVMFFIObjectDeleter>,
165    // private padding to ensure 8 bytes alignment
166    #[cfg(target_pointer_width = "32")]
167    __padding: u32,
168}
169
170impl TVMFFIObject {
171    pub fn new() -> Self {
172        Self {
173            combined_ref_count: AtomicU64::new(0),
174            type_index: 0,
175            __padding: 0,
176            deleter: None,
177        }
178    }
179}
180
181/// Second union in TVMFFIAny - 8 bytes
182#[repr(C)]
183#[derive(Copy, Clone)]
184pub union TVMFFIAnyDataUnion {
185    /// Integers
186    pub v_int64: i64,
187    /// Floating-point numbers
188    pub v_float64: f64,
189    /// Typeless pointers
190    pub v_ptr: *mut c_void,
191    /// Raw C-string
192    pub v_c_str: *const i8,
193    /// Ref counted objects
194    pub v_obj: *mut TVMFFIObject,
195    /// Data type
196    pub v_dtype: DLDataType,
197    /// Device
198    pub v_device: DLDevice,
199    /// Small string
200    pub v_bytes: [u8; 8],
201    /// uint64 repr mainly used for hashing
202    pub v_uint64: u64,
203}
204
205/// TVM FFI Any value - a union type that can hold various data types
206#[repr(C)]
207#[derive(Copy, Clone)]
208pub struct TVMFFIAny {
209    /// Type index of the object.
210    /// The type index of Object and Any are shared in FFI.
211    pub type_index: i32,
212    /// small string length or zero padding
213    pub small_str_len: u32,
214    /// data union - 8 bytes
215    pub data_union: TVMFFIAnyDataUnion,
216}
217
218impl TVMFFIAny {
219    /// create a new instance of TVMFFIAny that represents None
220    pub fn new() -> Self {
221        Self {
222            type_index: TVMFFITypeIndex::kTVMFFINone as i32,
223            small_str_len: 0,
224            data_union: TVMFFIAnyDataUnion { v_int64: 0 },
225        }
226    }
227}
228
229/// Byte array data structure used by String and Bytes.
230#[repr(C)]
231pub struct TVMFFIByteArray {
232    pub data: *const u8,
233    pub size: usize,
234}
235
236impl TVMFFIByteArray {
237    pub fn new(data: *const u8, size: usize) -> Self {
238        Self { data, size }
239    }
240    /// Convert the TVMFFIByteArray to a str view
241    ///
242    /// # Arguments
243    /// * `self` - The TVMFFIByteArray to convert.
244    ///
245    /// # Returns
246    /// * `&str` - The converted str view.
247    pub fn as_str(&self) -> &str {
248        unsafe { std::str::from_utf8_unchecked(std::slice::from_raw_parts(self.data, self.size)) }
249    }
250    /// Unsafe function to create a TVMFFIByteArray from a string
251    /// This function is unsafe as it does not check lifetime of the string
252    /// the caller must ensure that the string is valid for the lifetime of the TVMFFIByteArray
253    ///
254    /// # Arguments
255    /// * `data` - The string to create the TVMFFIByteArray from.
256    ///
257    /// # Returns
258    /// * `TVMFFIByteArray` - The created TVMFFIByteArray.
259    pub unsafe fn from_str(data: &str) -> Self {
260        Self {
261            data: data.as_ptr(),
262            size: data.len(),
263        }
264    }
265}
266
267/// Safe call type for function ABI
268pub type TVMFFISafeCallType = unsafe extern "C" fn(
269    handle: *mut c_void,
270    args: *const TVMFFIAny,
271    num_args: i32,
272    result: *mut TVMFFIAny,
273) -> i32;
274
275/// Function cell
276#[repr(C)]
277pub struct TVMFFIFunctionCell {
278    /// A C API compatible call with exception catching.
279    pub safe_call: TVMFFISafeCallType,
280    pub cxx_call: *mut c_void,
281}
282
283unsafe impl Send for TVMFFIFunctionCell {}
284unsafe impl Sync for TVMFFIFunctionCell {}
285
286#[repr(i32)]
287#[derive(Debug, Copy, Clone, PartialEq, Eq)]
288pub enum TVMFFIBacktraceUpdateMode {
289    kTVMFFIBacktraceUpdateModeReplace = 0,
290    kTVMFFIBacktraceUpdateModeAppend = 1,
291}
292
293/// Error cell used in error object following header.
294#[repr(C)]
295pub struct TVMFFIErrorCell {
296    pub kind: TVMFFIByteArray,
297    pub message: TVMFFIByteArray,
298    pub backtrace: TVMFFIByteArray,
299    pub update_backtrace: unsafe extern "C" fn(
300        self_ptr: *mut c_void,
301        backtrace: *const TVMFFIByteArray,
302        update_mode: i32,
303    ),
304}
305
306/// Shape cell used in shape object following header.
307#[repr(C)]
308pub struct TVMFFIShapeCell {
309    pub data: *const i64,
310    pub size: usize,
311}
312
313/// Field getter function pointer type
314pub type TVMFFIFieldGetter =
315    unsafe extern "C" fn(field: *mut c_void, result: *mut TVMFFIAny) -> i32;
316
317/// Field setter function pointer type
318pub type TVMFFIFieldSetter =
319    unsafe extern "C" fn(field: *mut c_void, value: *const TVMFFIAny) -> i32;
320
321/// Information support for optional object reflection
322#[repr(C)]
323pub struct TVMFFIFieldInfo {
324    /// The name of the field
325    pub name: TVMFFIByteArray,
326    /// The docstring about the field
327    pub doc: TVMFFIByteArray,
328    /// The metadata of the field in JSON string
329    pub metadata: TVMFFIByteArray,
330    /// bitmask flags of the field
331    pub flags: i64,
332    /// The size of the field
333    pub size: i64,
334    /// The alignment of the field
335    pub alignment: i64,
336    /// The offset of the field
337    pub offset: i64,
338    /// The getter to access the field
339    pub getter: Option<TVMFFIFieldGetter>,
340    /// The setter to access the field.
341    ///
342    /// When kTVMFFIFieldFlagBitSetterIsFunctionObj is NOT set (default),
343    /// this is a TVMFFIFieldSetter function pointer cast to *mut c_void.
344    /// When kTVMFFIFieldFlagBitSetterIsFunctionObj IS set,
345    /// this is a TVMFFIObjectHandle pointing to a FunctionObj.
346    ///
347    /// The setter is set even if the field is readonly for serialization.
348    pub setter: *mut c_void,
349    /// The default value or factory of the field, this field holds AnyView.
350    /// Valid when flags set kTVMFFIFieldFlagBitMaskHasDefault.
351    /// When kTVMFFIFieldFlagBitMaskDefaultFromFactory is also set,
352    /// this is a callable factory function () -> Any.
353    pub default_value_or_factory: TVMFFIAny,
354    /// Records the compile-time static type kind of the field.
355    pub field_static_type_index: i32,
356}
357
358/// Object creator function pointer type
359pub type TVMFFIObjectCreator = unsafe extern "C" fn(result: *mut TVMFFIObjectHandle) -> i32;
360
361/// Method information that can appear in reflection table
362#[repr(C)]
363pub struct TVMFFIMethodInfo {
364    /// The name of the field
365    pub name: TVMFFIByteArray,
366    /// The docstring about the method
367    pub doc: TVMFFIByteArray,
368    /// Optional metadata of the method in JSON string
369    pub metadata: TVMFFIByteArray,
370    /// bitmask flags of the method
371    pub flags: i64,
372    /// The method wrapped as ffi::Function, stored as AnyView
373    /// The first argument to the method is always the self for instance methods
374    pub method: TVMFFIAny,
375}
376
377/// Extra information of object type that can be used for reflection
378///
379/// This information is optional and can be used to enable reflection based
380/// creation of the object.
381#[repr(C)]
382pub struct TVMFFITypeMetadata {
383    /// The docstring about the object
384    pub doc: TVMFFIByteArray,
385    /// An optional function that can create a new empty instance of the type
386    pub creator: Option<TVMFFIObjectCreator>,
387    /// Total size of the object struct, if it is fixed and known
388    ///
389    /// This field is set optional and set to 0 if not registered.
390    pub total_size: i32,
391    /// Optional meta-data for structural eq/hash
392    pub structural_eq_hash_kind: i32,
393}
394
395/// Column array that stores extra attributes about types
396///
397/// The attributes stored in a column array that can be looked up by type index.
398/// Note that the TypeAttr behaves like type_traits so column T so not contain
399/// attributes from base classes.
400#[repr(C)]
401pub struct TVMFFITypeAttrColumn {
402    /// The data of the column, indexed by (type_index - begin_index).
403    pub data: *const TVMFFIAny,
404    /// The number of elements in the data array.
405    /// The column covers type indices [begin_index, begin_index + size).
406    pub size: i32,
407    /// The starting type index of the column data.
408    /// Lookup: if begin_index <= type_index < begin_index + size,
409    /// the entry is data[(type_index - begin_index) as usize].
410    pub begin_index: i32,
411}
412
413/// Runtime type information for object type checking
414#[repr(C)]
415pub struct TVMFFITypeInfo {
416    /// The runtime type index
417    /// It can be allocated during runtime if the type is dynamic
418    pub type_index: i32,
419    /// number of parent types in the type hierachy
420    pub type_depth: i32,
421    /// the unique type key to identify the type
422    pub type_key: TVMFFIByteArray,
423    /// `type_acenstors[depth]` stores the type_index of the acenstors at depth level
424    /// To keep things simple, we do not allow multiple inheritance so the
425    /// hieracy stays as a tree
426    pub type_acenstors: *const *const TVMFFITypeInfo,
427    /// Cached hash value of the type key, used for consistent structural hashing
428    pub type_key_hash: u64,
429    /// number of reflection accessible fields
430    pub num_fields: i32,
431    /// number of reflection acccesible methods
432    pub num_methods: i32,
433    /// The reflection field information
434    pub fields: *const TVMFFIFieldInfo,
435    /// The reflection method
436    pub methods: *const TVMFFIMethodInfo,
437    /// The extra information of the type
438    pub metadata: *const TVMFFITypeMetadata,
439}
440
441/// Mandatory header preceding each Object body. See `TVMFFIObjectAllocHeader`
442/// in `tvm/ffi/c_api.h`.
443#[repr(C)]
444pub struct TVMFFIObjectAllocHeader {
445    pub delete_space: Option<unsafe extern "C" fn(ptr: *mut c_void)>,
446}
447
448/// Custom allocator entry. See `TVMFFICustomAllocator` in `tvm/ffi/c_api.h`.
449#[repr(C)]
450pub struct TVMFFICustomAllocator {
451    pub allocate: Option<
452        unsafe extern "C" fn(
453            size: usize,
454            alignment: usize,
455            type_index: i32,
456            context: *mut c_void,
457        ) -> *mut c_void,
458    >,
459    pub context: *mut c_void,
460}
461
462unsafe extern "C" {
463    pub fn TVMFFIGetCustomAllocator() -> *mut TVMFFICustomAllocator;
464    pub fn TVMFFISetCustomAllocator(allocator: *mut TVMFFICustomAllocator) -> i32;
465
466    pub fn TVMFFITypeKeyToIndex(type_key: *const TVMFFIByteArray, out_tindex: *mut i32) -> i32;
467    pub fn TVMFFITypeRegisterAttr(
468        type_index: i32,
469        attr_name: *const TVMFFIByteArray,
470        attr_value: *const TVMFFIAny,
471    ) -> i32;
472    pub fn TVMFFIGetTypeAttrColumn(
473        attr_name: *const TVMFFIByteArray,
474    ) -> *const TVMFFITypeAttrColumn;
475    pub fn TVMFFIFunctionGetGlobal(
476        name: *const TVMFFIByteArray,
477        out: *mut TVMFFIObjectHandle,
478    ) -> i32;
479    pub fn TVMFFIFunctionSetGlobal(
480        name: *const TVMFFIByteArray,
481        f: TVMFFIObjectHandle,
482        can_override: i32,
483    ) -> i32;
484    pub fn TVMFFIFunctionCreate(
485        self_ptr: *mut c_void,
486        safe_call: TVMFFISafeCallType,
487        deleter: Option<unsafe extern "C" fn(*mut c_void)>,
488        out: *mut TVMFFIObjectHandle,
489    ) -> i32;
490    pub fn TVMFFIAnyViewToOwnedAny(any_view: *const TVMFFIAny, out: *mut TVMFFIAny) -> i32;
491    pub fn TVMFFIFunctionCall(
492        func: TVMFFIObjectHandle,
493        args: *const TVMFFIAny,
494        num_args: i32,
495        result: *mut TVMFFIAny,
496    ) -> i32;
497    pub fn TVMFFIErrorMoveFromRaised(result: *mut TVMFFIObjectHandle);
498    pub fn TVMFFIErrorSetRaised(error: TVMFFIObjectHandle);
499    pub fn TVMFFIErrorSetRaisedFromCStr(kind: *const i8, message: *const i8);
500    pub fn TVMFFIErrorCreate(
501        kind: *const TVMFFIByteArray,
502        message: *const TVMFFIByteArray,
503        backtrace: *const TVMFFIByteArray,
504        out: *mut TVMFFIObjectHandle,
505    ) -> i32;
506    pub fn TVMFFITensorFromDLPack(
507        from: *mut c_void,
508        require_alignment: i32,
509        require_contiguous: i32,
510        out: *mut TVMFFIObjectHandle,
511    ) -> i32;
512    pub fn TVMFFITensorToDLPack(from: TVMFFIObjectHandle, out: *mut *mut c_void) -> i32;
513    pub fn TVMFFITensorFromDLPackVersioned(
514        from: *mut c_void,
515        require_alignment: i32,
516        require_contiguous: i32,
517        out: *mut TVMFFIObjectHandle,
518    ) -> i32;
519    pub fn TVMFFITensorToDLPackVersioned(from: TVMFFIObjectHandle, out: *mut *mut c_void) -> i32;
520    pub fn TVMFFIStringFromByteArray(input: *const TVMFFIByteArray, out: *mut TVMFFIAny) -> i32;
521    pub fn TVMFFIBytesFromByteArray(input: *const TVMFFIByteArray, out: *mut TVMFFIAny) -> i32;
522    pub fn TVMFFIDataTypeFromString(str: *const TVMFFIByteArray, out: *mut DLDataType) -> i32;
523    pub fn TVMFFIDataTypeToString(dtype: *const DLDataType, out: *mut TVMFFIAny) -> i32;
524    pub fn TVMFFITraceback(
525        filename: *const i8,
526        lineno: i32,
527        func: *const i8,
528        cross_ffi_boundary: i32,
529    ) -> *const TVMFFIByteArray;
530    pub fn TVMFFIGetTypeInfo(type_index: i32) -> *const TVMFFITypeInfo;
531    pub fn TVMFFITestingDummyTarget() -> i32;
532}
tvm_ffi_sys/c_api.rs

tvm_ffi_sys/
c_api.rs
