新闻动态

简介

计算机系统结构领域的顶级国际会议ASPLOS 2020于3月16日正式召开。北京大学高能效计算与应用中心梁云研究员带领研究小组在ASPLOS 2020上发表了最新论文。该论文研究在异构系统中为张量计算自动优化与生成代码，研究提出了利用高层次编程抽象结合强化学习方法自动搜索优化空间的策略，自动为不同平台生成高性能算子实现。论文第一作者为梁云研究员指导的一年级博士研究生郑思泽，梁云研究员为通讯作者。

ASPLOS是CCF-A类会议，也是CSRanking收录的计算机系统结构领域的顶级会议之一，专注于编程语言、编译器、操作系统、网络和体系结构等领域及它们的交叉领域。

ASPLOS 2020收到超过四百篇投稿，但论文接收率仅为18.1%。

受新型冠状病毒肺炎疫情影响，ASPLOS 2020会议改为线上进行，采用录制视频报告的方式。

本文章已经在ASPLOS 2020进行发表。

论文题目为: FlexTensor: An Automatic Schedule Exploration and Optimization Framework for Tensor Computation on Heterogeneous System

作者: Size Zheng, Yun Liang, Shuo Wang, Renze Chen, Kaiwen Sheng

论文报告视频已上传至youtube：https://www.youtube.com/watch?v=Ee15mcGc8RA

研究背景与动机

张量计算在诸多领域有着重要意义，包括但不限于科学计算、数据分析、图像处理以及机器学习。在近些年火热的深度学习领域，几乎所有的神经网络都离不开张量计算。这些张量计算被抽象为一个个算子，如矩阵乘法(GEMM)，卷积(CONV)等，它们本质上是多重嵌套循环，非常适合在加速器（如GPU，FPGA）上进行加速。

但是在加速器上开发高性能的算子实现并非易事，往往需要程序员懂得算法优化和体系结构的细节，并能恰当地进行使用。为此，许多硬件厂商会提供高性能算子库给用户使用。典型的库包括Nvidia的cuBlas, cuDNN，以及Intel的MKL, MKL-DNN等。但是这种计算库往往并不开源，并且是高度定制化的，这就带来了两个问题：

1. 普通用户难以参与到新算子开发工作中，因为不清楚库中实现方式，难以将对已有算子的了解推广到新的算子。
2. 库的功能受限于厂商提供的支持，新的算子添加速度缓慢，难以跟上新的应用和算法发展的脚步，限制了应用和算法的进一步创新与实践。

为了解决这一问题，在先前研究工作中，学术界和工业界都提出了自动代码生成技术。利用高层次的计算描述来编写算子数学公式，再利用编译技术自动生成底层实现。该方向的几个重要工作包括PlaidML，Halide和TVM等。目前性能最佳的自动代码生成框架就是TVM，它沿用了Halide提出的计算描述与优化策略分离的理念，针对张量计算对多平台进行代码生成。

但是这些自动代码生成框架都面临这一个关键问题：如何进行优化决策。因为同一个算子在不同硬件上的优化方式是不同的，带来的性能差异也很大，仅凭编译器难以自动地进行决策。因此，这部分工作很大程度上仍然要程序员手工完成，即使是TVM本身，也只提供了一个半自动的优化工具AutoTVM帮助程序员进行优化，但是这个工具需要程序员手写一个优化模板，这本身也不是件轻松的事情。

优化决策

我们沿用Halide的习惯，称一个高层次的计算描述为compute，称优化策略为schedule。如图1是一个GEMM的compute，在lambda表达式中就是对于GEMM的数学公式描述。

图1. GEMM的compute

相应地，可以为这个compute提供一个schedule，如图2所示。一个完整的scehdule需要有一系列优化原语 (primitive) 以及它们的参数 (parameter) 组成。在这个例子里，展示的原语有tile (矩阵分块)，split (循环划分)，reorder (循环重排) 和vectorize (向量化)。而它们的参数包括目标循环，划分的大小 (bn和factor)以及循环顺序等。实际开发时，可选的原语和参数数量众多，而且不同的硬件还会使用不同的原语。对于程序员来说，在这段短小的程序里究竟如何选择不同的原语和参数，成为了难以解决的问题，只有不断地进行实际测试，才能选出一个较好的方案。

图2. GEMM的一种schedule

我们要解决的问题就是，如何自动地产生这些优化决策 (即原语序列和其参数的确定)，从而免去使用者手动确定优化决策的过程。