DOCSEARCHbyAlgolia
Submit Docs

Apache TVM FFI

repository·main·Indexed 19 days ago

https://github.com/apache/tvm-ffi

A framework-agnostic open ABI and FFI designed for high-performance interoperability between machine learning kernels, DSLs, frameworks like PyTorch and JAX, and various programming languages. It includes the tvm_ffi_orcjit addon for JIT-compiling functions to object files and loading them at runtime, as well as a CUBIN Launcher for executing CUDA kernels from CUBIN files.

Tokens
85.5K
Snippets
275
Records
349
Agent score
64%

What's inside apache-tvm-ffi

  1. Overview of TVM FFI Rust Packages

    main

    The Rust support for the tvm-ffi ABI is currently experimental and consists of three main crates designed to provide varying levels of abstraction over the raw C ABI:

    • tvm-ffi: Provides safe, ergonomic Rust bindings.
    • tvm-ffi-sys: Provides low-level, raw exposure of the C ABIs.
    • tvm-ffi-macros: Provides procedural macros (such as derive/object helpers and exported function helpers) used by tvm-ffi.

    The project prioritizes low-level, direct access to the ABI to ensure maximum efficiency and interoperability.

  2. Overview of Apache TVM FFI

    main

    Apache TVM FFI is an open, framework-agnostic, and minimal ABI (Application Binary Interface) and FFI (Foreign Function Interface) designed for machine learning systems. It provides a stable convention for interoperability between:

    • Kernel libraries: Allowing a single wheel to support multiple frameworks, Python versions, and languages.
    • Kernel DSLs: Providing a reusable ABI for JIT and AOT kernel exposure.
    • Frameworks and Runtimes: Serving as a uniform extension point for libraries like PyTorch, JAX, PaddlePaddle, and NumPy/CuPy.
    • ML Infrastructure: Enabling out-of-the-box bindings across Python, C++, and Rust.
    • Coding Agents: Providing a unified mechanism for shipping generated code into production.

  4. Generate Python type stubs using tvm-ffi-stubgen

    main

    TVM-FFI provides tvm-ffi-stubgen, a tool that generates Python type stubs from C++ reflection metadata. This enables IDE auto-completion and static type checking for registered global functions and classes.

    There are two primary ways to use this tool:

    1. CMake-based Generation (Recommended): Automatically runs stub generation after each build using the tvm_ffi_configure_target function.
    2. CLI-based Generation: Used for standalone usage or custom build systems.

  5. Understand the TVM-FFI ABI Core Principles

    main

    The TVM-FFI ABI is designed for machine learning systems with the following goals:

    • Minimal and efficient: Provides close-to-metal performance with simple structures.
    • Stability guarantee: Remains stable across compiler versions and is independent of host languages or frameworks.
    • Expressive for machine learning: Includes native support for tensors, shapes, and common ML data types.
    • Extensible: Uses a dynamic type registration system to support user-defined types.

    Authoritative specifications are found in the C headers:

    • tvm/ffi/c_api.h: Core ABI.
    • tvm/ffi/extra/c_env_api.h: Extra support features.

  6. Overview of tvm-ffi C++ Container Types

    main

    tvm-ffi provides several built-in container types for passing and storing collections of values compatible with the FFI. These containers are categorized by their mutability and semantics:

    TypeHeaderMutabilitySemantics
    Array<T>container/array.hImmutable (copy-on-write)Homogeneous sequence
    List<T>container/list.hMutable (shared reference)Homogeneous sequence
    Tuple<Ts...>container/tuple.hImmutable (copy-on-write)Heterogeneous fixed-size sequence
    Map<K,V>container/map.hImmutable (copy-on-write)Homogeneous key-value mapping
    Dict<K,V>container/dict.hMutable (shared reference)Homogeneous key-value mapping

  7. Overview of TVM-FFI Container Types

    main

    TVM-FFI provides five built-in container types for storing and exchanging collections of values across C++, Python, and Rust. All containers are heap-allocated, reference-counted objects that can be stored in tvm::ffi::Any and passed through the FFI boundary.

    Containers are categorized into two types based on their memory semantics:

    1. Immutable Containers: Use copy-on-write semantics. Mutations create a new backing object if the current one is shared.
    2. Mutable Containers: Use shared-reference semantics. Mutations happen in-place, and all handles to the object see the changes.
    TypeC++ ClassPython ClassMutabilitySemantics
    Arrayffi::Array<T>tvm_ffi.ArrayImmutableHomogeneous sequence (copy-on-write)
    Listffi::List<T>tvm_ffi.ListMutableHomogeneous sequence (shared-reference)
    Tupleffi::Tuple<Ts...>(backed by Array)ImmutableHeterogeneous fixed-size sequence
    Mapffi::Map<K, V>tvm_ffi.MapImmutableHomogeneous key-value (copy-on-write)
    Dictffi::Dict<K, V>tvm_ffi.DictMutableHomogeneous key-value (shared-reference)

  8. Overview of TVM-FFI export mechanisms

    main

    TVM-FFI provides three distinct mechanisms to expose functions and classes across C, C++, and Python, depending on your use case:

    1. C Symbols: Best for kernel libraries and compiler codegen. Functions are exported as __tvm_ffi_<name> in a shared library and loaded via tvm_ffi.load_module.
    2. Global Functions: Best for application-level APIs and cross-language callbacks. Functions are registered by a string name in a shared registry and can be retrieved from any language.
    3. Classes: Best for structured data. You define a C++ subclass of tvm::ffi::Object, which can then be used in Python as a dataclass with fields and methods.

  9. Understand Python and C++ object lifetime interoperability

    main

    TVM-FFI manages cross-language lifetimes using reference counting. Each Python tvm_ffi.Object instance holds a C handle (void*) to the underlying C++ object.

    Lifetime Behavior:

    1. Construction: When a Python object is created (e.g., obj = MyObject(...)), the C++ object is created and its reference count is set to 1.
    2. Aliasing: Creating a new Python variable pointing to the same object (e.g., obj2 = obj) creates a Python alias but does not change the C++ reference count.
    3. Deletion: Removing a Python reference (e.g., del obj) decrements the reference count. When the last Python reference is removed and the count reaches 0, the C++ object is destroyed.
    obj = MyObject(42, "test")    # C++ object created, C++ refcount = 1
obj2 = obj                    # Python alias created, C++ refcount unchanged
del obj                       # Python alias removed, C++ refcount unchanged
del obj2                      # Last Python reference gone, C++ refcount -> 0, object destroyed

  10. Load and execute CUDA kernels from CUBIN files

    main

    The CUBIN Launcher provides a way to load and execute CUDA kernels from CUBIN files using TVM-FFI. It uses the cubin_launcher.h header to wrap the CUDA Runtime or Driver API, providing RAII-based resource management via CubinModule and CubinKernel objects.

    It supports two primary backend APIs:

    • CUDA Runtime API (CUDA >= 12.8): Uses cudaLibraryLoadData(), cudaLibraryGetKernel(), and cudaLaunchKernel().
    • CUDA Driver API: Uses cuLibraryLoadData(), cuLibraryGetKernel(), and cuLaunchKernel().

    By default, the implementation uses the Runtime API if compiled with CUDA >= 12.8, falling back to the Driver API for older versions. You can force the API choice by defining the macro TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API before including the header (1 for Driver API, 0 for Runtime API).

  11. How Any and AnyView work for type-erased value passing

    main

    At the core of TVM-FFI is TVMFFIAny, a 16-byte tagged union used for type-erased value passing across language boundaries. It can hold either atomic POD (Plain Old Data) types or pointers to heap-allocated objects.

    Ownership Semantics

    In C++, these are wrapped by two classes with identical memory layouts but different ownership models:

    • Owning (tvm::ffi::Any): Reference-counted; manages the object's lifetime.
    • Borrowing (tvm::ffi::AnyView): A non-owning view; the caller must ensure the underlying data remains valid.

    To convert a borrowing tvm::ffi::AnyView to an owning tvm::ffi::Any, use the function TVMFFIAnyViewToOwnedAny.

    Runtime Type Index

    The type_index field determines the storage strategy:

    • Atomic POD types: (where type_index < kTVMFFIStaticObjectBegin) Stored inline in the payload union without heap allocation or reference counting.
    • Object types: (where type_index >= kTVMFFIStaticObjectBegin) Stored as pointers to heap-allocated, reference-counted TVM-FFI objects.

  12. Understand the low-level StructuralVisitor

    main

    The StructuralVisitor is a low-level traversal object intended for use within custom structural visit hooks or C++ integrations. Python users should generally use structural_walk instead.

    Key Methods

    • visitor.visit(value): Recursively visit a child value.
    • visitor.def_region_kind(): Inspect the current definition-region mode.
    • visitor.with_def_region_kind(kind, callback): Run a recursive visit under a temporary definition-region mode.

    Custom Visit Hooks

    Custom hooks are registered via the __s_visit__ type attribute. These hooks receive the active visitor and the current object, and are responsible for manually calling visitor.visit(child) on the object's structural children.