tvm_ffi_sys/
c_env_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::os::raw::c_char;
25
26use crate::c_api::TVMFFIObjectHandle;
27use crate::dlpack::DLTensor;
28
29// ----------------------------------------------------------------------------
30// Stream context
31// Focusing on minimalistic thread-local context recording stream being used.
32// We explicitly not handle allocation/de-allocation of stream here.
33// ----------------------------------------------------------------------------
34
35/// The type of the stream handle.
36pub type TVMFFIStreamHandle = *mut c_void;
37
38/// DLPack tensor allocator function type
39pub type DLPackManagedTensorAllocator = unsafe extern "C" fn(
40    prototype: *mut DLTensor,
41    out: *mut *mut c_void, // DLManagedTensorVersioned**
42    error_ctx: *mut c_void,
43    set_error: unsafe extern "C" fn(*mut c_void, *const c_char, *const c_char),
44) -> i32;
45
46unsafe extern "C" {
47    pub fn TVMFFIEnvSetStream(
48        device_type: i32,
49        device_id: i32,
50        stream: TVMFFIStreamHandle,
51        opt_out_original_stream: *mut TVMFFIStreamHandle,
52    ) -> i32;
53
54    pub fn TVMFFIEnvGetStream(device_type: i32, device_id: i32) -> TVMFFIStreamHandle;
55
56    pub fn TVMFFIEnvSetDLPackManagedTensorAllocator(
57        allocator: DLPackManagedTensorAllocator,
58        write_to_global_context: i32,
59        opt_out_original_allocator: *mut DLPackManagedTensorAllocator,
60    ) -> i32;
61
62    pub fn TVMFFIEnvGetDLPackManagedTensorAllocator() -> DLPackManagedTensorAllocator;
63
64    pub fn TVMFFIEnvCheckSignals() -> i32;
65
66    pub fn TVMFFIEnvRegisterCAPI(name: *const c_char, symbol: *mut c_void) -> i32;
67
68    pub fn TVMFFIEnvModLookupFromImports(
69        library_ctx: TVMFFIObjectHandle,
70        func_name: *const c_char,
71        out: *mut TVMFFIObjectHandle,
72    ) -> i32;
73
74    pub fn TVMFFIEnvModRegisterContextSymbol(name: *const c_char, symbol: *mut c_void) -> i32;
75
76    pub fn TVMFFIEnvModRegisterSystemLibSymbol(name: *const c_char, symbol: *mut c_void) -> i32;
77}
tvm_ffi_sys/c_env_api.rs

tvm_ffi_sys/
c_env_api.rs
