TVM 核心库架构

TVM 的 C++ 产物被拆成多个共享库，职责清晰、可按需部署。本文按依赖层次和编译/运行分工说明各库的功能与作用。

整体关系

库	何时需要	典型体积
libtvm_ffi	始终	小
libtvm_runtime	始终（执行）	中
libtvm_compiler	编译/优化 IR	大（含 LLVM）
libtvm_runtime_cuda	在 GPU 上跑 CUDA 算子	小

---

1. tvm-ffi（libtvm_ffi.so）

定位：最底层的开放 ABI 和 FFI 基础设施，已从 TVM 独立为 apache-tvm-ffi 项目，TVM 通过 3rdparty/tvm-ffi 引入。

核心能力

• 稳定 C ABI：跨语言、跨框架的统一调用约定
• 类型系统：Any / AnyView（类型擦除值）、Object / ObjectRef（引用计数对象）
• 函数注册：Function + 全局注册表（register_global_func / get_global_func）
• Tensor 互操作：基于 DLPack 与 PyTorch / JAX / CuPy 零拷贝交换
• Module 加载：load_module("xxx.so") 动态加载用户编译出的 kernel

在 TVM 中的角色

libtvm_runtime 和 libtvm_compiler 都链接 tvm_ffi_shared（见 CMakeLists.txt）。Python 侧 import tvm 时，最先加载的是 tvm_ffi 的 Cython 扩展；TVM 的 C++ 代码通过 #include <tvm/ffi/...> 使用同一套基础设施。

类比：tvm-ffi 像「操作系统内核接口」——对象模型、函数调用、模块加载都由它定义，TVM 和其他 kernel 库（FlashInfer、TileLang 等）都建在其上。

---

2. libtvm_runtime（libtvm_runtime.so）

定位：设备无关的通用执行运行时，负责「跑已经编译好的程序」。

包含的源码

CMake 中 RUNTIME_SRCS 主要来自：

• src/runtime/*.cc
• src/runtime/vm/*.cc
• src/runtime/memory/*.cc
• src/runtime/rpc/minrpc/*.cc（可选）

主要功能：

模块	作用
Module / PackedFunc	加载 .so 模块、调用编译产物
VM（Relax VM）	执行 Relax 编译后的字节码
Memory	内存池、工作区管理
RPC	远程设备调试/部署（可选）
Registry	全局函数/类型注册

关键设计

1. 与设备后端解耦：CUDA / ROCm / Vulkan 等设备 API 不在此库，而是独立 DSO（如 libtvm_runtime_cuda.so）。
2. RTLD_GLOBAL 加载：符号暴露给后续动态加载的模块（NVRTC kernel、用户 .so 等）。见 python/tvm/base.py。
3. 可单独部署：runtime-only wheel 只带 libtvm_runtime，适合纯推理场景（如 mlc-llm 部署端）。

类比：通用「虚拟机 + 加载器」——知道如何加载模块、调度函数，但不知道具体怎么操作 GPU。

---

3. libtvm_compiler（libtvm_compiler.so）

定位：编译器前端 + 中端 + 代码生成，负责「把 IR 编译成可执行代码」。

包含的源码（节选）

CMake 中 COMPILER_SRCS 和 CODEGEN_SRCS 覆盖：

• src/ir/、src/arith/、src/te/、src/tirx/、src/relax/
• src/driver/（tvm.compile() / tvm.build() 入口）
• src/target/、src/backend/cuda/codegen/、src/backend/rocm/codegen/ 等

主要功能：

模块	作用
IR（Relax / TIR / TIRx）	中间表示、变换 Pass
Driver	tvm.compile() / tvm.build() 入口
Codegen	生成 LLVM IR、CUDA C、Metal、Vulkan SPIR-V 等
LLVM 后端	CPU 代码生成、CUDA LLVM 路径（需 USE_LLVM=ON）
Arith / TE	表达式化简、张量表达式

关键设计

1. 依赖 LLVM（USE_LLVM=ON）：CPU codegen 和 LLVM 路径的 CUDA codegen 都需要。
2. RTLD_LOCAL 加载：编译器符号不污染全局命名空间，避免与 LLVM 等冲突。
3. CUDA 编译 vs CUDA 运行是分开的：
   • 编译（生成 .cu / PTX / cubin）：在 libtvm_compiler（src/backend/cuda/codegen/）
   • 运行（分配显存、启动 kernel）：在 libtvm_runtime_cuda

类比：GCC / Clang 编译器本身——读源码、优化、生成目标文件，但不负责在 GPU 上执行。

---

4. libtvm_runtime_cuda（libtvm_runtime_cuda.so）

定位：CUDA 设备运行时后端（USE_CUDA=ON 时构建），负责「在 NVIDIA GPU 上执行」。

包含的源码

见 cmake/modules/CUDA.cmake：

• src/backend/cuda/runtime/*.cc
• src/runtime/vm/cuda/*.cc

文件	作用
cuda_device_api.cc	CUDADeviceAPI：设备枚举、显存 alloc/free、H2D/D2H 拷贝、Stream、设备属性
cuda_module.cc	加载 cubin/PTX、启动 CUDA kernel
src/runtime/vm/cuda/*.cc	Relax VM 在 CUDA 上的 builtin 算子

cuda_device_api.cc 实现了 DeviceAPI 接口，例如 SetDevice、GetAttr、AllocDataSpace 等。

加载时机

不是 import tvm 时立即加载，而是在 CUDA backend 注册时按需加载（见 python/tvm/backend/cuda/__init__.py 的 register_backend()）。加载后会向全局注册表注册 device_api.cuda 等函数，使 tvm.cuda().exist 能检测到 GPU。

关键设计

1. 链接 libcuda + libcudart：需要 NVIDIA 驱动和 CUDA runtime。
2. 依赖 libtvm_runtime：复用通用 Module / PackedFunc 框架。
3. 不含编译器逻辑：不会调用 nvcc，只负责执行已编译产物。
4. 可交叉编译：编译 TVM 时不需要实际 GPU，只要有 CUDA toolkit。

类比：CUDA Driver API 的适配层——让 TVM runtime 能在 NVIDIA GPU 上分配内存、拷贝数据、启动 kernel。

---

5. libtvm_runtime_extra（补充）

虽然不在核心四件套之列，但编译时常遇到：libtvm_runtime_extra.so 是可选扩展运行时，包含 Disco 分布式、NCCL、NVSHMEM、cuDNN/cuBLAS/CUTLASS 等 contrib 模块。

• 源码来自 src/runtime/extra/ 及各 cmake/modules/contrib/ 模块
• USE_CUDA=ON 时会 PUBLIC 链接 libtvm_runtime_cuda
• 若启用了含 .cu 的 contrib（如 USE_NCCL、USE_THRUST、USE_CUTLASS），需正确设置 CMAKE_CUDA_ARCHITECTURES

---

完整生命周期：以 CUDA kernel 为例

1. 编译阶段（libtvm_compiler）：IR → CUDA 源码/PTX，不需要 GPU
2. 注册阶段（libtvm_runtime_cuda）：注册 CUDA 设备 API
3. 执行阶段（libtvm_runtime + libtvm_runtime_cuda）：加载模块、分配显存、启动 kernel

---

Python 加载顺序

import tvm_ffi          # FFI 基础设施（Cython 扩展）
import tvm
  ├─ load libtvm_runtime.so     (RTLD_GLOBAL，必须)
  ├─ load libtvm_compiler.so    (RTLD_LOCAL，runtime-only 模式跳过)
  └─ autoload backends
       └─ load libtvm_runtime_cuda.so  (按需，RTLD_LOCAL)
       └─ load libtvm_runtime_extra.so   (若存在，RTLD_LOCAL)

验证当前加载了哪些库：

import tvm
print(tvm.base._LOADED_LIBS.keys())
# dict_keys(['tvm_runtime', 'tvm_compiler', 'tvm_runtime_cuda', ...])

print("\n".join(f"{k}: {v}" for k, v in tvm.support.libinfo().items()))
# USE_CUDA: ON, USE_LLVM: ..., LLVM_VERSION: ...

---

部署场景对照

场景	需要的库
仅 CPU 推理	libtvm_ffi + libtvm_runtime
CUDA 推理（mlc-llm 部署）	上述 + libtvm_runtime_cuda [+ libtvm_runtime_extra]
本地编译 + CUDA 推理	全部四个核心库
交叉编译（如在 x86 上编译给 Jetson 用）	编译端需要 libtvm_compiler；运行端只需 runtime 系列

---

相关文档

• Install from Source — 从源码构建 TVM
• TVM FFI 文档 — FFI 独立项目说明
• CMakeLists.txt — 库拆分与链接关系
• cmake/modules/CUDA.cmake — CUDA 运行时构建