diff --git a/docs/tvm_ffi_integration.md b/docs/tvm_ffi_integration.md new file mode 100644 index 000000000..a0d6b8f5e --- /dev/null +++ b/docs/tvm_ffi_integration.md @@ -0,0 +1,170 @@ +# TVM-FFI 集成方案 + +## 概述 + +本文档描述了在 GraphNet 中引入 `tvm-ffi` 的集成方案。`tvm-ffi` 提供了统一的张量表示和跨框架互操作能力,以及标准化的编译器 ABI 接口,能够显著简化 GraphNet 对不同编译器后端的集成和维护。 + +## TVM-FFI 简介 + +### 什么是 TVM-FFI + +**TVM-FFI (Foreign Function Interface)** 是 Apache TVM 项目提出的开放应用二进制接口(ABI)和 FFI 标准,旨在解决机器学习系统生态割裂与互操作性难题。它通过定义稳定的 C ABI 和利用 DLPack 协议,实现了不同框架与底层编译器之间的高效连接。 + +### 核心特性 + +1. **开放的 ABI 标准** + - 定义统一的应用二进制接口(ABI),使不同编译器可以遵循同一套调用规范 + - 支持跨语言、跨平台的函数调用 + - 提供稳定的接口契约,降低版本兼容性问题 + +2. **零拷贝数据传递** + - 基于 DLPack 协议实现张量数据的零拷贝传递 + - 支持 PyTorch、JAX、NumPy 等框架间的无缝数据交换 + - 减少内存拷贝开销,提升性能 + +3. **统一的张量表示** + - 提供 `Tensor` 和 `TensorView` 等统一的数据结构 + - 标准化张量的元数据(形状、数据类型、设备等) + - 简化不同框架间的张量转换逻辑 + +4. **模块化设计** + - 支持 AOT(Ahead-of-Time)编译的模块加载 + - 提供 `call_tvm_ffi` 原语用于调用符合 ABI 的函数 + - 支持动态模块加载和运行时调用 + +### 技术原理 + +TVM-FFI 的核心技术栈包括: + +- **C ABI 层**:提供稳定的 C 语言应用二进制接口,确保跨平台兼容性 +- **DLPack 协议**:实现张量数据的零拷贝传递,支持多种框架 +- **FFI 机制**:通过 Foreign Function Interface 实现跨语言调用 +- **模块系统**:支持编译后的模块动态加载和执行 + +### 解决的问题 + +在机器学习系统生态中,存在以下痛点: + +1. **框架割裂**:PyTorch、JAX、TensorFlow 等框架各自为政,数据交换需要多次转换 +2. **编译器集成复杂**:每个编译器都有独特的接口,集成成本高 +3. **性能损失**:框架间数据转换带来额外的内存拷贝开销 +4. **维护困难**:多套接口标准导致代码重复和维护成本高 + +TVM-FFI 通过统一的 ABI 和 FFI 标准,有效解决了这些问题。 + +### 在 GraphNet 中的价值 + +GraphNet 作为张量编译器的基准测试框架,需要: + +- **测试多种编译器**:CINN、TorchInductor、TVM、TensorRT、XLA、BladeDISC 等 +- **处理多框架数据**:PyTorch、JAX 等框架的计算图 +- **标准化接口**:为不同编译器提供统一的调用方式 +- **性能优化**:减少数据转换开销,提高基准测试效率 + +引入 TVM-FFI 可以: +- ✅ 统一不同编译器的调用接口 +- ✅ 实现零拷贝的张量数据交换 +- ✅ 简化新编译器的集成流程 +- ✅ 提高基准测试的公平性和可重复性 + +## 集成价值 + +基于 TVM-FFI 的核心特性,在 GraphNet 中引入 TVM-FFI 可以带来以下具体价值: + +### 1. 统一张量表示与跨框架互操作 + +**现状问题**: +- GraphNet 需要处理来自 PyTorch、JAX 等不同框架的计算图 +- 不同框架的张量格式不统一,需要多次转换 +- 数据转换带来额外的内存拷贝和性能开销 + +**TVM-FFI 解决方案**: +- **零拷贝数据传递**:通过 DLPack 协议实现框架间的零拷贝张量传递,显著减少内存开销 +- **统一张量表示**:使用 `Tensor` 和 `TensorView` 标准化处理不同框架的张量,简化转换逻辑 +- **简化数据交换**:在不同编译器后端间传递张量时无需多次格式转换 + +**实际收益**: +- 减少基准测试中的数据传输开销 +- 提高多框架数据处理的效率 +- 降低数据转换带来的性能损失 + +### 2. 编译器集成的标准化接口 + +**现状问题**: +- GraphNet 需要测试多种编译器(CINN、TorchInductor、TVM、TensorRT、XLA、BladeDISC 等) +- 每个编译器都有独特的调用接口和集成方式 +- 新增编译器后端需要大量定制化代码 + +**TVM-FFI 解决方案**: +- **标准 ABI 接口**:所有符合 tvm-ffi ABI 的编译器都可以通过统一接口调用 +- **call_tvm_ffi 原语**:提供标准化的函数调用方式,无需了解各编译器的具体实现 +- **模块 API**:统一的模块加载和运行接口,支持 AOT 编译的模块 +- **跨语言支持**:通过 C ABI 实现跨语言调用,支持更多编译器生态 + +**实际收益**: +- 新编译器集成成本降低 50% 以上(预估) +- 统一的调用接口提高代码可维护性 +- 更公平的编译器性能对比(消除接口差异带来的性能影响) + +### 3. 简化新编译器后端的集成 + +**现状问题**: +- 每个新编译器都需要实现完整的 `GraphCompilerBackend` 接口 +- 需要处理框架特定的张量转换逻辑 +- 集成代码重复度高,维护成本大 + +**TVM-FFI 解决方案**: +- **ABI 标准化**:新编译器只需遵循 tvm-ffi ABI 即可接入 +- **统一适配层**:通过统一的张量适配器处理数据转换 +- **最小化集成代码**:只需实现基本的模块加载和函数调用逻辑 + +**实际收益**: +- 新编译器集成时间从数周缩短到数天 +- 减少重复代码,提高代码质量 +- 降低长期维护成本 + +### 4. 提升基准测试的公平性和可重复性 + +**额外价值**: +- **统一的调用路径**:所有编译器使用相同的调用接口,消除接口差异带来的性能偏差 +- **标准化的数据流**:统一的数据转换路径确保测试结果的公平性 +- **更好的可重现性**:标准接口使测试结果更容易在不同环境中重现 + +## 架构设计 + +### 集成层次 + +``` +┌─────────────────────────────────────────┐ +│ GraphNet Test Framework │ +├─────────────────────────────────────────┤ +│ Unified Tensor Exchange Layer │ ← 基于 tvm-ffi/DLPack +│ (tvm_ffi_tensor_adapter.py) │ +├─────────────────────────────────────────┤ +│ Compiler Backend Interface │ +│ (graph_compiler_backend.py) │ +├─────────────────────────────────────────┤ +│ TVM-FFI Backend │ ← 新增 +│ (tvm_ffi_backend.py) │ +├─────────────────────────────────────────┤ +│ Other Backends (TVM, XLA, etc.) │ +└─────────────────────────────────────────┘ +``` + +### 核心组件 + +1. **统一张量适配器** (`tvm_ffi_tensor_adapter.py`) + - 提供基于 DLPack 的张量转换工具 + - 支持 PyTorch、NumPy、JAX 等框架的张量互转 + +2. **TVM-FFI 后端** (`tvm_ffi_backend.py`) + - 实现 `GraphCompilerBackend` 接口 + - 支持调用符合 tvm-ffi ABI 的编译函数 + +3. **增强的 TVM 后端** (`tvm_backend.py`) + - 可选:使用 tvm-ffi 增强现有 TVM 后端 + + +## 参考资料 + +- [TVM-FFI ](https://github.com/apache/tvm-ffi)