AutoLALA: Automatic Loop Algebraic Locality Analysis for AI and HPC Kernels

Abstract: 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.

• The paper introduces AutoLALA, a tool that computes closed-form symbolic reuse distances and DMD metrics for affine loop kernels.
• It employs a DSL and polyhedral lowering to derive quasi-polynomial expressions parameterized by program sizes and cache configurations.
• Empirical results reveal distinctions between algorithms with similar operation counts, enabling cache-agnostic analysis for optimization.

AutoLALA: Automatic Loop Algebraic Locality Analysis for AI and HPC Kernels

Introduction

The AutoLALA system proposes a rigorous polyhedral-based toolchain for symbolic data locality analysis across modern AI and HPC loop kernels, focusing on quantifying data movement complexity—rather than merely operation count—at the level of program parameterization. By leveraging recent advances in symbolic reuse analysis and polyhedral computation, AutoLALA systematically computes closed-form quasi-polynomial expressions for reuse distance distributions and Data Movement Distance (DMD) metrics for affine loop programs. Its methodology decouples program-specific kernel features from machine-dependent caching specifics, enabling analysis that is parameterized by program problem size and cache configurations.

Methodology

Domain-Specific Language and Polyhedral Lowering

The AutoLALA DSL is purpose-built for concise and unambiguous specification of affine loop programs, with all relevant constraints (iteration bounds, memory access subscripts, conditionals) required to be affine in program parameters and iterator variables. This restriction guarantees the applicability of the polyhedral model, wherein iteration domains and access functions become parametric polyhedra and affine maps, respectively. The lowering pipeline constructs a lex-ordered timestamp space and an access map for each memory operation, representing all possible executions in integer polyhedral form.

Symbolic Locality Analysis

Building on the approach in Zhu et al. (Zhu et al., 10 Mar 2026), the system avoids stack simulation or recursive working set analysis, instead computing the reuse distance as the cardinality of the image of the reuse-interval window under the access map. Core steps include:

1. Construction of lexicographic order relations on timestamp space.
2. Identification of same-element access pairs via polyhedral composition.
3. Automated derivation of immediate predecessor/successor relations.
4. Reuse-interval computation as the cardinality of all timestamp points between predecessor and current access.
5. Computation of reuse distance as the number of distinct data elements accessed within the reuse interval, realized as a parametric polyhedral projection.
6. Extraction of symbolic multiplicity expressions for accesses at each reuse distance via Barvinok’s integer point counting.

This pipeline produces piecewise quasi-polynomial distributions of reuse distances and closed-form DMD expressions parameterized by program size, block size, associativity, and other symbolic parameters.

System Implementation

AutoLALA is implemented in Rust (edition 2024), robustly partitioned into three crates: a core analysis library, a command-line interface, and a web-based playground. Polyhedral computations are orchestrated via the Barvinok library (invoked through Rust bindings), supporting counting and projection operations essential to the analysis. The pipeline is designed for modularity and scalability, and handles arbitrary affine loop nests, conditionals, and statement interleaving.

Numerical Results and Empirical Validation

AutoLALA can distinguish between algorithms with identical operation count but differing memory traffic, such as naive versus tiled matrix multiplication, via their symbolic DMD formulas. The tool applies region-scaling filters to isolate asymptotically relevant contributions, discarding constant-size or parameter-pinned regions that do not scale with program size. By supporting parametric cache models (e.g., block size and set count) and producing formulas valid for all parameter choices, AutoLALA directly supports comparative analysis of algorithm/data structure variants and hardware-aware kernel optimization.

While the paper focuses on methodological exposition and tool demonstration, it explicitly inherits the correctness and precision guarantees of earlier symbolic analysis work (Zhu et al., 10 Mar 2026). AutoLALA faithfully implements an entirely symbolic, cache-independent analysis: the reuse distance distribution and associated DMD formulas do not require simulation, profiling, or any form of program execution.

Implications and Theoretical Significance

AutoLALA enables automated, symbolic understanding of the asymptotic and parametric locality properties of AI and HPC kernels—especially tensor contractions, stencils, and dense linear algebra. This facilitates:

• Algorithm selection and optimization: Analysts can compare fully symbolic DMD formulas to guide optimization decisions sensitive to target cache/memory configurations.
• Compiler and scheduling integration: The toolchain is designed for future integration into polyhedral IR-enabled compilers such as MLIR or Polygeist, enabling locality analysis of complex pipelines at compile time.
• Generality and extensibility: The tool is agnostic to loop dimensionality and is parameterized for input sizes, cache lines, and associativities, enabling rapid what-if analysis.

Limitations include the inapplicability to kernels with non-affine, data-dependent accesses, and indirect addressing. Theoretical extension to parallel loop analysis and further integration with tiling and scheduling transformations represents a logical future direction, aligning with recent work on symbolic parallel locality analysis [liu2024pluss].

Future Directions

Key suggested extensions include:

• Extension to symbolic thread counts for parallel programs.
• Automated kernel extraction from polyhedral IRs within MLIR.
• Broader support for tiling and hardware mapping transformations.
• Deeper integration with high-level numerical computing DSLs and AI frameworks for automated kernel resource estimation.

Conclusion

AutoLALA advances the state of symbolic data locality analysis for AI and HPC loop kernels, providing a robust, fully symbolic, and machine-parameterized toolchain for closed-form reuse characterization. Its integration of a DSL, polyhedral analysis, Barvinok-based counting, and practical filtering for scaling regions makes it applicable to a wide variety of affine codebases. AutoLALA’s outputs are directly useful for locality-driven algorithm design, kernel optimization, and hardware-software co-design, and establish a strong foundation for future generalizations to parallel and more complex program models.