自动分析AI与高性能计算中循环程序的数据局部性,优化内存访问效率。
AutoLALA: Automatic Loop Algebraic Locality Analysis for AI and HPC Kernels
- 基于符号化代数方法,直接计算数据重用距离,无需模拟堆栈。
- 可生成任意仿射循环嵌套的重用间隔闭式公式,支持矩阵乘等典型算子。
- 开源工具,提供命令行与网页界面,适合编译器开发者和性能优化研究者。
数据移动是现代计算系统的主要瓶颈。对于人工智能和高性能计算中常见的循环程序(如矩阵乘法、张量收缩、模板计算和einsum操作),数据在内存层次结构中的移动成本往往超过算术运算成本。本文提出AutoLALA,一个开源工具,用于分析仿射循环程序中的数据局部性。该工具接收用小型领域特定语言(DSL)编写的程序,将其降阶为多面体集合与映射,并生成重用距离和数据移动复杂度的闭式符号公式。AutoLALA实现了Zhu等人提出的完全符号化局部性分析,结合Smith等人提出的数据移动距离(DMD)框架。特别地,它通过访问空间在访问映射下的像来计算重用距离,避免了堆栈模拟和Denning的递归工作集公式。我们描述了DSL语法及其形式语义,多面体降阶管道如何通过仿射变换构建时间戳空间和访问映射,以及一系列Barvinok计数操作序列,用于推导符号化的重用间隔和重用距离分布。系统以Rust实现,分为三个模块化库,具备对Barvinok库的安全绑定。提供命令行接口和支持LaTeX渲染输出公式的交互式Web沙盒。该工具可处理任意仿射循环嵌套,涵盖张量收缩、einsum表达式、模板计算和一般多面体程序等负载。
原文摘要 · Abstract (English)
Data movement is the primary bottleneck in modern computing systems. For loop-based programs common in high-performance computing (HPC) and AI workloads, including matrix multiplication, tensor contraction, stencil computation, and einsum operations, the cost of moving data through the memory hierarchy often exceeds the cost of arithmetic. This paper presents AutoLALA, an open-source tool that analyzes data locality in affine loop programs. The tool accepts programs written in a small domain-specific language (DSL), lowers them to polyhedral sets and maps, and produces closed-form symbolic formulas for reuse distance and data movement complexity. AutoLALA implements the fully symbolic locality analysis of Zhu et al. together with the data movement distance (DMD) framework of Smith et al. In particular, it computes reuse distance as the image of the access space under the access map, avoiding both stack simulation and Denning's recursive working-set formulation. We describe the DSL syntax and its formal semantics, the polyhedral lowering pipeline that constructs timestamp spaces and access maps via affine transformations, and the sequence of Barvinok counting operations used to derive symbolic reuse-interval and reuse-distance distributions. The system is implemented in Rust as a modular library spanning three crates, with safe bindings to the Barvinok library. We provide both a command-line interface and an interactive web playground with LaTeX rendering of the output formulas. The tool handles arbitrary affine loop nests, covering workloads such as tensor contractions, einsum expressions, stencil computations, and general polyhedral programs.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。