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    pub cause_chain: TVMFFIObjectHandle,
305    pub extra_context: TVMFFIObjectHandle,
306}
307
308/// Shape cell used in shape object following header.
309#[repr(C)]
310pub struct TVMFFIShapeCell {
311    pub data: *const i64,
312    pub size: usize,
313}
314
315/// Field getter function pointer type
316pub type TVMFFIFieldGetter =
317    unsafe extern "C" fn(field: *mut c_void, result: *mut TVMFFIAny) -> i32;
318
319/// Field setter function pointer type
320pub type TVMFFIFieldSetter =
321    unsafe extern "C" fn(field: *mut c_void, value: *const TVMFFIAny) -> i32;
322
323/// Information support for optional object reflection
324#[repr(C)]
325pub struct TVMFFIFieldInfo {
326    /// The name of the field
327    pub name: TVMFFIByteArray,
328    /// The docstring about the field
329    pub doc: TVMFFIByteArray,
330    /// The metadata of the field in JSON string
331    pub metadata: TVMFFIByteArray,
332    /// bitmask flags of the field
333    pub flags: i64,
334    /// The size of the field
335    pub size: i64,
336    /// The alignment of the field
337    pub alignment: i64,
338    /// The offset of the field
339    pub offset: i64,
340    /// The getter to access the field
341    pub getter: Option<TVMFFIFieldGetter>,
342    /// The setter to access the field.
343    ///
344    /// When kTVMFFIFieldFlagBitSetterIsFunctionObj is NOT set (default),
345    /// this is a TVMFFIFieldSetter function pointer cast to *mut c_void.
346    /// When kTVMFFIFieldFlagBitSetterIsFunctionObj IS set,
347    /// this is a TVMFFIObjectHandle pointing to a FunctionObj.
348    ///
349    /// The setter is set even if the field is readonly for serialization.
350    pub setter: *mut c_void,
351    /// The default value or factory of the field, this field holds AnyView.
352    /// Valid when flags set kTVMFFIFieldFlagBitMaskHasDefault.
353    /// When kTVMFFIFieldFlagBitMaskDefaultFromFactory is also set,
354    /// this is a callable factory function () -> Any.
355    pub default_value_or_factory: TVMFFIAny,
356    /// Records the compile-time static type kind of the field.
357    pub field_static_type_index: i32,
358}
359
360/// Object creator function pointer type
361pub type TVMFFIObjectCreator = unsafe extern "C" fn(result: *mut TVMFFIObjectHandle) -> i32;
362
363/// Method information that can appear in reflection table
364#[repr(C)]
365pub struct TVMFFIMethodInfo {
366    /// The name of the field
367    pub name: TVMFFIByteArray,
368    /// The docstring about the method
369    pub doc: TVMFFIByteArray,
370    /// Optional metadata of the method in JSON string
371    pub metadata: TVMFFIByteArray,
372    /// bitmask flags of the method
373    pub flags: i64,
374    /// The method wrapped as ffi::Function, stored as AnyView
375    /// The first argument to the method is always the self for instance methods
376    pub method: TVMFFIAny,
377}
378
379/// Extra information of object type that can be used for reflection
380///
381/// This information is optional and can be used to enable reflection based
382/// creation of the object.
383#[repr(C)]
384pub struct TVMFFITypeMetadata {
385    /// The docstring about the object
386    pub doc: TVMFFIByteArray,
387    /// An optional function that can create a new empty instance of the type
388    pub creator: Option<TVMFFIObjectCreator>,
389    /// Total size of the object struct, if it is fixed and known
390    ///
391    /// This field is set optional and set to 0 if not registered.
392    pub total_size: i32,
393    /// Optional meta-data for structural eq/hash
394    pub structural_eq_hash_kind: i32,
395}
396
397/// Column array that stores extra attributes about types
398///
399/// The attributes stored in a column array that can be looked up by type index.
400/// Note that the TypeAttr behaves like type_traits so column T so not contain
401/// attributes from base classes.
402#[repr(C)]
403pub struct TVMFFITypeAttrColumn {
404    /// The data of the column, indexed by (type_index - begin_index).
405    pub data: *const TVMFFIAny,
406    /// The number of elements in the data array.
407    /// The column covers type indices [begin_index, begin_index + size).
408    pub size: i32,
409    /// The starting type index of the column data.
410    /// Lookup: if begin_index <= type_index < begin_index + size,
411    /// the entry is data[(type_index - begin_index) as usize].
412    pub begin_index: i32,
413}
414
415/// Runtime type information for object type checking
416#[repr(C)]
417pub struct TVMFFITypeInfo {
418    /// The runtime type index
419    /// It can be allocated during runtime if the type is dynamic
420    pub type_index: i32,
421    /// number of parent types in the type hierachy
422    pub type_depth: i32,
423    /// the unique type key to identify the type
424    pub type_key: TVMFFIByteArray,
425    /// `type_acenstors[depth]` stores the type_index of the acenstors at depth level
426    /// To keep things simple, we do not allow multiple inheritance so the
427    /// hieracy stays as a tree
428    pub type_acenstors: *const *const TVMFFITypeInfo,
429    /// Cached hash value of the type key, used for consistent structural hashing
430    pub type_key_hash: u64,
431    /// number of reflection accessible fields
432    pub num_fields: i32,
433    /// number of reflection acccesible methods
434    pub num_methods: i32,
435    /// The reflection field information
436    pub fields: *const TVMFFIFieldInfo,
437    /// The reflection method
438    pub methods: *const TVMFFIMethodInfo,
439    /// The extra information of the type
440    pub metadata: *const TVMFFITypeMetadata,
441}
442
443/// Mandatory header preceding each Object body. See `TVMFFIObjectAllocHeader`
444/// in `tvm/ffi/c_api.h`.
445#[repr(C)]
446pub struct TVMFFIObjectAllocHeader {
447    pub delete_space: Option<unsafe extern "C" fn(ptr: *mut c_void)>,
448}
449
450/// Custom allocator entry. See `TVMFFICustomAllocator` in `tvm/ffi/c_api.h`.
451#[repr(C)]
452pub struct TVMFFICustomAllocator {
453    pub allocate: Option<
454        unsafe extern "C" fn(
455            size: usize,
456            alignment: usize,
457            type_index: i32,
458            context: *mut c_void,
459        ) -> *mut c_void,
460    >,
461    pub context: *mut c_void,
462}
463
464unsafe extern "C" {
465    pub fn TVMFFIGetCustomAllocator() -> *mut TVMFFICustomAllocator;
466    pub fn TVMFFISetCustomAllocator(allocator: *mut TVMFFICustomAllocator) -> i32;
467
468    pub fn TVMFFITypeKeyToIndex(type_key: *const TVMFFIByteArray, out_tindex: *mut i32) -> i32;
469    pub fn TVMFFITypeRegisterAttr(
470        type_index: i32,
471        attr_name: *const TVMFFIByteArray,
472        attr_value: *const TVMFFIAny,
473    ) -> i32;
474    pub fn TVMFFIGetTypeAttrColumn(
475        attr_name: *const TVMFFIByteArray,
476    ) -> *const TVMFFITypeAttrColumn;
477    pub fn TVMFFIFunctionGetGlobal(
478        name: *const TVMFFIByteArray,
479        out: *mut TVMFFIObjectHandle,
480    ) -> i32;
481    pub fn TVMFFIFunctionSetGlobal(
482        name: *const TVMFFIByteArray,
483        f: TVMFFIObjectHandle,
484        can_override: i32,
485    ) -> i32;
486    pub fn TVMFFIFunctionCreate(
487        self_ptr: *mut c_void,
488        safe_call: TVMFFISafeCallType,
489        deleter: Option<unsafe extern "C" fn(*mut c_void)>,
490        out: *mut TVMFFIObjectHandle,
491    ) -> i32;
492    pub fn TVMFFIAnyViewToOwnedAny(any_view: *const TVMFFIAny, out: *mut TVMFFIAny) -> i32;
493    pub fn TVMFFIFunctionCall(
494        func: TVMFFIObjectHandle,
495        args: *const TVMFFIAny,
496        num_args: i32,
497        result: *mut TVMFFIAny,
498    ) -> i32;
499    pub fn TVMFFIErrorMoveFromRaised(result: *mut TVMFFIObjectHandle);
500    pub fn TVMFFIErrorSetRaised(error: TVMFFIObjectHandle);
501    pub fn TVMFFIErrorSetRaisedFromCStr(kind: *const i8, message: *const i8);
502    pub fn TVMFFIErrorCreate(
503        kind: *const TVMFFIByteArray,
504        message: *const TVMFFIByteArray,
505        backtrace: *const TVMFFIByteArray,
506        out: *mut TVMFFIObjectHandle,
507    ) -> i32;
508    pub fn TVMFFIErrorCreateWithCauseAndExtraContext(
509        kind: *const TVMFFIByteArray,
510        message: *const TVMFFIByteArray,
511        backtrace: *const TVMFFIByteArray,
512        cause_chain: TVMFFIObjectHandle,
513        extra_context: TVMFFIObjectHandle,
514        out: *mut TVMFFIObjectHandle,
515    ) -> i32;
516    pub fn TVMFFITensorFromDLPack(
517        from: *mut c_void,
518        require_alignment: i32,
519        require_contiguous: i32,
520        out: *mut TVMFFIObjectHandle,
521    ) -> i32;
522    pub fn TVMFFITensorToDLPack(from: TVMFFIObjectHandle, out: *mut *mut c_void) -> i32;
523    pub fn TVMFFITensorFromDLPackVersioned(
524        from: *mut c_void,
525        require_alignment: i32,
526        require_contiguous: i32,
527        out: *mut TVMFFIObjectHandle,
528    ) -> i32;
529    pub fn TVMFFITensorToDLPackVersioned(from: TVMFFIObjectHandle, out: *mut *mut c_void) -> i32;
530    pub fn TVMFFIStringFromByteArray(input: *const TVMFFIByteArray, out: *mut TVMFFIAny) -> i32;
531    pub fn TVMFFIBytesFromByteArray(input: *const TVMFFIByteArray, out: *mut TVMFFIAny) -> i32;
532    pub fn TVMFFIDataTypeFromString(str: *const TVMFFIByteArray, out: *mut DLDataType) -> i32;
533    pub fn TVMFFIDataTypeToString(dtype: *const DLDataType, out: *mut TVMFFIAny) -> i32;
534    pub fn TVMFFITraceback(
535        filename: *const i8,
536        lineno: i32,
537        func: *const i8,
538        cross_ffi_boundary: i32,
539    ) -> *const TVMFFIByteArray;
540    pub fn TVMFFIGetTypeInfo(type_index: i32) -> *const TVMFFITypeInfo;
541    pub fn TVMFFITestingDummyTarget() -> i32;
542}
tvm_ffi_sys/c_api.rs

tvm_ffi_sys/
c_api.rs
