MLIR的Systolic Array映射与优化:一次硬件调试的血泪史
去年做AI加速器项目,在FPGA上跑一个8x8的脉动阵列,MLIR生成的代码死活跑不对。波形抓出来一看,数据流时序全乱套了——PE(处理单元)之间的数据传递总是差一拍,有些PE甚至吞掉了本该传递的激活值。当时盯着Vivado的时序报告,脑子里只有一个念头:MLIR的Systolic Array映射,远没有文档里写的那么“优雅”。
从“拍数”说起:脉动阵列的时序本质
脉动阵列的核心就一句话:数据像心跳一样在PE之间“泵送”。每个PE在时钟上升沿做三件事——从左边邻居拿数据、从上面邻居拿权重、做乘加运算、把结果往右边和下面传。听起来简单,但MLIR的affine dialect在生成循环嵌套时,默认的调度策略会把这种“流水线”拆成顺序执行。
我踩的第一个坑:用affine.for直接写脉动阵列的循环,MLIR默认给每个循环迭代分配独立的时钟周期。结果8x8的阵列,数据从左上角传到右下角,硬生生走了64个周期——每个PE都在等前一个PE算完。这根本不是脉动阵列,这是串行计算。
正确的做法是用scf.for配合pipeline属性,或者直接操作asyncdialect。MLIR里有个pipelinepass,能把循环体里的操作重叠起来。但注意,这个pass默认只做软件流水线,不会自动感知硬件PE的物理连接。你得手动告诉它:第i行第j列的PE,它的输入依赖第i-1行和第j-1列的输出。