TVM FFI（三）：多语言绑定与工具链

这是 TVM FFI 系列第三篇。前两篇讲了核心骨架和数据层，这一篇把剩下的生态拼图补全：多语言绑定、JIT 编译、工程工具和运行时机制。

Python SDK 全景

Python 侧不只是简单的 C 接口包装——它是一套完整的运行时，提供自动类型转换、运算符重载和跨语言异常。

自动类型转换

_convert.py 里的 convert() 函数是核心——当 Python 值传给 FFI 函数时，自动做如下映射：

import tvm_ffi

# Python list/tuple → ffi.Array
a = tvm_ffi.convert([1, 2, 3])       # → Array<IntImm>

# Python dict → ffi.Map
m = tvm_ffi.convert({"a": 1, "b": 2}) # → Map<String, IntImm>

# Python str/bytes → ffi.String/Bytes（零拷贝）
s = tvm_ffi.convert("hello")          # → String
b = tvm_ffi.convert(b"world")         # → Bytes

# Python callable → ffi.Function
f = tvm_ffi.convert(lambda x: x + 1)  # → Function

对 PyTorch / NumPy 张量的转换也是自动的：

import torch
import tvm_ffi

# torch.Tensor → ffi.Tensor（DLPack 零拷贝）
x = torch.randn(3, 4)
ffi_x = tvm_ffi.convert(x)   # 内部走 DLPack，数据不动

# numpy.ndarray 同理
import numpy as np
a = np.array([1, 2, 3])
ffi_a = tvm_ffi.convert(a)

运算符重载

C++ 侧定义的算术运算（Add、Sub、Mul 等），Python 侧也自动支持：

expr_a = tvm.tirx.Add(tvm.tirx.Var("x"), 1)
expr_b = tvm.tirx.Mul(expr_a, 2)

# 等价于
expr_b = (tvm.tirx.Var("x") + 1) * 2

_dunder.py 负责定义每个 IR 节点的 __add__、__mul__、__getitem__ 等方法，IR 构造体验和纯 Python 代码一样自然。

跨语言异常

C++ 里抛出的异常会被 Python 端正常捕获：

// C++ 侧
TVM_FFI_THROW(ValueError) << "Expected at least 2 arguments";

# Python 侧
try:
    mod.bad_func()
except ValueError as e:
    print(e)  # "Expected at least 2 arguments"

tvm_ffi/error.py 维护了 C++ 异常类型到 Python 异常类型的映射表。反向亦然——Python 回调中抛出的异常也能被 C++ 端捕获。

Cython 胶水层

cython/core.pyx 是 Python 和 C++ 之间的实际执行通道。它处理：

TVMFFIAny 的打/拆包——Python 整/浮点数 ↔ C TVMFFIAny
引用计数管理——进入 FFI 边界时 IncRef，离开时 DecRef
DLTensor 构造——从 torch.Tensor 的 __dlpack__() 提取底层指针
packed function 调用——参数打包成 TVMFFIAny[]，调用函数指针，解包返回值

这层是性能关键路径，必须用 Cython（编译成 C）而非纯 Python。

注册表

全局函数注册表是 TVM 的核心发现机制：

import tvm_ffi

# 注册函数
@tvm_ffi.register_global_func("my_package.norm")
def compute_norm(x):
    return (x * x).sum().sqrt()

# 之后任何地方都可以通过名字查到
fn = tvm_ffi.get_global_func("my_package.norm")
result = fn(tensor)

# 也可以列出所有注册的函数
for name in tvm_ffi.list_global_func_names():
    print(name)

JIT 编译系统

这是 TVM FFI 最独特的特性之一——Python 字符串里写 C++/CUDA，即时编译为可调用的 Module。

load_inline 的原理

import tvm_ffi.cpp

mod = tvm_ffi.cpp.load_inline(
    name="my_kernels",
    cpp_sources=r"""
        void add_one(tvm::ffi::TensorView x, tvm::ffi::TensorView y) {
            for (int i = 0; i < x.size(0); ++i)
                ((float*)y.data_ptr())[i] = ((float*)x.data_ptr())[i] + 1;
        }
    """,
    functions=["add_one"],
)

extension.py 在背后做的事情：

生成 CMake 项目——根据 sources 和 functions 生成 CMakeLists.txt
哈希缓存——sources + flags 做 SHA256，避免重复编译
调用 Ninja 构建——增量编译，只编译修改过的部分
dlopen 加载——Module.load_from_file() 动态加载产物
符号发现——扫描 .so 中的 __tvm_ffi_<name> 符号，生成函数表

整个过程对用户是一个函数调用，但背后是一套完整的构建系统。

CUDA NVRTC 路径

import tvm_ffi.cpp

mod = tvm_ffi.cpp.load_inline(
    name="gpu_kernels",
    cuda_sources=r"""
        __global__ void VecAddKernel(float* a, float* b, float* c, int n) {
            int i = blockIdx.x * blockDim.x + threadIdx.x;
            if (i < n) c[i] = a[i] + b[i];
        }
        void vec_add(TensorView a, TensorView b, TensorView c) {
            int n = a.size(0);
            cudaStream_t stream = TVMFFIEnvGetStream(a.device());
            VecAddKernel<<<(n+255)/256, 256, 0, stream>>>(
                (float*)a.data_ptr(), (float*)b.data_ptr(), (float*)c.data_ptr(), n);
        }
    """,
    functions=["vec_add"],
)

nvrtc.py 封装了 NVRTC 库调用，在运行时编译 PTX。和 CPU 路径一样——用户不需要碰 CMake 或编译器命令行。

Rust 绑定

tvm-ffi/rust/ 下有三个 crate，设计思路和 Python 侧镜像：

tvm-ffi-sys

最底层，用 libloading 动态加载 libtvm_ffi 共享库，生成所有 C ABI 函数的 raw FFI 绑定。

tvm-ffi

安全的 Rust API，提供和 C++/Python 一致的类型：

use tvm_ffi::{Module, Tensor, Function, Result, into_typed_fn};

fn main() {
    // 加载 .so
    let lib: Module = Module::load_from_file("add_one_cpu.so").unwrap();

    // 查找函数
    let add_one: Function = lib.get_function("add_one_cpu").unwrap();

    // 转换为类型安全的包装
    let typed_add_one = into_typed_fn!(
        add_one,
        Fn(&Tensor, &Tensor) -> Result<()>
    );

    // 调用
    let x = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], &[4]).unwrap();
    let y = Tensor::from_slice(&[0.0f32; 4], &[4]).unwrap();
    typed_add_one(&x, &y).unwrap();
    // y.data_as_slice::<f32>() → &[2.0, 3.0, 4.0, 5.0]
}

关键的是 into_typed_fn! 宏——它在编译时将类型擦除的 packed function 包装为带类型的 Rust 函数，参数和返回值的打包/拆包完全在后台处理。

tvm-ffi-macros

过程宏，让 Rust struct 可以注册为 TVM FFI Object：

#[tvm_ffi_macros::object(type_key = "my.Type")]
struct MyType {
    value: i64,
    name: String,
}

效果和 C++ 侧的 TVM_FFI_DECLARE_OBJECT_INFO + RegisterReflection() 等价。

Stubgen：类型标注自动生成

这是一个被低估的工具。C++ 侧通过反射注册了每个字段的类型和每个方法的签名。stubgen 读取这些信息，自动生成 Python .pyi stub 文件：

uv run tvm-ffi-stubgen python

效果——IDE 能为 FFI 对象提供完整的自动补全：

from tvm import tirx

expr = tirx.Add(...)
# IDE 知道 expr 有 .a, .b 字段，类型都是 PrimExpr
# IDE 知道 tirx.Var 的构造函数接受 name_hint: str 和 dtype: DataType

不是运行时内省——是编译时类型标注，和手写的没有区别。

Addons

torch_c_dlpack_ext

PyTorch ≤ 2.9 的 DLPack 实现和标准有微小偏差。这个 addon 提供兼容层，让老版本 PyTorch 也能零拷贝传张量：

pip install torch-c-dlpack-ext

安装后，torch.Tensor 自动获得正确的 __dlpack__ / __dlpack_device__ 方法。

tvm_ffi_orcjit

基于 LLVM ORC JIT 引擎的模块加载器。和 load_inline 不同——它不生成 .so 文件，而是直接在内存中 JIT 编译：

// 直接内存编译，不落盘
auto mod = tvm_ffi::orcjit::CompileModule(cpp_sources, functions);

适合不需要缓存、快速迭代的场景。

Stream 管理

GPU kernel 需要感知当前的 CUDA stream——否则结果可能和框架的计算不同步。TVM FFI 通过环境感知机制解决：

# C++ 侧 kernel 中自动获取当前 stream
# TVMFFIEnvGetStream(device_type, device_id)
# → 返回调用方的 CUDA stream 句柄

# Python 侧自由切换 stream 上下文
from tvm_ffi import use_torch_stream

with torch.cuda.stream(my_stream):
    with use_torch_stream(my_stream):
        # 在此上下文中调用的 FFI kernel 自动获取正确的 stream
        mod.cuda_kernel(x, y)

C++ kernel 调用 TVMFFIEnvGetStream(x.device()) 时，得到的 stream 和 Python 侧 torch.cuda.current_stream() 一致。不需要显式传入——这在 kernel 库开发中是巨大的便利。

Access Path：值来源追踪

一个有趣的调试/优化工具。AccessPath 记录一个值在 IR 树中的完整路径：

auto mismatch = StructuralEqual::GetFirstMismatch(expr_a, expr_b);
if (mismatch) {
    // mismatch->lhs → "/func.body[2]/body/indices[0]"
    // mismatch->rhs → "/func.body[2]/body/indices[0]"
    // 精确知道是哪棵树的哪个字段不相等
}

对于大型 IR 模块——包含数千个节点的 IRModule——这个机制能在秒级定位到不匹配的具体位置。

全系列总结

三篇文章覆盖了 TVM FFI 的完整面貌：

篇章	内容
第一篇	对象系统、类型擦除、Packed Function、DLPack 零拷贝
第二篇	容器类型、Any/AnyView、结构相等性、反射、Dataclass
第三篇	Python/Rust 绑定、JIT 编译、stubgen、addons、stream 管理

贯穿始终的设计哲学很简单：用稳定的 C ABI 做”通用语言”，在各语言之上提供类型安全的原生体验。C++ 程序员用 TVM_FFI_DLL_EXPORT_TYPED_FUNC，Python 程序员用 mod.add_one(x, y)，Rust 程序员用 into_typed_fn!——三种体验，一个 ABI。