[Paper] Tenstorrent Wormhole上的Stencil计算

Source: arXiv - 2605.07599v1

概述

本文评估了 Tenstorrent Wormhole——一种基于 RISC‑V 的 AI 数据流加速器——在运行经典 HPC 核心：二维 5 点模板（2‑D 5‑point stencil）时的表现。通过将模板重新构想为一系列逐元素矩阵操作（Axpy）或矩阵乘法（MatMul），作者将其性能、能耗以及瓶颈与传统 CPU 实现进行比较。

关键贡献

• 两种新颖的映射策略用于 AI 加速器上的 stencil 工作负载：
  1. Axpy – 将 stencil 分解为元素级子矩阵操作的流水线。
  2. MatMul – 将 stencil 重构为密集矩阵乘法，以利用硬件的 GEMM 引擎。
• 全面的性能分析，将 Wormhole 上的计算时间与主机端开销（PCIe 传输、设备初始化、预处理）分离。
• 能效分析表明，对于大规模问题，Axpy 实现的功耗低于 CPU 基准，即使其墙钟时间较慢。
• 识别出架构和软件瓶颈（内存带宽、PCIe 延迟、缺乏原生 stencil 原语），并给出具体的硬件‑软件协同设计建议，以使 AI 加速器更适合 HPC。

方法论

1. 基准选择 – 作者使用标准的二维 5‑点模板（在热扩散、CFD 和图像处理代码中常见）。
2. 内核重新设计 –
   • Axpy: 模板更新 out[i,j] = a*in[i,j] + b*in[i-1,j] + c*in[i+1,j] + d*in[i,j-1] + e*in[i,j+1] 被拆分为五个独立的逐元素矩阵加法/乘法，能够自然映射到 Wormhole 的向量 ALU 上。
   • MatMul: 通过填充和重塑输入网格，将模板表达为稀疏矩阵‑向量乘积，经过适当的分块后，加速器可以将其视为密集 GEMM。
3. 实现堆栈 – 内核使用 Tenstorrent 的 SDK（基于 Python 的 API）编写，并编译为片上数据流 ISA。主机代码在 x86‑64 CPU 上运行，负责通过 PCIe 进行数据移动。
4. 基线 – 一个高度优化的多线程 CPU 版本（使用 OpenMP 和缓存分块）作为参考。
5. 收集的指标 – 端到端运行时间、加速器独立计算时间、PCIe 传输量、功耗（通过板载传感器），以及求解能耗。

结果与发现

要点：一旦数据驻留在加速器上，Wormhole 可以像现代 CPU 那样快速计算 stencil。主要的慢点来自通过 PCIe 传输数据以及一次性的设备初始化。Axpy 变体虽然整体更慢，但在大规模问题上能效更佳。

实际影响

• Accelerator‑first HPC pipelines – 对于在多个 stencil 传递中重复使用相同数据的工作负载（例如时间步进模拟），将数据保留在芯片上可以抵消 PCIe 的惩罚，使 AI 加速器具备竞争力。
• Energy‑constrained edge supercomputing – Axpy 的每次求解能耗更低，表明 AI 芯片在远程或嵌入式 HPC 节点（电源预算是主要限制）中具有吸引力。
• Software stack considerations – 开发者需要考虑数据移动成本；使用统一内存或 NVMe 直接分阶段可以缩小主机侧的差距。
• Algorithm redesign – 将传统内核重新构造为匹配加速器优势的形式（例如 GEMM‑friendly）可以释放隐藏的性能，如矩阵乘法映射所示。
• Tooling – 本研究强调了需要能够区分主机侧和设备侧开销的分析工具，这对实现真实的性能预算至关重要。

限制与未来工作

• PCIe瓶颈 – 当前的Wormhole平台依赖相对较慢的 PCIe Gen3 链接；更新的互连技术（CXL、PCIe Gen4/5）有望显著提升端到端时间。
• 内存层次结构 – 片上 SRAM 容量有限；更大的模板会溢出到片外 DRAM，导致的延迟在论文中未得到充分探讨。
• 单精度聚焦 – 实验使用 FP32；混合精度或整数模板变体（在某些 CFD 代码中常见）尚未测试。
• 可扩展性 – 研究仅限于单个加速器；多节点扩展、集体通信以及工作负载划分仍是未解之题。
• 软件成熟度 – Tenstorrent SDK 仍在发展中；更丰富的原语（原生模板操作、更好的 DMA 调度）可能降低手动内核重构的需求。

作者提出的未来方向 包括更紧密的 CPU‑加速器集成、对 halo 交换的硬件支持，以及能够自动将模板 DSL 翻译为 AI 芯片最优数据流图的编译器扩展。

作者

• Lorenzo Piarulli
• Daniele De Sensi

论文信息

• arXiv ID: 2605.07599v1
• 分类: cs.DC, cs.ET
• 出版日期: 2026年5月8日
• PDF: Download PDF