Related Work
- 基于Graph IR的XLA、DLVM:现有的基于静态图表示的深度学习编译器的主要问题是图的语义过于高层了,不利于下面的优化。如果希望支持不同的后端,需要各个都手动实现lower阶段,从而生成较为底层的LLVM代码,带来了巨大的工作量。
- Tensor comprehension:基于多面体变换的黑盒搜索优化,CUDA only。
- TACO:实现了通用的CPU代码生成技术,用于生成任意的(包含稀疏)线性代数计算。
Automating Optimization
AutoTVM由两个模块组成:schedule explorer用于产生新的可能会有性能提升的代码,然后用一个基于ML的耗时评估模型来预测这个代码的性能,这样不断迭代搜索,直到得到一个满意的性能结果。注意,这里并不是产生一个schedule就上硬件跑,而是先经过模型评估,直接剪枝砍掉那些看起来不会有性能提升的结果,然后真实运行的是那些看起来能提升性能的结果。并且在真的硬件上跑完之后,要用相应的数据更新性能评估模型,从而实现完整的自动优化loop。
Schedule Space Specification
首先的问题是我们如何描述一个算法的schedule?这里指的就是算法经过TVM的DSL描述以后,再使用schedule原语进行循环变换的相关API。例如下面的split(循环拆分):
A = tvm.placeholder((m,), name='A')
B = tvm.compute((m,), lambda i: A[i]*2, name='B')
s = tvm.create_schedule(B.op)
xo, xi = s[B].split(B.op.axis[0], factor=32)
print(tvm.lower(s, [A, B], simple_mode=True))
这些变换能够保证不影响程序的正确执行,当然这就要求应用变换的时候需要满足一定的前置条件。即便如此,可能产生的组合也是极大的,因此我们需要在其中进行搜索优化。
Cost Model
接下来的问题是如何评估一个schedule的效率?我们需要一个硬件模型。这里使用的cost model是个数学统计模型而不是性能模拟器这样的东西,主要的motivation就是我们无法忍受每次都上真实硬件评估性能的开销,更不可能对各种硬件做软件层面的复杂精细建模——更慢。所以本质上相当于先在这里套一个剪枝。
既然是用统计模型,就会有两个问题:模型结构、优化目标与特征抽取。对于模型结构来说,主要是得快——如果比物理硬件还慢就没有意义了。目标函数不需要设置为预测精确的执行时间(个人猜测这样可能不利于优化),而是设置为预测执行耗时的相对比例,能够用于不同schedule之间的对比即可。于是这里他们选取了速度比较给力的XGBoost作为模型。使用的特征包括访存总量、cache利用率、以及所有应用到的循环变换所组成得one-hot向量等等。另外也有尝试一种暴力的TreeRNN实现,可以直接把代码的AST作为输入,从而避免了人工的feature engineering,两者其实效果相似,但是前者更快。
总体上,使用ML的cost model进行搜索的效率显然是远远大于暴力搜索的,也比遗传算法之类的启发搜索算法效果上好很多。
Schedule Exploration
当我们有了schedule以及cost model,就可以开始我们的搜索过程了。注意在初始状态下,cost model完全没有被训练过,因此需要随机挑选一些schedule先跑起来看看再说。那么最后一个问题就是新的schedule如何生成呢?当抽取出代码特征以后,这本质上其实是高维特征空间上的非凸优化问题,可以用模拟退火策略来搞定,具体就不再赘述了。