动手写一个简单版的谷歌TPU-指令集

时间:2022-09-05 11:41:46

深度学习飞速发展过程中,人们发现原有的处理器无法满足神经网络这种特定的大量计算,大量的开始针对这一应用进行专用芯片的设计。谷歌的张量处理单元(Tensor Processing Unit,后文简称TPU)是完成较早,具有代表性的一类设计,基于脉动阵列设计的矩阵计算加速单元,可以很好的加速神经网络的计算。本系列文章将利用公开的TPU V1相关资料,对其进行一定的简化、推测和修改,来实际编写一个简单版本的谷歌TPU,以更确切的了解TPU的优势和局限性。

动手写一个简单版的谷歌TPU系列目录

谷歌TPU概述和简化

TPU中的脉动阵列及其实现

神经网络中的归一化和池化的硬件实现

TPU中的指令并行和数据并行

Simple TPU的设计和性能评估

SimpleTPU实例:图像分类

拓展

TPU的边界(规划中)

重新审视深度神经网络中的并行(规划中)

TPU V1定义了一套自己的指令集,虽然在介绍处理器时,往往会先谈指令集架构,但此处却把它放到了最后,这主要基于两个原因;其一在于个人的对处理器不太了解,这也是主要原因,其二在于公开资料中并没有TPU指令集的细节和TPU微架构的描述。从数据流和计算单元出发对TPU进行分析固然容易很多,但如果想理解TPU的设计思想,依旧需要回到其架构设计上进行分析。这一部分内容有些超出了我现有的能力,不当之处还请多多指正。

本文主要探讨从架构设计上看,TPU时如何做高性能和高效能的设计。高性能的多来自于并行,因此本文分别讨论了指令并行和数据并行的设计方法。由于论文中并未描述TPU指令集的具体设计,除特别说明外,本文关于TPU指令集的探讨均为推测;另外,SimpleTPU的指令设计并不系统/完整,此处仅阐明设计中的几种基本思想。

1. TPU的指令集

TPU的指令集采用CISC设计,共计有十多条指令,主要的五条指令包括

  1. Read_Host_Memory 将数据从CPU的内存中读取到TPU的Unified Buffer上
  2. Read_Weights 将weight从内存中读取到TPU的 Weight FIFO 上.
  3. MatrixMultiply/Convolve 执行卷积或矩阵乘法操作.
  4. Activate 执行人工神经网络中的非线性操作和Pooling操作(如有)
  5. Write_Host_Memory 将结果从Unified Buffer写回CPU内存.

从给出的五条指令可以看出,TPU的指令集设计和通用处理器有很大的不同。指令需要显示指定数据在内存和片上buffer之间搬移的过程。而执行指令(MatrixMultiply)直接指定了Buffer的地址,指令上并不能看到一系列通用寄存器。这是因为TPU本质上还是一个专用的处理芯片,其高性能和高效能都是建立在失去一定灵活性的前提下的。为了获得更高的性能,可以采用一系列的常规方法进行设计,包括

  • 指令并行,即一次性处理更多指令,让所有执行单元高效运行
  • 数据并行,即一次性处理多组数据,提高性能

后文会针对这两点做进一步描述,并简单讨论TPU设计中的更多其他的优化方法和方向。

2. 指令并行

2.1 Simple TPU中的流水线

为了提高吞吐率和时钟频率,处理器通常使用流水线设计,经典的五级流水线设计一般如下所示

 

clk0

clk1

clk2

clk3

clk4

clk5

clk6

clk7

instruction 0

IF

ID

EX

MEM

WB

     

instruction 1

 

IF

ID

EX

MEM

WB

   

instruction 2

   

IF

ID

EX

MEM

WB

 

instruction 3

     

IF

ID

EX

MEM

WB

其中,IF指取指(insturction fetch),ID指指令译码(instruction decode),EX指执行(Execute),MEM指内存读写(Memory Access),WB指写回寄存器(Write back)。采用流水线设计可以提高性能,如果不采用流水线设计,那么instruction1需要在clk5才能开始进行IF,严重影响其性能;如果在同一周期完成IF/ID/EX/MEM/WB的功能,由于逻辑极其复杂,会严重影响工作频率。

TPU论文中介绍其采用四级流水线设计,Simple TPU中采用了两级流水线,来完成控制过程。

 

clk0

clk1

clk2

clk3

clk4

clk5

clk6

clk7

instruction 0

IF&ID

EX

           

instruction 1

 

IF&ID

EX

         

instruction 2

   

IF&ID

EX

       

instruction 3

     

IF&ID

EX

     

也认为Simple TPU内部有四级流水线,这是因为在实际执行过程中,包括了读取寄存器,执行和写回三个部分,这三个部分是流水设计的。

2.2 超长指令字(VLIW)

如前文所述,Simple TPU中有两个基本的计算单元——矩阵乘法阵列和池化计算单元。除此之外,还有一些没有显式描述的执行单元,譬如载入和存储。在这一前提下,即使TPU的指令流水线做得再好,每条指令占有的周期数也不可能小于1。如果其他执行单元的执行周期数很小,此时总会有一些执行单元处于闲置状态,处理器的瓶颈会出现在指令上。为了解决这一问题,很直接的想法时每个周期发射多条指令(另一个方法时让执行单元的执行时间变长,Simple TPU通过向量体系结构设计也有这一处理)。

由于TPU的专用性,以及计算过程中不存在跳转和控制的原因,采用VLIW设计多发射处理器似乎是一个很适合的方式。在这一设计下,指令发射结构时固定的,而且所有的冒险可以由编译器事先检测并处理,这很大程度可以降低硬件实现的复杂度。在Simple TPU中借鉴了VLIW的思想进行设计,如下所示(示意图)

动手写一个简单版的谷歌TPU-指令集

其中各个字段具体描述如下

  • model mask 指定了当前指令运行的模块
  • load weight 指定了从内存将weight读取到SRAM的指令
  • load act. & mac & store result 指定了将操作数(act.)读取到寄存器,乘加阵列计算以及将结果写回到存储器的过程
  • set weight 指定了将操作数(weight)读取到计算阵列寄存器的过程
  • load act. & pooling& store result field指定了将操作数(act.)读取到寄存器,完成pooling和归一化计算以及将结果写回到存储器的过程

VLIW的设计放弃了很多的灵活性和兼容性,同时将很多工作放到软件完成,但依旧适合在TPU这样的偏专用的处理器中使用。Simple TPU中没有对数据冲突、依赖进行任何处理,软件需要事先完成分析并进行规避。在这一设计下一条指令可以调度最多四个模块同时工作,效率得到了提升。

3. 卷积计算中的数据并行

3.1 单指令多数据(SIMD)

单指令多数据,故名思意是指在一条指令控制多组数据的计算。显然,TPU core的设计中采用了这样一种数据并行的方式——一条instruction控制了256*256个乘加计算单元(MatirxMultiply/Convolve)。根据指令流和数据流之间的对应关系,可以将处理器分为以下几个类别

  • SISD,单指令流单数据流,顺序执行指令,处理数据,可以应用指令并行方法
  • SIMD,单指令流多数据流,同一指令启动多组数据运算,可以用于开发数据级并行
  • MISD,多指令流单数据流,暂无商业实现
  • MIMD,多指令流多数据流,每个处理器用各种的指令对各自的数据进行操作,可以用在任务级并行上,也可用于数据级并行,比SIMD更灵活

由于TPU应用在规则的矩阵/卷积计算中,在单个处理器内部的设计上,SIMD是数据并行的最优选择。SIMD有多种实现方式,根据给出的描述(MatirxMultiply/Convolve指令接受B*256输入,输出B*256个结果),TPU中应该采用了类似向量体系结构的设计方法。

3.2 向量体系结构

基本单元-矩阵乘法阵列所述,计算单元完成矩阵乘法计算,即向量计算。以《计算机体系结构 : 量化研究方法》给出的例子为例,如需计算

for(int i=;i<N;i++)
y[i] += a*x[i];

以MIPS为例,对于一般的标量处理器和向量处理器而言,执行的指令序列如下所示

动手写一个简单版的谷歌TPU-指令集

最大的不同在于向量处理器大幅的减小了指令的数目,缩减了指令带宽。同时,简单的MIPS指令中可能存在互锁的情况,会降低性能,而这一现象在向量处理器中则不存在。

对于卷积神经网络中的卷积操作而言,计算可以表示为(已忽略bias)

for(int i=;i<M;i++){
for(int j=;j<N;j++){
for(int k=;k<K;k++){
for(int c=;c<C;c++){
for(int kw=;kw<KW;kw++){
for(int kh=;kh<KH;kh++){
result(i,j,k) += feature(i+kw,j+kh,c)*w(k,kw,kh,c);
}
}
}
}
}
}

由于KW和KH可能为1(即卷积核的宽度和高度),而weight在计算过程中认为是固定在计算阵列内部的,因此调整循环顺序后有

for(int kw=;kw<KW;kw++){
for(int kh=;kh<KH;kh++){
for(int k=;k<K;k++){
for(int i=;i<M;i++){
for(int j=;j<N;j++){
for(int c=;c<C;c++){
result(i,j,k) += feature(i+kw,j+kh,c)*w(k,kw,kh,c);
}
}
}
}
}
}

其中第一二层循环通过指令进行控制,第三层循环在计算阵列中以256并行度进行计算,指令调度;第4-6层循环按向量处理器的设计思路进行设计,通过一条指令完成三层循环的计算。为了完成循环的计算,需要设置三个向量长度寄存器,另外,由于向量在SRAM中的地址并不连续,还需要设定三个不同的步幅寄存器。参考 基本单元-矩阵乘法阵列的代码,具体为

    short ubuf_raddr_step1;
short ubuf_raddr_step2;
short ubuf_raddr_step3;
short ubuf_raddr_end1;
short ubuf_raddr_end2;
short ubuf_raddr_end3

采用这样的设计,SimpleTPU中一条指令可以完成大量数据的计算,提高了数据并行度。这些数据会并行的进入到计算阵列中完成计算(可以认为是多条车道)。由于SimpleTPU中数据的读取延时是固定的(指从SRAM),因此向量化的设计较一般处理器还更为简单。

根据谷歌论文中的描述,TPU中有repeat fileld,但MatrixMultiply/Convolve指令长度有限,因此可能只有一组或两组向量长度寄存器和步幅寄存器,但设计思路应该类似。

4. 其他

从谷歌论文中的参数来看,TPU具有极高单位功耗下性能。这一部分来自于其内核设计,正如之前的文章中所描述的

  • 采用了INT8数据类型进行计算
  • 采用了脉动阵列优化计算
  • 没有采用缓存,没有分支跳转,预测和数据冲突处理(编译器完成)

而从本文的内容可以看出,TPU还采用了简单的指令集设计+SIMD+向量体系结构+VLIW来进一步优化单位功耗下性能;除此之外,在V2/V3中google更进一步,还利用多核和多处理器设计进一步提高了性能。

参考

Jouppi, Norman P. , et al. "In-Datacenter Performance Analysis of a Tensor Processing Unit." the 44th Annual International Symposium IEEE Computer Society, 2017.

JohnL.Hennessy, and DavidA.Patterson. Computer architecture : a quantitative approach = 计算机体系结构 : 量化研究方法/ 5th ed.

动手写一个简单版的谷歌TPU-指令集的更多相关文章

  1. 动手写一个简单版的谷歌TPU

    谷歌TPU是一个设计良好的矩阵计算加速单元,可以很好的加速神经网络的计算.本系列文章将利用公开的TPU V1(后简称TPU)相关资料,对其进行一定的简化.推测和修改,来实际编写一个简单版本的谷歌TPU ...

  2. 动手写一个简单版的谷歌TPU-矩阵乘法和卷积

    谷歌TPU是一个设计良好的矩阵计算加速单元,可以很好的加速神经网络的计算.本系列文章将利用公开的TPU V1相关资料,对其进行一定的简化.推测和修改,来实际编写一个简单版本的谷歌TPU.计划实现到行为 ...

  3. 动手写一个简单的Web框架(模板渲染)

    动手写一个简单的Web框架(模板渲染) 在百度上搜索jinja2,显示的大部分内容都是jinja2的渲染语法,这个不是Web框架需要做的事,最终,居然在Werkzeug的官方文档里找到模板渲染的代码. ...

  4. 动手写一个简单的Web框架(Werkzeug路由问题)

    动手写一个简单的Web框架(Werkzeug路由问题) 继承上一篇博客,实现了HelloWorld,但是这并不是一个Web框架,只是自己手写的一个程序,别人是无法通过自己定义路由和返回文本,来使用的, ...

  5. 动手写一个简单的Web框架(HelloWorld的实现)

    动手写一个简单的Web框架(HelloWorld的实现) 关于python的wsgi问题可以看这篇博客 我就不具体阐述了,简单来说,wsgi标准需要我们提供一个可以被调用的python程序,可以实函数 ...

  6. 自己动手写一个简单的MVC框架(第一版)

    一.MVC概念回顾 路由(Route).控制器(Controller).行为(Action).模型(Model).视图(View) 用一句简单地话来描述以上关键点: 路由(Route)就相当于一个公司 ...

  7. 自己动手写一个简单的(IIS)小型服务器

    因为第一次在博客园发表随笔,不太会用,这个笔记是我之前在印象笔记中写好的,然后直接copy过来,有兴趣自己做一个IIS服务器的小伙伴们可以参照下面的流程做一次,也可以叫我要源代码,不过要做完,我觉得花 ...

  8. 手写一个简单版的SpringMVC

    一 写在前面 这是自己实现一个简单的具有SpringMVC功能的小Demo,主要实现效果是; 自己定义的实现效果是通过浏览器地址传一个name参数,打印“my name is”+name参数.不使用S ...

  9. 自己动手写一个简单的MVC框架(第二版)

    一.ASP.NET MVC核心机制回顾 在ASP.NET MVC中,最核心的当属“路由系统”,而路由系统的核心则源于一个强大的System.Web.Routing.dll组件. 在这个System.W ...

随机推荐

  1. 对于angularJS的一点思考

    已经找好工作近两周了,入职基本上还算顺利,自己两年来的挑灯夜战也算是有了收获,于是这两周基本上是按部就班的工作,没有学习什么新技术.在上个公司的时候,同事在项目中使用angularJs,之前他也没有接 ...

  2. GitHub使用教程

    一直以来都想使用Git来管理自己平时积累的小代码,就是除了工作之外的代码了.有时候自己搞个小代码,在公司写了,就要通过U盘或者网盘等等一系列工具进行Copy,然后回家才能继续在原来的基础上作业.Cop ...

  3. WordPress菜单函数wp&lowbar;nav&lowbar;menu&lpar;&rpar;详细介绍

    导航菜单函数wp_nav_menu()进行详细的说明. 1.wp_nav_menu()函数介绍: worpdress发展到3.0以后增加了一个自定义菜单函数wp_nav_menu(),使得wordpr ...

  4. &lbrack;C&num;&rsqb;Main&lpar;String&lbrack;&rsqb; args&rpar;参数输入问题

    Main函数是程序的入口点,它是入口点,那它的参数,又是怎样来的呢?首先写个简单的测试程序看看args到底是什么? class Program { static void Main(string[] ...

  5. ExtJS4&period;2 Ext&period;grid&period;panel Store更改后刷新表格

    //////////////////////// // Prepare store //////////////////////// // prepare fields and columns var ...

  6. wdcp php5&period;3添加pdo&lowbar;mysql模块

    先查看探针: pdo没有支持mysql.导致了PHpwind以及thinkphp框架的一些运用了pdo进行mysql操作的程序无法运行. php5.3默认是封装了pdo_mysq的.那么就没必要单独下 ...

  7. Struts 之 通配符 路径匹配 常量用法 配置默认值

    Struts 框架学习 Action的开发的几种方式 方式1 : 继承ActionSupport     如果使用Struts校验功能,必须继承此类 方式2 : 实现Action接口 方式3 :不继承 ...

  8. Robot Framework学习笔记(八)------ride标签使用

    一.edit标签使用 1.导入库 点击 Edit 标签页右侧的"Library"按钮,来添加库.在添加库之前,首先库已经在 Python 下进行了安装.如,添加"Sele ...

  9. python学习之路网络编程篇(第二篇)

    新课程知识的引入:python作用域 #python中无块级别作用域 if 1 == 1 : name = 'alex' print(name) for i in range(10): name = ...

  10. html&plus;css&plus;jq随记

    随便写个博客吧,记录一下自己的历程,今天忽然用自己好久不用的jq还做项目,并且从零开始搭建,让自己慌乱不已啊!遇到了如下问题 1.ios端点击闪屏的问题,解决办法如下 在body上添加 -webkit ...