谷歌 TPU 是一个设计良好的矩阵计算加速单元, 可以很好的加速神经网络的计算. 本系列文章将利用公开的 TPU V1(后简称 TPU)相关资料, 对其进行一定的简化, 推测和修改, 来实际编写一个简单版本的谷歌 TPU. 计划实现到行为仿真为止, 仅为更确切的了解 TPU 的优势和局限性, 暂无在 FPGA 等硬件上进一步实现的计划.
系列目录
谷歌 TPU 概述和简化
基本单元 - 矩阵乘法阵列
基本单元 - 归一化和池化(待发布)
TPU 中的 Instruction (待完成)
SimpleTPU 实例: (计划中)
拓展
TPU 的边界(规划中)
重新审视深度神经网络中的并行(规划中)
1. TPU 设计分析
人工神经网络中的大量乘加计算 (譬如三维卷积计算) 大多都可以归纳成为矩阵计算. 而之前有的各类处理器, 在其硬件底层完成的是一个 (或多个) 标量 / 向量计算, 这些处理器并没有充分利用矩阵计算中的数据复用; 而 Google TPU V1 则是专门针对矩阵计算设计的功能强大的处理单元. 参考 Google 公开的论文 In-Datacenter Performance Analysis of a Tensor Processing Unit,TPU V1 的结构框图如下所示
结构框图中最受瞩目的是巨大的 Matrix Multiply Unit, 共计 64K 的 Mac 可以在 700MHz 的工作频率下提供 92T int8 Ops 的性能. 这样一个阵列进行矩阵计算的细节将会在基本单元 - 矩阵乘法阵列进行更进一步的阐述. TPU 的设计关键在于充分利用这一乘加阵列, 使其利用率尽可能高.
结构图中其他的部分基本都是为尽可能跑满这个矩阵计算阵列服务的, 据此有以下设计
Local Unified Buffer 提供了 256*8b@700MHz mailto:256×8b@700MHz 的带宽(即 167GiB/s,0.25Kib*700/1024/1024=167GiB/s), 以保证计算单元不会因为缺少 Data in 而闲置;
Local Unified Buffer 的空间高达 24MiB, 这意味着计算过程的中间结果几乎无需和外界进行交互, 也就不存在因为数据带宽而限制计算能力的情况;
Matrix Multiply Unit 中每个 Mac 内置两个寄存器存储 Weight, 当一个进行计算时另一个进行新 Weight 的载入, 以掩盖载入 Weight 的时间;
30GiB/s 的带宽完成 256*256Weight 的载入需要大约 1430 个 Cycles, 也就意味着一组 Weight 至少需要计算 1430Cycles, 因此 Accumulators 的深度需要为 2K(1430 取 2 的幂次, 论文中给出的数值是 1350, 差异未知);
由于 Mac 和 Activation 模块之间需要同时进行计算, 因此 Accumulators 需要用两倍存储来进行 pingpang 设计, 因此 Accumulators 中存储的深度设计为 4k;
因此从硬件设计上来看, 只要 TPU ops/Weight Byte 达到 1400 左右, 理论上 TPU 就能以接近 100% 的效率进行计算. 但在实际运行过程中, 访存和计算之间的调度, 读写之间的依赖关系(譬如 Read After Write, 需要等写完才能读), 指令之间的流水线和空闲周期的处理都会在一定程度影响实际的性能.
为此, TPU 设计了一组指令来控制其访问存和计算, 主要的指令包括
- Read_Host_Memory
- Read_Weights
- MatrixMultiply/Convolve
- Activation
- Write_Host_Memory
所有的设计都是为了让矩阵单元不闲下来, 设计希望所有其他指令可以被 MatrixMultiply 指令所掩盖, 因此 TPU 采用了分离数据获取和执行的设计(Decoupled-access/execute), 这意味着在发出 Read_Weights 指令之后, MatrixMultiply 就可以开始执行, 不需要等待 Read_Weight 指令完成; 如果 Weight/Activation 没有准备好, matrix unit 会停止.
需要注意的是, 一条指令可以执行数千个周期, 因此 TPU 设计过程中没有对流水线之间的空闲周期进行掩盖, 这是因为由于 Pipline 带来的数十个周期的浪费对最终性能的影响不到 1%.
关于指令的细节依旧不是特别清楚, 更多细节有待讨论补充.
2. TPU 的简化
实现一个完整的 TPU 有些过于复杂了, 为了降低工作量, 提高可行性, 需要对 TPU 进行一系列的简化; 为做区分, 后文将简化后的 TPU 称为 SimpleTPU. 所有的简化应不失 TPU 本身的设计理念.
TPU 中为了进行数据交互, 存在包括 PCIE Interface,DDR Interface 在内的各类硬件接口; 此处并不考虑这些标准硬件接口的设计, 各类数据交互均通过 AXI 接口完成; 仅关心 TPU 内部计算的实现, 即下图红框所示.
由于 TPU 的规模太大, 乘法器阵列大小为 256*256, 这会给调试和综合带来极大的困难, 因此此处将其矩阵乘法单元修改为 32*32, 其余数据位宽也进行相应修改, 此类修改包括
Resource | TPU | SimpleTPU |
Matrix Multiply Unit | 256*256 | 32*32 |
Accumulators RAM | 4K*256*32b | 4K*32*32b |
Unified Buffer | 96K*256*8b | 16K*32*8b |
由于 Weight FIFO 实现上的困难(难以采用 C 语言描述), Weight 采用 1K*32*8b 的 BRAM 存放, Pingpang 使用;
由于 Matrix Multiply Unit 和 Accumulators 之间的高度相关性, SimpleTPU 将其合二为一了;
由于 Activation 和 Normalized/Pool 之间的高度相关性, SimpleTPU 将其合二为一了(TPU 本身可能也是这样做的), 同时只支持 RELU 激活函数;
由于并不清楚 Systolic Data Setup 模块到底进行了什么操作, SimpleTPU 将其删除了; SimpleTPU 采用了另一种灵活而又简单的方式, 即通过地址上的设计, 来完成卷积计算;
由于中间结果和片外缓存交互会增加 instruction 生成的困难, 此处认为计算过程中无需访问片外缓存;(这也符合 TPU 本身的设计思路, 但由于 Unified Buffer 大小变成了 1/24, 在这一约束下只能够运行更小的模型了)
由于 TPU V1 并没有提供关于 ResNet 中加法操作的具体实现方式, SimpleTPU 也不支持 ResNet 相关运算, 但可以支持 channel concate 操作;(虽然有多种方式实现 Residual Connection, 但均需添加额外逻辑, 似乎都会破坏原有的结构)
简化后的框图如下所示, 模块基本保持一致
3. 基于 Xilinx HLS 的实现方案
一般来说, 芯片开发过程中多采用硬件描述语言(Hardware Description Language), 譬如 Verilog HDL 或者 VHDL 进行开发和验证. 但为了提高编码的效率, 同时使得代码更为易懂, SimpleTPU 试图采用 C 语言对硬件底层进行描述; 并通过 HLS 技术将 C 代码翻译为 HDL 代码. 由于之前使用过 Xilinx HLS 工具, 因此此处依旧采用 Xilinx HLS 进行开发; 关于 Xilinx HLS 的相关信息, 可以参考高层次综合(HLS)- 简介, 以及一个简单的开发实例利用 Xilinx HLS 实现 LDPC 译码器.
虽然此处选择了 Xilinx HLS 工具, 但据我所了解, HLS 可能并不适合完成这种较为复杂的 IP 设计. 尽管 SimpleTPU 已经足够简单, 但依旧无法在一个函数中完成所有功能, 而 HLS 并不具有函数间相对复杂的描述能力, 两个模块之间往往只能是调用关系或者通过 FIFO Channel 相连. 但由于 HLS 易写, 易读, 易验证, 此处依旧选择了 HLS, 并通过一些手段规避掉了部分问题. 真实应用中, 采用 HDL 或者 HDL 结合 HLS 进行开发是更为合适的选择.
按规划之后将给出两个关键计算单元的实现, 以及控制逻辑和指令的设计方法; 最后将给出一个实际的神经网络及其仿真结果和分析.
来源: https://www.cnblogs.com/sea-wind/p/10993958.html