AI与HPC交叉领域的Stencil计算

中文摘要 AI 辅助

TinyTC和OpenAI Triton等张量编译器最初是为AI工作负载开发的，但相同的分块和内存抽象可以应用于实现面向科学和工业应用的高阶stencil。我们针对一个具有边界条件的8阶25点声学stencil，在规模要求苛刻的网格上进行了演示，目标平台为GPGPU，并将硬件专用的TinyTC实现与可移植的PyTorch/Triton实现进行了比较。评估的目标平台包括Intel B70、B580、GPU MAX 1550、NVIDIA A100/RTX6000 Blackwell/H100和AMD MI325x。例如，在Battlemage B580上，TinyTC在随机初始化下达到15.6 Gpts/s，而PT/Triton为13.5 Gpts/s；而零初始化运行由于硬件内存压缩可达35.8 Gpts/s。通过roofline和内存层次分析，我们表明——正如预期——性能主要受带宽限制，并且编译器管理的L1/LSC缓存可以有效地替代程序员管理的共享内存暂存，适用于此类stencil。总体而言，结果将TinyTC定位为Intel硬件上的性能导向路径，而PyTorch/Triton则作为跨厂商HPC stencil开发的强可移植性/生产力基线。

英文摘要

Tensor compilers such as TinyTC and OpenAI Triton were originally developed for AI workloads, but the same tiling and memory abstractions can be applied to implement efficient high-order stencils for scientific and industrial applications. We demonstrate this for an 8th-order, 25-point acoustic stencil with boundary conditions over an a demanding-sized grid, targeting GPGPUs, where we compare the hardware-specialized TinyTC implementation with a portable PyTorch/Triton implementation. The target platforms for evaluation include Intel B70, B580, GPU MAX 1550, NVIDIA A100/RTX6000 Blackwell/H100, and AMD MI325x. For instance, on Battlemage B580 TinyTC reaches 15.6 Gpts/s versus 13.5 Gpts/s for PT/Triton under random initialization, while zero-initialized runs reach up to 35.8 Gpts/s due to hardware memory compression. Using roofline and memory-hierarchy profiling, we show that -as expected- performance is predominantly bandwidth-limited and that compiler-managed L1/LSC caching can effectively replace programmer-managed shared-memory staging for this stencil class. Overall, the results position TinyTC as the performance-oriented path on Intel hardware and PyTorch/Triton as a strong portability/productivity baseline for cross-vendor HPC stencil development.

发表机构

• Intel Corporation（英特尔公司）
• TotalEnergies EP Research and Technology US（道达尔能源EP美国研究与技术）

机构由 AI 辅助整理，请以论文原文为准。