论文解读 | 自动编译框架AMOS
软件到硬件资源的自动映射是编译器的一项重要技术，也是业界研究的一个热门方向。本文主要观点来自发表在计算机体系结构领域顶级会议ISCA2022上的一篇论文——AMOS: Enabling Automatic Mapping for Tensor Computations On Spatial Accelerators with Hardware Abstraction，希望可以给读者一些启发。
1. 摘要
如今计算设备逐渐向专用化方向发展，加速器设备（Spatial hardware accelerators）中有多层级专用的计算和存储资源，已经在深度学习、科学计算、数据挖掘等领域显示出较高的性能优势。软件开发者需要使用硬件专用的指令集才能发挥出加速器的计算性能。然而，开发者普遍缺少对这类加速器设备和对应的专用指令集的理解，比如其硬件抽象表示、性能优化空间及自动化探索优化空间等。现有的编译器通常使用手工方式和基于模板方式进行代码性能优化，其中手工方式开发成本较高，基于模板方式可能会导致得到次优解。
本文提出一个用于计算加速器的自动编译框架AMOS。框架的核心是硬件抽象，不仅可以准确描述硬件指令行为，而且可以统一定义计算任务到硬件的映射。基于硬件抽象，作者使用性能代价模型在所有可能的映射空间中自动探索最优映射。总体来看，该编译框架使用硬件抽象作为编译器IR，探索计算映射和存储映射，最终可为多种硬件后端生成高性能代码。实验表明AMOS在TensorCore上比手工优化性能提升2.5倍，在Intel CPU AVX-512指令集支持的芯片上比TVM性能提升1.37倍，在移动端的Mali GPU上比AutoTVM性能提升25.04倍。
1. 简介
作者将计算加速器的软件到硬件的映射方式划分为两类：硬件可感知和指令集可感知。
硬件可感知：编译器知晓硬件架构细节信息，比如PE（processing element）数量、PE间互联方式，这样软硬件映射过程可以转化为一个在给定硬件约束下的优化问题。这种映射方式具有较高的能效比，但是牺牲了一定的灵活性（当搜索空间发生变化后，映射过程需要重新生成）。代表性编译器有Triton、CoSA、SRAR、HASCO。
指令集可感知：映射过程对硬件架构无感知，只根据对外暴露的指令集进行软硬件映射。XLA、TVM、Ansor、AKG、Tiramisu等均归为指令集可感知。本文聚焦于指令集可感知的映射。
指令集可感知方式通过提供张量计算的若干内置函数(Intrinsics)提高了可编程性，但仍面临两方面挑战：
给定算法下，如何生成最优指令。比如，二维卷积的7层循环映射到Tensor Core上，映射方式有35种。现有的编译器如TVM、ISA Mapper等主要依赖基于模板或手工编写方式，这样可能会得到次优解。
理想的映射方法对不同加速器是统一的，然而不同加速器提供的指令集相差较大，在给定算法下，无法统一开发。比如，Tensor Core使用不同的内置函数来描述矩阵加载/存储/相乘操作，在Mali GPU中只需使用一条内置函数
arm_dot就可以完成数据加载和执行。
本文设计了一个在指令集之上的抽象层来解决如上两方面挑战。抽象层包括计算抽象和存储抽象两部分，可以统一描述不同加速器的指令集，实现了统一的自动映射流程。另外，作者设计了一个性能评估模型来评价所有可能解。最终，AMOS将映射问题抽象后转化为一个迭代优化问题。
2. 背景和动机
2.1 专用计算加速器
专用计算加速器通常分为多个层级，最内层是有多个PE组成的计算阵列，中间层是由多个计算阵列组成的subcore，最外层是由多个subcore组成的core以及core间共享的全局存储器，如图1(a)所示。专用计算加速器通常嵌入在通用计算架构类芯片中，如在CPU或GPU中嵌入专用计算加速器。Nvidia GPU中的Tensor Core，Mali GPU中的dot计算模块，华为Ascend NPU中的cube模块都被看做专用计算架构。
图 1
硬件厂商一般都会为加速器设计专门的指令，这类指令主要分计算型指令和访存型指令两种。图1(b)中AVX-512指令集中__mm512_add_pd是一个向量加指令，__mm512_loadu_ph是一个向量数据加载指令。Tensor Core中mma_sync是一个矩阵乘指令，load_maxtrix_sync是一个矩阵数据加载指令。
2.2 现有的映射流程
表1中列出来业界领先的编译器，可以分为两类：
硬件可感知映射：编译器知晓PE互联和存储配置等信息，因此编译器可以将软件层面上的计算任务lower到每个PE和存储单元，这一过程主要应用软硬协同设计，即在特定硬件资源约束下最大化资源利用率。由于映射搜索空间巨大，这类编译器通常使用线性规划来实现软件到硬件资源的映射。
指令集可感知映射：空间加速器都对外提供专用指令集，为了生成高性能代码，软件开发者一般需要编写调度模板，并且在给定的调度空间内搜索出最佳调度。例如TVM提供了一个
tensorize接口原语，开发者需要手动实现该接口函数；作为多面体优化编译器的代表，AKG依赖多面体模型和模板来进行软硬件映射。
表 1
2.3 动机
表2对比了深度学习编译器XLA在典型模型上的映射效果，虽然XLA中使用了大量复杂模板，但是运行典型模型时也只有少量算子可以映射到NVIDIA V100 Tensor Core上。例如，ShuffleNet中的深度可分离卷积和分组卷积是普通二维卷积的两个变种，内部同样是大量乘加操作，原理上讲是适合在Tensor Core上运行的，但是XLA却无法将这两类算子映射到Tensor Core上。原因是因为XLA中手工编写的模板无法成功匹配到这两类算子。
表 2
手工编写模板有两个限制：第一，手工编写的模板依赖于显式地基于指令集实现代码，这种方式人力成本高、开发周期长，因此只能支持少量部分算子的映射。第二，手工编写的模板适用性受限，例如对算子稍加改动（例如改变数据排布方式）后，就可能导致模板匹配失败。为解决如上问题，本文设计了一个无模板不依赖设备厂商第三方加速库（如cuDNN/cuTLASS）的自动映射编译流程。
3. 原理
图2展示了AMOS的编译流程，输入是由Python编写的高层次代码和基于计算和访存抽象出的硬件指令组成，然后生成不同的映射组合。一个映射组合由计算映射和访存映射组成：
计算映射是将算子转为计算指令映射到空间加速器上；
访存映射通过访存指令描述数据的访问和存储行为。
整个映射生成过程分为2个阶段，第一阶段是将循环迭代映射到一个虚拟加速器上，第二阶段是加入硬件资源和指令集限制，调整以前的映射过程最终映射到物理加速器上。
以Conv2D为例进行说明，首先将循环迭代空间映射为指令集迭代空间。在下图中循环迭代维度n,p,q映射至指令集的i1维度，循环迭代维度k映射至指令集的i2维度，循环迭代维度c,r,s映射至指令集的r1维度。然后，将卷积操作转化为矩阵乘操作，完成计算和访存代码到指令集的映射。第一阶段映射完成，在该阶段不考虑硬件资源和指令集约束。
第二阶段中，实际的硬件加速器只能计算特定shape的数据，文中称之为问题大小限制。针对存储容量限制，每种存储器只有有限的存储空间，因此文中将张量数据切分为片，然后多次执行load/store来完成访存操作。这样基地址和地址偏移量会存在多次更新。
如上操作可能会生成一些非法的映射，因此文中也提出一种映射结果验证方法用来过滤掉非法映射。映射结果验证方法描述如下，核心思想是通过分析生成的0、1构成的矩阵得出合法映射。
在探索最优映射方面，文中构建了一个代价模型的数学描述（只考虑时延指标）。由于空间加速器都是分层次设计，系统的整体时延等于最上层硬件抽象层的时延，每层的时延是计算、读取、写入操作时延的最大值。
4. 实现
实现上，硬件抽象层通过IR抽象描述，文中基于TVM的tensorize接口新增了2个IR节点（计算节点和访存节点）：计算节点IR中的第一个属性是张量数据，第二个属性是计算表达式Expr，第三个属性是指令迭代空间；访存节点IR中的第一个属性是目的地址数据，第二个属性是描述数据驻留的存储器类型（global、shared、register），第三个属性是源地址数据和加载索引。在代码生成时将新增的IR节点插入到AST中。
5. 效果
AMOS在绝大部分测试项中均超过PyTorch，加速比0.91到10.42。Bert模型中有大量矩阵乘操作，矩阵乘操作已经在多个第三方库中充分优化了多年，即便如此，在batchsize为16时AMOS仍然可以达到第三方库90%的性能。在ShuffleNet中可达到10.42的性能加速比，原因是该模型中存在大量深度可分离卷积和分组卷积，第三方库没有充分将这两类卷积支持到Tensor Core中。同时，AMOS与TVM相比也有明显的性能优势。
同时，AMOS具有良好的泛化性，可支持多种加速器设备，如Tensor Core、Intel CPU的AVX-512矢量单元、Mali GPU的点乘单元，作者在这些加速器设备上也做了大量实验，在这些加速器设备上与TVM相比均有一定的性能提升。
最后给我们团队的开源项目打个广告，Adlik 是中兴通讯贡献的深度学习推理工具，已获得 Linux AI 基金会支持，目前仍在持续完善中，期待大家的支持和关注：Github Adlik；另外，我们也在做深度学习编译器方向的前沿技术研发工作，欢迎感兴趣的小伙伴加入我们~~
