GPTPU: Accelerating Applications using Edge Tensor Processing Units(SC 21′)

作者：Kuan-Chieh Hsu and Hung-Wei Tseng @University of California, Riverside

本文介绍了一种在Edge TPU上进行通用计算加速的开源框架。

论文为什么要做这个事情？

• 理论可行性：NN加速器理论上是一种张量处理器，因此对于任何以张量为输入、输出的应用都有潜在的加速效果；
• NN加速器的局限性：现有的商用NN加速器只向用户开放了AI或机器学习专用的调用接口，对通用领域没有相关设计；
• 框架的必要性：NN加速器几乎没有开放其硬件设计的细节，对于程序员来说，很难仅仅通过调用接口来加速通用的算法。因此，设计一个计算框架成为了提升可编程型的必要工作。

多年之前，GPU只是用来着色和渲染的工具，有CUDA和OpenCL之后，GPU凭借其大规模并行能力，面向了矢量计算领域。本工作和TPU的关系与CUDA与GPU的关系异曲同工。

这个事情之前别人是怎么做的，存在什么问题？

① 背景知识

• TPU：通过创建对张量执行操作的脉动阵列，来加速机器学习领域中的神经网络任务的领域特定硬件。
• Edge TPU：谷歌开发的谷歌云TPU的精简版。相比于谷歌云TPU，它具有公共可得性、且其开放了部分后端C++代码、有更好的每瓦特性能（2TOPS/W vs. 0.36TOPS/W），因而被本工作选用为加速器。除谷歌之外，在TPU领域仍有非常多的工作，但是有绝大部分的TPU并不生产，用来学术研究；或者其性能和功耗表现较差，没有被本工作选用。然而，这些TPU都以张量为输入和输出，且其接口大同小异，所以本文提出的框架应该具备可移植性。

② 前人工作*的问题

本工作和前人工作有本质上的不同。

1. 以往的NPU只可以加速与之前训练到的神经网络模型相匹配的算法， 而GPETPU可以将张量运算映射到TPU上来加速任何用户定义的算法；
2. NPU只能产生近似于矩阵运算的效果，但GPETPU可以使用Edge TPU实现确切的矩阵操作结果；
3. NPU被神经网络模型大致的结果所限制了精确度，GPETPU可以通过对原始输入的不同部分不断迭代达到所需的精度。
4. 有一些工作关注了ASICS或者稀疏矩阵压缩，在这里并不属于同一范围内的问题。

(*: Neural acceleration for general-purpose approximate programs,2012; Neural acceleration for GPU throughput processors,2015)

该论文解决了什么难点，难点存在的原因是什么，作者是如何发现这些难点的？

设计一个在TPU上的通用计算框架有以下的几个难点：

1. 领域专用性：NN加速器为神经网络任务而进行优化，直接将传统的算法映射到加速器上会产生sub-optimal问题；
2. 容错性的考虑：NN加速器往往会牺牲一定的精度换来面积或者功耗，在设计通用计算时需要做一定的调整；
3. 底层不可知性：现存的框架对于神经网络加速器的软硬件接口介绍很少，普通的应用因此而调整参数或者数据格式将造成严重的性能开销；
4. 算法适用性：通用领域优化过的算法对于张量计算不再适用，需要重新设计基于张量的优化算法。

针对以上难点，作者各做了哪些优化，优化背后的思想是什么？

分析Edge TPU的特点

作者采用了以下的原型机对TPU的特点进行了探索：

TPUs: built two quad-EdgeTPU PCIe expansion cards using QNAP QM2-4P-384A, containing 4× M.2 Edge TPUs with M.2 slots connected to a PCIe switch (totally 8x TPUs)

HOST: AMD Ryzen 3700X, 4.4GHz

Cache: 32MB LLC

Main memory: 64GB DDR

Storage: NVMe SSD

分析Edge TPU指令

文章分析了Edge TPU的一些指令的性能表现。由于像TOPS(tera op per second)、IPS(inference per second)这些常用的衡量标准都是面向神经网络的，所以本文建立了一些新的衡量标准RPS(results per second)。 RPS计算了某条指令产生结果数量的效率，测量得到的结果是：

从上表我们可以得到以下结论：

• conv2D的RPS非常高，这是因为卷积操作的使用频率高，TPU进行了特别的优化；
• 不同的指令RPS相差很大；
• OPS和RPS不是强相关的

分析Edge TPU采用的数据和模型格式

Edge TPU 指令的输入一般是两种数据：

1. 待推断的张量数据集，
2. TFLite编译生成的模型数据。

对于通用计算的输入来说，必须要把其中一个张量数据转换为TFLite的模型格式，底层的TPU才能够接收。 而由于直接采用TFLite进行编译的延迟过长，对于除机器学习之外的应用不可忍受。 因此，通过对TFLite改变不同的数据输入、维度以及值，该工作解码了TFlite编译生成的文件。主要包括四个部分：

1. 头部(Header)：每个模型头部都是有120byte的用于识别的模型头，最后4byte描述了数据段的大小。
2. 数据段(Data Section)：这里用行优先的方式存放了二进制编码的8比特整数，超出8比特将使用一个缩放因子进行缩放。
3. 额外数据段(Metadata Section)：存放了数据段的维度和缩放系数。
4. TFLite编译生成的模型采用小端存储的方式进行数据存放。

GPETPU整体架构

整个GPETPU可以分为如下几个部分：

• 前端开发了OpenCtpu方便编程，以及与之配套的编译器。
• 后端为API库和运行时系统，他们的作用是连接用户程序与EdgeTPU接口。

OpenCtpu: 编程接口（前端）

作用

描述TPU任务，协调异构计算与数据交换。

特点

• 将应用程序和设备使用的控制放在主机端；
• 需要程序员显示指定TPU运算的输入输出缓冲区；
• 提供了让程序员描述计算任务的函数

用法

1. 编写描述计算的核函数：注意，在描述计算时，用invoke函数可以显式地表示这里需要采用TPU进行加速。除了这些显式的操作之外，在和函数中进行矩阵加减乘除等操作的时候也会默认使用TPU进行加速；
2. 准备TPU运算核函数的输入输出缓冲区；
3. 以任务的形式使核函数进队列 ；
4. 适时进行同步：这里在调度时系统的设计中，任务之间是可以并行的，所以需要在恰当的时候进行同步操作来获得结果。

GPETPU的运行时系统

任务调度

GPETPU的运行时任务调度策略是一种基于数据流的算法，由前端任务队列OPQ和后端指令队列IQ组成。OPQ的每一项包含任务ID，任务操作符，输入、输出、量化参数等必要的信息。IQ的每一项包含任务ID、TPU操作指令码、输入输出位置等。

OPQ

1. 主机端调用enqueue API时，运行时系统新建一个taskID,为即将进入队列的kernel函数做准备，并且开始执行该kernel函数。
2. 在kernel函数中，若触发了TPU操作，则阻塞当前的kernel函数，用当前taskID创建一个OPQ任务，并在OPQ的参数准备好后进入OPQ队列。
3. OPQ任务可以并行在Tensorizer上 。

IQ

Tensorizer将OPQ任务转换为TPU指令，将数据转换为TPU可以接受的形式，并将其送至IQ上。若有一些IQ项有相同的输入、量化标志、taskID，但输出位置不同，IQ便将他们调度到同一个TPU上操作，以减少数据转移的开销以及数据格式的转换。其他采用FCFS方式。

Tensorizer

作用

作用主要是把用户要求的操作转换成对应的TPU指令，在这个过程中需要对数据也进行裁剪和缩放以适应TPU的特征，并且要仿照TFLite对数据进行包装。

将操作转换为TPU指令

• 将任务数据块划分为TPU指令可以达到最优性能的尺寸大小；
• Pair-wise操作(add): TPU指令在每个子块上操作后，整合运算结果得到输出；
• Matrix-wise操作(max): TPU指令在每个子块上操作后，生成CPU代码来集成最终结果。相较于迭代法，减少了数据的转移；
• 算术操作(conv2D/FC): 类似于矩阵乘的分块算法，将运算分为PxQ大小的块，先用TPU指令再生成CPU代码来集成出最后的结果。

将数据转换为适用于TPU的格式

Tensorizer将指令的输入数据缩放到定点数范围之内，并将这些数字包装成TPU可接受的模型或矩阵上。 缩放系数的计算：

S与kernel函数上的操作顺序、操作总数、input的范围相关。

Tensorizer的开销

Tensorizer还需要借助前文解码的TFLite生成的模型格式，仿照TFLite将一个张量包装成模型的形式。 实验测得：基于C的Tensorizer包装的速度比基于python的TFLite快1500倍。此时延可以与对EdgeTPU之间的数据转移相抵消。

为GPETPU优化应用

Tensorizer完成了Task级别的优化，但是要完成一个应用，采用什么样的操作符进行计算仍然需要探究。我们以矩阵乘法为例：

1. 用全连接操作符完成矩阵乘法非常简单，因为全连接本质上是一种矩阵-向量乘法，所以很容易实现，但是根据本工作前文对TPU指令的量化，全连接操作的RPS非常低。
2. 除此之外我们还可以使用卷积操作实现矩阵乘法，这时候需要数据进行一定程度的变形。但是他的RPS非常好。

上图是使用全连接和卷积两种操作的效果，对比对象是CPU上使用OpenBLAS库进行计算。在只使用全连接算子时，其速度还不如CPU；但用卷积操作则可以加速超过两倍。 这个实验告诉我们，采用GPETPU进行编程时仍然需要注意算子的选择，可以依据之前测量的RPS进行选择。

为了进一步方便程序员的编程，GPETPU将一些领域内最常用的算法通过优化封装在标准库里，可以直接调用。领域涵盖了图计算、线性代数、物理仿真、模式识别以及金融领域。用户在前端像调用TPU基本操作一样调用这些算子就可以实现这些基本操作。

实验结果

实验平台

同上。

单核性能比较：TPU vs. CPU

Baseline是用单核CPU进行应用实现的方案。可以看到TPU平均可以实现2.46倍的加速；且能耗要比CPU低45%。

GPETPU的扩展性

其次，使用多个TPU进行并行计算也有很好的扩展性。左边这个图的baseline仍然是使用单核CPU的计算速度。可以看到采用8核CPU的平均速度提升大概只在2.7倍左右，但采用8核TPU则可以提升13倍左右。 随着TPU数量的增多，大部分应用的速度几乎呈线性提升。

与GPU的比较

最后这张图展示了GPU与TPU的性能表现对比以及相对能耗。仍然是对比了单核的CPU。 虽然TPU的性能不如GPU，但是能耗非常低，相对于CPU可以平均节省40%的能量。 综合考虑能耗和性能，8x Edge TPU可以超出baseline 46%，强于GPU。

自己的思考

借助GPETPU框架，Edge TPU在边缘计算以及嵌入式设备中有较好的应用前景。类似的设计方法在各种加速器设计上都可以参考。

赏

大哥大嫂过年好！

支付宝

微信