Nexus Machine架构：稀疏矩阵计算与主动消息优化

1. Nexus Machine架构设计背景与核心挑战

稀疏矩阵计算和图形处理等不规则工作负载在现代机器学习、科学计算和网络分析中越来越普遍。这类计算通常表现出三个典型特征：非零元素分布不均匀、内存访问模式不可预测、计算密度变化大。传统架构（如CPU/GPU）在处理这类问题时面临显著效率瓶颈——根据研究数据，GPU在稀疏矩阵乘法(SpMSpM)上的利用率通常不足30%，主要由于线程发散和内存访问冲突。

1.1 不规则工作负载的硬件加速困境

当前硬件加速方案主要面临三重挑战：

• 内存墙问题：稀疏数据结构的间接寻址导致缓存命中率低下。例如在CSR格式的稀疏矩阵中，行指针和列索引的额外开销可能占存储空间的30%以上
• 计算利用率低：静态调度的计算单元（如TPU的脉动阵列）难以适应动态变化的计算密度。当矩阵稀疏度达到90%时，传统阵列中超过90%的乘法单元处于闲置状态
• 通信开销大：分布式系统中数据重分配带来的网络传输可能消耗50%以上的执行时间。在图形处理中，顶点度数差异可达数个数量级，导致严重的负载不均衡

1.2 主动消息机制的历史沿革

主动消息(Active Messages)概念最早由加州大学伯克利分校在1992年提出，其核心思想是将计算逻辑嵌入消息本身。当消息到达目标节点时，携带的处理程序自动执行，减少显式同步和数据搬运。Nexus Machine的创新点在于：

1. 将AM与粗粒度可重构阵列(CGRA)结合，形成统一的计算-通信范式
2. 支持消息在传输路径上的任意空闲PE执行（称为"途中计算"）
3. 通过编译器静态分析与运行时动态调度的协同，实现计算任务的自动分解与分配

关键洞见：传统AM多用于集群通信，而Nexus Machine将其下沉到芯片级互连网络，使每个PE既是计算单元又是路由节点，形成真正的数据流架构

2. 架构设计详解

2.1 整体硬件组成

Nexus Machine采用4×4 PE阵列的实例化设计，每个PE包含：

• 计算部分：4×4 INT16乘加阵列，支持8/16/32位精度可配置
• 存储层次：
  • 1KB SRAM（64B/cycle带宽）
  • 1KB AM队列（70位/entry，深度可调）
• 网络接口：
  • 动态路由控制器（支持XY/转向模型）
  • 拥塞感知电路（基于缓冲区水位监测）

全局架构参数：

• 峰值频率：588MHz（22nm FDSOI工艺）
• 片外内存带宽：4.7GB/s（AXI4接口）
• 能效比：194MOPS/mW（稀疏工作负载下）

2.2 关键创新组件

2.2.1 主动消息处理单元

AM的硬件实现包含三个核心模块：

1. 消息扫描器(Scanner)：并行检测输入消息中的有效载荷，识别操作数位置。采用SIMD比较电路，单周期可处理16个地址匹配
2. 依赖解析器：构建动态数据流图(DFG)，通过标签匹配确定操作间的生产者-消费者关系
3. 执行上下文池：维护16个硬件线程上下文，支持乱序发射但顺序提交

AM格式示例：

| PE_ID(4b) | Opcode(6b) | Op1_Addr(12b) | R1_PE(4b) | R1_Addr(12b) | | R2_PE(4b) | R2_Addr(12b) | Immediate(16b) |

2.2.2 非相似性感知数据分区

针对稀疏矩阵的存储优化算法（Algorithm 1）包含以下步骤：

1. 访问模式分析：统计每行非零元素访问的存储体集合$L_i$
2. 相似度计算：对任意两行$r_i,r_j$，计算$d(i,j)=|L_iΔL_j|$（对称差集大小）
3. 聚类分配：将相似度高的行分配到同一PE，差异大的分散到不同PE

该算法可将存储体冲突降低40-60%，实测在85%稀疏度的矩阵上，PE利用率仍能保持72%以上。

2.3 内存子系统优化

2.3.1 分布式SRAM组织

与传统CGRA的边缘共享存储不同，Nexus Machine采用全分布式设计：

• 每个PE独占1KB SRAM，通过NoC互连形成统一地址空间
• 支持三种访问模式：
  • 本地访问（1周期延迟）
  • 近邻直接访问（2-3周期）
  • 全局路由访问（4-8周期，带拥塞规避）

2.3.2 带宽-容量权衡

通过实验发现（图16）：

• 当片内存储≥256KB时，稀疏度为85%的SpMSpM带宽需求趋于稳定
• 极端稀疏度(95%)下，采用值/坐标分离存储可节省37%带宽
• 对输出密集型操作（如SDDMM），采用压缩传输协议可提升有效带宽2.1倍

3. 软件栈与编译流程

3.1 前端编译器设计

基于LLVM的编译器流程包含四个关键阶段：

1. 数据流图提取：通过程序切片识别关键计算子图
2. 稀疏模式识别：分析访存模式，标注规律性（如CSR/ELL格式）
3. AM生成：将循环体转换为AM序列，插入路由提示
4. 配置内存分配：静态调度确定PE功能单元配置

编译示例（SpMV内核）：

// 原始循环 for(int i=0; i<row_ptr[n]; ++i) { y[row[i]] += val[i] * x[col[i]]; } // 转换后的AM序列 AM_LOAD(row_ptr, PE0); AM_LOAD(col, PE1); AM_COMPUTE(val, x, y); // 自动展开为PE间流水线

3.2 运行时管理系统

轻量级运行时(约50KB代码)负责：

• 动态负载均衡：监控PE利用率，当差异超过阈值时触发AM迁移
• 拥塞控制：基于Hop-by-hop背压机制调节注入速率
• 错误恢复：通过AM校验和实现细粒度回滚

实测运行时开销仅占执行时间的0.3-1.2%，比传统CGRA的任务调度器轻量5-8倍。

4. 性能评估与对比

4.1 实验设置

• 测试平台：Cadence Genus综合结果 + 周期精确模拟器
• 对比基线：
  • 脉动阵列（Google TPU类似设计）
  • 传统CGRA（HyCube架构）
  • TIA（触发指令架构）
• 工作负载：
  • 稀疏类：SpMSpM、SpMV、SDDMM
  • 稠密类：MatMul、Conv
  • 图计算：BFS、PageRank

4.2 关键结果

4.2.1 性能对比

在稀疏矩阵乘法(SpMSpM)上：

• 相比脉动阵列：1.8倍加速（稀疏度70%时）
• 相比传统CGRA：1.35倍加速
• 途中计算占比：最高达63%（S4类工作负载）

能效比表现：

• 稀疏工作负载：平均194MOPS/mW
• 稠密工作负载：仍保持87MOPS/mW（牺牲部分效率换取灵活性）

4.2.2 面积开销分析

总芯片面积增加17.3%，但换来稀疏场景下2.1倍的能效提升。

5. 应用场景与优化技巧

5.1 典型使用模式

5.1.1 稀疏神经网络推理

以Pruned ResNet-50为例：

1. 权重矩阵划分：按输出通道分组，保持PE间负载均衡
2. 激活数据分布：采用行循环划分，提高AM重用率
3. 动态跳过：当输入激活为零时，AM自动终止传播

实测在70%稀疏度下，端到端延迟降低41%，能耗减少33%。

5.1.2 图形分析

BFS实现的优化点：

• 顶点分配：按度数排序后交替分配（高-低-高模式）
• AM批处理：将同一层次的多个更新合并为单个AM
• 优先级路由：将深度较小的消息标记为高优先级

在infect-dublin数据集上，相比传统CGRA提升吞吐量2.7倍。

5.2 编程实践建议

1. 数据布局优化：

   • 对主要张量使用非相似性分区
   • 次要张量采用就近放置原则
   • 静态已知的稀疏模式可预编译为AM模板

2. AM设计原则：

   • 单个AM应包含4-16个基本操作
   • 避免在AM中引入条件分支
   • 对频繁使用的AM启用PE本地缓存

3. 调试技巧：

   • 使用模拟器的波形视图追踪AM传播路径
   • 在AM头部添加唯一ID便于日志关联
   • 对跨PE依赖插入显式同步标记

6. 局限性与未来方向

当前架构存在以下待改进点：

• 精度限制：INT16处理单元不适合某些科学计算场景
• 扩展性瓶颈：超过8×8阵列时，NoC延迟成为性能瓶颈
• 编译复杂度：自动AM生成对不规则循环的支持仍需完善

可能的演进方向包括：

1. 引入混合精度支持（FP8/INT32可配置）
2. 采用光互连降低大规模阵列通信延迟
3. 开发面向AM的领域特定语言(DSL)