# 前言

本系列文章将从代码和流程图入手,详细介绍 TVM AI 编译器的编译流程。本文章为第一篇,取名为 “必要知识”,约等于 “大杂烩”。食之无味,弃之可惜。但基于我个人的主观喜好和偏见,这些知识值得在整个代码讲解之前声明。

作为初学者,错误在所难免,还望不吝赐教。

# 工程结构

TVM 主要由 C++ 语言和 Python 语言实现。

  • a.TVM C++ 后端核心代码 位于 /home/user/tvm/src
    C++ 后端核心代码的功能是生成 libtvm.so libtvm_runtime.so,分别是模型编译过程中的需要的 libtvm.so 库,和模型部署运行时需要的 libtvm_runtime.so 库。对应的头文件位于 /home/user/tvm/include.

    • 1. 其中 src/runtime 是模型部署相关 CPP 文件
      • src/runtime/graph_executor 最基本的图执行器
      • src/runtime/module.cpp 运行时 Module 基类
      • src/runtime/contrib BYOC 第三方设备运行及实现 包括 dnnl cblas cudnn ethosu 等等
      • srd/runtime/cuda CUDA 运行时 Module
      • 还有其他设备的运行时 Module 如 opencl hexagon metal 等,以及一些接口文件
    • 2.src/relay 是 TVM 高级中间表示 IR 的相关 CPP 文件
    • 3.src/relax TVM 正在发展的高级中间表示 IR,功能上等同于 Relay
    • 4.src/te TVM 的 TE 表达式
    • 5.src/topi 算子的计算定义和后端调度
    • 6.src/tir TVM 低级中介表示 IR,接近于硬件表示
    • 7.src/auto_shedule TVM 自动优化调度相关
  • b. TVM Python 前端代码 位于 /home/user/tvm/python
    Python 前端代码为 C++ API 和执行编译的驱动代码,提供 Python 绑定。

表面上看 TVM 工程应当是 Python 前端调用 C++ 后端生成的 libtvm.so libtvm_runtime.so ,实际上 TVM 支持 Python 和 C++ 的相互调用。例如一些算子的 compute 计算公式和 schedule 调度位于 Python 代码中,由 C++ 进行调用。

在编译阶段,开发者编写 Python 代码,调用 TVM Python 前端,而 TVM 的 Python 前端又依赖于 C++ 后端生成的 libtvm.so libtvm_runtime.so 动态库。

部署运行阶段即可以编写 Python 来执行,也可以直接编译 C++ 为可执行文件,参考 TVM C++ 部署

# 编译

编译过程可以参考 TVM 英文文档和中文文档。文档中内容全面,这里只摘必要步骤,是最简单的部署安装方式。

1. 编译 libtvm.so libtvm_runtime.so

mkdir build  # 下载工程后进入工程目录 创新文件夹
cp cmake/config.cmake ./build
    # 将 cmake 配置文件拷贝到 build 目录,可以更改需要的额配置选项
cd build
cmake ..
    
make -j32

2. 配置临时运行环境

个人喜欢用临时环境运行项目,可以避免 .bashrc 文件中添加过多环境路径;特别是开发人员有多个 TVM 项目的时候,可以避免项目的混淆。

# 指定 tvm 库路径
export PATH=/home/user/tvm/build:$PATH
# 指定 tvm Python 路径
export PYTHONPATH=/home/user/tvm/python:${PYTHONPATH}
# 运行
python3 ***.py

# TVM 基类

大概从事工程项目开发的人员对 Object, ObjectPtr, ObjectRef 等写法相对熟悉,但是没有相关经验的算法研究人员,还是先认识一下 TVM 中的这三个基类。

TVM 中大部分类都继承自三大基础类 Object, ObjectPtr, ObjectRef。想多了解一下的话,可以学习博客 TVM: Object, ObjectPtr, ObjectRef (以 tir.sch 为例)

可能比较粗显的总结一下:TVM 中以 "Node" 结尾的类是 Object 的派生类,去掉 "Node" 的类是 ObjectRef 的派生类,Object 的派生类负责 "做事",ObjectRef 的派生类方便程序员引用,ObjectPtr 是桥梁。但总之这三者是一个东西(刚学时我就被这搞晕了:( )。

举例说明如何使用 TVM 的三大基础类构造对象:

class Anode: public tvm::runtime::Object{   
    public:
        Anode(){};
        ~Anode(){};
        std::string a(){return "instance";}
};
class A: public tvm::runtime::ObjectRef{       
    public:    
    TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(A, tvm::runtime::ObjectRef, Anode);
}; 
int main(){    
   tvm::runtime::ObjectPtr<Anode> n = tvm::runtime::make_object<Anode>();  // 桥梁 Ptr = make_object (本体 node)    
   A ref(std::move(n));  // 指针
   ref->a();          // Anode 指针,可以引用 a ()
   //ref.get ()->a (); // 基类 Object 指针,error Object has no member a ()
}

Object 应用例子

可以再看一下 TVM 工程中的例子,这段代码创建了本体 CallNode ,通过 make_object 生成桥梁 ObjectPtr<CallNode> ,最后赋值给指针 Calldata_ 字段。

# Python 和 C++ 相互调用

TVM 实现了 FFI (Foreign Function Interface) 机制,目标是使得任意语言下定义的函数都可以被任意其他语言调用。

宏 TVM_REGISTER_GLOBAL

/*!
 * \brief Register a function globally.
 * \code
 *   TVM_REGISTER_GLOBAL("MyPrint")
 *   .set_body([](TVMArgs args, TVMRetValue* rv) {
 *   });
 * \endcode
 */
 #define TVM_REGISTER_GLOBAL(OpName) \
  TVM_STR_CONCAT(TVM_FUNC_REG_VAR_DEF, __COUNTER__) = ::tvm::runtime::Registry::Register(OpName)

在 C++ 后端同过 “TVM_REGISTER_GLOBAL” 这个宏将函数注册为全局函数,该全局函数的类型为 PackedFunc。

TVM_REGISTER_GLOBAL("tvm.graph_executor_factory.create")  // 注册为全局函数 PackedFunc
    .set_body([](TVMArgs args, TVMRetValue* rv) {      
        *rv = Module(exec);
    });

在 Python 前端可以通过 get_global_func () 函数来获取全局函数 PackedFunc 并执行

fcreate = get_global_func("tvm.graph_executor_facttory.create")
self.module = fcreate(graph_json_str, libmod, libmod_name, *args)

在 C++ 端也可以通过名字获取 全局函数 PackedFunc

PackedFunc pf = it.GetFunction(name)

python 端有大量如下类似代码,将有相同后缀的 PackedFunc 注册到 Python 端

"""FFI APIs for tvm.ir"""
import tvm._ffi
tvm._ffi._init_api("ir", __name__)

Python 端获取并执行 PackedFunc

_ffi_api.Module_UpdateGlobalInfo(self, name, global_info)

# 后记

本博客目前以及可预期的将来都不会支持评论功能。各位大侠如若有指教和问题,可以在我的 github 项目 或随便一个项目下提出 issue,并指明哪一篇博客,我看到一定及时回复!

Edited on

Give me a cup of [coffee]~( ̄▽ ̄)~*

XianMu WeChat Pay

WeChat Pay

XianMu Alipay

Alipay