当前位置:首页>python>TVM FFI 深度剖析:Python 如何调用 C++

TVM FFI 深度剖析:Python 如何调用 C++

  • 2026-10-11 06:05:21
TVM FFI 深度剖析:Python 如何调用 C++

从 Apache TVM 源码逐层拆解 Foreign Function Interface 的完整架构


TVM 的核心编译器、runtime、内存管理全部由 C++ 实现,而用户接口层是 Python。两者之间通过 tvm-ffi[1] 这个独立的 FFI(Foreign Function Interface)库连接。要理解 TVM 的工作方式,必须先理解这个桥梁。

整体架构图


Layer 6 — 物理根基:编译产物(多库拆分)

最底层的基础 — TVM 的 C++ 代码在编译期(cmake --build 阶段)就编译好了。

编译过程

#   Step 1: CMake 配置(根据 config.cmake 选择后端)cmake .. -DUSE_LLVM=ON -DUSE_CUDA=ON#   Step 2: 编译cmake --build . --parallel#   Step 3: 产物(不是单一 libtvm.so,而是拆分为三个库)build/lib/libtvm_compiler.dylib   #  编译器核心(含所有 IR/Pass/Operator/Codegen)build/lib/libtvm_runtime.dylib    #  轻量级 runtime(仅推理执行所需)build/lib/libtvm_ffi.dylib        #  FFI 桥接库(Python ↔ C++ 的函数注册/调用)build/lib/libtvm_ffi_static.a     #  FFI 静态库(用于 MicroTVM 等嵌入式场景)

为什么拆成三个库?

库
用途
场景
libtvm_compiler
完整编译器栈(IR/Pass/Codegen)
开发 / 训练 / 调优
libtvm_runtime
最小化推理运行时
生产部署(体积小)
libtvm_ffi
C ABI + C++17 API + Python 绑定
所有跨语言功能
libtvm_ffi_static.a
FFI 静态库
MicroTVM 嵌入式

Python 侧加载(tvm_ffi Cython 扩展)

#   base.pxi 通过 cdef extern 声明 libtvm_ffi.dylib 中的 C API 符号:#   int TVMFFIFunctionGetGlobal(TVMFFIByteArray* name, TVMFFIObjectHandle* out)#   int TVMFFIFunctionSetGlobal(TVMFFIByteArray* name, TVMFFIObjectHandle f, int override)#   int TVMFFIObjectDecRef(TVMFFIObjectHandle obj)#   ... 等等

当 Python 执行 import tvm_ffi 时,Cython 编译生成的 C 扩展自动 load,进而可以调用 libtvm_compiler.dylib 和 libtvm_runtime.dylib 中暴露的所有 C++ API。

关键理解:C++ 不是在 Python 调用时才编译的,而是在 pip install tvm 或 cmake --build 时就已经编译成 build/lib/ 下的 .dylib/.so 了。Python 只是通过 tvm-ffi 的 Cython 扩展访问这些已存在的库。


Layer 5 — C++ API 实现层:全局注册表 + 类型擦除

FFI 在 C++ 侧的"接线板",由 tvm-ffi 库提供(3rdparty/tvm-ffi/src/ffi/)。两个核心机制:

全局函数注册表

tvm-ffi 在 C++ 侧维护一个全局的函数表:

  全局注册表(tvm::ffi 内部)┌──────────────────────────────────────────────────┐│  "ir.conv2d"                          → Function ││  "relax.op.nn.matmul"                 → Function ││  "tirx.build"                         → Function ││  "ffi.FunctionListGlobalNamesFunctor" → Function ││  "ffi.Array"                          → Function ││  ...                                             │└──────────────────────────────────────────────────┘

C++ 侧注册

//  方式1: 通过 tvm::ffi C++ APItvm::ffi::Function::RegisterGlobal("ir.conv2d")    .set_body_typed([](Expr data, Expr weight, ...) {        return conv2d(data, weight, ...);    });//  方式2: 通过反射注册表自动注册 (tvm::ffi::reflection)// 在静态初始化阶段,注册表条目自动调用 TVMFFIFunctionSetGlobalFromMethodInfo

Python 侧注册

import tvm_ffi@tvm_ffi.register_global_func("my.echo")def echo(x):    return x#  内部调用 TVMFFIFunctionSetGlobal("my.echo", packed_func, override=0)

双向 FFI — 之后 Python 注册的函数可以直接被 C++ 或其他语言调用。

静态初始化:当 libtvm_compiler.dylib 被 dlopen 加载时,所有 C++ static 初始化代码自动执行,注册表就已填满。Python 侧不需要任何手动注册。

类型擦除的 Function 对象

tvm-ffi 的核心类型是 Function,定义在 include/tvm/ffi/function.h:

  任何 callable 都可以包装成 Function:    C++ lambda:  通过 make_function_wrapper    C 函数指针:  通过 TVMFFIFunctionCreate    Python 函数: 通过 _convert_to_ffi_func (Cython)  统一的调用约定:    (const AnyView* args, int32_t num_args, Any* rv) → void  其中 AnyView 是类型擦除的值容器:    int, float, string, ObjectRef, DLTensor*, ...

Layer 4 — Cython C-API 桥接层:零拷贝穿越

最精细、性能最关键的一层。位于 3rdparty/tvm-ffi/python/tvm_ffi/cython/。

为什么用 Cython 而不是纯 ctypes?

维度
ctypes ❌
Cython ✅
性能
每次调用都有 Python↔C 开销
编译成 C 扩展,接近原生速度
类型安全
运行时检查,弱
编译期检查,强
GIL
默认持有 GIL
可精细控制 nogil
对象管理
手动管理 C++ 生命周期
自动 ref-counting

关键文件与职责

  cython/ 目录结构├──   base.pxi          ← cdef extern: 声明 C API 符号├──   core.pyx          ← Cython 主模块: 导入时触发加载├──   function.pxi      ← Function 类 + 查找/注册逻辑├──   object.pxi        ← CObject 类: 句柄 + ref-counting├──   tensor.pxi        ← Tensor/DLTensor: DLPack 协议└──   tvm_ffi_python_helpers.h ← TVMFFIPyFuncCall C 实现

base.pxi — C 函数声明

#  这不是实现,只是声明 "libtvm_ffi.dylib 中有这些 C API 符号"cdef extern from "tvm/ffi/c_api.h":    intTVMFFIFunctionGetGlobal(TVMFFIByteArray* name, TVMFFIObjectHandle* out) nogil    int TVMFFIFunctionSetGlobal(TVMFFIByteArray* name, TVMFFIObjectHandle f, intoverride) nogil    int TVMFFIObjectDecRef(TVMFFIObjectHandle obj) nogil    ...#  Python helper (在 tvm_ffi_python_helpers.h 中实现):cdef extern from "tvm_ffi_python_helpers.h":    int TVMFFIPyFuncCall(void* func_handle, PyObject* py_arg_tuple,                         TVMFFIAny* result, int* c_api_ret_code,                         int release_gil, const DLPackExchangeAPI** c_ctx_dlpack_api)    ...

function.pxi — Function 类(实际调用入口)

cdef class Function(CObject):    cdef int c_release_gil    def __call__(self, *args:Any) -> Any:        cdef TVMFFIAny result        cdef int c_api_ret_code        cdef const DLPackExchangeAPI* c_ctx_dlpack_api = NULL        #   Step 1: 初始化返回值        result.type_index = kTVMFFINone        result.v_int64 = 0        #   Step 2: 调用 Python helper —— PyObject* 直接传入,C 层完成转换        TVMFFIPyFuncCall(            (<CObject>self).chandle, <PyObject*>args,            &result, &c_api_ret_code,            self.release_gil,            &c_ctx_dlpack_api        )        #   Step 3: 检查错误码        if c_api_ret_code == 0:            return make_ret(result, c_ctx_dlpack_api)        if c_api_ret_code == -2:            raise raise_existing_error()        error = move_from_last_error()        raise error.py_error()

Function.__call__() = 穿越点 — Python 调用在此进入 libtvm_ffi.dylib 地址空间,路由到 libtvm_compiler.dylib 中的 C++ 函数。

关键性能设计:TVMFFIPyFuncCall 接受 PyObject* 元组并在 C 层遍历转换(PyArgDispatcher),整个参数列表只需 一次 语言边界穿越。

_get_global_func 实现

def _get_global_func(name: str, allow_missing: bool):    cdef TVMFFIObjectHandle chandle    cdef ByteArrayArg name_arg = ByteArrayArg(c_str(name))    #  调用 C API 查找全局函数    CHECK_CALL(TVMFFIFunctionGetGlobal(name_arg.cptr(), &chandle))    if chandle != NULL:        ret = Function.__new__(Function)        (<CObject>ret).chandle = chandle  #  包装 C++ 句柄        return ret    if allow_missing:        return None    raise ValueError("Cannot find global function %s" % name)

_register_global_func 实现

def _register_global_func(name: str, pyfunc, override: bool) -> Function:    if not isinstance(pyfunc, Function):        pyfunc = _convert_to_ffi_func(pyfunc)  #  Python → Function    CHECK_CALL(TVMFFIFunctionSetGlobal(name_arg.cptr(), pyfunc.chandle, override))    return pyfunc

对象引用计数

双向绑定 — Python GC 自动触发 C++ 释放。

cdef class CObject:    cdef void* chandle    def __dealloc__(self):        if self.chandle != NULL:            CHECK_CALL(TVMFFIObjectDecRef(self.chandle))  #   C++ refcount--            self.chandle = NULL

Layer 3 — 注册表层:get_global_func() 与 init_ffi_api()

核心文件:3rdparty/tvm-ffi/python/tvm_ffi/registry.py[2] — Python 代码中最常打交道的入口。

get_global_func() — 按名查找

def get_global_func(name, allow_missing=False):    """按名称从 C++ 全局注册表中查找函数"""    return core._get_global_func(name, allow_missing)    #  core._get_global_func 调用 TVMFFIFunctionGetGlobal    #  返回 Function 对象(Python 侧的 C++ Function 包装)

register_global_func() — Python → C++ 注册

def register_global_func(func_name, f=None, override=False):    """将 Python 函数注册到 C++ 全局注册表"""    #  内部调用 core._register_global_func    #  → TVMFFIFunctionSetGlobal(name, packed_func, override)

双向 FFI — C++ 也可以反过来调用 Python 注册的函数(常见于 callback)。

register_object() — Python 类 ↔ C++ 类型绑定

@register_object("tvm.ir.GlobalVar")class GlobalVar(Object):    pass

绑定后的效果:

•C++ 返回的 GlobalVar 对象自动包装成 Python GlobalVar•Python 子类自动获得所有 C++ 反射字段的 property 访问

list_global_func_names() — 列出所有函数

def list_global_func_names() -> list[str]:    """列出 C++ 侧注册的所有函数名"""    name_functor = get_global_func("ffi.FunctionListGlobalNamesFunctor")()    num_names = name_functor(-1)    return [name_functor(i) for i in range(num_names)]

init_ffi_api() — 批量绑定(关键!)

def init_ffi_api(namespace: str, target_module_name: str | None = None) -> None:    target_module_name = target_module_name if target_module_name else namespace    #  去掉可选的 "tvm." 前缀    if namespace.startswith("tvm."):        prefix = namespace[4:]    else:        prefix = namespace    target_module = sys.modules[target_module_name]    #  遍历所有已注册的全局函数    for name in list_global_func_names():        if not name.startswith(prefix):            continue        fname = name[len(prefix) + 1:]   #  提取函数短名        if fname.find(".") != -1:         #  跳过子命名空间            continue        f = get_global_func(name)        setattr(f, "__name__", fname)        setattr(target_module, fname, f)  #  挂载为 Python 模块属性

源码验证:3rdparty/tvm-ffi/registry.py:305-352[3]

⚠️ 注意:这是 eager(立即)加载,不是懒加载。init_ffi_api 在 import 时立即遍历所有全局函数名并 setattr。

实际调用示例

#  python/tvm/ir/_ffi_api.py (实际源码)import tvm_ffitvm_ffi.init_ffi_api("ir", __name__)#  namespace="ir" → prefix="ir" → 挂载所有 "ir.*" 函数
from tvm.ir import _ffi_api#  _ffi_api.conv2d → 直接是 Function 对象(已在 import 时挂载)result = _ffi_api.conv2d(data, weight)

Layer 2 — 模块级批量绑定层

每个 TVM 子模块一个 _ffi_api.py(共 49 个_ffi_*.py),按模块组织 C++ 函数入口。

python/tvm/├──   ir/_ffi_api.py              # init_ffi_api("ir", __name__)├──   ir/_ffi_analysis_api.py     # init_ffi_api("ir.analysis", __name__)├──   ir/_ffi_transform_api.py    # init_ffi_api("transform", __name__)├──   ir/_ffi_instrument_api.py   # init_ffi_api("instrument", __name__)├──   relax/_ffi_api.py           # init_ffi_api("relax", __name__)├──   tirx/_ffi_api.py            # init_ffi_api("tirx", __name__)├──   s_tir/_ffi_api.py           # init_ffi_api("s_tir", __name__)├──   driver/_ffi_api.py          # init_ffi_api("driver", __name__)├──   target/_ffi_api.py          # init_ffi_api("target", __name__)├──   arith/_ffi_api.py           # init_ffi_api("arith", __name__)├──   runtime/_ffi_api.py         # init_ffi_api("runtime", __name__)├──   runtime/_ffi_node_api.py    # init_ffi_api("node", __name__)├──   support.py                  # init_ffi_api("support", __name__)├──   topi/cpp/generic.py         # init_ffi_api("topi.generic", ...)├──   topi/cpp/cuda.py            # init_ffi_api("topi.cuda", ...)├──   topi/cpp/x86.py             # init_ffi_api("topi.x86", ...)└──   ...(共 49 个)

设计优点:C++ 侧新增函数注册后,Python 侧零改动即可调用 — namespace 前缀匹配即自动挂载。


Layer 1 — 公共 API 层

用户可见的入口。python/tvm/ffi.py 直接重导出 tvm_ffi 全部 API。

#  python/tvm/ffi.py (实际源码)"""Redirects to tvm_ffi"""from tvm_ffi import *
#  python/tvm/__init__.py (实际源码)from tvm_ffi import register_object, register_global_func, get_global_func

用户典型用法:

import tvm#  方式 1: 按名查找f = tvm.get_global_func("relax.op.nn.conv2d")result = f(data, weight)#  方式 2: 通过 _ffi_api(更常用)from tvm.ir import _ffi_apiresult = _ffi_api.conv2d(data, weight)

⏱ 完整调用时序图

以一个真实调用 _ffi_api.conv2d(data, weight) 为例:

═══════════════════════════════════════════════════════════════════════  阶段 A: Import 绑定(仅执行一次)═══════════════════════════════════════════════════════════════════════t=0     Python: import tvm.ir._ffi_api        │       ↓ init_ffi_api("ir", __name__)t=1     Python: list_global_func_names() → 遍历 C++ 注册表        │       → get_global_func("ir.conv2d") → setattr "conv2d"        │       → get_global_func("ir.matmul")  → setattr "matmul"        │       → ... (所有 "ir.*" 一次性挂载)        │       ↓ 就绪: _ffi_api.conv2d 已是 Function 对象═══════════════════════════════════════════════════════════════════════  阶段 B: 函数调用(每次调用执行)═══════════════════════════════════════════════════════════════════════t=2    Python: _ffi_api.conv2d(data, weight)       │        ↓ 进入 Cythont=3    Cython: Function.__call__(self, data, weight)       │        result.type_index = kTVMFFINone       │        ↓ TVMFFIPyFuncCall(...)t=4      C:   PyArgDispatcher 遍历 args 元组,逐元素转换       │       Python int→C int, Python Tensor→DLTensor*, ...       │       ↓ 构建 AnyView[] 传给 C++t=5     C++:  tvm::ffi::Function::Call(args, &ret)       │        → 解包 AnyView 参数       │        → 调用注册的 lambda: conv2d(Expr, Expr, ...)       │        → 返回结果到 Any* ret       │        ↓t=6      C:   返回 ret_code=0t=7    Cython: make_ret(result, dlpack_api)       │        → type_index switch: Tensor→DLPack, Object→PyClass, ...       │        ↓t=8    Python: result = ...  

关键性能特点

特点
说明
仅绑定一次
Import 阶段执行,之后 _ffi_api.xxx 是模块属性,不再查表
一次穿越
所有参数打包成 PyObject* 在 C 层批量转换,只穿越语言边界一次
零拷贝
ObjectRef 传递 void* chandle 裸指针,无序列化
原生速度
Cython 编译为 C 扩展模块,调用路径极短
按类型分发
make_ret()
 通过 type_index switch 精确创建 Python 对象

总结

层
位置
职责
关键 API
6
build/lib/
编译产物(三个库拆分)
cmake build
5
tvm-ffi/src/ffi/
C++ 注册表 + 类型擦除
Function::RegisterGlobal
4
cython/function.pxi
Cython 桥接,穿越语言边界
TVMFFIPyFuncCall
3
tvm_ffi/registry.py
Python 注册表 API
get_global_func
, init_ffi_api
2
python/tvm/*/_ffi_api.py
批量绑定 (49 文件)
init_ffi_api("ir", __name__)
1
python/tvm/__init__.py
公共 API 入口
from tvm_ffi import ...

一句话概括:libtvm_compiler + libtvm_runtime + libtvm_ffi 三个库在编译期就已包含所有 C++ 代码和静态初始化的注册表,Python 通过 tvm_ffi 独立 Cython 扩展以原生速度跨过语言边界,init_ffi_api 提供了 zero-boilerplate 的自动函数绑定,用户通过 tvm.get_global_func() 或 _ffi_api.xxx() 透明调用 C++ 函数,全程零序列化开销。


•基于 Apache TVM 源码逐行走读验证编写*

References

[1] tvm-ffi: 3rdparty/tvm-ffi/[2]: 3rdparty/tvm-ffi/python/tvm_ffi/registry.py[3] 3rdparty/tvm-ffi/registry.py:305-352: 3rdparty/tvm-ffi/python/tvm_ffi/registry.py#L305-L352

最新文章

随机文章