导读
7月6日,后摩智能联合智东西公开课策划推出的「存算一体大算力AI芯片在线研讨会」顺利完结。东南大学电子科学与工程学院副研究员司鑫、后摩智能联合创始人&芯片研发副总裁陈亮、后摩智能联合创始人&产品推出副总裁信晓旭三位主讲人参与了本次在线研讨会并进行了主题分享。
陈亮博士的演讲主题为《从硬件架构到软件工具链,存算一体大算力AI芯片的创新与实践》。他首先以FSD 和Tenstorrent 两个范例的AI处理器架构为例,先容了传统处理器架构的不敷,之后环绕存算单元Macro、层次化的AI核设计、软件工具链、软件栈、编程模型等方面对后摩智能存算一体大算力AI芯片进行了深入讲解。错过直播的朋友,可以点击“阅读原文”不雅观看回放。

本文为陈亮博士的主讲回顾:
大家好,我是后摩智能联合创始人&芯片研发副总裁陈亮,很高兴跟大家做本日的分享,我讲解的主题是《从硬件架构到软件工具链,存算一体大算力AI芯片的创新与实践》,紧张从以下三个部分展开先容:
第一部分会讲下范例的AI处理器架构,这里会环绕特斯拉FSD 和Tenstorrent 芯片做先容;
第二部分是后摩AI处理器的架构设计,紧张分为存算单元Macro 的设计考虑、层次化的AI 核的设计,还有基于存算一体的AI核设计过程中的一些工程化考虑;
第三部分是软件工具链方面,我们知道每个NPU,每个AI核都有自己的软件工具链、编译器等,我会先容下后摩智能AI处理器的软件栈和编程模型。
一、范例的AI处理器架构
第一个范例的AI处理器架构是特斯拉的FSD。特斯拉FSD 发布于2017年旁边,是一个非常简洁、高效的设计,由于它是特斯拉专用的AI处理器,紧张的组成部分包括一个96×96的MACs 阵列、一个片内32MB的SRAM、一个非常简洁的指令集,指令集里面包括了两个DMA 指令、三个点乘指令、一个scale 即标量指令、一个eltwise 指令和一个stop 指令。如果不算stop 指令,FSD 只有7条打算指令。
上图的下半部分是它的架构图和版图,从它的架构图和版图来看,它具有非常简洁的设计,带来的好处是非常高效。当时在打算效率、能效比等方面,FSD 要比英伟达的GPU 高几个量级。这也是由于运用的场景不一样,它是一个专用的AI处理器,特斯拉自己的算法运行在特定的处理器上。FSD 架构的紧张设计者是吉姆·凯勒,他是处理器领域架构设计的大牛,简洁设计也跟他个人的风格有关,他本人是一个非常崇尚极简主义的架构师。
第二个范例的AI处理器架构是吉姆·凯勒的其余一个作品,这是他近两年加盟的一个新创业公司Tenstorrent,Warmhole 是2021年他们公司一款芯片的名字。我们认为Tenstorrent Warmhole,还有常听说过的GraphCore,它们都属于近存打算的范畴。
从这两个范例的AI处理器架构可以看出,经由一段韶光的演化,传统存储和打算分离的架构已经演进成了近存打算。那近存打算是指什么呢?从上图可以看出全体芯片架构的概况。左下角大的方块,是它的芯片,芯片内部由多少个Tensix core 组成,这些小方块便是Tensix core,这些Tensix core 之间是用Network-on-chip,即NoC 连接在一起的。多少个二维的Tensix core 阵列通过NoC 连接在一起,组成一个芯片,芯片间又通过Ethernet 把芯片连接在一起,这样在芯片间再做一个2Dmesh 扩展。
从打算的角度来看,多少个Tensix core 可以同时完成一个大的Tensor 打算,一个Tensor 可以分解成多少个mini-Tensor,每个mini-Tensor 可以分布在不同的Tensix core 上运行。Tensix core 看起来也是相对简洁的设计,包括了一个CPU matrix vector Engine,再加上近存部分,便是它的SRAM,每一个Tensix core 里有一小块SRAM。除此之外,为了增强数据共享,还增加了Router 和Packet Manager 的处理单元,来担保数据在不同的小Tensix core 之间,还有芯片之间做数据的传输和互联。
从全体架构来看,可以看到它的打算是一个分布式的打算,它的存储也是一个分布式的存储。每个Tensix core 里边的SRAM 组成了片内相称大的分布式存储资源。有了打算,有了存储,其余一个关键的问题便是数据传输,它利用了Network-on-chip(NoC)的数据传输办理方案。
从设计角的度来看,也是一个非常简洁的设计,基本上可以认为把一个Tensix core 和NoC 设计好之后,不断的实行copy-paste 就可以组成一个大的芯片。因此,它有极好的可扩展性。Warmhole 是Tenstorrent 公司2021年的产品,前几年的产品还非常小,只是一个4×4的Tensix core阵列,现在已经发展成一个非常大的阵列。
从2017年到2021年,可以看到AI处理器的架构已经从传统打算和存储分离的架构,演化成了一个近存的架构。之前司鑫老师也讲过,后摩智能所做的是更进一步把存储和打算完备领悟在一起,而不但是一个近存打算。
二、后摩智能AI处理器架构
下面更详细先容下后摩的AI处理器设计。从Macro 到Cluster,我们认为这是一种分布式打算和集中式打算的折衷,是一种trade off。如上图所示,最右边这是一个Macro,便是刚才司鑫老师讲的一个存算单元,由多少个Macro 组成一个Macro Group array。Macro Group array 在Tile 里面是Tensor Engine 最主要的一个打算单元,Tile 同时又是AI Core 里一个主要的组成部分。Tile 内部除了Tensor Engine 以外,还包括了CPU、Special Function Unit(SFU)、Vector Processor(VP)、还有Shared Memory&Controller。Tile 有点类似于上面讲到的特斯拉FSD Core,包括了非常大的算力,每个Macro Group 可以供应4TOPS 算力,以是它类似一个大算力的Core。
多少个Tile 又可以组成一个AI Core。AI Core 里边除了Tile 之外,还包括了像神经网络处理或者AI打算里的前处理或后处理的处理单元、LDST 单元,还有一个Bus Node 单元,Bus Node 用来在Tile 之间或者Core 之间做数据共享和数据路由的掌握来源。
在SOC层面,多少个AI Core 又可以组成一个AI Core Clueter,通过总线将多少个AI Core 连接在一起。以是这是一个层次化的设计,多少个Macro 可以组成一个Macro Group,Macro Group 又组成了一个Tile,多少个Tile 可以组成一个AI Core,多少个AI Core 可以组成一个AI Core Clueter,那我们的算力就可以从最小的1个Macro Group,4TOPS 算力不断叠加,终极可以在1个SoC 里面实现几百TOPS 算力。
接下来我会以这4个层次为根本,详细先容后摩智能的存算一体大算力AI芯片设计。首先是Macro 存算单元,刚刚司鑫老师也有讲到,存算是有不同的路线可以选择的,包括了一些非易失落存储的工艺,还有SRAM 工艺,我们采取的方案是SRAM 工艺。SRAM 方案也有仿照和数字两种实现办法,我们把数字存算称为CIMD,仿照存算称为CIMA。
上图的表格中比拟了数字和仿照CIM 实现的一些特点。从打算事理上来看,CIMD 是数字逻辑,CIMA 采取了诸如电流分压、时域延时,或者是电荷共享、电容耦合的打算办法;在运算精度方面,CIMD 可以达到8比特乃至更高的比特,比如16比特,但CIMA 的精度会低一些,由于它受限于ADC 的精度,还有它须要把仿照旗子暗记转换成数字旗子暗记,如果想达到比较高的能效比,一样平常来说只能实现小于即是4比特的打算精度;从工艺上来看,CIMD 与前辈工艺相匹配,不会受到PVT 颠簸的影响,但CIMA 与前辈工艺无法匹配,它须要非常繁芜的校正模块,由于仿照电路会随意马虎受滋扰;在可靠性方面,数字电路抗滋扰能力是比较强的,而仿照电路随意马虎受到温度、噪声等成分的影响,以是可靠性会差一些。
从PPA 的角度来看,CIMD 的算力密度和能效比是比较高的,我们的数据在不同前辈工艺下,它的能效比可以达到几十TOPS 乃至上百TOPS。由于是纯数字电路,以是CIMD 是没有打算偏差的。仿照电路的算力密度也可以做得比较高,能效比乃至会比数字电路更高,但是它是有打算偏差的,而且其打算精度会比较低。如果在同等精度下的话,仿照的存算电路并不会比数字的存算电路有太大的上风。但在比较低精度的情形下,仿照电路的能效比会更高一些,但是它的打算偏差、精度的问题,还有校正、噪声、温度等这些影响是比较难办理的。所往后摩智能的第一代落地量产的产品因此CIMD 为根本的。
讲完最底层的CIMD,接下来看看若何把CIMD 组成一个可以用来打算的单元。我们把多少个CIMD Macro 组成一个Macro Group,数据因此数据流stream 的办法流入Macro,结果同样以流的办法流出。这里的一个好处是,在数据流入Macro 过程中,如果feature 数据有很多0,是可以节省一部分功耗的,乃至节省打算韶光。这里天然的可以支持feature 稀疏化的效果,不像有些AI处理器里讲稀疏化,比如Orin 的稀疏化是对weight 做稀疏,须要做重新的设计、演习,而我们的打算单元针对feature 可以做稀疏化的加速。
但须要考虑的问题是打算须要在Macro 上映射。上图中举了一个例子,我们的打算单元Macro 会组成一个3×3的阵列,如果是有4个3×3的阵列,可以把它类似组成一个Systolic Array,类似TPU 脉动阵列的形式,数据feature 可以从阁下的Tile 里边或者是Tile SRAM 里的Multi Bank Shared Memory 流入到Macro Group 里,结果同样可以流出Macro Group。
Macro Group 之间的结果可能是一个Partial sum,Partial sum 做一个加法之后,可以写到当前Tile的Shared Memory 里面去,也有可能通过Partial sum bus 流到临近Tile 的打算单元里面。
除了传统的AI NPU 或者AI Core,还要考虑一个利用率的问题,利用率上除了韶光利用率,还要考虑空间利用率。比如,上面的例子为什么是一个3×3的构造,由于大多数卷积神经网络中的卷积核,最常用的是3×3 kernel,以是用一个3×3的阵列是可以最高效的打算3×3的卷积。但是除了3×3以外,还有5×5,7×7或者1×1等的卷积规格,这时若何把一个5×5、7×7的一个kernel映射到3×3的阵列里,这是一个非常难的问题,须要大家仔细考虑怎么把空间上的利用率用满,这也是在用Macro 设计时须要重点关注或考虑的问题。
Macro Group 可以作为Tensor Engine 里一个最主要的单元,那Tensor Engine 和其他的一些掌握单元或者打算单元,组成一个Tile。从上图可以看出,Tile 里包括了CPU,CPU 紧张用来做掌握,通过一个指令分发单元把不同的指令分发给Tensor Engine、Vector Engine,Special Function Unit,还有一个多通道的DMA,以及一个Switch。Switch 起到了一个路由的浸染,用来在不同的Tile 之间传输数据。我们自己设计了一个数据的传输总线,通过传输和路由,还有多播的机制,让数据可以在不同的Tile,乃至在不同的AI Core 之间进行共享和传播,这样可以极大的提高带宽的利用率,减少数据和memory 之间的传输。
其余,还须要考虑的是Tile 之间的同步问题。比如两个Tile 之间须要共同完成一个打算时,那Tile0 和Tile1 之间怎么进行同步,这也须要在设计中考虑。
再上一层便是Core level。Core 是由多少个Tile 组成,上图中有4个Tile 的示例了,那Tile 之间以什么样的形式去做拓扑,这也是在做Core level 设计时须要第一个考虑的问题。上图是一个环形的拓扑,4个Tile 组成了1个环。如果是Tile 很多时,也可以是2D Mesh 这样的一个 拓扑形式。
上图的Core level 除了4个Tile 以外,包括一些AI打算里的前处理、后处理的处理单元,Load/Store 等处理单元。Core 里边会有其余的Bus Node 处理单元,Bus Node 用来将不同的Core 之间连接在一起,让不同的Core 之间可以直接传输数据。类似于英伟达最新H100 的架构,H100 架构里面也有类似的设计,它叫DSMEM,是Tensor 之间直接传输数据,而不用通过global memory,也是一个节省数据带宽,节省DDR 带宽很有效的办法。
右边的图是第一个Core 的版图,可以看到这里边包括了Macro Group、shared memory、CPU,还有Special Function Unit 等等,灰色部分便是Core level 里面的前后处理、Load/Stop等处理单元,还有主要的数据传输和互联,这是第一代Core 设计的版图,里面包括了4个Tile。
除了架构的设计以外,基于存算一体的AI 处理器设计,还有很多工程化的问题须要考虑,由于它跟传统的数字电路已经有比较大的差别,这些问题都是我们在实际工程当中碰着的问题。
第一个须要考虑的问题是SI 和PI。SI 是旗子暗记完全性,PI 是电源完全性,由于每个Macro 供应的算力很大,有4TOPS。如果有几百TOPS 算力,这几百TOPS 算力在同一韶光运转,那对电源和旗子暗记完全性是非常大的寻衅。
除此之外,Macro DFT 该怎么做?如果不做DFT,实际上是无法实现大规模量产的,我们对Macro 做了非常多DFT 干系的设计,紧张包括MBIST 和Repair,Repair 是修复由于面历年夜了往后可能会有default,我们就须要把它修复。我们做的修复电路里可以测试和修复SRAM bit cell 电路,由于存算本身是由SRAM bit cell 加上一些逻辑电路组成的。同时,还有诊断功能,与传统的SRAM 比较,须要开拓自己的lvlib,由于它已经跟传统的SRAM 行为已经有不一样的地方,没有标准的工具可以支持。
其余在CIM 打算模式下,我们设计了定制化的Macro BIST。定制化BIST 须要自己设计BIST RTL,同时可以测试Macro 内部的打算表,以是我们是把Macro 分为SRAM mode 和CIM 两种模式下进行BIST 和Repair 设计。
三、软件工具链
上面紧张讲到硬件设计方面的的考虑,如果把硬件设计类比成人的身体,那有了身体还须要有灵魂,而软件工具链便是我们的灵魂。很多人会问基于存算电路的软件工具链、软件生态是否会与传统电路的AI处理器有不一样或者有不兼容的地方?
从软件的角度来看,险些已经看不到底层的存算电路的存在,它已经在AI处理器里边,通过AI处理器的架构设计把它屏蔽掉。而从软件工具链的角度来看,并不须要过多的考虑底层的存算电路,以是软件工具链须要考虑的事情与用传统的数字电路去做NPU 没有实质上的差异。
我们的软件工具链紧张供应了两个开拓工具:一个是算子的开拓工具,一个是模型的开拓工具。对付大多数用户而言,如果没有分外的自定义算子开拓需求,可以用模型开拓工具,这里包括了一个加速的算子库,算子库有很丰富的算子,可以cover 80%-90%的算子需求,除了极个别的自定义算子不包括在算子库里边。模型开拓SDK 里面还包括了推理引擎、Graph IR、设备内存的分配,设备内存有片内memory 和片外memory 的分配器,还有一个图优化器,以及运行时的一些东西。
如果是高等的用户,可以开拓自己的自定义算子,我们也供应算子的开拓工具SDK,这里边包括了编程模型方面,基于CUDA 的扩展措辞,叫后摩 Data parallel language 或者叫hardware data parallel language(HDPL措辞),还有 Schedule Language,最底层是标准的C++。
编译器方面包括了HDPL 编译器,底层C/C++编译器。工具链方面也供应了丰富的工具链,包括了debugger、调试器、汇编和反汇编的工具、Objdump 工具,还有一个HM profeiler,用它来可以方便的调试CIM。
除此之外,我们也对推理引擎方面有支持。我们的编译器可以让用户从开拓类似GPU 的一些程序里,无缝的过渡到AI处理器的开拓过程中来,由于我们是一个类CUDA 的编程模型。上层对接的推理引擎可以是多种多样的,包括百度的PaddlePaddle、ONNX、TensorFlow、MXNet 等。这些工具产生的模型,经由量化工具会翻译成一个Relay IR 的中间表达。这个中间表达之后会通过Tensor graph 优化器来做优化,Tensor 优化器做了哪些事情呢?包括自动算子的领悟,优化带宽的瓶颈,自动流水的分配机制,可以充分利用AI Core 的硬件并行性,以及设备内存SRAM 和DDR 的自动分配,层间的调度优化等,来提高硬件利用率。
再下一层是IPU graph runtime,即图的运行时。运行时最下层会调用IPU 各种各样的资源。IPU 资源包括了AI Core,便是AI写处理器,还有一些扩展的打算资源,大多数的打算可以通过卷积在CIM 里来实现。还有一些其他的分外操作,比如最常用的pooling 等类似的操作,可以在Special Function Unit 来实现。如果有自定义的算子还可以通过HDPL 措辞的接口,在不同的打算扩展单元里实现。我们的编程模型也是针对数据并行的模型,上图有一个大略的例子,它与CUDA 编程办法是非常靠近的。
以上便是本次分享的紧张内容,感谢大家的不雅观看。