尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

数据流AI芯片算子映射实战:从35%到82%利用率的关键方法

数据流AI芯片算子映射实战:从35%到82%利用率的关键方法 前阵子一个做算法移植的同事找我拿着画好的计算图问我为什么在数据流AI芯片上只跑出了35%的利用率。他说这个模型每个算子我都核对了在硬件上连得好好的像玩连连看一样一对接就过可片上buffer总是满DMA一直在抢总线PE阵列一半时间在等数据。我看了眼他的映射表问题不在算子本身而在于他理解的“数据流”是纸面上计算图里那条箭头不是硬件里真正流动的、带生命周期和位置的token。这篇文章就想聊聊数据流AI芯片上算法到硬件的正确映射方式以及那些让大家把算法和硬件玩成连连看的陷阱。适合做AI芯片编译、推理引擎、嵌入式AI加速的工程师看也适合刚接触硬件加速的算法同学避开几个血压升高的瞬间。1. 计算图连线 ≠ 数据流大多数人挂在第一道解释上1.1 你在连连看里连的“线”到底是数据还是依赖第一次接触数据流架构的人几乎都会干同一件事把网络结构图画出来算子是方块数据依赖是箭头然后照着箭头把硬件上的通道、DMA、缓冲区连起来。连完了以为万事大吉结果仿真一跑带宽爆炸、buffer溢出、PE空转性能比串行CPU还难看。这里头有个根本性误读计算图里的箭头表示的是“计算依赖”它只说明某个算子必须等前一个算子算完才能开始。而硬件上的数据流连接需要回答的是另外四个问题数据从哪个存储来、在哪个时刻到达、以什么粒度传输、在哪个PE上被消费。换句话说计算图上的线是逻辑硬件上的线是物理时序和空间路径。你把逻辑线直接当物理线去连就是在硬件上玩连连看。更麻烦的是计算图里一个tensor往往被多个算子消费比如同一个feature map既进分类头又进回归头。脚本看起来是一根线分叉硬件上却意味着需要复制指针、复制数据或者设计多播路由。很多芯片没有真正的多播能力那就只能在buffer里存两份或者反复从DRAM读两遍。连线的“分叉”在硬件上不是免费的甚至可能是最贵的部分之一。所以第一步不是连图而是给每个中间tensor做一个生命周期表谁产生它、谁消费它、它活多久、能不能全程待在片上。这个表没拉清楚之前任何映射方案都是盲人摸象。1.2 Sobel算子那个案例流程图骗了我的性能拿一个非常经典的例子讲Sobel边缘检测。教科书实现通常是三个步骤先做高斯模糊再算Sobel梯度最后算幅值和方向。在GPU上这么写没问题因为GPU的缓存层级能兜住一部分中间数据而且kernel之间的同步开销相对可控。但到了数据流AI芯片上这套流程按计算图映射会暴露一个致命伤每一级输出都是一整幅中间图像都会被写回DRAM然后下一个算子再从DRAM读出来。我当时做的是一个1024x1024灰度图三个中间结果是16bit整型。原本的期望是只要读一遍输入图、写一遍输出图实际按“三步走”映射高斯输出写一遍DRAM、Sobel输出再写一遍、中间还要分别读两遍额外增加了2GB写入和2GB读取的流量。对于一个边缘检测这种计算量很小的任务这个搬运开销直接把加速器的性能拖没了。后来改成数据流写法让高斯窗口、Sobel窗口和幅值计算卷在一个循环里在行缓冲上做滑动窗口每一行像素进入片上之后依次经过三个算子中间结果只在寄存器或者行缓冲里存几个cycle整图只访问一次DRAM。效果立竿见影算子数量没变计算的乘法数没变DMA流量少了接近4倍总延迟少了60%多。这个例子告诉你算法流程图是给人做逻辑理解用的不是给硬件做数据安排用的。在数据流芯片上能不能把整条链路的中间tensor压在片内比算子本身优化什么指令、用多少MAC更重要。1.3 先画生产者-消费者距离再谈算子映射教大家一个我自己用了很久的习惯拿到一个待映射模型不急着打开工具链先在白纸上画一张“生产者-消费者距离图”。对每个中间tensor标注生产者是哪个算子、第一个消费者是哪个算子、中间跨越了几级存储如果跨越DRAM就标一个红圈。凡是红圈数量多的地方就是性能黑洞所在地。要回答的核心问题有四个一个tensor被几个消费者复用如果多路复用片内能放几份引用还是必须复制从生产者输出第一个有效数据到消费者真正需要它中间有多少时间窗口这个窗口够不够做一次完整DMA传输如果按tile切分计算这个tensor的最小切片能不能放到目标buffer里如果不能DMA的切块策略是什么切块会不会破坏数据复用这四个问题回答完你基本已经知道这个模型在这颗芯片上能不能跑瓶颈在哪。之后再去画PE绑定、选stationary模式才叫映射。否则就是在玩连连看——方块连对了数据流却全是反的。2. 带宽账本没有花在算力上先算搬运再排流水2.1 用Roofline之前先手算一下搬运/计算比在数据流芯片上判断一个算子或者一个fused block究竟受算力限制还是受搬运限制一定要先算一个数搬运/计算比。公式很简单设某个计算块的总计算量为 CMAC次数从DRAM到片上需要搬运的数据量为 M字节搬运/计算比 r M / C。这个r越低说明算得越划算r越高说明你很多时间都在给数据搬家。拿3x3卷积举例。输入是 H x W x Ci输出是 H x W x Co。计算量大约是 H x W x Co x 9 x Ci。如果输入feature map只从DRAM读一次M ≈ H x W x Ci那么 r ≈ 1 / (9 x Co)。Co一般上百r非常小属于纯算力受限优化重点是把MAC流水排满。但如果你按算子分成两步走中间结果写回DRAM再读回来搬运量就变成中间feature map的读写也就是 2 x H x W x Co。此时 r ≈ 2 / (9 x Ci)。Ci如果也是64、128还好但如果Ci很小比如第一层卷积只有3通道输入r会非常大性能直接被数据搬运锁死。所以在排流水线之前先把每个候选fused block的r值算一遍。凡是r明显偏高的优先考虑改tile让数据在片上多活一段时间而不是优化算子内部循环。这个步骤花不了十分钟却决定了后面调度策略的大方向。2.2 选择stationary模式比选择PE摆放更影响性能数据流加速器里有个很容易被忽略的决策数据驻留模式。所谓stationary就是哪个数据“赖在”PE里不动其他数据流经它。常规有三种模式核心思想数据复用适合场景典型代表Weight Stationary权重提前加载进PE输入数据流过来权重复用小kernel、权重占主导的卷积TPU类脉动阵列Output Stationary累加器留在PE内输入数据和权重不断流入流出输出部分和复用大通道数、需要大量累加的场景常见于脉动/近存储阵列Row Stationary一行数据沿PE链传播同时做多个窗口计算输入和权重混合复用卷积/稀疏结构Eyeriss类很多人把PE摆放、片上网络拓扑调了半天性能没起色其实是stationary模式选错了。比如一维1x1卷积权重本来就小你非要用weight stationary把权重在每个PE里存一份输入数据广播来广播去复用没吃到DRAM压力也没减小。正确做法是output stationary累加器固定在PE里输入卷进来权重按通道顺序搬入部分和始终不离开PE。反之如果是深度卷积Depthwise Conv每个channel独立一个权重核权重很小不用多路复用这时row stationary可能更合适让输入沿PE阵列走一行多窗口共享同一块数据。选stationary不能凭感觉。做法是画出最内层循环的数据访问模式看哪个数据在时间上被反复使用哪个数据流式经过然后选择驻留重复使用的那一个。一个映射方案如果连最内层循环的驻留数据都想不清楚那说明还没到写连接的阶段。2.3 一次加载、一次输出、部分和不出PE的一组设计账讲一个实际tile设计账本。假设一个卷积层输入 224x224x64输出 224x224x1283x3 kernel8bit定点。我打算把输入按16行、整宽度224切成一个tile在片上做整个convrelupool的融合。输入tile大小16行 x 224宽 x 64通道 x 1字节 229376字节加上卷积halo的2行约232KB。片上buffer如果只有256KB这已经快满了不能同时再放权重。所以我把权重分成两份每份64个输出通道算出第一个半段的输出之后释放权重buffer再加载后半段权重。这样片上buffer分配是输入232KB 部分和累加器16x224x128x1字节 458KB显然放不下。于是再把空间tile切小按4行输出、全通道累加部分和只有4x224x128114KB加上输入halo的约58KB再加一半权重buffer约36KB总共200KB左右勉强装进256KB。整个过程你会发现决定tile size的不是算子的边界而是片上存储容量。你把算子挨个排开没用必须让整条链路的buffer账本合得上。这种账我在paper里叫“数据驻留计划”。没有这个计划后面连线再漂亮也只是把问题从DMA转移到了cache替换策略。3. 真正的映射技巧tile怎么切、循环怎么换、流水怎么排3.1 切块不是按算子切而是按内存层次切很多人的第一版映射是“按算子切块”卷积算一块池化算一块激活算一块。这种切法在CPU/GPU上问题不大因为L2缓存够大算子间中间结果大概率能留在cache里。但数据流芯片的片上存储通常很小几百KB到几MB按算子切会导致每个算子都要把中间结果吐到DRAM然后下一个算子再读回来搬运账瞬间爆炸。正确切法是按内存层次切把多个算子包进同一个tile让tile内部的中间数据全部留在片上。比如Conv3x3 - BatchNorm - ReLU - MaxPool这一条链按H方向切一个16行tile输入要读162行经过Conv和BNReLU后中间feature map还是16行或者14行MaxPool之后变成8行。最后只有8行输出要写DRAM中间所有数据都在片上行缓冲和部分和暂存器里流动。计算tile size的公式很朴素tile内所有中间buffer之和 片上可用buffer同时还要给double buffer留余量。注意通道方向也得考虑。如果一个tile内的通道数太多输出feature map的buffer爆掉就得按通道再切一层。经验值是把tile的片内buffer占用控制在总buffer的60%左右剩下40%留给DMA地搏双缓冲、流水线缓冲和意外峰值。我曾经因为贪心把tile开到90%结果流水线一跑起来反压信号频繁触发性能反而掉了15%。3.2 空间展开和时间复用的边界数据流加速器常用的大招是空间展开把循环维度展开到不同的PE上一个cycle算多个输出。比如64x64的PE阵列理论上可以把输出通道和空间位置各展开一部分。但空间展开不是白拿的它要求每个输出对应的中间数据或权重同时存在片上否则PE算到一半发现数据还没来只能stall。相反时间复用让一个PE分时算多个输出数据复用效率高但吞吐率低。最理想的映射是在空间展开和数据复用之间找平衡点。给你一个简单例子。输出有256通道每个通道一个3x3卷积PE阵列为64x64。如果我把256个输出通道全部空间展开需要256组权重同时驻留每组3x3x输入通道个数x64个PE远超片上buffer。所以必须时间复用。我通常做法是把输出通道切成co_tile32组一次展开32个输出通道对应64x64阵列中每个PE算一个3x3窗口的一个输出通道再通过时间循环把32组串起来。空间展开度越高每个cycle的MAC利用率越容易拉满但buffer压力线性增长。经验法则是优先展开那些数据复用机会大的维度比如输出通道和kernel窗口对数据复用差的维度比如batch用时间复用。整个映射的性能不是由最擅长展开的维度决定而是由所有维度里搬运瓶颈最严重的那一个决定。3.3 流水线里有气泡buffer里有死锁算子间流水是压榨数据流芯片性能的关键但流水线不是自动来的它需要仔细切stage。一个stage的吞吐率由它的瓶颈决定。比如前一个算子每cycle产出一个32x32的tile下一个算子一次只能消费16x16的tile下游就会周期性停滞产生气泡。我调流水线的时候一定会给每个stage画一个“生产-消费速率表”。一个stage的输出速率大于下一个stage的输入速率就得在中间加缓冲。缓冲大小至少要覆盖一次传输的突发长度加上排队时间否则就会丢数据或者反压。公式经验buffer size 带宽 x 流水线往返延迟。比如DMA峰值带宽128B/cycleDMA的启动延迟200cycle那么至少需要25600字节的缓冲才能掩盖延迟低于这个数的大概率会在高负载时露出气泡。另一个更头疼的问题是死锁。常见死锁模型是A算子要等B算子的同步token而B在等A的空buffer。两边都认为自己没有错握手动作永远完成不了。排查死锁我有两个土办法一是在片上网络每个端口放一个“buffer full超过N cycle”的计数器看哪个端口先卡满二是把同步token的命令排序打印出来确认是否存在循环等待。基本上所有死锁都能在依赖图里找到一个环解决办法就是打破这个环比如预留一个空buffer或者改变连接顺序。4. 动态控制流数据流芯片最不舒服的角落4.1 把if变成select把for变成固定trip count数据流芯片不像CPU有PC寄存器它没有“跳转”的概念指令执行是由数据token到达触发的。所以算法里那些自然写的if、while、动态break到了数据流上没有一个能原样映射。最简单可靠的策略是if变select。两个分支都算条件作为选择信号选其中一个输出。代价是两条分支的算力、功耗都会真实消耗所以必须评估分支概率如果一条分支很少触发那用select就很不划算。另一个策略是把分支放到host CPU上做只把可预测的密集计算放上加速器。for循环必须把trip count变成编译期常数。动态循环次数比如while循环直到误差收敛在静态数据流里几乎映射不动。工程上常见的做法是设置一个最大迭代次数循环跑到这个次数就强制结束或者把循环体“展开”成一个固定的数据流图。如果循环次数真的太动态那只能在host端分多次下发每完成一次迭代回主机判断是否继续这个往返延迟一般会让人放弃这条路线。所以做算法部署时提前把模型里的动态shape、动态循环剪枝干净比在芯片上调调度策略重要得多。很多模型压缩、剪枝算法最后没有上板不是算力不够而是动态控制流破坏了数据流硬件的静态调度假设。4.2 静态数据流与动态数据流选型前先认清约束数据流架构本身也有流派。静态数据流Static Dataflow的操作数token在编译期就确定了路由硬件开销低调度确定性强但表达能力弱循环次数、分支都必须静态化。动态数据流Dynamic Dataflow会给token加tag运行时靠匹配tag来把数据送到对应算子表达能力更强能处理动态依赖但tag匹配逻辑、存储管理开销非常大。类型触发方式硬件开销表达能力适合场景静态数据流编译期确定token无tag路由低循环次数静态、无动态分支深度学习推理、固定形状计算动态数据流运行期用tag匹配高支持动态循环和分支通用数据流计算、图计算市面上绝大多数面向AI推理的数据流芯片走的是静态或半静态路线。所谓半静态就是允许运行期切换tile大小、动态batch但不允许在计算图内部出现任意分支。因此写kernel的人必须清楚一个循环里如果出现依赖运行数据的提前exit硬件轻则多耗周期重则直接不支持。老老实实把动态内容拆出来放到预处理或者后处理里比所有花活都管用。4.3 反压、死锁与ready/valid握手三个真实故障调数据流硬件最常碰到的三个故障全和握手有关。第一个是反压风暴。当片上网络某个端口buffer满向上游发送反压上游算子不能继续发数据只能stall。如果多个算子共享同一个DMA通道一个算子拥堵会把整条链路都堵死。典型现象是PE利用率不高但busy信号全亮。第二个是死锁。两个算子互相等对方的buffer形成循环等待握手信号永远完成不了。我印象最深的一次是调试一个带两个分支的splitconcat结构A分支占满了到B分支的中间bufferB分支的输出又要回到A分支作为输入结果两边谁都不让片上一个token都走不动。查了整整两天最后发现是A分支产生数据太快把公共池占满了而B分支需要的数据又被A占着。解决方式很简单给两个分支各分配独立buffer池禁止共享。第三个是valid信号丢失。跨时钟域或者跨模块握手时ready信号更新得比valid慢导致发送方丢了一个cycle的有效数据。这个问题非常隐蔽最终结果只是偶发算错几个像素但排除起来极其痛苦。我的经验是仿真时打开数据面上的byte-level trace拿一个已知结果做差分比对一旦出现不匹配立刻看对应数据的valid/ready波形九成能定位到握手竞争。5. 一次检测头映射复盘从35%利用率到82%的过程5.1 最初“连连看”映射三个典型症状前阵子我接手一个检测模型在数据流AI芯片上的部署具体是YOLO系列的一个检测头结构不算复杂几个3x3卷积、1x1卷积、split成分类和回归两个分支各自走sigmoid和线性输出最后concat回来。第一版映射是典型的“连连看”式照着计算图把每个算子分别映射到不同的计算单元中间feature map全部经过全局buffer暂存。上板之后瓶颈一览无遗DMA带宽占用97%PE利用率只有35%。原因是每个算子都要从DRAM读输入、写输出带宽被中间tensor的读写吃掉了。编译好的调度经常因为buffer溢出而stall。因为split和concat产生了大量数据副本全局buffer被无谓地占用。分支结构引发的bank冲突频繁两个分支同时写同一块buffer的不同地址导致访问被串行化。后来我把每层产生、消费的数据量拉出来看发现真正必要的DRAM流量只有实际流量的四分之一。也就是说四分之三的搬运全是计算图连线方式造成的多余开销。这印证了前面说的沿着逻辑箭头连连出来的是带宽黑洞。5.2 按数据流重写的步骤和调度伪代码重写的时候我没有动网络结构只改了映射策略。核心思想一句话把整个检测头当成一个大的嵌套循环按输入空间位置切tiletile内部所有算子融合起来让中间结果只活在片上。步骤很清晰把feature map按H和W切成16x16的tile卷积的halo多读2行2列。把通道方向分成co_tile32份保证输出累加器能塞进片上buffer。把所有pointwise操作1x1卷积、sigmoid直接融合到前一个算子的输出循环里不产生独立feature map。Split不再复制数据改成两个分支共享同一份读指针Concat不再拷贝靠DMA的scatter地址直接写最终输出。调度伪代码大概长这样注意这不是某个厂商的API只是表达思路# 伪代码示意数据流调度结构 for y0 in range(0, H, tile_h): # 空间tile外层 for x0 in range(0, W, tile_w): im_buf dma_load(input, y0-1:y0tile_h1, x0-1:x0tile_w1) # fused convreluconv for co0 in range(0, Co, co_tile): acc conv3x3(im_buf, w[co0]) # 输出驻留累加器 acc conv1x1(im_buf, w2[co0]) # 同一tile内复用 acc activate(acc) # split分支共享同一份acc不拷贝 cls_out sigmoid(acc[:co_cls]) reg_out acc[co_cls:co_reg] dma_store_scatter(out, cls_out, offset_cls) dma_store_scatter(out, reg_out, offset_reg)这个调度的关键点是co_tile决定了PE阵列能做多少空间展开tile_h和tile_w决定了halo读入和输出缓冲的大小。我最终选的tile16x16co_tile32依赖片上buffer刚好卡在80%以内同时给double buffer留了余量。重写之后的结果DRAM流量降到了原来的四分之一PE利用率从35%提到了82%端到端延迟缩短了2.3倍。最明显的变化是片上网络几乎不再反压DMA带宽占用降到53%空出来的带宽可以给其他业务用。5.3 值得保留的工具链经验仿真器、mapper和片上计数器复盘这次优化工具链用得最值的地方有三个。第一个是cycle approximate simulator。在改映射之前先用仿真器跑一遍能看到每个算子的cycle分布在哪些阶段是等DMA、等buffer还是等PE。很多问题在仿真阶段就能暴露不用反复上板试。强烈建议先建立“DRAM流量/计算量”的r值表格再跑仿真否则连改方向都找不到。第二个是自动化mapper。数据流芯片厂商都会提供一个搜索工具把循环顺序、tile大小、PE绑定、buffer分配作为搜索变量。我一开始还想手写映射表后来发现搜索空间大得离谱。用类似NSGA-II或粒子群这类多目标优化算法去搜能在几小时内找到一组pass的候选方案再人工挑一挑比手调快得多。注意不要只盯着算力目标优化目标里必须有DMA带宽占用和buffer占用这两个约束否则mapper会给你一个算力很满但带宽爆炸的方案。第三个是片上profiler计数器。真机调试时硬件自带的各种counter特别重要。我最常看三个计数器含义正常范围参考PE stall cycle由于数据未就绪导致的PE空转周期占比低于20%DMA wait cycleDMA请求被buffer占满而等待的时间占比越低越好NoC buffer full count片上网络各端口buffer满的次数每百万cycle低于100次如果PE stall超过20%先查DMA wait和NoC buffer full哪个高就先优化哪个。数据流芯片的调试就是一个不断看计数器、改调度、再验证的循环。一旦养成这个习惯所谓“硬件上的连连看陷阱”基本就不太会踩了。最后再分享一个小技巧每次跑完一组性能数据顺手把当时的映射参数存下来包括tile大小、co_tile、buffer分配、循环顺序以及对应的DMA流量和PE利用率。下次遇到类似模型直接拿这些参数当初始值去搜索能省掉非常多前期摸索的时间。这个习惯帮我处理过的模型越多后面新模型的收敛就越快。数据流芯片的映射没有银弹但有方法论核心就一句话让数据在正确的时间出现在正确的地点而不是让算法逻辑图上的方块用线连得好看。
返回列表