tvm_ffi/macros.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// rexport paste under macro namespace so downstream do not need to specify dep
20pub use paste;
21// ----------------------------------------------------------------------------
22// Macros for error handling
23// ----------------------------------------------------------------------------
24
25/// Macro gto get the name of the function
26///
27/// # Usage
28/// Usage: function_name!()
29#[macro_export]
30macro_rules! function_name {
31 () => {{
32 // dummy function to get the name of the function
33 fn f() {}
34 fn type_name_of<T>(_: T) -> &'static str {
35 std::any::type_name::<T>()
36 }
37 let name = type_name_of(f);
38 // remove the f() from the name
39 &name[..name.len() - 3]
40 }};
41}
42
43/// Resolves a registered global [`Function`](crate::function::Function) by name
44/// and caches it for the lifetime of the process (lock-free after the first
45/// lookup). Each call site gets its own cache and the macro evaluates to a
46/// `&'static Function`. Panics — naming the function — if it is not registered.
47///
48/// Usage: `cached_global_func!("ffi.Map").call_packed(args)`
49#[macro_export]
50macro_rules! cached_global_func {
51 ($name:literal) => {{
52 static FUNC: std::sync::LazyLock<$crate::function::Function> =
53 std::sync::LazyLock::new(|| {
54 $crate::function::Function::get_global($name).unwrap_or_else(|_| {
55 panic!(concat!("global function `", $name, "` is not registered"))
56 })
57 });
58 &*FUNC
59 }};
60}
61
62/// Implement zero-copy conversions from a derived object reference to one of
63/// its base object-reference types.
64///
65/// # Safety
66///
67/// Generated bindings must ensure that `target` is a registered base of
68/// `source` and that every `source` satisfies the target reference invariants.
69/// An incorrect declaration makes the generated safe conversion unsound.
70#[macro_export]
71macro_rules! impl_object_upcast {
72 ($($source:ty => $target:ty),+ $(,)?) => {
73 $(
74 impl ::std::convert::From<$source> for $target {
75 #[inline]
76 fn from(value: $source) -> Self {
77 let data = <$source as $crate::object::ObjectRefCore>::into_data(value);
78 // SAFETY: The macro declaration promises that `target` is
79 // a registered base of `source`. Both references retain
80 // the same allocation and only change its static view.
81 let data = unsafe {
82 $crate::object::ObjectArc::from_raw(
83 $crate::object::ObjectArc::into_raw(data).cast::
84 <<$target as $crate::object::ObjectRefCore>::ContainerType>(),
85 )
86 };
87 // SAFETY: The macro declaration promises both the
88 // container inheritance relation and every additional
89 // invariant imposed by the target reference view.
90 unsafe {
91 <$target as $crate::object::ObjectRefCore>::from_data(data)
92 }
93 }
94 }
95
96 impl ::std::convert::From<&$source> for $target {
97 #[inline]
98 fn from(value: &$source) -> Self {
99 value.clone().into()
100 }
101 }
102
103 )+
104 };
105}
106
107/// Check the return code of the safe call
108///
109/// # Arguments
110/// * `ret_code` - The return code of the safe call
111///
112/// # Returns
113/// * `Result<(), Error>` - The result of the safe call
114/// Macro to check safe calls and automatically update traceback with file/line info
115///
116/// Usage: check_safe_call!(function(args))?;
117#[macro_export]
118macro_rules! check_safe_call {
119 ($expr:expr) => {{
120 let ret_code = $expr;
121 if ret_code == 0 {
122 Ok(())
123 } else {
124 let error = $crate::error::Error::from_raised();
125 Err(error)
126 }
127 }};
128}
129
130/// Create a new error with file/line info attached
131///
132/// This macro automatically appends file/line info to the traceback
133///
134/// # Arguments
135/// * `error_kind` - The kind of the error
136/// * `msg` - The message of the error
137/// * `args` - The posisble format arguments
138///
139/// # Returns
140/// * `Result<(), Error>` - The result of the safe call
141#[macro_export]
142macro_rules! bail {
143 ($error_kind:expr, $fmt:expr $(, $args:expr)* $(,)?) => {{
144 let context = format!(
145 " File \"{}\", line {}, in {}\n",
146 file!(),
147 line!(),
148 $crate::function_name!()
149 );
150 return Err($crate::error::Error::new($error_kind, &format!($fmt $(, $args)*), &context));
151 }};
152}
153
154/// Create a new error with file/line info attached
155///
156/// This macro automatically appends file/line info to the traceback
157///
158/// # Arguments
159/// * `kind` - The kind of the error
160/// * `msg` - The message of the error
161/// * `args` - The posisble format arguments
162///
163/// # Returns
164/// * `Result<(), Error>` - The result of the safe call
165#[macro_export]
166macro_rules! ensure {
167 ($cond:expr, $error_kind:expr, $fmt:expr $(, $args:expr)* $(,)?) => {{
168 if !$cond {
169 $crate::bail!($error_kind, $fmt $(, $args)*);
170 }
171 }};
172}
173
174/// Attach a context to a result if it is error
175///
176/// This macro automatically appends file/line info to the traceback
177///
178/// # Arguments
179/// * `error` - The error to attach the context to
180/// * `msg` - The message of the error
181///
182/// # Returns
183/// * `Result<(), Error>` - The result of the safe call
184#[macro_export]
185macro_rules! attach_context {
186 ($error:expr) => {{
187 match $error {
188 Ok(value) => Ok(value),
189 Err(error) => {
190 let context = format!(
191 " File \"{}\", line {}, in {}\n",
192 file!(),
193 line!(),
194 $crate::function_name!()
195 );
196 Err(Error::with_appended_backtrace(error, &context))
197 }
198 }
199 }};
200}
201
202// ----------------------------------------------------------------------------
203// Macros for any definitions
204// ----------------------------------------------------------------------------
205
206// implements try from any for all integer types
207/// Macro to implement `TryFrom<AnyView>` and `TryFrom<Any>` for a list of types
208#[macro_export]
209macro_rules! impl_try_from_any {
210 ($($t:ty),* $(,)?) => {
211 $(
212 impl<'a> TryFrom<$crate::any::AnyView<'a>> for $t {
213 type Error = $crate::error::Error;
214 #[inline(always)]
215 fn try_from(
216 value: $crate::any::AnyView<'a>
217 ) -> Result<Self, Self::Error> {
218 type TryFromTemp = $crate::any::TryFromTemp<$t>;
219 return TryFromTemp::try_from(value).map(TryFromTemp::into_value);
220 }
221 }
222
223 impl TryFrom<$crate::any::Any> for $t {
224 type Error = $crate::error::Error;
225 #[inline(always)]
226 fn try_from(
227 value: $crate::any::Any
228 ) -> Result<Self, Self::Error> {
229 type TryFromTemp = $crate::any::TryFromTemp<$t>;
230 return TryFromTemp::try_from(value).map(TryFromTemp::into_value);
231 }
232 }
233 )*
234 };
235}
236
237/// Macro to implement `TryFrom<AnyView>` and `TryFrom<Any>` for generic types like `Option<T>`
238#[macro_export]
239macro_rules! impl_try_from_any_for_parametric {
240 ($generic_type:ident<$param:ident>) => {
241 impl<'a, $param: AnyCompatible> TryFrom<$crate::any::AnyView<'a>>
242 for $generic_type<$param>
243 {
244 type Error = $crate::error::Error;
245 #[inline(always)]
246 fn try_from(value: $crate::any::AnyView<'a>) -> Result<Self, Self::Error> {
247 type TryFromTemp<T> = $crate::any::TryFromTemp<$generic_type<$param>>;
248 return TryFromTemp::<T>::try_from(value).map(TryFromTemp::<T>::into_value);
249 }
250 }
251
252 impl<$param: AnyCompatible> TryFrom<$crate::any::Any> for $generic_type<$param> {
253 type Error = $crate::error::Error;
254 #[inline(always)]
255 fn try_from(value: $crate::any::Any) -> Result<Self, Self::Error> {
256 type TryFromTemp<T> = $crate::any::TryFromTemp<$generic_type<$param>>;
257 return TryFromTemp::<T>::try_from(value).map(TryFromTemp::<T>::into_value);
258 }
259 }
260 };
261}
262
263/// Macro to implement IntoArgHolder for a list of types
264#[macro_export]
265macro_rules! impl_into_arg_holder_default {
266 ($($t:ty),*) => {
267 $(
268 impl $crate::function_internal::IntoArgHolder for $t {
269 type Target = $t;
270 fn into_arg_holder(self) -> Self::Target {
271 self
272 }
273 }
274 impl<'a> $crate::function_internal::IntoArgHolder for &'a $t {
275 type Target = &'a $t;
276 fn into_arg_holder(self) -> Self::Target {
277 self
278 }
279 }
280 )*
281 };
282}
283
284/// Macro to implement ArgIntoRef for a list of types
285#[macro_export]
286macro_rules! impl_arg_into_ref {
287 ($($t:ty),*) => {
288 $(
289 impl $crate::function_internal::ArgIntoRef for $t {
290 type Target = $t;
291 fn to_ref(&self) -> &Self::Target {
292 &self
293 }
294 }
295 impl<'a> $crate::function_internal::ArgIntoRef for &'a $t {
296 type Target = $t;
297 fn to_ref(&self) -> &Self::Target {
298 &self
299 }
300 }
301 )*
302 }
303}
304
305// ----------------------------------------------------------------------------
306// Macros for function definitions
307// ----------------------------------------------------------------------------
308
309/// Macro to export a typed function as a C symbol that follows the tvm-ffi ABI
310///
311/// # Arguments
312/// * `$name` - The name of the function
313/// * `$func` - The function to export
314///
315/// # Example
316/// ```rust
317/// use tvm_ffi::*;
318///
319/// fn add_one(x: i32) -> Result<i32> { Ok(x + 1) }
320///
321/// tvm_ffi_dll_export_typed_func!(add_one, add_one);
322/// ```
323#[macro_export]
324macro_rules! tvm_ffi_dll_export_typed_func {
325 ($name:ident, $func:expr) => {
326 $crate::macros::paste::paste! {
327 // `#[no_mangle]` is required so the symbol is preserved in a
328 // `cdylib` and matches the `__tvm_ffi_<name>` naming convention
329 // that `ffi.Module.load_from_file.<format>` looks up via
330 // `GetSymbolWithSymbolPrefix`. Without it, the linker strips the
331 // function from the output `.so`.
332 //
333 // Using plain `#[no_mangle]` (rather than `#[unsafe(no_mangle)]`,
334 // which would require rustc >= 1.82) keeps the crate buildable
335 // on older toolchains. Edition-2024 callers will see a
336 // deprecation warning, which is harmless.
337 //
338 // The path-qualified `$crate::tvm_ffi_sys::…` reference (rather
339 // than a bare `tvm_ffi_sys::…`) lets downstream crates use the
340 // macro without having to add `tvm-ffi-sys` to their own
341 // `[dependencies]`.
342 #[no_mangle]
343 pub unsafe extern "C" fn [<__tvm_ffi_ $name>](
344 _handle: *mut std::ffi::c_void,
345 args: *const $crate::tvm_ffi_sys::TVMFFIAny,
346 num_args: i32,
347 result: *mut $crate::tvm_ffi_sys::TVMFFIAny,
348 ) -> i32 {
349 let packed_args =
350 std::slice::from_raw_parts(args as *const $crate::any::AnyView, num_args as usize);
351 let ret_value = $crate::function_internal::call_packed_callable($func, packed_args);
352 match ret_value {
353 Ok(value) => {
354 *result = $crate::any::Any::into_raw_ffi_any(value);
355 0
356 }
357 Err(error) => {
358 $crate::error::Error::set_raised(&error);
359 -1
360 }
361 }
362 }
363 }
364 };
365}
366
367///-----------------------------------------------------------
368/// into_typed_fn
369///
370/// Converts a generic `Function` into a typed function with compile-time
371/// argument count and type checking. This macro provides a convenient way
372/// to create type-safe wrappers around TVM functions.
373///
374/// # Arguments
375/// * `$f` - The function identifier to convert
376/// * `$trait` - The trait type (typically `Fn`)
377/// * `($t0, $t1, ...)` - The argument types
378/// * `$ret_ty` - The return type
379///
380/// # Example
381/// ```rust
382/// use tvm_ffi::*;
383///
384/// let func = Function::from_typed(|x: i32, y: i32| -> Result<i32> { Ok(x + y) });
385/// let typed_func = into_typed_fn!(func, Fn(i32, &i32) -> Result<i32>);
386/// let result = typed_func(10, &20).unwrap(); // Returns 30
387/// assert_eq!(result, 30);
388/// ```
389/// Note that the `into_typed_fn!` macro can specify arguments to be passed either
390/// by reference or by value in the argument list.
391/// We recommend passing by reference for ObjectRef types such as Tensor.
392/// Since the ffi mechanism requires us to pass arguments by reference.
393///
394/// # Supported Argument Counts
395/// This macro supports functions with 0 to 8 arguments.
396///-----------------------------------------------------------
397#[macro_export]
398macro_rules! into_typed_fn {
399 // Case for 0 arguments
400 ($f:expr, $trait:ident() -> $ret_ty:ty) => {{
401 let _f = $f;
402 move || -> $ret_ty { Ok(_f.call_tuple_with_len::<0, _>(())?.try_into()?) }
403 }};
404 // Case for 1 argument
405 ($f:expr, $trait:ident($t0:ty) -> $ret_ty:ty) => {{
406 let _f = $f;
407 move |a0: $t0| -> $ret_ty {
408 use $crate::function_internal::IntoArgHolderTuple;
409 let tuple_args = (a0,).into_arg_holder_tuple();
410 Ok(_f.call_tuple_with_len::<1, _>(tuple_args)?.try_into()?)
411 }
412 }};
413 // Case for 2 arguments
414 ($f:expr, $trait:ident($t0:ty, $t1:ty) -> $ret_ty:ty) => {{
415 let _f = $f;
416 move |a0: $t0, a1: $t1| -> $ret_ty {
417 use $crate::function_internal::IntoArgHolderTuple;
418 let tuple_args = (a0, a1).into_arg_holder_tuple();
419 Ok(_f.call_tuple_with_len::<2, _>(tuple_args)?.try_into()?)
420 }
421 }};
422 // Case for 3 arguments
423 ($f:expr, $trait:ident($t0:ty, $t1:ty, $t2:ty) -> $ret_ty:ty) => {{
424 let _f = $f;
425 move |a0: $t0, a1: $t1, a2: $t2| -> $ret_ty {
426 use $crate::function_internal::IntoArgHolderTuple;
427 let tuple_args = (a0, a1, a2).into_arg_holder_tuple();
428 Ok(_f.call_tuple_with_len::<3, _>(tuple_args)?.try_into()?)
429 }
430 }};
431 // Case for 4 arguments
432 ($f:expr, $trait:ident($t0:ty, $t1:ty, $t2:ty, $t3:ty) -> $ret_ty:ty) => {{
433 let _f = $f;
434 move |a0: $t0, a1: $t1, a2: $t2, a3: $t3| -> $ret_ty {
435 use $crate::function_internal::IntoArgHolderTuple;
436 let tuple_args = (a0, a1, a2, a3).into_arg_holder_tuple();
437 Ok(_f.call_tuple_with_len::<4, _>(tuple_args)?.try_into()?)
438 }
439 }};
440 // Case for 5 arguments
441 ($f:expr, $trait:ident($t0:ty, $t1:ty, $t2:ty, $t3:ty, $t4:ty) -> $ret_ty:ty) => {{
442 let _f = $f;
443 move |a0: $t0, a1: $t1, a2: $t2, a3: $t3, a4: $t4| -> $ret_ty {
444 use $crate::function_internal::IntoArgHolderTuple;
445 let tuple_args = (a0, a1, a2, a3, a4).into_arg_holder_tuple();
446 Ok(_f.call_tuple_with_len::<5, _>(tuple_args)?.try_into()?)
447 }
448 }};
449 // Case for 6 arguments
450 ($f:expr, $trait:ident($t0:ty, $t1:ty, $t2:ty, $t3:ty, $t4:ty, $t5:ty) -> $ret_ty:ty) => {{
451 let _f = $f;
452 move |a0: $t0, a1: $t1, a2: $t2, a3: $t3, a4: $t4, a5: $t5| -> $ret_ty {
453 use $crate::function_internal::IntoArgHolderTuple;
454 let tuple_args = (a0, a1, a2, a3, a4, a5).into_arg_holder_tuple();
455 Ok(_f.call_tuple_with_len::<6, _>(tuple_args)?.try_into()?)
456 }
457 }};
458 // Case for 7 arguments
459 ($f:expr, $trait:ident($t0:ty, $t1:ty, $t2:ty, $t3:ty, $t4:ty, $t5:ty, $t6:ty)
460 -> $ret_ty:ty) => {{
461 let _f = $f;
462 move |a0: $t0, a1: $t1, a2: $t2, a3: $t3, a4: $t4, a5: $t5, a6: $t6| -> $ret_ty {
463 use $crate::function_internal::IntoArgHolderTuple;
464 let tuple_args = (a0, a1, a2, a3, a4, a5, a6).into_arg_holder_tuple();
465 Ok(_f.call_tuple_with_len::<7, _>(tuple_args)?.try_into()?)
466 }
467 }};
468 // Case for 8 arguments
469 ($f:expr, $trait:ident($t0:ty, $t1:ty, $t2:ty, $t3:ty, $t4:ty, $t5:ty, $t6:ty, $t7:ty)
470 -> $ret_ty:ty) => {{
471 let _f = $f;
472 move |a0: $t0, a1: $t1, a2: $t2, a3: $t3, a4: $t4, a5: $t5, a6: $t6, a7: $t7| -> $ret_ty {
473 use $crate::function_internal::IntoArgHolderTuple;
474 let tuple_args = (a0, a1, a2, a3, a4, a5, a6, a7).into_arg_holder_tuple();
475 Ok(_f.call_tuple_with_len::<8, _>(tuple_args)?.try_into()?)
476 }
477 }};
478}