Canonicalization of Batched Einstein Summations for Tuning Retrieval
Fuente:
arXiv
Saved in:
| Main Authors: | , |
|---|---|
| Format: | Preprint |
| Published: |
2026
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866917207846420480 |
|---|---|
| author | Kulkarni, Kaushik Klöckner, Andreas |
| author_facet | Kulkarni, Kaushik Klöckner, Andreas |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2601_12220 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | Canonicalization of Batched Einstein Summations for Tuning Retrieval Kulkarni, Kaushik Klöckner, Andreas Mathematical Software Distributed, Parallel, and Cluster Computing 15-04, 65Y15 G.4; G.1.3 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. |
| title | Canonicalization of Batched Einstein Summations for Tuning Retrieval |
| topic | Mathematical Software Distributed, Parallel, and Cluster Computing 15-04, 65Y15 G.4; G.1.3 |
| url | https://arxiv.org/abs/2601.12220 |