Canonicalization of Batched Einstein Summations for Tuning Retrieval

TL;DR

This work addresses the lack of a canonical representation for batched Einstein summations, a common pattern in scientific computing where multiple contractions share data. It introduces a graph-based canonicalization pipeline that maps batched einsums to a unique normal form and extends the primitive to functional batched einsums to enable fusion of surrounding computations. The Feinsum framework stores transformation plans and maps functional batched einsums to idealized forms, enabling reuse of optimized subprograms across real workloads. Experimental results show meaningful speedups over state-of-the-art approaches (e.g., a geomean of $4.7\times$ on TCCG and FEM solver, and up to $2.1\times$ over JAX with memory-bound cases reaching $3.6\times$), while keeping canonicalization overheads modest (less than $1\,\mathrm{ms}$). The work demonstrates practical impact by enabling near-roofline performance for memory-bound tensor contractions and provides a reproducible open-source pipeline for canonicalization and transformation retrieval.

Abstract

We present an algorithm for normalizing \emph{Batched Einstein Summation} expressions by mapping mathematically equivalent formulations to a unique normal form. Batches of einsums with the same Einstein notation that exhibit substantial data reuse appear frequently in finite element methods (FEM), numerical linear algebra, and computational chemistry. To effectively exploit this temporal locality for high performance, we consider groups of einsums in batched form. Representations of equivalent batched einsums may differ due to index renaming, permutations within the batch, and, due to the commutativity and associativity of multiplication operation. The lack of a canonical representation hinders the reuse of optimization and tuning knowledge in software systems. To this end, we develop a novel encoding of batched einsums as colored graphs and apply graph canonicalization to derive a normal form. In addition to the canonicalization algorithm, we propose a representation of einsums using functional array operands and provide a strategy to transfer transformations operating on the normal form to \emph{functional batched einsums} that exhibit the same normal form; crucial for fusing surrounding computations for memory bound einsums. We evaluate our approach against JAX, and observe a geomean speedup of $4.7\times$ for einsums from the TCCG benchmark suite and an FEM solver.

Figures (9)

• Figure 1: Batched einsum $e_1$.
• Figure 2: Batched einsum $e_2$.
• Figure 4: The Feinsum approach to achieve separation of concerns. Computational scientist expresses his workload as a functional batched einsum and retrieves optimizations from the database. Transformation writer implements transformations for idealized batched einsums and records their performance in the database. The normal form is integral in allow the reuse of transformation knowledge from the performance engineer to the high-level programmer.
• Figure 5: Arithmetic intensities (AI) for the TCCG benchmark suite with double-precision operands. The horizontal dashed lines indicate saturation AIs of modern accelerators, above which the workloads become compute-bound.
• Figure 6: Observed Speedups for batched einsums executed using Feinsum relative to equivalent implementation using JAX/.

Theorems & Definitions (32)

• Definition 1
• Definition 2: Array Domain
• Definition 3: Array Codomain
• Definition 4: Index
• Definition 5: Index List
• Definition 6: Length of an Index List
• Definition 7: Einstein Summation ("einsum")
• Example 1: Einsum
• Definition 8: Equality of Index Lists
• Definition 9: Equality of einsums