Skip to main content

2.4 Flexible PE Array

4.1 节的 Flexible PE Array 是 PanGaea 加速器的计算核心。它的“灵活”体现在：同一套 12×12 的处理单元（PE）阵列，可以通过动态配置，分别支持线性链、WFA-Affine 对齐和 WFA-Edit 对齐三种截然不同的计算模式。 这样就不需要为每个流水线阶段单独设计专用硬件，从而节省面积、提高利用率，并减少阶段切换时的上下文开销。

下面分几层来解释。

1. 为什么需要“灵活”的 PE 阵列？​

S2G 流水线中，不同阶段的计算模式差异很大：

  • Linear Chaining（线性链）：稀疏动态规划，每个 PE 评估一个锚点对 (i,j)(i,j)，计算间隙惩罚，然后归约找最优 jj。
  • WFA-Affine（仿射间隙波前对齐）：密集动态规划，计算 E1,E2,F1,F2,HE_1, E_2, F_1, F_2, H 五个矩阵，需要从多个 penalty level 读取数据。
  • WFA-Edit（编辑距离波前对齐）：也是密集 DP，但代价模型更简单，只需三个矩阵，从 s−1s-1 层读取。

如果为每种模式都做专用硬件，面积会很大，而且每个阶段只用到一部分，利用率低。
所以 PanGaea 设计了一个可重构 PE 阵列：底层共享算术单元，上层通过多路复用器（MUX）切换数据通路和参数，不用的部分电源门控。这样一套硬件就能分时处理三种模式。

2. PE 阵列总体结构​

  • 规模：12×12=14412 \times 12 = 144 个 PE。
  • 共享算术单元：每个 PE 内部有加法器、减法器、比较器、绝对值、最小值、最大值、多路复用器等。
  • 模式特定数据通路：根据当前模式，MUX 选择不同的输入、操作和输出路径。
  • 电源门控：当前模式不用的逻辑块被关掉，降低功耗。
  • 局部 scratchpad：PE 阵列直接连接一块局部 SRAM，用来存 WFA 的 DP 矩阵，访问延迟只有 1 个周期（从 GLB 读要 3 个周期），这对 WFA 的高频访问至关重要。

3. 三种模式详解​

Mode 1：Linear Chaining（线性链）​

任务：对固定的锚点 ii，遍历所有可能的 j<ij < i，计算链分数，找最优 j∗j^*。

公式（论文公式 7）：

fi=max⁡j<i{fj+score(i,j)−penalty(Δx,Δy)}f_i = \max_{j< i}\{f_j + \mathrm{score}(i,j) - \mathrm{penalty}(\Delta x,\Delta y)\} penalty=cgap⋅∣Δx−Δy∣+cskip⋅min⁡(Δx,Δy)+0.5⋅log⁡2(∣Δx−Δy∣+1)\mathrm{penalty} = c_{\mathrm{gap}}\cdot |\Delta x - \Delta y| + c_{\mathrm{skip}}\cdot \min (\Delta x,\Delta y) +0.5\cdot \log_2(|\Delta x - \Delta y| + 1)

其中 Δx=xi−xj\Delta x = x_i - x_j，Δy=yi−yj\Delta y = y_i - y_j。

PE 内部实现：

  • 两个 32 位减法器算 Δx,Δy\Delta x, \Delta y；
  • 绝对值、最小值单元算惩罚项；
  • 一个 256 项的查找表（2 KB，存在 scratchpad）算 log⁡2\log_2；
  • 两个定点乘法器乘 cgapc_{\mathrm{gap}} 和 cskipc_{\mathrm{skip}}；
  • 比较器检查线性条件 yi>yjy_i > y_j 和带宽限制 ∣Δx−Δy∣≤B|\Delta x - \Delta y| \leq B。

数据流（4.2 节详述）：

  • 锚点数据广播到所有 PE；
  • 每个 PE 独立算一个 (i,j)(i,j) 对的分数；
  • 然后通过 PE 网格做垂直-水平归约：每列先垂直比较找最大，再把各列最大值水平传到 PE0PE_0，得到全局最大分数和最佳 jj。
  • 归约树直接嵌入 PE 的输出混合逻辑，不单独做专用树，节省约 15% 面积，减少 3 周期延迟。

Mode 2A：WFA-Affine（仿射间隙波前对齐）​

任务：用两层仿射间隙模型做波前对齐，短间隙用 (o1,e1)(o_1, e_1)，长间隙用 (o2,e2)(o_2, e_2)。

公式（论文公式 8）：

E1s,d=max⁡(Hs−(o1+e1),d−1,E1s−e1,d−1)E2s,d=max⁡(Hs−(o2+e2),d−1,E2s−e2,d−1)F1s,d=max⁡(Hs−(o1+e1),d+1,F1s−e1,d+1)+1F2s,d=max⁡(Hs−(o2+e2),d+1,F2s−e2,d+1)+1Hs,d=max⁡(Hs−x,d+match(q,r),E1,E2,F1,F2)\begin{array}{rl} E_{1s,d} =& \max (H_{s - (o_1 + e_1),d - 1},E_{1s - e_1,d - 1})\\ E_{2s,d} =& \max (H_{s - (o_2 + e_2),d - 1},E_{2s - e_2,d - 1})\\ F_{1s,d} =& \max (H_{s - (o_1 + e_1),d + 1},F_{1s - e_1,d + 1}) + 1\\ F_{2s,d} =& \max (H_{s - (o_2 + e_2),d + 1},F_{2s - e_2,d + 1}) + 1\\ H_{s,d} =& \max (H_{s - x,d} + \mathrm{match}(q,r),E_1,E_2,F_1,F_2) \end{array}

PE 输入：查询碱基 qb[7:0]q_b[7:0]、参考碱基 rb[7:0]r_b[7:0]，以及从 scratchpad 读出的五个 DP 值 H,E1,E2,F1,F2H, E_1, E_2, F_1, F_2。

PE 内部实现：

  • 8 位比较器判断 match；
  • 四个并行 max 单元算 E1,E2,F1,F2E_1, E_2, F_1, F_2；
  • 两个增量器给 FF 矩阵加 1；
  • 五路 max 树选最优状态。

关键点：每周期每个 PE 需要 9 次 scratchpad 读，因为要访问多个 penalty level 的矩阵。所以 scratchpad 必须多端口。

数据流（4.2 节）：

  • 波前沿对角线有依赖：计算 (s,d)(s,d) 需要 (s−δ,d±1)(s-\delta, d\pm 1) 的值。
  • 相邻对角线共享大量数据，所以 PE 排成 2D 网格，邻居之间水平交换：
    • 每个 PE 把 EE 矩阵发给右邻居；
    • 把 FF 矩阵发给左邻居；
    • 通过 All-reduce 通信传播。

Mode 2B：WFA-Edit（编辑距离波前对齐）​

任务：计算编辑距离对齐，每个操作代价统一为 +1+1，用于图 WFA 跨变异泡。

公式（论文公式 9）：

Es,d=max⁡(Hs−1,d+1,Es−1,d−1+1)Fs,d=max⁡(Hs−1,d+1,Fs−1,d+1+1)Hs,d=max⁡(Hs−1,d+mismatch(q,r),E,F)\begin{array}{rl} E_{s,d} =& \max (H_{s - 1,d} + 1,E_{s - 1,d - 1} + 1)\\ F_{s,d} =& \max (H_{s - 1,d} + 1,F_{s - 1,d + 1} + 1)\\ H_{s,d} =& \max (H_{s - 1,d} + \mathrm{mismatch}(q,r),E,F) \end{array}

特点：

  • 只需三个 DP 矩阵，比 Mode 2A 少 40% 存储；
  • 每周期只需 3 次 scratchpad 读，且都来自 s−1s-1 层；
  • 与 Mode 2A 共享算术单元：match 检测器、加法器、MAX 树复用，只通过 MUX 选择仿射参数或统一 +1+1 常数。
  • 这种共享让 PE 面积比复制一份硬件节省 23%。

4. 灵活性的体现​

  • 共享算术单元：加法、比较、max 等基础单元三种模式都用。
  • MUX 动态切换：根据模式选择输入源、操作类型、输出路径。
  • 电源门控：不用的数据通路断电，降低功耗。
  • 统一 PE 网格：同一套 144 个 PE 分时处理链、对齐，减少阶段切换开销。
  • 局部 scratchpad：为 WFA 提供低延迟、多端口访问，支持密集 DP。

5. 与整体架构的配合​

  • GLB：存储 anchors、图拓扑、中间结果。PE 阵列从 GLB 取数据，但 WFA 矩阵放在局部 scratchpad。
  • 数据流映射：
    • Mode 1 用广播-单播-归约；
    • Mode 2A/2B 用对角线依赖 + 水平交换 + All-reduce。
  • GPF 引擎：独立处理图遍历，PE 阵列专注链和对齐。
  • Traceback：对齐完成后，回溯 DP 矩阵生成 CIGAR。

6. 一句话总结​

Flexible PE Array 是一个 12×12 的可重构计算阵列，通过共享算术单元和动态 MUX，在三种模式间切换：线性链的稀疏 DP、WFA-Affine 的密集仿射 DP、WFA-Edit 的编辑距离 DP。它用一套硬件覆盖了 S2G 流水线的主要计算模式，既节省面积又提高利用率，是 PanGaea 实现跨阶段融合的关键。