ARTICLE DETAIL

资讯详情

深耕编程入门与网站建设的一线实战洞察。

从零设计AI加速器:脉动阵列与矩阵乘法硬件实现

从零设计AI加速器:脉动阵列与矩阵乘法硬件实现 1. 为什么我想从零设计一款AI加速器1.1 一个念头当通用处理器不再“通用”我最早动这个念头是在跑一个中等规模的卷积网络推理时。当时用一块普通的桌面级CPU跑单张图片的推理时间在几百毫秒量级风扇狂转功耗也不好看。换到GPU上确实快了很多但随之而来的是一整套生态依赖驱动版本、运行时库、显存管理稍有不慎就报“设备能力不匹配”或者“设备已移除”之类的错误。那段时间我反复在想一个问题如果我只是想跑一个固定的、结构已知的神经网络为什么非要背上一整套通用计算的包袱这个疑问其实就是AI加速器存在的根本理由。CPU是“什么都能干”的通用选手它的设计目标是分支预测、乱序执行、大缓存为的是应对不可预测的通用程序。GPU是“大量重复计算”的并行选手几千个核心一起上但它依然要兼顾图形渲染、科学计算等一大堆任务。而NPU神经网络处理单元这类AI加速器目标非常纯粹把神经网络里最核心的运算——矩阵乘加——用最省电、最省面积的方式做到极致。所以“从零设计一款自己的AI加速器”对我而言不是一个炫技的项目而是一次把“矩阵运算”这件事从算法一路打通到硬件的完整实践。它适合有一定数字逻辑基础、懂一点神经网络、又想搞清楚“芯片到底怎么算矩阵”的人。哪怕你最后不流片只在仿真器里跑通收获也远比调通一个现成框架大得多。1.2 先想清楚加速器到底加速的是什么在动手之前必须把这个问题回答清楚否则设计出来的东西就是空中楼阁。神经网络推理抛开激活函数、归一化这些“边角料”计算量的绝对大头是矩阵乘法和卷积。而卷积在数学上可以展开成矩阵乘法im2col全连接层本身就是矩阵乘法注意力机制里的QKV计算也是矩阵乘法。可以说谁把矩阵乘法做快做省谁就抓住了AI加速的命脉。一个矩阵乘法 C A × B假设A是M×KB是K×N那么输出C是M×N总共需要M×N×K次乘加运算。这个立方级的复杂度就是加速器要啃的硬骨头。CPU靠SIMD指令一次算几个数GPU靠成千上万个线程分摊而专用加速器靠的是二维甚至三维的乘加阵列让几百上千个乘法器同时工作一个时钟周期就完成一大片计算。理解了这一点整个设计的主线就清晰了围绕矩阵乘法设计数据怎么进来、怎么在阵列里流动、结果怎么出去。这也是我后面所有章节要展开的核心。2. 整体架构设计从算法到硬件的映射2.1 顶层思路以矩阵乘法单元为核心我的整体架构可以用一句话概括一个以脉动阵列为核心的矩阵乘法引擎配上灵活的片上存储和简单的控制逻辑。为什么选脉动阵列Systolic Array而不是别的结构这是有讲究的。脉动阵列最早由H.T. Kung提出它的精髓在于数据像心跳一样有节奏地在处理单元PE之间流动每个PE只负责一次乘加然后把结果传给下一个。这样做的好处是数据复用率极高——一个从左边流入的数据会被阵列里一整行的PE依次使用一个从上边流入的数据会被一整列的PE依次使用。相比每个PE都去访问存储器脉动阵列把访存次数降到了最低而访存恰恰是能耗大户。我做过一个粗略的估算在同样的工艺下一次片上存储访问的能耗大约是一次乘加运算的几十倍甚至上百倍。所以“让数据多流动、少访存”是加速器省电的第一原则。脉动阵列天然符合这个原则这就是我选它的核心理由。当然脉动阵列也有缺点它适合规则的、尺寸较大的矩阵运算对于稀疏的、不规则的运算效率会下降。但考虑到主流神经网络的结构相对规整这个取舍是值得的。2.2 关键参数怎么定阵列尺寸、数据位宽、频率架构定了接下来是一堆参数要拍板。这些参数不是拍脑袋定的每一个背后都有计算和权衡。阵列尺寸。我最终选了16×16的PE阵列也就是256个乘法累加单元。为什么是16而不是8或者328×8只有64个PE算力偏弱32×32有1024个PE面积和布线压力陡增对于个人项目来说仿真都跑不动。16×16是一个甜点算力够看面积可控而且16正好是2的幂地址计算和分块都很方便。实测在仿真里16×16阵列处理一个256×256的矩阵乘法分块调度起来逻辑清晰不会太复杂。数据位宽。这是最纠结的地方。神经网络对精度其实没那么敏感推理时8位定点INT8往往就够了精度损失很小。但为了通用性和验证方便我最终选择了16位定点作为主数据通路同时保留8位模式。16位的好处是动态范围够大不容易溢出调试时心里有底8位的好处是面积和带宽减半。实际做的时候我把乘法器设计成可配置的先跑通16位再切8位对比。工作频率。这个在FPGA原型或者仿真阶段其实不是最关键的但我还是定了一个目标100MHz。理由很简单100MHz在大多数FPGA和ASIC工艺下都容易达到时序收敛压力小而且算力已经够用。算一下256个PE每个PE每周期一次乘加100MHz下就是256×100M 25.6 GMAC/s每秒256亿次乘加。这个数字放在几年前是相当可观的放到现在虽然不算顶尖但对于理解原理、验证设计来说绰绰有余。2.3 存储层次为什么片上缓存比什么都重要前面说了访存是能耗大户所以存储层次的设计直接决定加速器的效率。我的方案是三级外部存储DDR→ 全局缓存Global Buffer→ 本地寄存器PE内。外部存储容量大但慢且费电只用来存放原始数据和最终结果。全局缓存是片上的一块SRAM用来暂存当前正在计算的数据块。PE内的寄存器最小最快存放当前周期要用的操作数。这里有个关键设计数据分块Tiling。因为矩阵可能很大片上缓存装不下整个矩阵所以要把大矩阵切成小块一块一块地送进阵列计算。分块的大小要和阵列尺寸匹配——16×16的阵列一次吃16×16的数据块最合适。分块策略直接影响到数据复用的次数进而影响访存量。我后面会专门讲分块怎么算。3. 核心模块拆解与实操要点3.1 处理单元PE加速器的心脏PE是整个加速器最小的计算单元它的结构决定了整个阵列的能力。一个最基础的PE需要做三件事接收数据、执行乘加、传递数据。我的PE设计包含一个乘法器、一个累加器、几个输入寄存器、以及用于阵列互联的旁路通路。数据从左边进来记为A从上边进来记为B两者相乘后累加到本地累加器同时A向右传、B向下传。这样经过若干个周期后整个阵列就完成了一次矩阵乘法的“波前”计算。这里有个实操细节值得说累加器的位宽一定要留够。假设输入是16位那么乘积是32位如果K维度是256累加结果可能达到32840位。我一开始只给了32位累加器结果在跑大K值的时候溢出了输出全是错的。后来加到48位才稳。这个坑很典型位宽估算一定要按最坏情况来宁可多几位不能少一位。另一个细节是流水线。乘法器本身有延迟如果每个周期都要出结果就得在PE里插入流水线寄存器。我最初没加流水线频率上不去加了之后虽然单个PE的延迟增加了但吞吐率上来了整体反而更快。这就是典型的“用延迟换吞吐”。3.2 数据流控制让数据有节奏地流动脉动阵列最迷人的地方就是数据流。但要让数据“有节奏”地流动控制逻辑必须精确到每个周期。我的做法是A矩阵的数据从左边界逐行注入B矩阵的数据从上边界逐列注入两者在阵列中相遇并相乘。为了让不同位置的数据在同一时刻对齐注入的时间要错开——第i行的A数据要延迟i个周期注入第j列的B数据要延迟j个周期注入。这个“错拍”是脉动阵列正确工作的关键错一个周期结果就全乱。我踩过的坑是一开始没做错拍所有数据同时注入结果只有对角线上的PE算对了其他全错。后来画了个时序图把每个周期每个PE该收到什么数据都标出来才把错拍逻辑理清楚。强烈建议动手前先画时序图比直接写代码高效得多。3.3 片上缓存与地址生成全局缓存我用的是单口SRAM模型读写不能同时。为了不让读写冲突拖慢整体我采用了双缓冲Double Buffering一块缓存用于当前计算另一块用于预取下一块数据。这样计算和加载可以重叠阵列不用等数据。地址生成模块负责把逻辑上的矩阵坐标转换成SRAM的物理地址。这里有个技巧把矩阵按行优先存储地址 行号 × 行宽 列号。分块的时候只需要计算每个块的起始地址和步长就能连续读取。我一开始用复杂的多维地址计算逻辑又乱又容易错后来统一成线性地址加偏移清爽多了。提示地址生成模块一定要单独验证。我见过太多项目计算单元没问题结果全错在地址算错上。写个简单的测试把每个块的地址打印出来和手工计算对比能省下大量调试时间。4. 完整实操流程从仿真到跑通一个矩阵乘法4.1 开发环境与工具选型个人做硬件设计工具选型很关键。我的选择是Verilog做RTL设计Verilator做仿真Python做数据生成和结果比对。为什么用Verilator而不是商业仿真器因为它快、免费、跨平台而且和C/Python的联合仿真很方便。对于个人项目仿真速度直接决定你能迭代多少次。Verilator编译出来的模型跑起来比很多商业工具还快这点在调试大矩阵时优势明显。Python这边我用NumPy生成随机矩阵把数据按定点格式量化后写入文件仿真跑完再把结果读回来和NumPy的结果比对。整个流程自动化改一次设计跑一次回归几分钟就能知道对错。4.2 分块策略与参数计算假设要计算C A × BA是M×KB是K×N阵列是16×16。分块的核心思想是把M、N、K三个维度都按16切分。具体来说把A按行切成M/16块按列切成K/16块把B按行切成K/16块按列切成N/16块。计算时外层循环遍历M和N的块内层循环遍历K的块每个内层循环做一次16×16×16的阵列计算累加到结果块上。这里有个参数要算清楚每个块的计算周期数。一个16×16的阵列完成一次16×16×16的乘加需要大约3×16 48个周期数据注入、计算、排空各占一部分。如果K方向有K/16个块那么一个输出块需要(K/16)×48个周期。这个估算帮我判断整体性能也帮我决定缓存要开多大。我实测下来一个256×256×256的矩阵乘法用16×16阵列总共大约需要(256/16)×(256/16)×(256/16)×48 ≈ 16×16×16×48 ≈ 196608个周期。在100MHz下大约2毫秒。对比CPU单核可能要几十毫秒加速比还是很明显的。4.3 仿真验证怎么确认结果是对的验证是硬件设计的生命线。我的验证分三层第一层是单元测试单独测PE。给一个PE喂固定的A和B看输出是不是A×B。这一步确保最基本的乘法累加没问题。第二层是小阵列测试比如4×4阵列跑4×4矩阵。规模小出错了容易定位。我会把每个PE的输出都dump出来和手工计算对比。第三层是全阵列随机测试16×16阵列跑随机矩阵和NumPy结果比对。这一步跑通基本就说明设计正确了。我踩过的一个坑是定点量化的舍入方式。NumPy算的是浮点我硬件里是定点两者比对时会有微小误差。一开始我用严格相等判断结果全是“错误”。后来改成允许一定误差范围比如相对误差小于1%才通过。这个经验很重要定点验证要比对误差不能比绝对值。5. 常见问题与排查技巧实录5.1 结果全错或者部分错怎么定位这是最常见的问题。我的排查顺序是先看数据注入对不对再看阵列互联对不对最后看累加和输出对不对。数据注入错通常是错拍逻辑写错了。解决办法是把注入时序打印出来逐周期核对。阵列互联错通常是PE之间的连线接反了比如A传给了下面而不是右边。这个要靠仔细检查例化时的端口连接。累加错多半是位宽不够或者清零时机不对。我整理了一个速查表遇到问题按这个顺序查基本都能定位现象可能原因排查方法结果全为0数据没注入或使能没拉高检查注入逻辑和使能信号只有对角线对错拍逻辑缺失检查注入延迟结果偏小累加器位宽不够溢出扩大累加器位宽结果随机错时序竞争或复位问题检查复位和时钟域误差略大定点量化舍入调整舍入策略放宽比对5.2 仿真跑得太慢怎么办大矩阵仿真慢是常态。我的优化手段有三个一是减小测试规模先用小矩阵验证逻辑最后再跑大矩阵二是用Verilator的优化选项开-O3能快不少三是减少dump波形波形文件巨大且拖慢仿真只在必要时开。还有一个技巧把长时间不变的部分做成C模型只对改动部分做RTL仿真。比如存储器的行为可以用C模拟只仿真计算阵列。这样能大幅提速。5.3 定点精度的取舍经验定点设计最头疼的就是精度。我的经验是输入8位、累加32位以上基本能覆盖大多数推理场景。如果发现精度不够优先加累加器位宽而不是加输入位宽因为输入位宽翻倍会让乘法器面积翻四倍而累加器位宽增加代价小得多。另外量化时的舍入方式也有讲究。直接截断误差大四舍五入好一些随机舍入最好但硬件复杂。个人项目用四舍五入就够了。6. 后续可以怎么扩展这套设计跑通之后其实还有很多可以玩的方向。比如加入激活函数单元把ReLU、Sigmoid做进流水线这样就能端到端跑一个完整的网络层。再比如支持稀疏矩阵跳过零元素能进一步提升效率。还可以做多阵列级联把几个16×16阵列拼起来算力线性增长。我自己下一步想尝试的是把卷积直接映射到阵列上不做im2col展开而是设计专门的卷积数据流。这样能省掉展开带来的存储开销对卷积网络更友好。这个方向有点挑战但想清楚了会很有意思。最后分享一个我在整个项目里体会最深的心得硬件设计里时序图比代码重要验证比设计重要位宽估算比什么都重要。把这三件事做好从零设计一个AI加速器并没有想象中那么遥不可及。
返回列表