4. 数据流 (Dataflow) 优化

做AI芯片优化这么多年,我越来越觉得数据流是个被低估的环节。很多人一上来就调算子、改量化,结果发现瓶颈根本不在计算,而在数据搬运。说白了,数据流决定了数据怎么在存储和计算单元之间流动,这直接影响了芯片的能效和吞吐。

今天咱们就聊聊三种经典的数据流策略:权值固定、输入固定、输出固定,以及行固定和脉动阵列的映射策略。嗯,这些概念听起来有点抽象,但我会结合项目经验讲清楚。

为什么数据流这么重要?

先问个问题:一个MAC(乘累加)操作,真正花在计算上的能量是多少?我告诉你,大概只有0.1pJ左右。但从DRAM里取一次数据,要花200倍以上的能量。你想想看,如果数据流设计不好,大部分能量都浪费在搬运上了。

我在做某款边缘AI芯片时,一开始没太在意数据流,结果芯片跑ResNet-50时功耗直接超标。后来重新设计了数据流,功耗降了40%多。从那以后,我每次做架构设计,第一件事就是画数据流图。

三大经典数据流策略

这三种策略,说白了就是决定哪个数据在计算阵列里"待着不动",哪个数据"到处跑"。

1. 权值固定(Weight Stationary)

权值固定的思路很简单:把权重存在PE(处理单元)里不动,输入数据流过来,和权重做乘累加,结果再传出去。

核心特点:

  • 权重只在开始时加载一次,后续不再移动
  • 输入数据在PE阵列中广播或逐行传递
  • 部分和结果在PE间传递累加

我个人习惯在卷积层用权值固定。为什么?因为卷积的权重复用率很高,一个权重要和多个输入做计算。权值固定正好利用了这一点,减少了权重的搬运次数。

我的经验: 在跑MobileNet这类深度可分离卷积时,权值固定效果特别好。因为深度卷积的权重通道数少,每个权重被复用的次数多,权值固定的优势就体现出来了。

2. 输入固定(Input Stationary)

输入固定正好反过来:输入数据待在PE里不动,权重流过来做计算。

核心特点:

  • 输入数据固定在PE中
  • 权重在PE阵列中广播或传递
  • 部分和结果在PE间传递累加

什么时候用输入固定?全连接层。我记得有一次做BERT的推理优化,全连接层的权重矩阵特别大,如果每次都要从外部搬权重,带宽根本扛不住。用输入固定后,把输入激活值固定在PE里,权重流进来,效果立竿见影。

注意: 输入固定对PE的本地存储要求比较高。如果输入特征图很大,PE里的buffer可能装不下。我曾经踩过这个坑,后来不得不把输入切分成tile来处理。

3. 输出固定(Output Stationary)

输出固定是把部分和结果留在PE里,等所有累加完成后一次性写回。

核心特点:

  • 部分和结果固定在PE中累加
  • 权重和输入数据都流经PE
  • 只有最终结果才写回全局存储

输出固定有个好处:减少了部分和的搬运次数。你想想看,如果每次乘累加都把部分和写回再读出来,那得多浪费带宽。输出固定让部分和在PE内部累加,只有最终结果才写出去。

我在做YOLOv5的检测头优化时,发现输出固定特别适合。因为检测头的输出通道数少,每个输出像素需要累加的乘累加次数多,输出固定正好能减少部分和的搬运。

三种策略对比

策略 数据复用方式 适合场景 主要瓶颈
权值固定 权重复用 卷积层(权重复用率高) PE本地存储大小
输入固定 输入复用 全连接层(输入复用率高) 输入buffer容量
输出固定 部分和累加 输出通道少的层 累加器精度

行固定(Row Stationary)与脉动阵列映射

行固定是Eyeriss架构里提出的一种数据流策略,它把卷积操作按行分解,让数据在PE阵列中以行为单位流动。

说白了,行固定是权值固定和输入固定的一个折中方案。它既考虑了权重的复用,也考虑了输入的复用,还兼顾了部分和的累加。

行固定的核心思想

行固定把卷积核的一行和输入特征图的一行映射到一个PE上。每个PE负责计算一个输出像素的部分和,然后通过PE间的传递完成累加。

行固定的优势:

  • 同时利用了权重复用和输入复用
  • 减少了数据搬运次数
  • 对PE阵列的规模要求相对灵活

脉动阵列的映射策略

脉动阵列是另一种经典的数据流架构。数据在PE阵列中像心跳一样有节奏地流动,每个PE只和相邻的PE通信。

我做过一个脉动阵列的加速器,用的是权值固定的映射方式。具体来说:

// 脉动阵列权值固定映射伪代码
for (int i = 0; i < num_rows; i++) {
    for (int j = 0; j < num_cols; j++) {
        // 权重预加载到PE
        PE[i][j].weight = weight_matrix[i][j];
    }
}

// 输入数据从左侧流入
for (int t = 0; t < num_cycles; t++) {
    for (int i = 0; i < num_rows; i++) {
        // 输入数据从左向右传递
        input_data = (i == 0) ? new_input : PE[i-1][j].input;
        PE[i][j].input = input_data;
        
        // 乘累加计算
        PE[i][j].partial_sum += input_data * PE[i][j].weight;
        
        // 部分和从下向上传递
        if (i > 0) {
            PE[i-1][j].partial_sum += PE[i][j].partial_sum;
        }
    }
}

嗯,这里要注意:脉动阵列的映射策略不是固定的。你可以根据数据流的特点选择不同的映射方式。比如:

  • 权值固定映射:权重预加载,输入和部分和流动
  • 输入固定映射:输入预加载,权重和部分和流动
  • 输出固定映射:部分和驻留,权重和输入流动

避坑指南: 我曾经在脉动阵列里用了输出固定映射,结果发现部分和的精度损失很大。后来排查发现,是因为部分和在PE里累加的次数太多,超出了累加器的动态范围。解决办法是定期做一次部分和的归一化,或者用更高精度的累加器。

如何选择数据流策略?

说实话,没有一种数据流策略是万能的。选择哪种策略,取决于你的模型结构、芯片架构和性能目标。

我一般会按这个思路来选:

  1. 看数据复用率:权重复用率高用权值固定,输入复用率高用输入固定
  2. 看PE阵列规模:PE多的话,行固定和脉动阵列更合适
  3. 看带宽瓶颈:如果外部带宽是瓶颈,优先减少数据搬运次数
  4. 看精度要求:输出固定要注意累加器的精度问题

最后说一句:数据流优化是个系统工程,不能只看局部。我建议你在做架构设计时,先把数据流图画出来,看看数据是怎么流动的,瓶颈在哪里。很多时候,一个简单的数据流调整,就能带来意想不到的性能提升。

核心总结:

  • 权值固定:权重不动,输入和部分和流动
  • 输入固定:输入不动,权重和部分和流动
  • 输出固定:部分和不动,权重和输入流动
  • 行固定:按行分解,兼顾三种复用
  • 脉动阵列:数据有节奏流动,映射策略灵活
数据流优化策略对比 权值固定 权重:驻留PE不动 输入:广播/传递 部分和:PE间传递 适合:卷积层 输入固定 输入:驻留PE不动 权重:广播/传递 部分和:PE间传递 适合:全连接层 输出固定 部分和:驻留PE累加 权重:流经PE 输入:流经PE 适合:输出通道少 行固定(Row Stationary) 按行分解卷积操作,同时利用权重复用、输入复用和部分和累加 适合:灵活配置的PE阵列 脉动阵列(Systolic Array) 数据有节奏地在PE间流动,支持权值固定/输入固定/输出固定映射 适合:规则计算、高吞吐场景