AI时代进击的CPU们
科技
科技 > 互联网 > 正文

毫无疑问,GPU和AI加速器才是AI时代算力的最大提供者。训练AI模型需要GPU的超大算力,但一块GPU越来越存不下今天的模型,互联多块GPU成了刚需,各种高端的网卡,DPU也诞生了。推理AI模型也需要非常高的算力,同时对存储带宽的要求也越来越高。即便是端侧,这几年移动GPU和NPU的进步,AI的部署也大多从CPU侧迁移到更大算力,更高能效的模块中。一切都在向着异构 + 专有结构的方向发展。

CPU这个曾经的通用计算王者现在面临很尴尬的处境:往往沦为大型AI负载的调度器和数据搬运工具,还要承担加速器处理不了的长尾需求,被吐槽拖慢了整体速度。然而各家CPU厂商总要生存,他们近几年也陆陆续续在自家CPU上引入AI的加速功能,以期能在这个言必称AI的市场里获得机会。算力密度上的天生弱势,导致CPU厂商更多在推理侧发力。

AI运算的二八法则

还好AI的计算规律也符合二八法则:从深度神经网络开始流行,到现在的大语言模型,最大的计算模块基本就是卷积,矩阵乘法一类的计算密集型操作,它们在绝大多数AI模型中大约能占据八成的计算规模,剩下几十到上千种算子占据二成左右。

深度学习中的卷积运算,可以转换成矩阵乘法的模式进行计算,如下面两篇文章展示:M1 Max初相遇,快快乐乐写卷积,“远超”理论浮点峰值。

矩阵乘法是由乘法和加法这两种基本运算所组成,且成对儿出现。我们假设相乘的两个矩阵分别是 𝑚×𝑘 维的𝐴(记作 𝐴𝑚×𝑘),以及 𝑘×𝑛 维的𝐵 (记作 𝐵𝑘×𝑛),结果矩阵就是 𝑚×𝑛 维矩阵 𝐶(记作 𝐶𝑚×𝑛)。那么矩阵乘法的操作就是:𝐶𝑚×𝑛+=𝐴𝑚×𝑘×𝐵𝑘×𝑛。这个计算过程包含了 𝑚×𝑛×𝑘 次乘法,和同样数量的加法。所以高性能计算和AI里统计矩阵乘法的运算次数就是乘法和加法的总次数:2×𝑚×𝑛×𝑘 ,用这个数字除以矩阵乘法的计算时间,就可以得到我们通常表示算力大小的FLOPS。如果这个值高于 109 ,就可以表示成GFLOPS;如果高于 1012 ,就可以表示成TFLOPS... 这个指标就是计算矩阵乘法或者卷积等AI核心操作的绝对性能指标。

人们习惯于把复杂的问题拆成简单的步骤组合,或者庞大的任务拆成很多粒度更小的任务组合。矩阵乘法或者卷积运算,往往就是拆成多种粒度的分块的矩阵乘法,最终分解到硬件的计算单元上去执行并行计算。以nvidia为代表的GPU架构,最早在17年加入了TensorCore这种矩阵乘法指令,相比早先的多核单元内SIMT运算,直接提升了一个数量级的性能;多数AI加速器也是以某种形式的Matrix结构作为主要的计算单元(例如脉动阵列),支持卷积和矩阵乘法类的运算,也能取得很优秀的性能。唯独CPU,因为功能定位通用多元,芯片设计复杂以及尽量不破坏传统体系结构的原因,早年一直通过对SIMD向量指令的扩展来支持这些算子,直到最近两年才开始加入类似Matrix结构的专用DSA。

本文也从这个角度简单回顾下各家CPU厂商的AI能力构建历史。

x86:成也生态,败也生态

即便2023年,nvidia的营收也才突破600亿美元,intel年年600亿美元,市值仅有nvidia的零头,AMD的一半... 事实上,intel这10年来也一直在不停地加大对AI的投入:Xeon Phi,Nevana Sys,Movidius,Altera FPGA,Habana以及自研架构的GPU。但这些似乎都不那么成功,往往难产,最终被砍掉;或者推迟,等真正发布的时候已经不见优势。

所以这十年来,intel的战略就是通过Xeon服务器CPU去和GPU打,在x86里面塞入了大量的重型负载。比如AVX512指令集,就是为对抗GPU的高性能计算能力诞生的。这个指令集十分的庞大,除了将前面所有的指令集升级到512位以外,还加入了大量新的特性,尤其是AI。

下面通过四个部分介绍下intel对于x86的AI能力构建的历史过程:

基本的单精度浮点向量运算

CPU最早是通过单精度的向量指令的方式去拼凑出整个矩阵乘法的操作。无论是训练还是推理,32位单精度浮点操作都足够用。

早年的SSE指令支持128位的单精度向量乘法和向量加法,也就是一条指令只能执行4个float类型的乘法,或者加法,指令算力密度很低;AVX指令的引入,增加到256位,算是提升了一倍;因为矩阵乘法,卷积,FFT这些操作,乘法和加法总是成对出现,所以AVX2增加了乘加指令,相当于一条指令执行了 𝑉𝑐+=𝑉𝑎×𝑉𝑏 ,绑定了成对的乘法和加法,将单指令算力密度再提升一倍;最后就是AVX512F指令的引入,将乘加指令延长至512位。从SSE到AVX512F,单一指令的计算密度提升到了8倍。指令操作如下图所示:

图片

从SSE到AVX512F四种单精度向量乘加指令

按照关于sgemm_hsw的一点解释说明,x86使用AVX/FMA指令实现寄存器分块矩阵乘法的示意图如下:

图片

矩阵乘法kernel的FMA指令寄存器分块示意

这里实际使用的是向量-标量乘加,拼凑出一个矩阵乘法外积的运算。x86并没有直接从寄存器读取标量,和另一个向量做乘加的指令,所以这里需要额外的broadcast类指令帮忙把内存(实际在L1 cache)中的标量数据广播给一个向量寄存器的所有通道。

低比特定点运算

单精度指令对于AI推理来讲,一般过于冗余,很多算法模型都可以将参数精度压缩到8bit定点。Intel也在AVX512系列指令集中,单独加入了针对int8和int16点乘的AVX512-VNNI指令集。

DP4A和DP2A类指令最早是nvidia在Pascal架构GPU中引入,每个CUDA Core有一个针对int8的dp4a单元。存储在32位寄存器中的int8向量,与另一个int8向量做点积运算(即DP),中间结果和结果32位寄存器中的数据累加到一起,存回结果寄存器。DP2A与之类似,只是把8bit数据换成16位的数据,如int16,fp16或者bf16。后两类的结果寄存器也从32位整型,换成fp32浮点。

Intel的AVX512-VNNI指令就是DP4A/DP2A类的指令,如下图所示:

图片

Intel VNNI指令集中使用的dp4a/dp2a类指令

AVX512-VNNI支持uint8乘以int8,以及int16乘以int16两种规格。在512bit寄存器上分成16组dp4a或dp2a单元。这样相比AVX-512f的单精度乘加指令,乘加计算密度又有了4倍吞吐。因为AVX512指令集资源消耗过重,只在云端上放了完整的支持,桌面和移动端从12代酷睿开始,都取消了AVX512的支持。为了应对这种尴尬的局面,Intel引入了AVX-VNNI指令,只支持256bit向量的VNNI。同时,AVX512-VNNI中的256bit指令,与AVX-VNNI的功能虽然完全相同,但使用不一样的指令编码,需要加上{vex}前缀。12代以后的桌面和移动端intel处理器,开始支持这个指令集。同时12代开始的酷睿又引入了大小核的设置,gracemont的小核心也开始加入AVX-VNNI的支持,只是int8吞吐量并没有单精乘加的4倍,只有2倍。然后intel发现u8s8和s16s16的方式似乎不够用,又找补了AVX-VNNI-INT8和AVX-VNNI-INT16两种新指令集,将u8u8,s8s8,s8u8,u16s16,s16u16,u16u16这六种组合方式也加入了进来,但还不知道会用在哪款CPU上。

从SSE到AVX512-VNNI,单一指令的乘加计算吞吐量增大到32倍。

半精度指令

Intel引入VNNI的时候,更多的是将CPU定位为推理算力。为了和nvidia更全面地竞争,intel也想在CPU上做训练的支持,因此又引入了两种半精度浮点的支持:bfloat16和fp16。

同样也是在AVX512指令集系列里做的扩充,分别叫做AVX512-bf16和AVX512-fp16。Intel最早在第三代Xeon Icelake SP服务器上引入的AVX512-bf16,在第四代Xeon Sapphire Rapids引入了AVX512-fp16。有趣的是,前者提供的是dp2a类型的指令,且除了和fp32的互转,没有提供更多的运算支持;后者提供的是标准的向量乘加指令,且有类似fp32的各类计算操作支持。两者的吞吐量,理论上都是fp32的两倍,补充了VNNI指令只支持定点的缺陷。

至此,Intel在现有SIMD寄存器体系下对AI的支持已经完成。我们总结一下,这个体系下intel主要依靠增加单条指令的计算密度来提高吞吐,大致分成两个方向:一个是增加向量的长度,从128bit,到512bit;另一个是压缩参数位数,从32bit浮点,到8bit定点。这两个方向都是线性增长,且都存在扩展上限,所以intel并没有满足于此,抛开现有SIMD体系,走上了在CPU核内增加特定AI DSA架构的路线。

专用的DSA——Intel AMX

AMX是Intel针对矩阵乘法设计的全新DSA,与SIMD指令的最大区别是,它抛开了现有SIMD的寄存器结构,引入了全新的TILE二维寄存器文件。下图简单描绘了AMX的结构,每个处理器核心都拥有这样一个结构:

图片

Intel AMX基本结构

灰色框中标注的就是它独有的TILE寄存器体系,当前的版本(Sapphire Rapids),TILE寄存器的数量是8,支持int8和bf16两种数据类型。TILE的结构,其实就是一个二维 16×16 的方阵。方阵的元素是个32bit的数据,可以存4个int8,2个bf16或者1个int32/fp32。AMX的核心计算,就是将两个TILE中存储的矩阵相乘,累加到第三个矩阵中。下图展示了AMX_INT8的TMUL计算的矩阵乘法规模:

图片

AMX_INT8中TMUL操作示意图

每个TILE寄存器都是 16×16×4 字节的容量,在int8的case里,相当于计算了 𝐶16×16+=𝐴16×64×𝐵64×16 这样规模的矩阵乘,其中每四个int8通过dp4a操作累加到一个int32。bf16的操作类似,只是每两个bf16通过dp2a操作累加到fp32,计算吞吐比int8小一倍。

这个过程的具体实现,是一个标准的脉动阵列。bf16的计算指令是TDPBF16PS,int8的计算指令有四个:TDPBSSD,TDPBSUD,TDPBUSD,TDPBUUD,就是int8和uint8的四种组合。每条计算指令启动一个矩阵乘法的脉动数据流。脉动阵列的时空流水较复杂,这里不详细描述,有非常多的文章介绍这个结构的原理。除了核心的TMUL计算指令,AMX还提供了load和store数据的基础指令,对TILE清零的指令,以及配置计算规模和数据类型的TILECFG类指令。下面的表格里给出各个指令的延迟吞吐:

图片

AMX指令延迟和吞吐

其中TDP类核心计算指令,吞吐周期是16,这意味着每隔16个周期可以发起一条新的数据无关的脉动流水计算指令;延迟周期是52,表明这条脉动指令52周期后会计算完成,得到所有的C矩阵 16×16 个结果。根据我们以前讲到测试指令吞吐的方式,需要至少 ⌈5216⌉=4 条无依赖的指令排在一起,才能掩盖单个脉动流水的延迟。下面的代码可以测试AMX计算峰值:

tilezero %tmm0

tilezero %tmm1

tilezero %tmm2

tilezero %tmm3

tilezero %tmm4

tilezero %tmm5

.amx.int8.mm.s32s8s8:

tdpbssd %tmm4, %tmm5, %tmm0

tdpbssd %tmm4, %tmm5, %tmm1

tdpbssd %tmm4, %tmm5, %tmm2

tdpbssd %tmm4, %tmm5, %tmm3

sub $0x1, %rdi

jne .amx.int8.mm.s32s8s8

ret

具体用AMX优化矩阵乘法或者卷积还要考虑很多因素,比如cache上数据排布的转换。我们用上面代码的写法,在一台双路64核(每路32核)的Xeon Sapphire Rapids服务器上测试了AMX的峰值,结果如下:

图片

Xeon Sapphire Rapids 64核服务器AMX峰值测试

图左是单核的测试结果,图右是64核的测试结果。单核可以做到7+Tops,这在CPU的历史上是独孤求败的存在;64核可以做到接近400Tops的峰值,可以媲美很多云端大算力AI加速器的水平,且这还是最低端的一款SPR,intel还有单路56核的高端版本。

AMX这种专用的DSA,做到了SIMD里最高吞吐指令VNNI的8倍。而且从intel的文档和材料里看,这个结构的扩展性很强:从结构图里提到的accelerator,当前版本还只有一个,后续版本可以扩展出更多;TILE寄存器的容量和个数也都可以扩展;还可以增加更多的低精度数据类型。扩展维度非常多。至此,intel想利用CPU去抗衡GPU的布局正式完成,Xeon CPU正式长成一颗大算力的怪物,intel甚至为它配上了HBM的版本。

关于生态

我们都知道intel在推出64位处理器时候犯下的错误:2001年Intel推出安腾架构处理器(IA-64),对32位的x86(IA-32)采取了完全不兼容的策略,使用了VLIW的指令集体系结构。在后来的一系列发展过程中,固然有VLIW架构在通用计算本身的性能问题,但是不兼容先前的x86和所有的x86软件,导致用户软件迁移的成本巨大,安腾的市场一直做不起来,最终早早退出历史舞台。同时AMD在03年抢先推出了x86的64位扩展,命名为amd64,并兼容运行32位的x86程序,市场大范围跟进。这导致不久以后intel不得不开始兼容amd64,甚至到现在为止,很多操作系统发行版对于x64的名字依然叫做amd64。

x86生态在PC和通用服务器市场如此重要,导致intel后续在移动互联网,高性能计算和AI这些新战场的初期,都过于迷信x86生态,错过了快速迭代,占领用户心智的机会。移动互联网要求高度的可定制化,非常低的成本,arm有天然的先发适配优势,且最初由苹果主导,x86没有掀起任何水花;高性能计算和AI对单点性能,互联扩展性有极强的要求,nvidia GPU顺应了时代,快速进步,并推出完整的解决方案,致使最初intel基于Xeon Phi的解决方案很快败下阵来。intel在长达五六年的时间内没有拿出任何有竞争力的高性能计算加速方案,将这个市场拱手让给nvidia。

Intel现在x86处理器中塞入了过多重量级的计算负载,固然是对抗nvidia的无奈之举。然而本身也使自身的枷锁越来越沉重,甚至影响了自己的server市场:AMD和arm都有各自的专用高性能和AI方案,不需要CPU过度负担,且通用单核性能早已追平intel,在server端轻轻松松就可以堆到128核甚至更多,市场进展势如破竹。intel却很难堆更多的核,甚至不得不启用自家小核(E核)方案去做堆核方案,竞争力相当堪忧。Intel需要尽快明确并坚持自己在AI领域的专用加速方案,并与自家x86生态配合,为x86减负增效,回归本来的定位,否则连目前利润最高的服务器市场都会被对手蚕食殆尽。

arm:RISC还是CISC?

arm是Advanced RISC Machine的简称,一直以来标榜自己是RISC处理器。我们来看看arm对RISC的定义What is RISC? – Arm®:

A Reduced Instruction Set Computer is a type of microprocessor architecture that utilizes a small, highly-optimized set of instructions rather than the highly-specialized set of instructions typically found in other architectures. RISC is an alternative to the Complex Instruction Set Computing (CISC) architecture and is often considered the most efficient CPU architecture technology available today.

然而当我们翻开arm的架构指令集文档Arm Architecture Reference Manual,并和intel的同类型文档做个对比:

图片

arm和intel指令集文档页数对比

emmm... RISC的R改成Regular好了...

arm AI指令集的SIMD演进

回归正题,本节简要介绍一下arm在SIMD指令体系下对AI的支持。

arm早年的SIMD指令使用neon/asimd这一套体系,这是一套仅有128bit长度的向量指令集。neon的指令设计相比x86的向量指令大同小异,但更高效紧凑:neon除了支持标准的向量-向量乘法,加法和乘加以外,还支持选择向量中的某条lane(arm术语,即向量中的某个位置的标量值),和另一条向量的所有lane做乘加,即标量-向量乘加。这样的指令设计,相比x86的broadcast + 向量-向量乘加的指令组合,节省了指令数量和读取数据的次数。

arm从armv8.0开始支持fp32和fp64的向量乘法,加法和乘加运算指令(asimd/fp);从armv8.2开始增加了fp16的支持(asimdhp),以及int8的dp4a指令支持(asimddp)。然后arm在neon现有的寄存器体系之上,又增加了针对bf16和int8的一种矩阵乘法指令模式mmla,其中针对bf16的指令集就叫bf16,针对int8的指令集叫做i8mm,这两个指令集在armv8.6中引入。如下图所示:

图片

mmla类指令运算示意

这个系列的指令设计非常巧妙,输入向量 𝑉𝑎 分成两个64位向量 𝑎1 和 𝑎2 , 𝑉𝑏 也分成两个64位向量 𝑏1 和 𝑏2 ,根据数据类型,这四个向量可以是8个int8,或者4个bf16。两两一组做向量点积,得到四个32位结果(int8点积结果int32,bf16点积结果fp32)。累加到结果寄存器 𝑉𝑐 的四个32位值里。由于输入向量的每个值都复用两次,所以这条指令的吞吐,比对应数据类型的dp4a或者浮点乘加有一倍的吞吐提升。mmla和标量-向量计算类指令,都是x86目前所没有的设计。

arm在服务器端通过增加向量长度扩展指令吞吐,为此引入了SVE/SVE2这样一套新的SIMD指令,这两种指令集并不兼容neon,但是功能是类似的,最大的变化是支持更长的向量(从128bit到2048bit)。在移动端则是通过增大SIMD单元发射数量来扩展指令吞吐,而向量长度保持128bit不变。比如Cortex-A57/A72有一个SIMD单元;从A76开始增加到两个;Cortex-X1系列开始增加到4个;最近发布的最新Cortex-X925更是一举扩展到6个。同时配合越来越低比特精度数据,以及乘加密度更高的指令设计,arm在AI领域的进步神速。

SME——外积矩阵乘DSA

苹果自家的apple silicon芯片在2020年正式发布,在CPU上引入了名为的AMX矩阵乘法DSA,与intel的AMX同名,但是完全不同的设计。arm也在2021年引入了新的指令集SME,它包含了SVE/SVE2,同时为矩阵乘法引入了类似苹果AMX的DSA。

SME里提供的矩阵乘法指令基于外积计算(outer product),如下图所示:

图片

向量外积计算矩阵乘法

向量外积相当于做一次 𝑘=1 的矩阵乘法:TILE_C是结果矩阵C的一个分块( 𝑀×𝑁 维);VEC_A是其在矩阵A中对应的水平M行条带分块中的一个列向量;VEC_B是其在矩阵B中对应的垂直N列条带分块中的一个行向量;这两个向量分别读到向量寄存器中,进行一次外积计算,乘累加到TILE_C,就完成了一次外积运算;然后沿着A和B当前的条带分块的K方向不断读取新的向量,形成计算流水线,最终计算完整个TILE_C的值,写回内存。

SME为这个操作提供了全新的Z Array寄存器体系和ZA Tiles分块方式,并通过新指令提供对A和B矩阵条带分块的数据重排方案。实际上我们观察到,不管是x86还是arm的SIMD方案,也都是通过构造矩阵外积的寄存器分块方式进行计算的。只是它们提供的指令计算粒度,只能做标量和短向量的乘加,通过多条指令的组合,构造一个完整的外积矩阵乘法。而SME提供的外积指令,只需要一条指令即可做一个完整且大规模的外积乘法,这就是DSA对指令密度的提升。

向量外积有哪些特色呢:首先一次外积计算相当于 𝑀×𝑁 个乘加操作阵列,这些乘加操作之间理论上无数据依赖,可同时并行流水(实际上电路的布线可能会限制这个特点,M和N不能大规模扩展);其次,向量外积是个非常简单的顺序Streaming读取重排后的矩阵A和B,同时与计算操作进行重叠;最后,向量外积的读取和计算的比例最低,是 𝑀+𝑁/𝑀×𝑁 ,对于读取带宽的需求有很大降低。

SME的外积指令相比mmla指令,虽然都是矩阵乘法指令,但前者对读取数据的利用率更高,VEC_A的每个数据复用了N次,VEC_B的每个数据复用了M次;而mmla指令对读取数据只复用了两次。前者是建立在使用更庞大的寄存器文件扩大了矩阵乘分块实现的,后者的好处是在现有SIMD寄存器体系下就可以轻松提升一倍的吞吐。

arm CPU的AI能力是移动端生态建设的基石

最近AI手机,AI PC刮起了一阵旋风,让各种大算力的移动GPU和NPU出尽了风头,然而arm的CPU仍然占有重要的地位。在手机端,目前有苹果,高通的adreno,arm自己的mali,以及imagination等四种主要的移动GPU平台;PC上有nvidia,AMD和intel等三种GPU平台。虽然有OpenCL/Vulkan等计算API可以写一套代码适配大部分GPU,但不同的GPU架构要优化到位,还有很多特化的优化要去做。NPU的生态就更破碎了,不同NPU的开发接口,编程能力完全不同,很难用一种统一的方式支持所有主流的移动NPU。这对于手机和PC上的App开发者简直是噩梦。arm架构的CPU虽然也有代际差异,但基本反映在支持的AI标准指令集的多寡。只要把本文提到的几种指令集都支持好,对不同芯片使用运行时识别指令集的方式支持本地最新的AI指令,就能把最通用的AI能力集成到自家App里。有了这基础通用的AI能力,可以再适配不同的GPU或者NPU方案,丰俭由人。

其它CPU架构:家家有本难念的经

其它种类的CPU在高性能的场景,如手机移动端,桌面PC或者服务器端暂时都没有很大的使用量,更多还是用在IoT,MCU等场景,这些场景对AI也有很大的需求。

risc-v:繁荣还是混乱?

risc-v是目前业界除x86和arm以外非常火爆的指令集架构,它是由开放社区制定的开源指令集标准,任何组织都可以免费使用其标准构筑自己的实现。不过risc-v是一种很新的架构,许多标准,例如vector extension这样的SIMD指令标准,是最近一两年才正式发布。在这前后有很多公司发布过自己的标准,造成了一些混乱。官方正式版的vector extension简称rvv1.0,是一种很现代化的SIMD指令标准。它提供了针对fp64/fp32/fp16的基本乘加指令,对于AI的支持有个很好的基础。有些risc-v厂商在这个标准之外设计了自己的AI指令集,如进迭时空的K1芯片(为志坚兄做个广告~),通过测试可以看出AI能力有很大的提高:

图片

SpacemiT K1芯片带AI加速指令的四个核峰值测试

这在很多IoT场景甚至不用额外添加NPU了,而且与CPU本身的标量和SIMD指令配合,可以做到长尾算子的高效执行,某些情况下比CPU + NPU的组合还要好。

国产指令集——loongarch

loongarch是龙芯2020年左右以mips为蓝本发布的全新自研指令集标准,成为国产自主指令集的独苗。loongarch目前支持LSX和LASX两种SIMD指令,分别是128bit和256bit向量长度。遗憾的是,由于种种原因,龙芯暂未公布loongarch SIMD指令的标准文档,我只hack出了基础的fp32和fp64的标准向量乘加指令,还不知道有没有针对AI的指令。龙芯的生态最近几年面临一些官司,同时由于一些历史原因,ABI还分成旧世界和新世界两种,可谓是一出道就陷入了泥潭。23年3A6000的发布,在性能上已经让龙芯进入可用的程度,希望龙芯可以早日统一规格标准,快速适配国产生态环境。

写在最后

本文是个综述性的科普文章,把市面上现有流行的CPU架构对AI的支持简要过了一遍,技术上蜻蜓点水,也在关注一些历史逻辑和未来方向,观点或多有偏颇。希望这些内容对于初学者,或者关注相关领域的朋友有所帮助。

亲爱的凤凰网用户:

您当前使用的浏览器版本过低,导致网站不能正常访问,建议升级浏览器

第三方浏览器推荐:

谷歌(Chrome)浏览器 下载

360安全浏览器 下载