专栏首页海风TPU中的指令并行和数据并行

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

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

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

  • 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=0;i<N;i++)
    y[i] += a*x[i];

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

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

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

for(int i=0;i<M;i++){
    for(int j=0;j<N;j++){
        for(int k=0;k<K;k++){
            for(int c=0;c<C;c++){
                for(int kw=0;kw<KW;kw++){
                    for(int kh=0;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=0;kw<KW;kw++){
    for(int kh=0;kh<KH;kh++){
        for(int k=0;k<K;k++){
            for(int i=0;i<M;i++){
                for(int j=0;j<N;j++){
                    for(int c=0;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.

本文参与腾讯云自媒体分享计划,欢迎正在阅读的你也加入,一起分享。

我来说两句

0 条评论
登录 后参与评论

相关文章

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

    人工神经网络中的大量乘加计算(譬如三维卷积计算)大多都可以归纳成为矩阵计算。而之前有的各类处理器,在其硬件底层完成的是一个(或多个)标量/向量计算,这些...

    sea-wind
  • Simple TPU的设计和性能评估

    在TPU中的脉动阵列及其实现中介绍了矩阵/卷积计算中的主要计算单元——乘加阵列(上图4),完成了该部分的硬件代码并进行了简单的验证;在 神经网络中的归一...

    sea-wind
  • 神经网络加速器应用实例:图像分类

    在Simple TPU的设计和性能评估中,一个神经网络加速器的硬件雏形已经搭建完成了;在https://github.com/cea-wind/Simp...

    sea-wind
  • 谷歌云TPU服务正式全面开放:「AlphaGo背后的芯片」进入商用化

    机器之心
  • 【深度】基于论文,对谷歌 TPU 的最全分析和专业评价

    【新智元导读】本文以 Google 最新公开的 TPU 论文《在数据中心中对张量处理器进行性能分析》的译本为基础,对该论文及 TPU 进行了评价。 源起 2...

    新智元
  • 【重磅】谷歌公布TPU细节之后,AI业界怎么看?

    AI科技评论按:我们对于谷歌的TPU并不陌生,正是它支撑了AlphaGo强大快速的运算力,但谷歌一直未曾披露其细节,使得TPU一直保有神秘感。 美国当地时间4月...

    AI科技评论
  • 如何薅羊毛 | PyTorch终于能用上谷歌云TPU,推理性能提升4倍

    现在PyTorch官方已经在Github上给出示例代码,教你如何免费使用谷歌云TPU训练模型,然后在Colab中进行推理。

    昱良
  • 谷歌开放TPU应对英伟达GPU挑战

    谷歌宣布将以“有限数量”向谷歌云客户开放张量处理器(Tensor Processing Unit,简称TPU)服务,按时收费,每小时成本6.50美元。 谷歌宣布...

    人工智能快报
  • Kaggle竞赛硬件如何选择?不差钱、追求速度,那就上TPU吧

    图 1:在 Kaggle Notebook 中可以免费使用 CPU、GPU 和 TPU。

    机器之心
  • 谷歌硬件工程师揭秘,TPU为何会比CPU、GPU快30倍?

    在谷歌发布TPU一年后,这款机器学习定制芯片的神秘面纱终于被揭开了。 昨日,谷歌资深硬件工程师Norman Jouppi刊文表示,谷歌的专用机器学习芯片TPU处...

    AI研习社

扫码关注云+社区

领取腾讯云代金券