Function and Module

Function and Module#

TVM-FFI provides a unified and ABI-stable calling convention that enables cross-language function calls between C++, Python, Rust, and other languages. Functions are first-class TVM-FFI objects.

This tutorial covers defining, registering, and calling TVM-FFI functions, exception handling, and working with modules.

Glossary#

TVM-FFI ABI, or “Packed Function”. TVMFFISafeCallType

A stable C calling convention where every function is represented by a single signature, which enables type-erased, cross-language function calls. This calling convention is used across all TVM-FFI function calls at the ABI boundary. See Stable C ABI for a quick introduction.

TVM-FFI Function. tvm_ffi.Function, tvm::ffi::FunctionObj, tvm::ffi::Function

A reference-counted function object and its managed reference, which wraps any callable, including language-agnostic functions and lambdas (C++, Python, Rust, etc.), member functions, external C symbols, and other callable objects, all sharing the same calling convention.

TVM-FFI Module. tvm_ffi.Module, tvm::ffi::ModuleObj, tvm::ffi::Module

A namespace for a collection of functions, loaded from a shared library via dlopen (Linux, macOS) or LoadLibraryW (Windows), or statically linked to the current executable.

Global Functions and Registry. tvm_ffi.get_global_func() and tvm_ffi.register_global_func()

A registry is a table that maps string names to Function objects and their metadata (name, docs, signatures, etc.) for cross-language access. Functions in the registry are called global functions.

Common Usage#

TVM-FFI C Symbols#

Shared library. Use TVM_FFI_DLL_EXPORT_TYPED_FUNC to export a function as a C symbol that follows the TVM-FFI ABI:

static int AddTwo(int x) { return x + 2; }

TVM_FFI_DLL_EXPORT_TYPED_FUNC(/*ExportName=*/add_two, /*Function=*/AddTwo)

This creates a C symbol __tvm_ffi_<ExportName> in the shared library, which can then be loaded and called via tvm_ffi.load_module():

import tvm_ffi

mod = tvm_ffi.load_module("path/to/library.so")
result = mod.add_two(40)  # -> 42

System library. For symbols bundled in the same executable, use TVMFFIEnvModRegisterSystemLibSymbol() to register each symbol during static initialization within a TVM_FFI_STATIC_INIT_BLOCK. See tvm_ffi.system_lib() for a complete workflow.

Global Functions#

Register a global function. In C++, use tvm::ffi::reflection::GlobalDef to register a function:

#include <tvm/ffi/tvm_ffi.h>

static int AddOne(int x) { return x + 1; }

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef()
      .def("my_ext.add_one", AddOne, "Add one to the input");
}

The TVM_FFI_STATIC_INIT_BLOCK macro ensures that registration occurs during library initialization. The registered function is then accessible from Python by the name my_ext.add_one.

In Python, use the decorator tvm_ffi.register_global_func() to register a global function:

import tvm_ffi

@tvm_ffi.register_global_func("my_ext.add_one")
def add_one(x: int) -> int:
    return x + 1

Retrieve a global function. After registration, functions are accessible by name.

In Python, use tvm_ffi.get_global_func() to retrieve a global function:

import tvm_ffi

# Get a function from the global registry
add_one = tvm_ffi.get_global_func("my_ext.add_one")
result = add_one(41)  # -> 42

In C++, use tvm::ffi::Function::GetGlobal() or tvm::ffi::Function::GetGlobalRequired() to retrieve a global function:

ffi::Function func = ffi::Function::GetGlobalRequired("my_ext.add_one");
int result = func(41);  // -> 42

Create Functions#

From C++. An tvm::ffi::Function can be created via tvm::ffi::Function::FromTyped() or tvm::ffi::TypedFunction’s constructor.

// Create type-erased function: add_type_erased
ffi::Function add_type_erased = ffi::Function::FromTyped([](int x, int y) {
  return x + y;
});

// Create a typed function: add_typed
ffi::TypedFunction<int(int, int)> add_typed = [](int x, int y) {
  return x + y;
};

// Convert a typed function to a type-erased function
ffi::Function generic = add_typed;

From Python. Any Python Callable is automatically converted to a tvm_ffi.Function at the ABI boundary. The example below demonstrates that in my_ext.bind:

import tvm_ffi

@tvm_ffi.register_global_func("my_ext.bind")
def bind(func, x):
  assert isinstance(func, tvm_ffi.Function)
  return lambda *args: func(x, *args)  # converted to `tvm_ffi.Function`

def add_x_y(x, y):
  return x + y

func_bind = tvm_ffi.get_global_func("my_ext.bind")
add_y = func_bind(add_x_y, 1)  # bind x = 1
assert isinstance(add_y, tvm_ffi.Function)
print(add_y(2))  # -> 3

tvm_ffi.convert() explicitly converts a Python callable to tvm_ffi.Function:

import tvm_ffi

def add(x, y):
  return x + y

func_add = tvm_ffi.convert(add)
print(func_add(1, 2))

Function#

Calling Convention#

All TVM-FFI functions ultimately conform to the TVMFFISafeCallType signature, which provides a stable C ABI for cross-language calls. The C calling convention is defined as:

int tvm_ffi_c_abi(
  void* handle,           // Resource handle
  const TVMFFIAny* args,  // Input arguments (non-owning)
  int32_t num_args,       // Number of input arguments
  TVMFFIAny* result       // Output argument (owning, zero-initialized)
);

Input arguments. The input arguments are passed as an array of tvm::ffi::AnyView values, specified by args and num_args.

Output argument. The output argument result is an owning tvm::ffi::Any that the caller must zero-initialize before the call.

Important

The caller must zero-initialize the output argument result before the call.

Return value. The ABI returns an error code that indicates:

  • Error code 0: Success

  • Error code -1: Error occurred, retrievable with TVMFFIErrorMoveFromRaised()

  • Error code -2: Very rare frontend error

Hint

See Any for more details on the semantics of tvm::ffi::AnyView and tvm::ffi::Any.

This design is called a packed function, because it “packs” all arguments into a single array of type-erased tvm::ffi::AnyView, and further unifies calling convention across all languages without resorting to JIT compilation.

More specifically, this mechanism enables the following scenarios:

  • Dynamic languages. Well-optimized bindings are provided for, e.g. Python, to translate arguments into packed function format, and translate return value back to the host language.

  • Static languages. Metaprogramming techniques, such as C++ templates, are usually available to directly instantiate packed format on stack, saving the need for dynamic examination.

  • Cross-language callbacks. Language-agnostic tvm::ffi::Function makes it easy to call between languages without depending on language-specific features such as GIL.

Performance Implications. This approach is in practice highly efficient in machine learning workloads.

  • In Python/C++ calls, we can get to microsecond level overhead, which is generally similar to overhead for eager mode;

  • When both sides of calls are static languages, the overhead will go down to tens of nanoseconds.

Note

Although we found it less necessary in practice, further link time optimization (LTO) is still theoretically possible in scenarios where both sides are static languages with a known symbol and linked into a single binary. In this case, the callee can be inlined into caller side and the stack argument memory can be passed into register passing.

Layout and ABI#

tvm::ffi::FunctionObj stores two call pointers in TVMFFIFunctionCell:

  • safe_call: Used for cross-ABI function calls; intercepts exceptions and stores them in TLS.

  • cpp_call: Used within the same DSO; exceptions are thrown directly for better performance.

See Function for the C struct definition.

Important

TVMFFIFunctionCall() is the idiomatic way to call a tvm::ffi::FunctionObj in C, while safe_call or cpp_call remain low-level ABIs for fast access.

Conversion with Any. Since tvm_ffi.Function is a TVM-FFI object, it follows the same conversion rules as any other TVM-FFI object. See Object Conversion with Any for details.

Throw and Catch Errors#

TVM-FFI gracefully handles exceptions across language boundaries without requiring manual error code management.

Important

Stack traces from all languages are properly preserved and concatenated in the TVM-FFI Stable C ABI.

Python. In Python, raise native Exception instances or derived classes. TVM-FFI catches these at the ABI boundary and converts them to tvm::ffi::Error objects. When C++ code calls into Python and a Python exception occurs, it propagates back to C++ as a tvm::ffi::Error, which C++ code can handle appropriately.

C++. In C++, use tvm::ffi::Error or the TVM_FFI_THROW macro:

#include <tvm/ffi/error.h>

void ThrowError(int x) {
  if (x < 0) {
    TVM_FFI_THROW(ValueError) << "x must be non-negative, got " << x;
  }
}

The TVM_FFI_THROW macro captures the current file name, line number, stack trace, and error message, then constructs a tvm::ffi::Error object. At the ABI boundary, this error is stored in TLS and the function returns -1 per the TVMFFISafeCallType calling convention.

Hint

A detailed implementation of such graceful handling behavior can be found in TVM_FFI_SAFE_CALL_BEGIN / TVM_FFI_SAFE_CALL_END macros.

Compiler developers commonly need to look up global functions in generated code. Use TVMFFIFunctionGetGlobal() to retrieve a function by name, then call it with TVMFFIFunctionCall(). See Function for C code examples.

Exception#

This section describes the exception handling contract in the TVM-FFI Stable C ABI. Exceptions are first-class citizens in TVM-FFI, and this section specifies:

  • How to properly throw exceptions from a TVM-FFI ABI function

  • How to check for and propagate exceptions from a TVM-FFI ABI function

When a TVM-FFI function returns a non-zero code, an error occurred. An ErrorObj is stored in thread-local storage (TLS) and can be retrieved with TVMFFIErrorMoveFromRaised().

  • Error code -1: Retrieve the error from TLS, print it, and release via TVMFFIObjectDecRef().

  • Error code -2: A rare frontend error; consult the frontend’s error mechanism instead of TLS.

To raise an error, use TVMFFIErrorSetRaisedFromCStr() to set the TLS error and return -1. For chains of calls, simply propagate return codes - TLS carries the error details.

See Exception for C code examples.

Modules#

A tvm_ffi.Module is a namespace for a collection of functions that can be loaded from a shared library or bundled with the current executable. Modules provide namespace isolation and dynamic loading capabilities for TVM-FFI functions.

Shared Library#

Shared library modules are loaded dynamically at runtime via dlopen (Linux, macOS) or LoadLibraryW (Windows). This is the most common way to distribute and load compiled functions.

Export functions from C++. Use TVM_FFI_DLL_EXPORT_TYPED_FUNC to export a function as a C symbol that follows the TVM-FFI ABI:

#include <tvm/ffi/tvm_ffi.h>

static int AddTwo(int x) { return x + 2; }

// Exports as symbol `__tvm_ffi_add_two`
TVM_FFI_DLL_EXPORT_TYPED_FUNC(add_two, AddTwo);

Load and call from Python. Use tvm_ffi.load_module() to load the shared library:

import tvm_ffi

# Load the shared library
mod = tvm_ffi.load_module("path/to/library.so")

# Access functions by name
result = mod.add_two(40)  # -> 42

# Alternative: explicit function retrieval
func = mod.get_function("add_two")
result = func(40)  # -> 42

Build and load from source. For rapid prototyping, tvm_ffi.cpp.load() compiles C++/CUDA source files and loads them as a module in one step:

import tvm_ffi.cpp

# Compile and load in one step
mod = tvm_ffi.cpp.load(
    name="my_ops",
    cpp_files="my_ops.cpp",
)
result = mod.add_two(40)

Essentially, tvm_ffi.cpp.load() is a convenience function that JIT-compiles the source files and loads the resulting library as a tvm_ffi.Module.

System Library#

System library modules contain symbols that are statically linked to the current executable.

This technique is useful when you want to simulate dynamic module loading behavior but cannot or prefer not to use dlopen or LoadLibraryW (e.g., on iOS). Functions are statically linked to the executable as a system library module. Symbols can be registered via TVMFFIEnvModRegisterSystemLibSymbol() and looked up via tvm_ffi.system_lib().

Register symbols in C/C++. Use TVMFFIEnvModRegisterSystemLibSymbol() to register a symbol during static initialization:

#include <tvm/ffi/c_api.h>
#include <tvm/ffi/extra/c_env_api.h>

// A function following the TVM-FFI ABI
static int add_one_impl(void*, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) {
  TVM_FFI_SAFE_CALL_BEGIN();
  int64_t x = reinterpret_cast<const tvm::ffi::AnyView*>(args)[0].cast<int64_t>();
  reinterpret_cast<tvm::ffi::Any*>(result)[0] = x + 1;
  TVM_FFI_SAFE_CALL_END();
}

// Register during static initialization
// The symbol name follows the convention `__tvm_ffi_<prefix>.<name>`
TVM_FFI_STATIC_INIT_BLOCK() {
  TVMFFIEnvModRegisterSystemLibSymbol(
      "__tvm_ffi_my_prefix.add_one",
      reinterpret_cast<void*>(add_one_impl)
  );
}

Access from Python. Use tvm_ffi.system_lib() to get the system library module:

import tvm_ffi

# Get system library with symbol prefix "my_prefix."
# This looks up symbols prefixed with `__tvm_ffi_my_prefix.`
mod = tvm_ffi.system_lib("my_prefix.")

# Call the registered function
func = mod.add_one  # looks up `__tvm_ffi_my_prefix.add_one`
result = func(10)  # -> 11

Note

The system library is intended for statically linked symbols that exist for the entire program lifetime. For dynamic loading with the ability to unload, use shared library modules instead.

Further Reading#

Contents