清华大学高鸣宇:基于Halide调度实现高效能的DNN加速

计算机体系结构领域国际顶级会议每次往往仅录用几十篇论文,录用率在20%左右,难度极大。国内学者在顶会上开始发表论文,是最近十几年的事情。

ASPLOS与HPCA是计算机体系结构领域的旗舰会议。其中ASPLOS综合了体系结构、编程语言、编译、操作系统等多个方向,HPCA则主要针对高性能体系结构设计。过去的三十多年里,它们推动了多项计算机系统技术的发展,RISC、RAID、大规模多处理器、Cluster架构网络存储、机器学习加速器等诸多计算机体系结构领域的重大技术突破,都最早发表在这两个会议之上。

2020年4月12日上午,北京智源人工智能研究院和北京大学高能效计算与应用中心联合主办了“AI芯片体系架构和软件专题报告会”,五位学者结合在2020年计算机体系结构顶级会议(ASPLOS和HPCA)中发表的最新研究成果,针对AI芯片和体系结构领域的几个关键挑战:芯片加速、能耗和反对抗攻击等,详细介绍了他们的最新思考和解决方案。超过1900名观众在线观摩了学者们的报告。报告主题分别是:

  • 本次报告会主席、智源青年科学家和北京大学研究员梁云《FlexTensor: An Automatic Schedule Exploration and Optimization  Framework for Tensor Computation on Heterogeneous System》

  • 美国南加州大学PhD Candidate骆沁毅《Prague: High-Performance Heterogeneity-Aware Asynchronous Decentralized Training》

  • 清华大学交叉信息研究院助理教授高鸣宇《Interstellar: Using Halide's Scheduling Language to Analyze DNN Accelerators》

  • 智源青年科学家、中国科学院计算技术研究所副研究员陈晓明《Communication Lower Bound in Convolution Accelerators》

  • 中国科学院信息工程研究所研究员和信息安全国家重点实验室副主任侯锐《DNN Guard: An Elastic Heterogeneous Architecture for DNN Accelerator against Adversarial Attacks》

本次活动的演讲内容都会整理成文字,陆续发布。今天,我们将介绍清华大学交叉信息研究院助理教授高鸣宇的《Interstellar: Using Halide's Scheduling Language to Analyze DNN Accelerators》。

 

高鸣宇,博士毕业于美国斯坦福大学电子工程系。主要研究成果包括针对大数据应用的近数据处理架构软硬件系统,高密度低功耗可重构计算芯片,及专用神经网络芯片的调度优化。已发表多篇国际顶级学术会议(ISCA、ASPLOS、HPCA、PACT等)论文,授权多项专利。曾获得IEEE Micro 2016年度计算机系统结构最佳论文奖(Top Picks)、欧洲HiPEAC论文奖、福布斯中国评选2019年“30 Under 30”等荣誉。

本次演讲中,高鸣宇带来了他发表在ASPLOS 2020上的最新作品——Interstellar。Interstellar是高鸣宇在Stanford求学期间与他人合作提出的基于Halide调度语言对DNN加速器进行分析的工具。在演讲中,高鸣宇从DNN加速计算中通用的7层循环计算入手,阐述了DNN加速器的Blocking,Resource Allocation和Dataflow三个维度,并使用循环变换描述设计空间,进而利用Halide中的调度语言进行实现。思路清新完整,令人豁然开朗。以下是智源社区编辑整理的讲座内容,让我们一起跟随高鸣宇开启DNN 加速器的星际探测之旅吧!

 

整理:智源社区蔡晓圣、孙超、张驰

01 DNN在硬件及软件层面的高需求

伴随着深度神经网络(DNN)正在包括自动驾驶、人脸识别、人脸检索和图像识别在内的各领域得到更加广泛地应用,业界对DNN加速器的需求也在不断增加。鉴于此,不论学术界还是工业界,都对DNN加速器开展了大量的研究,包括Google的TPU、华为麒麟NPU, MIT的Eyeriss以及UCLA在FPGA上的架构创新等。

 

图1: DNN加速器硬件成果

同时,我们在关注底层硬件架构的基础上,也同样需要把目光聚焦在上层软件层面,探索如何调度以及更好地使用底层硬件,从而最大化算力的资源使用。目前上层软件主要根据CNN等卷积神经网络相关算法进行调度,那么何为CNN算法?如图2所示,CNN是由多个2D的input features maps(ifmaps)产生多个输出的output features maps(ofmaps),在每一对ifmaps和ofmaps之间对应一对2D的filters,同时由于ifmaps和ofmaps是3D的,每一对均对应一对filters,所以filters实际上存在4D的张量结构,即除了每一个2D的filters之外,可以认为所有的ifmaps的一组和相同数量的filters产生一个ofmaps,而每一个ofmaps都对应一组filters。

 

图2: CNN计算形式

同时,在input和output之间均作一个Nb大小的batching,则input、output和filters三者均会形成一种4D结构(2维图像结构以及1维通道channel维度和1维batch维度),如果用循环的形式表示这种框架,则为七层循环架构,如图3所示。可以看到CNN的架构较为复杂,而所有neural network硬件的实现,均为针对这种七层循环结构进行的合理调度,使其最大化利用计算的便携性和硬件资源。

 

图3: 卷积层中的七层循环架构

 

02 DNN加速器设计空间

将神经网络的加速空间罗列出来(如图4所示),可以分为3个维度:

1. Resource Allocation(资源配置)为硬件方面的维度,比如采用多少并行的PE处理单元,采用寄存器大小或SRAM缓存大小参数的设计等。该维度主要作用为提供硬件资源,为软件调度提供约束。

2. On-chip的dataflow调度,主要描述片上数据在多个单元之间流动。

3. Off-chip维度,将DRAM迁移到SRAM里,进行片上运算,我们称为Blocking,其中涉及的选择为如何将大的数据集划分为小的数据块,然后将这些数据块分别放到片上进行计算。这其中就涉及如何将数据切分成小的数据块。

 

图4: DNN加速器的3D设计空间

那么,在Blocking维度,具体应如何将数据切分成小的数据块?刚开始(如图5所示),所有的input、filters、output均存放在Main Memory里,如果不考虑效率问题,则直接将片外的内存搬移至片上的cores,如ACK单元或CPU、GPU计算模块进行计算。但是很明显,这种方式计算并不高效,会导致所有数据在两者之间来回转换,而比较有效的方式为进行数据划分,即将数据切分成小块,每次只把数据的一部分放入Global Buffer里,从而最大化片上重用性,通过片上cores和global buffer之间的片上传输,避免高昂的线外传输。但在这个过程中也需要考虑很多问题,比如如何划分数据块大小,尤其是针对input和output数据块大小的划分等。

 

图5: 不同的Blocking方式

在dataflow维度,当将各类数据放到SRAM等片上global-buffer里面,如何将片上数据传输到并行的处理单元,以及如何在处理单元之间进行最大化的共享同样需要很多技巧和策略。图6列出了一些常用的数据流传输算法:

左侧第一个为简单的adder tree结构,所有的PE在input和weight之间进行乘法运算,之后进行加和得到某一个输出结果;

左侧第二个为MIT的Eyeriss架构所采用的unrolling simultaneously片上算法,即将weight按行放到PE里,output每一行按列放到PE的列里,input按照对角线的形式进行传输,从而避免PE分别从global-buffer获取数据,使得数据可以在PE之间传输,减少数据读出操作;

最右侧同样为常用架构,利用fully connected layer,将input channel和output channel分别放到PE列和PE行里,input从上到下进入PE array,output从左到右走出PE array,将其划分成systolic array的方式,目前的Google TPU即采用这种方式。

总而言之,Dataflow解决的主要是将哪些数据放在PE上运算及PE运算哪些数据的问题。

 

图6: 不同的dataflow架构

在Resource Allocation维度,涉及到Memory Hierarchy中包含多少层级,以及每一层的合适大小问题。在这个问题中涉及到一个很核心的tradeoff,即每一层的存储空间如用SRAM在片上实现buffer或者缓存,如果缓存的大小更大,则可以在片上缓存更多的数据,从而进一步减少线外的数据访问,进而减少数据访问的能耗。但同时,如果片上缓存做大,会导致对片上缓存的每一次访问能耗消耗更大,即片上访问会更昂贵。这里就涉及到减少线外访问次数和降低片上访问代价的折中处理。因此需要合理设计片上缓存大小,使其既能buffer一定的数据,同时又能避免增大片上访问的代价。

图7: 不同的Resource Allocation之间的权衡

 

03 选择 Halide的相关分析

通过循环变换描述设计空间,可以使用任何与循环变换兼容的编译系统来生成硬件设计。而在这项工作中,Halide是一个比较合适的选择。Halide是一款基于Loop-based计算,能够很好地描述调度选择的工具。

 

那么,选择Halide的具体价值如下:

 

1、Halide是比较紧密的调度空间。

(1)基于循环的数据流分类法可以和Halide现有的调度语言很好地关联在一起,比如前文提到的循环展开、循环并行、循环划分以及交换等,这些在Halide现有的Primitive(原语)里面都得到了支持。

 

(2)Halide是可以把早期算法和调度选择分开的工具之一。算法与调度表之间进行解耦的好处是,在给定算法的情况下只需要去调整它的Schedule Choice,不用重写算法实现不同的Schedule。这样,调整Schedule、探索设计空间时也不会担心因为重写算法而导致计算的正确性会发生变化,比如,Accidentally会改变算法的正确性,因为它可以改变循环的方式。

 

2、Halide可以重复使用编译框架

(1)Halide使用非常短小的代码就能进行高性能图像处理。

 

(2)Halide现在已经有了非常好的编译环境和编译框架,这不仅在学术界,而且在工业界都被广泛应用,比如现在已经可以支持很多CPU或者GPU等后端的代码生成。我们可以重复利用已有的编译环境以及对现有语法的分析,通过这些现有的工具来提供我们调度选择上需要的信息,比如数据的大小等等。

 

Halide 调度表原语

 

我们还会用到很多Halide原语,比如:

 

(1)Loop Block Computation有Loop split和Loop reorder。split和reorder可以把一个Loop划分为多个层次化的循环,或者改变循环之间的顺序。

 

(2)Resource Allocation中有in这种原语,它可以表明数据可以放在缓存里面,然后在这一级里增加一个缓存的层次。

 

(3)Compute and store granularity(计算和存储的粒度)可以用compute_at或者store_at两个原语,结合前面create another level of memory 中的in这种原语,三者之间可以很好地指定硬件上面应该在哪个地方分配一个存储的空间,包括指定存储Buffer大小是多少。

 

(4)Parallel computation (duplicate units)的unroll是把一个顺序循环展开,成为多个计算单元上面并行处理的循环。

 

片上的Dataflow,它可以利用很多变形的计算原语,比如unroll,所以现有的Halide Primitive是和前面所探讨的设计空间的维度紧密结合的,并且二者可以非常好地关联起来。

 

运用实例逐步完成生成设计的过程

高鸣宇团队没有具体讨论每个原语的具体语法和含义,而是通过一个例子给出现有的设计。

 

首先,从主内存中可以获取所有数据,每次搬运数据到片上Process Element进行计算,第一步先把数据进行划分,比如通过Split和Reorder把大块的数据划分成四个小块数据,以便把每一个小块搬运到片上的计算单元中去计算。

 

其次,划分之后还需要在片上分配出一个一个buffer,这里用到in和compute_at来创建 buffer空间。这就表明某一部分的数据应该被存放在这样的Buffer里面,而且每次是在buffer里取出数据进行计算。

 

最后,为了提高吞吐量,这里需要一种并行计算的方法,通过重载Halide展开原语可以实现这一点。此时进行给定多个PE的操作,比如首先进行Loop Unroll的模式,在xi这个Loop上分成四份,这样会提供四个PE,并且计算四个不同的Input,都会从Buffer当中取出数据。我们也可以针对新的设计进一步引入新的原语,比如systolic表明数据并不都是从Buffer读取,数据从一个PE中读取后,通过systolic在不同的PE之中进行流动。Accelerate也是新引入的原语,表明计算的是哪一部分的数据。

 

可以看出这三步就是前面提到的三个维度设计空间的策略选择:第一步是Blocking的维度,第二步是Buffer Resource Allocation,第三步是片上Dataflow这个维度。

 图8: 描述微框架的调度语言示例

 

不同数据流的调度语言描述

上面已经显示了基于循环的数据流分类法自然符合Halide调度语言,接下来高鸣宇团队使用Halide展开原语来生成不同的数据流。前文提到用Loop Unroll的方式描述Dateflows这个维度,Input Channel在C维度上进行变形,所以相当于将C Loop进行unroll,然后这里的unroll(c)是针对在二维PE中分别unroll两个不同的Loop,一个是Filter按行传递——unroll(fy),一个是Feature Map按行纵向传递——unroll(y),因此可以表达成为两个不同的Unroll Primitive。那么后面两个维度,Input和Output Channel也可以Unroll CK两个循环。

 图9: 描述不同数据流的调度语言示例

 

DNN加速器设计的Halide调度语言描述

将Halide调度原语映射到3维,可以使用相应的调度原语沿不同的轴进行优化。Resource allocation 使用in和compute_at分配存储的buffer,dataflow 使用unroll描述数据在哪个平面并行,blocking 使用split/reorder划分数据,表示数据访问的顺序。这样就构成了用Halide的原语去描述设计空间的方法。

 

在图10中可以看到Highlight的两个平面, dataflow和blocking所在的Mapping plane是在一个特定硬件的结构下如何决定片上和片外软件的调度策略,另一个平面就是Optimization Plane,这个平面在后期会有更多的关注,这个平面本质上是固定在某一种片上Dataflow的选择方式,例如固定硬件的模板化设计,如何选择最好的Blocking策略以及最好的Memory Buffer Allocation,因此相比起Mapping Plane和之前纯粹运用调度的优化来提升性能,Optimization Plane其实对整个Accelerator Efficiency有更大的影响,因而通过结合硬件Resource Allocation这种方式优化加速器会有更好的表现。

 图10: 将Halide调度原语映射到3维

 

04 分析模型与结论

尽管可以使用Halide将硬件工具链用于生成和合成硬件,但通常设计周期较长,所以为了快速探索大型设计空间,高鸣宇团队创建了一个分析模型来分析、搜索空间。这样可以单独分析每个维度的影响,通过这样的Analytical Cost Model直接计算最后各种设计带来的Energy Performance的表现。所有Layer Configuration和Build Cost Model作为Input,其中Cost Model代表每次训练后memory所需要的Energy和Communication Cost是多少,以及Communication Cost在每个PE上是多少,这些都是根据真实的Hardware结果,和现有研究成果进行验证综合而成的结果,具体细节可以参阅高鸣宇的论文。

 

下面是关于高鸣宇团队所探索的性能表现以及他们所找到的Insight。他们首先固定了以AlexNet为主的Workload, PE采用16×16大小,16Bit的精度,搜索Global Buffer大小的变化以及register file的一个变化,来探索这些硬件的参数对于最终的性能和最终Energy的影响有多大,以及包括前面设计空间的三个维度在内的不同的Dataflow和不同的Blocking策略对性能的影响。

 

1、首先比较一下Dateflow。Dataflow横轴采用的是不同Dateflow的方式,这里是用之前提到的Unroll Loop进行表示,比如,CK在这里表示的是C和K两个Loop,还有X和Y表示Filter和Feather Map两个2D维度。蓝点表示只利用了一个512B的Register File,红点表示我们用512Bit的Register File,每个PE会单独从Global Buffer读取数据,类似于broadcast这样的传递方式,这种情况下其实可以看到不同的Dataflow的方式对Energy的影响是非常小的,横轴不同的点带来的energy基本都是非常相似,我们需要从Global Buffer取出数据的次数不一样,总体来讲对于整体Energy的影响是非常小的。

 图11: 不同Dateflow的能量消耗

不同Dataflow的表现会有明显的不同,大概会带来20%-25%左右的PE Utilization上面的差别,这种带宽或者吞吐量上面的差别来自于不同的Workflow对于固定大小的二维PE Array上面的映射。这是因为可能会出现除不尽的情况,例如把16×16 Feature Map映射到3×3的空间, 16÷3余下的1就是额外的一次PE映射的一个Pass处理,这个时候相对于一个满的3×3的处理,它的PE Utilization就会下降,也就是运算达不到峰值的吞吐量。但如果把一个大的Feature Map拆分成许多小的Feature Map,每次映射一个小的Feature Map,或者反过来,对一些小的Feature Map映射到一个大的Array上面,这样一次性映射多个不同的Feature Map,比如3×3映射到16×16PE Array可以每次映射25个Feature Map,最大化地利用PE Size。在这种情况下,只要两个数据互相之间是比较匹配的,除不尽的情况会比较有限,Utilization总体来讲就可以达到80%以上。

图12: 不同Dataflow的PE Utilization的表现

2、下面来看Blocking的影响。这里把不同的Blocking带来的不同Energy做了分布曲线,由图可以看出曲线会受到Long Tail的影响——很大一部分blocking策略距离真正最优的Blocking选择是比较远的,大概可以带来2倍以上或者1.5-2倍以上的Energy浪费,在这种情况下比起On-chip Dataflow,Blocking的选择其实是有一个更大的Energy的影响,这也是为什么比起Dataflow更需要关注片外的Blocking选择。这方面的问题在于,片上的dataflow选择主要影响的是array,就是如何在PE ARRAY里面进行数据传递产生的energy cost,在一些常见的Configuration 或者网络层里面,这一部分Energy Cost是非常小的,其实就是Wire上面数据传输的Cost,比起访问Global Buffer来说是比较小的;或者从另一个角度来讲,对于一些这种Fully Configuration,有很多片外的数据传递,这样的片外传递成为对energy cost影响最大的一个地方,因而Dataflow Choice对整体的影响是比较小的。

 

图13: 不同的Blocking产生的Energy分布曲线

那么,既然Communication Cost要更小,高鸣宇团队选择了进一步优化Memory Allocation和改变RF size、buffer size这两种方式。Memory Buffer大小和层次对Energy的影响比较明显,比如比起512bit的Register File,使用64bit这种很小的Register File,Energy会小很多,因为64Bit Energy File每次访问数据的代价和能量的开销是比较小的,可以降低总体的Energy Cost。这里他们采用64Bit和512Bit两种不同的PE Register大小能耗的比较,可以看到采用更小的Register File可以很大程度上降低总体的能量。除了每一层的Buffer大小,他们还可以进行更大更深的层次,比如两级Register File,这会进一步带来大概20%能耗的提升。

 

最后高鸣宇团队使用分析模型,创建了自动优化器,联合穷举搜索最佳内存层次结构和整个网络的最佳hierarchy,通过使用获得的观察结果,修复了数据流以减少搜索空间。为了获得结果,所有层都使用相同的架构,并且具有最佳的hierarchy,可以看到相比之前的结果有一个明显的提升,最佳的Configuration是采用256KB 的Global Buffer和128 B + 16 B two-level Register File结构。

 

05 总结

 

现阶段各类DNN加速器层出不穷,高鸣宇团队从最为通用的7层计算循环入手,将DNN加速器设计空间划分为Blocking,Resource Allocation和Dataflow三个维度,并采用循环变换统一描述设计空间,进而利用Halide中的调度语言进行实现。该工作的亮点包括:

 

第一、Design Space可以通过Loop Transformation的方式描述,利用片上片外的方式去描述数据的选择,数据的传递。

第二、直接利用现有的Halide描述语言表达Space。

第三、许多不同的数据流可以达到相似且接近最佳的效率。

第四、优化内存层次结构对能源效率有很大影响,同时片上的Dataflow传统上大家可能以为比较关键,但是对整体Energy的影响是有限的。

 

Q&A

 

Q:如果Memory资源和Loop循环的次数不是倍数关系,有余数的情况怎么处理?

A:在现有的工作里假定最后一层有余数的情况,如果用一个并没有用满PE Array的情况去运行,便会影响我们总体Throughput(吞吐量)的下降,更好的处理方式可能是在硬件上进行改动,比如PE Array是不是可以动态地重构成不同的大小,或者动态划分成更小的单元,以适用于处理稍微有点小的余数,达到最大化的PE Utilization。

 

Q:Mapping 是针对Gem还是针对Convolution?

A:我们是统一处理的,不管是Gem还是Convolution都可以或多或少映射到前面的七层循环方式。在Convolution这种方式的情形下,它的某些循环loop bound是1,但可能循环的数量没有了,循环的数量就少一点,总体来说都是基于循环的计算,通过Halide调度的方式是可以比较统一地进行处理。

 

点击阅读原文,进入智源社区参与更多讨论。

Published by

风君子

独自遨游何稽首 揭天掀地慰生平