Reading Note: TVM

Overview

• TVM 是一个端到端的深度学习编译器，旨在自动优化深度学习模型的可执行代码，使得能够高效运行在不同的硬件上。（也就是说只关心部署，不关心训练）

• TVM 的主要目标：
• 性能可移植性（performance portability）：只需要制定 high level 的模型 specification，在多个硬件上都能取得良好性能，无需针对每个硬件都做手动编写和调优专门的算子库。
• 自动调优：通过自动搜索和调优降低人工调优的工作量

Core Design

• 作为一个端到端的深度学习编译器，TVM 接受来自各框架的深度学习模型规约描述，首先转换为计算图表示。

• 图级别优化：基于计算图表示可以做图级别的优化
• 算子融合：将多个小算子融合成一个整体，减少中间数据存储和内存访问，从而提高整体性能。
• 数据布局优化：根据目标硬件的内存层次结构和数据访问特点，将数据转换成更高效的布局（例如：将二维卷积数据划分成块状），以提高内存局部性。

• 算子级别优化：在进行图级别的优化之后，要产生最后的可执行代码，还需要有优化后的计算图中的对应 kernel 实现。TVM 通过自动调优自动生成需要的算子实现（在深度学习框架中，算子融合之后要产生代码还是要有事先实现好的 kernel 才可以）。

• TVM 借鉴了 Halide 的“计算（compute）”与“调度（schedule）”分离思想：
• 张量表达式语言(tensor expression language) 用于描述算子（例如卷积、矩阵乘法等）的计算规则，而不涉及具体的循环结构和执行细节。
• 嵌套并行与协作加载：支持在多层次线程（如 GPU 的线程组和线程）中实现数据协作加载(cooperative fetching)，共享内存的使用，从而提高内存带宽利用率。
• 张量化（Tensorization）：类似于向量化，但针对张量运算，将大规模计算分解为调用硬件特定的微内核指令（例如 8×8 的矩阵乘法）。
• 显式内存延迟隐藏：通过虚拟线程的方式，将高层并行程序自动降级为带有低级同步指令的单一指令流，从而在硬件流水线中重叠内存访问和计算操作。

• 自动调优和机遇机器学习的 cost model:
• 通过抽取循环程序的特征（如内存访问次数、数据复用比率、以及循环注解信息）来构建特征向量。
• 使用梯度提升树（如 XGBoost）或 TreeRNN 模型预测一个调度配置的性能表现，重点是预测相对顺序（用 ranking objective，比如 bpr）而非绝对运行时间。
• 模型在不断进行硬件实验后迭代更新，从而提高预测准确性，进而引导调度搜索算法（如模拟退火）更快地找到效果好的配置。

Summary

• 计算与调度分离：tensor expression 定义算子做什么，这对应一个基础的实现；调度原语定义了我们可以如何变换程序，带来了什么可以调的参数，这定义了我们在搜索空间里如何探索。

• 基于 ML 的 cost model：需要考虑速度和准确率的 trade-off，因为要频繁调用。