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/// Check the return code of the safe call
63///
64/// # Arguments
65/// * `ret_code` - The return code of the safe call
66///
67/// # Returns
68/// * `Result<(), Error>` - The result of the safe call
69/// Macro to check safe calls and automatically update traceback with file/line info
70///
71/// Usage: check_safe_call!(function(args))?;
72#[macro_export]
73macro_rules! check_safe_call {
74 ($expr:expr) => {{
75 let ret_code = $expr;
76 if ret_code == 0 {
77 Ok(())
78 } else {
79 let error = $crate::error::Error::from_raised();
80 Err(error)
81 }
82 }};
83}
84
85/// Create a new error with file/line info attached
86///
87/// This macro automatically appends file/line info to the traceback
88///
89/// # Arguments
90/// * `error_kind` - The kind of the error
91/// * `msg` - The message of the error
92/// * `args` - The posisble format arguments
93///
94/// # Returns
95/// * `Result<(), Error>` - The result of the safe call
96#[macro_export]
97macro_rules! bail {
98 ($error_kind:expr, $fmt:expr $(, $args:expr)* $(,)?) => {{
99 let context = format!(
100 " File \"{}\", line {}, in {}\n",
101 file!(),
102 line!(),
103 $crate::function_name!()
104 );
105 return Err($crate::error::Error::new($error_kind, &format!($fmt $(, $args)*), &context));
106 }};
107}
108
109/// Create a new error with file/line info attached
110///
111/// This macro automatically appends file/line info to the traceback
112///
113/// # Arguments
114/// * `kind` - The kind of the error
115/// * `msg` - The message of the error
116/// * `args` - The posisble format arguments
117///
118/// # Returns
119/// * `Result<(), Error>` - The result of the safe call
120#[macro_export]
121macro_rules! ensure {
122 ($cond:expr, $error_kind:expr, $fmt:expr $(, $args:expr)* $(,)?) => {{
123 if !$cond {
124 $crate::bail!($error_kind, $fmt $(, $args)*);
125 }
126 }};
127}
128
129/// Attach a context to a result if it is error
130///
131/// This macro automatically appends file/line info to the traceback
132///
133/// # Arguments
134/// * `error` - The error to attach the context to
135/// * `msg` - The message of the error
136///
137/// # Returns
138/// * `Result<(), Error>` - The result of the safe call
139#[macro_export]
140macro_rules! attach_context {
141 ($error:expr) => {{
142 match $error {
143 Ok(value) => Ok(value),
144 Err(error) => {
145 let context = format!(
146 " File \"{}\", line {}, in {}\n",
147 file!(),
148 line!(),
149 $crate::function_name!()
150 );
151 Err(Error::with_appended_backtrace(error, &context))
152 }
153 }
154 }};
155}
156
157// ----------------------------------------------------------------------------
158// Macros for any definitions
159// ----------------------------------------------------------------------------
160
161// implements try from any for all integer types
162/// Macro to implement `TryFrom<AnyView>` and `TryFrom<Any>` for a list of types
163#[macro_export]
164macro_rules! impl_try_from_any {
165 ($($t:ty),* $(,)?) => {
166 $(
167 impl<'a> TryFrom<$crate::any::AnyView<'a>> for $t {
168 type Error = $crate::error::Error;
169 #[inline(always)]
170 fn try_from(
171 value: $crate::any::AnyView<'a>
172 ) -> Result<Self, Self::Error> {
173 type TryFromTemp = $crate::any::TryFromTemp<$t>;
174 return TryFromTemp::try_from(value).map(TryFromTemp::into_value);
175 }
176 }
177
178 impl TryFrom<$crate::any::Any> for $t {
179 type Error = $crate::error::Error;
180 #[inline(always)]
181 fn try_from(
182 value: $crate::any::Any
183 ) -> Result<Self, Self::Error> {
184 type TryFromTemp = $crate::any::TryFromTemp<$t>;
185 return TryFromTemp::try_from(value).map(TryFromTemp::into_value);
186 }
187 }
188 )*
189 };
190}
191
192/// Macro to implement `TryFrom<AnyView>` and `TryFrom<Any>` for generic types like `Option<T>`
193#[macro_export]
194macro_rules! impl_try_from_any_for_parametric {
195 ($generic_type:ident<$param:ident>) => {
196 impl<'a, $param: AnyCompatible> TryFrom<$crate::any::AnyView<'a>>
197 for $generic_type<$param>
198 {
199 type Error = $crate::error::Error;
200 #[inline(always)]
201 fn try_from(value: $crate::any::AnyView<'a>) -> Result<Self, Self::Error> {
202 type TryFromTemp<T> = $crate::any::TryFromTemp<$generic_type<$param>>;
203 return TryFromTemp::<T>::try_from(value).map(TryFromTemp::<T>::into_value);
204 }
205 }
206
207 impl<$param: AnyCompatible> TryFrom<$crate::any::Any> for $generic_type<$param> {
208 type Error = $crate::error::Error;
209 #[inline(always)]
210 fn try_from(value: $crate::any::Any) -> Result<Self, Self::Error> {
211 type TryFromTemp<T> = $crate::any::TryFromTemp<$generic_type<$param>>;
212 return TryFromTemp::<T>::try_from(value).map(TryFromTemp::<T>::into_value);
213 }
214 }
215 };
216}
217
218/// Macro to implement IntoArgHolder for a list of types
219#[macro_export]
220macro_rules! impl_into_arg_holder_default {
221 ($($t:ty),*) => {
222 $(
223 impl $crate::function_internal::IntoArgHolder for $t {
224 type Target = $t;
225 fn into_arg_holder(self) -> Self::Target {
226 self
227 }
228 }
229 impl<'a> $crate::function_internal::IntoArgHolder for &'a $t {
230 type Target = &'a $t;
231 fn into_arg_holder(self) -> Self::Target {
232 self
233 }
234 }
235 )*
236 };
237}
238
239/// Macro to implement ArgIntoRef for a list of types
240#[macro_export]
241macro_rules! impl_arg_into_ref {
242 ($($t:ty),*) => {
243 $(
244 impl $crate::function_internal::ArgIntoRef for $t {
245 type Target = $t;
246 fn to_ref(&self) -> &Self::Target {
247 &self
248 }
249 }
250 impl<'a> $crate::function_internal::ArgIntoRef for &'a $t {
251 type Target = $t;
252 fn to_ref(&self) -> &Self::Target {
253 &self
254 }
255 }
256 )*
257 }
258}
259
260// ----------------------------------------------------------------------------
261// Macros for function definitions
262// ----------------------------------------------------------------------------
263
264/// Macro to export a typed function as a C symbol that follows the tvm-ffi ABI
265///
266/// # Arguments
267/// * `$name` - The name of the function
268/// * `$func` - The function to export
269///
270/// # Example
271/// ```rust
272/// use tvm_ffi::*;
273///
274/// fn add_one(x: i32) -> Result<i32> { Ok(x + 1) }
275///
276/// tvm_ffi_dll_export_typed_func!(add_one, add_one);
277/// ```
278#[macro_export]
279macro_rules! tvm_ffi_dll_export_typed_func {
280 ($name:ident, $func:expr) => {
281 $crate::macros::paste::paste! {
282 // `#[no_mangle]` is required so the symbol is preserved in a
283 // `cdylib` and matches the `__tvm_ffi_<name>` naming convention
284 // that `ffi.Module.load_from_file.<format>` looks up via
285 // `GetSymbolWithSymbolPrefix`. Without it, the linker strips the
286 // function from the output `.so`.
287 //
288 // Using plain `#[no_mangle]` (rather than `#[unsafe(no_mangle)]`,
289 // which would require rustc >= 1.82) keeps the crate buildable
290 // on older toolchains. Edition-2024 callers will see a
291 // deprecation warning, which is harmless.
292 //
293 // The path-qualified `$crate::tvm_ffi_sys::…` reference (rather
294 // than a bare `tvm_ffi_sys::…`) lets downstream crates use the
295 // macro without having to add `tvm-ffi-sys` to their own
296 // `[dependencies]`.
297 #[no_mangle]
298 pub unsafe extern "C" fn [<__tvm_ffi_ $name>](
299 _handle: *mut std::ffi::c_void,
300 args: *const $crate::tvm_ffi_sys::TVMFFIAny,
301 num_args: i32,
302 result: *mut $crate::tvm_ffi_sys::TVMFFIAny,
303 ) -> i32 {
304 let packed_args =
305 std::slice::from_raw_parts(args as *const $crate::any::AnyView, num_args as usize);
306 let ret_value = $crate::function_internal::call_packed_callable($func, packed_args);
307 match ret_value {
308 Ok(value) => {
309 *result = $crate::any::Any::into_raw_ffi_any(value);
310 0
311 }
312 Err(error) => {
313 $crate::error::Error::set_raised(&error);
314 -1
315 }
316 }
317 }
318 }
319 };
320}
321
322///-----------------------------------------------------------
323/// into_typed_fn
324///
325/// Converts a generic `Function` into a typed function with compile-time
326/// argument count and type checking. This macro provides a convenient way
327/// to create type-safe wrappers around TVM functions.
328///
329/// # Arguments
330/// * `$f` - The function identifier to convert
331/// * `$trait` - The trait type (typically `Fn`)
332/// * `($t0, $t1, ...)` - The argument types
333/// * `$ret_ty` - The return type
334///
335/// # Example
336/// ```rust
337/// use tvm_ffi::*;
338///
339/// let func = Function::from_typed(|x: i32, y: i32| -> Result<i32> { Ok(x + y) });
340/// let typed_func = into_typed_fn!(func, Fn(i32, &i32) -> Result<i32>);
341/// let result = typed_func(10, &20).unwrap(); // Returns 30
342/// assert_eq!(result, 30);
343/// ```
344/// Note that the `into_typed_fn!` macro can specify arguments to be passed either
345/// by reference or by value in the argument list.
346/// We recommend passing by reference for ObjectRef types such as Tensor.
347/// Since the ffi mechanism requires us to pass arguments by reference.
348///
349/// # Supported Argument Counts
350/// This macro supports functions with 0 to 8 arguments.
351///-----------------------------------------------------------
352#[macro_export]
353macro_rules! into_typed_fn {
354 // Case for 0 arguments
355 ($f:expr, $trait:ident() -> $ret_ty:ty) => {{
356 let _f = $f;
357 move || -> $ret_ty { Ok(_f.call_tuple_with_len::<0, _>(())?.try_into()?) }
358 }};
359 // Case for 1 argument
360 ($f:expr, $trait:ident($t0:ty) -> $ret_ty:ty) => {{
361 let _f = $f;
362 move |a0: $t0| -> $ret_ty {
363 use $crate::function_internal::IntoArgHolderTuple;
364 let tuple_args = (a0,).into_arg_holder_tuple();
365 Ok(_f.call_tuple_with_len::<1, _>(tuple_args)?.try_into()?)
366 }
367 }};
368 // Case for 2 arguments
369 ($f:expr, $trait:ident($t0:ty, $t1:ty) -> $ret_ty:ty) => {{
370 let _f = $f;
371 move |a0: $t0, a1: $t1| -> $ret_ty {
372 use $crate::function_internal::IntoArgHolderTuple;
373 let tuple_args = (a0, a1).into_arg_holder_tuple();
374 Ok(_f.call_tuple_with_len::<2, _>(tuple_args)?.try_into()?)
375 }
376 }};
377 // Case for 3 arguments
378 ($f:expr, $trait:ident($t0:ty, $t1:ty, $t2:ty) -> $ret_ty:ty) => {{
379 let _f = $f;
380 move |a0: $t0, a1: $t1, a2: $t2| -> $ret_ty {
381 use $crate::function_internal::IntoArgHolderTuple;
382 let tuple_args = (a0, a1, a2).into_arg_holder_tuple();
383 Ok(_f.call_tuple_with_len::<3, _>(tuple_args)?.try_into()?)
384 }
385 }};
386 // Case for 4 arguments
387 ($f:expr, $trait:ident($t0:ty, $t1:ty, $t2:ty, $t3:ty) -> $ret_ty:ty) => {{
388 let _f = $f;
389 move |a0: $t0, a1: $t1, a2: $t2, a3: $t3| -> $ret_ty {
390 use $crate::function_internal::IntoArgHolderTuple;
391 let tuple_args = (a0, a1, a2, a3).into_arg_holder_tuple();
392 Ok(_f.call_tuple_with_len::<4, _>(tuple_args)?.try_into()?)
393 }
394 }};
395 // Case for 5 arguments
396 ($f:expr, $trait:ident($t0:ty, $t1:ty, $t2:ty, $t3:ty, $t4:ty) -> $ret_ty:ty) => {{
397 let _f = $f;
398 move |a0: $t0, a1: $t1, a2: $t2, a3: $t3, a4: $t4| -> $ret_ty {
399 use $crate::function_internal::IntoArgHolderTuple;
400 let tuple_args = (a0, a1, a2, a3, a4).into_arg_holder_tuple();
401 Ok(_f.call_tuple_with_len::<5, _>(tuple_args)?.try_into()?)
402 }
403 }};
404 // Case for 6 arguments
405 ($f:expr, $trait:ident($t0:ty, $t1:ty, $t2:ty, $t3:ty, $t4:ty, $t5:ty) -> $ret_ty:ty) => {{
406 let _f = $f;
407 move |a0: $t0, a1: $t1, a2: $t2, a3: $t3, a4: $t4, a5: $t5| -> $ret_ty {
408 use $crate::function_internal::IntoArgHolderTuple;
409 let tuple_args = (a0, a1, a2, a3, a4, a5).into_arg_holder_tuple();
410 Ok(_f.call_tuple_with_len::<6, _>(tuple_args)?.try_into()?)
411 }
412 }};
413 // Case for 7 arguments
414 ($f:expr, $trait:ident($t0:ty, $t1:ty, $t2:ty, $t3:ty, $t4:ty, $t5:ty, $t6:ty)
415 -> $ret_ty:ty) => {{
416 let _f = $f;
417 move |a0: $t0, a1: $t1, a2: $t2, a3: $t3, a4: $t4, a5: $t5, a6: $t6| -> $ret_ty {
418 use $crate::function_internal::IntoArgHolderTuple;
419 let tuple_args = (a0, a1, a2, a3, a4, a5, a6).into_arg_holder_tuple();
420 Ok(_f.call_tuple_with_len::<7, _>(tuple_args)?.try_into()?)
421 }
422 }};
423 // Case for 8 arguments
424 ($f:expr, $trait:ident($t0:ty, $t1:ty, $t2:ty, $t3:ty, $t4:ty, $t5:ty, $t6:ty, $t7:ty)
425 -> $ret_ty:ty) => {{
426 let _f = $f;
427 move |a0: $t0, a1: $t1, a2: $t2, a3: $t3, a4: $t4, a5: $t5, a6: $t6, a7: $t7| -> $ret_ty {
428 use $crate::function_internal::IntoArgHolderTuple;
429 let tuple_args = (a0, a1, a2, a3, a4, a5, a6, a7).into_arg_holder_tuple();
430 Ok(_f.call_tuple_with_len::<8, _>(tuple_args)?.try_into()?)
431 }
432 }};
433}