脉动阵列架构(一):脉动阵列的基本原理、数据流方式

各位同学,今天我们来聊聊脉动阵列。说实话,我第一次接触这个概念是在做一款边缘AI芯片的时候。当时团队在争论用哪种架构做矩阵乘法,有人提议用脉动阵列,我第一反应是——这玩意儿不是几十年前的东西吗?后来真香了。

脉动阵列,英文叫Systolic Array。名字听着挺玄乎,其实说白了就是一种让数据像血液一样在计算单元之间规律流动的架构。每个处理单元(PE)只做简单的事——乘加运算,然后传给邻居。嗯,就这么简单。

3.1 脉动阵列的基本原理

想象一下,你有一排工人,每个人只负责把零件加工一下,然后传给下一个人。这就是脉动阵列的核心思想。每个PE(Processing Element)只做一件事:D_out = D_in + A_in × B_in

我习惯把脉动阵列比作流水线,但比流水线更灵活。流水线是单向的,脉动阵列可以是二维的,数据可以从多个方向流入。

核心特点:

  • 局部互联:每个PE只和邻居通信,没有全局连线
  • 规则结构:所有PE一模一样,布局布线非常规整
  • 高吞吐:每个时钟周期,每个PE都在干活,没有闲置
  • 低功耗:没有长距离走线,没有复杂的控制逻辑

你想想看,为什么脉动阵列在深度学习加速器里这么火?因为矩阵乘法太适合它了。每个PE算一个乘加,数据流过去,结果就出来了。我在项目中遇到过一个问题:一开始用通用处理器做矩阵乘法,功耗高得吓人。换成脉动阵列后,同样的算力,功耗降了70%。

3.2 三种数据流方式

脉动阵列的数据流方式,说白了就是数据怎么喂进去、怎么流出来。主要有三种:权重固定、输入固定、输出固定。我一个个说。

3.2.1 权重固定(Weight Stationary)

这种方式,权重存在PE里不动,输入数据流过去,输出结果也流过去。

具体怎么做?每个PE预先加载好权重W。然后输入数据从左边流进来,部分和从上面流下来。每个PE算完,把结果传给下面的PE。

我的经验:权重固定适合卷积层。因为卷积核的权重复用率很高,一个权重要和多个输入做乘加。我在做YOLO加速器时,卷积层全部用的权重固定,效果很好。

伪代码长这样:

// 权重固定脉动阵列
// 每个PE预先加载权重W[i][j]
for t = 0 to N-1:
    for i = 0 to M-1:
        for j = 0 to M-1:
            // 输入从左边来,部分和从上边来
            input = input_stream[t][i][j]
            partial_sum = top_pe[i][j].output
            // 乘加运算
            result = partial_sum + input * W[i][j]
            // 传给下面的PE
            bottom_pe[i][j].input = result
            // 输入传给右边的PE
            right_pe[i][j].input = input

优点很明显:权重不用来回搬,省带宽。缺点呢?如果网络层切换,权重要重新加载,有延迟。

3.2.2 输入固定(Input Stationary)

这种方式反过来。输入数据存在PE里不动,权重流过去。

每个PE预先加载好输入数据。然后权重从上面流下来,部分和从左边流过来。每个PE算完,把结果传给右边的PE。

注意:输入固定适合全连接层。因为全连接层的输入向量是固定的,权重矩阵在变化。我曾经在做一个推荐系统模型时,全连接层特别多,用输入固定比权重固定快了30%。

伪代码:

// 输入固定脉动阵列
// 每个PE预先加载输入数据I[i][j]
for t = 0 to N-1:
    for i = 0 to M-1:
        for j = 0 to M-1:
            // 权重从上边来,部分和从左边来
            weight = top_pe[i][j].output
            partial_sum = left_pe[i][j].output
            // 乘加运算
            result = partial_sum + I[i][j] * weight
            // 传给右边的PE
            right_pe[i][j].input = result
            // 权重传给下面的PE
            bottom_pe[i][j].input = weight

说白了,输入固定和权重固定是对称的。一个存权重,一个存输入。选哪个?看你的数据复用模式。

3.2.3 输出固定(Output Stationary)

这种方式,输出结果存在PE里不动。输入和权重都流过去。

每个PE里存的是累加结果。输入从左边来,权重从上边来。每个PE算完,结果累加到本地寄存器里。等所有数据算完,再把结果读出去。

避坑指南:我曾经在做一个语音识别加速器时,用了输出固定。结果发现PE里的累加寄存器位宽不够,导致精度损失。后来我把累加器从16位扩到32位,问题才解决。所以,输出固定要注意累加器的位宽设计。

伪代码:

// 输出固定脉动阵列
// 每个PE初始化为0
for t = 0 to N-1:
    for i = 0 to M-1:
        for j = 0 to M-1:
            // 输入从左边来,权重从上边来
            input = left_pe[i][j].output
            weight = top_pe[i][j].output
            // 乘加运算,结果累加到本地
            PE[i][j].accumulator += input * weight
            // 输入传给右边的PE
            right_pe[i][j].input = input
            // 权重传给下面的PE
            bottom_pe[i][j].input = weight

输出固定的好处是:不用频繁读写结果,减少带宽压力。坏处是:PE里要存累加结果,面积会大一些。

3.3 三种方式的对比

我整理了一个表格,方便你对比:

特性 权重固定 输入固定 输出固定
存储对象 权重 输入 输出(累加结果)
流动数据 输入、部分和 权重、部分和 输入、权重
适用场景 卷积层(权重复用高) 全连接层(输入复用高) 输出通道多的层
带宽需求 低(权重不动) 低(输入不动) 中(输入权重都动)
PE复杂度 中(需要累加器)
切换延迟 高(重加载权重) 高(重加载输入) 低(只需清空累加器)

你想想看,实际项目中怎么选?我个人的习惯是:先看你的模型里卷积层多还是全连接层多。卷积层多,用权重固定;全连接层多,用输入固定。如果两者都有,可以考虑混合架构。

3.4 小结

脉动阵列的核心就三点:

  • 每个PE只做乘加,数据在邻居间流动
  • 三种数据流方式,各有各的适用场景
  • 选哪种,看你的数据复用模式和带宽预算

嗯,今天就先聊到这儿。下一节我们会深入脉动阵列的硬件实现细节,包括PE的微架构、数据调度、以及如何解决数据冲突。到时候我会分享一个我在流片中踩过的坑——数据冒险导致的计算错误,那叫一个惨。

课后思考:如果让你设计一个支持三种数据流方式可配置的脉动阵列,你会怎么做?PE里需要加哪些控制逻辑?