Latent Idiom Recognition for a Minimalist Functional Array Language using Equality Saturation

TL;DR

LIAR addresses idiom recognition for high-performance libraries by encoding programs and library idioms in a minimalist functional array IR and applying equality saturation. The approach uses a target-independent IR with a small set of core rewrite rules plus target-specific extractors and idiom rules to expose library calls such as BLAS and PyTorch functions. Evaluation on PolyBench kernels shows a geometric mean speedup of 1.46x over reference C kernels (excluding gemver) and demonstrates how idioms evolve over saturation steps toward complete library coverage. The results indicate that automated idiom discovery using equality saturation is robust to input variations and adaptable to multiple libraries, reducing the need for hand-crafted analyses.

Abstract

Accelerating programs is typically done by recognizing code idioms matching high-performance libraries or hardware interfaces. However, recognizing such idioms automatically is challenging. The idiom recognition machinery is difficult to write and requires expert knowledge. In addition, slight variations in the input program might hide the idiom and defeat the recognizer. This paper advocates for the use of a minimalist functional array language supporting a small, but expressive, set of operators. The minimalist design leads to a tiny sets of rewrite rules, which encode the language semantics. Crucially, the same minimalist language is also used to encode idioms. This removes the need for hand-crafted analysis passes, or for having to learn a complex domain-specific language to define the idioms. Coupled with equality saturation, this approach is able to match the core functions from the BLAS and PyTorch libraries on a set of computational kernels. Compared to reference C kernel implementations, the approach produces a geometric mean speedup of 1.46x for C programs using BLAS, when generating such programs from the high-level minimalist language.

Figures (15)

• Figure 1: Expression is converted to e-graph , which is subsequently saturated. In this example, only rule is applied: $\texttt{x / N} \rightarrow \texttt{x $\gg$ $\log_2\texttt{N}$}$. From saturated e-graph , expression is selected by an extractor that prefers bitwise shift.
• Figure 2: Overview of LIAR, the proposed technique. Vector sum expression is converted to e-graph . Equality saturation applies a set of rules to e-graph . These rules consist of target-independent language semantics and target-specific idioms. Rule application yields updated e-graph . From e-graph , a target-specific extractor chooses expression . The only target-specific components are the target idioms and extractor, both of which are highlighted in blue.
• Figure 3: The grammar describing the minimalist IR. $\overline{\texttt{e}}$ indicates zero or more instances of e. N is a compile-time integer constant and f is an anonymous function. The set of available named functions depends on the problem domain. For example, gemm is a named function when targeting BLAS but not when targeting PyTorch.
• Figure 4: Reduction semantics for the minimalist IR.
• Figure 5: Eight rewrite rules that capture the relationships between build, array access, tuple construction, and tuple deconstruction, $\lambda$-abstraction and $\beta$-reduction.