面向深度学习算子的循环不变式外提算法

摘要

TVM是一个深度学习编译器,支持将TVM的领域专用语言即张量表达式定义的算子编译生成目标平台的代码,并在高级中间表示TVM IR上进行一系列优化。张量表达式对算子执行循环变换,产生与循环迭代变量相关的复杂表达式的计算,在多层嵌套循环内这些计算包含了大量的循环不变式。然而,传统的循环不变量外提技术不能判断不变量外提是否能带来额外收益,无法发现操作数顺序不同的循环不变表达式,不能处理嵌套的条件表达式,并且与目标平台编译器优化存在冲突等。由于这些挑战,传统的循环不变量外提算法无法直接用于深度学习编译器的优化,提出了一种融合深度学习代价函数和启发式策略的循环不变量外提算法。该算法基于深度学习编译器的高层中间表示,通过调整操作数顺序和简化嵌套条件表达式等方法规范化表达式。为了衡量优化的收益,在结合TVM IR和目标平台的特点的基础上,提出了一个新的面向深度学习的不变式外提代价指标函数。在开源编译器TVM 0.7版本上,通过新增优化遍的形式,具体实现了所介绍的算法以及代价函数。为评测算法的有效性,在Tesla P4的图形处理器(GPU)平台上对TVM TOPI的测试算子集中27个典型算子不同输入规模的511个测例进行了测试。实验结果表明47.6%的算子性能得到提升,最大加速比大于40.0%。