数据流优化：三大固定策略与脉动阵列映射

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

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

为什么数据流这么重要？

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

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

三大经典数据流策略

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

1. 权值固定（Weight Stationary）

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

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

2. 输入固定（Input Stationary）

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

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

3. 输出固定（Output Stationary）

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

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

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

三种策略对比

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

行固定（Row Stationary）与脉动阵列映射

行固定是Eyeriss架构里提出的一种数据流策略，它把卷积操作按行分解，让数据在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;
        }
    }
}

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

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

如何选择数据流策略？

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

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

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

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