Automatic Generation of Efficient Linear Algebra Programs

The level of abstraction at which application experts reason about linear algebra computations and the level of abstraction used by developers of high-performance numerical linear algebra libraries do not match. The former is conveniently captured by high-level languages and libraries such as Matlab and Eigen, while the latter expresses the kernels included in the BLAS and LAPACK libraries. Unfortunately, the translation from a high-level computation to an efficient sequence of kernels is a task, far from trivial, that requires extensive knowledge of both linear algebra and high-performance computing. Internally, almost all high-level languages and libraries use efficient kernels; however, the translation algorithms are too simplistic and thus lead to a suboptimal use of said kernels, with significant performance losses. In order to both achieve the productivity that comes with high-level languages, and make use of the efficiency of low level kernels, we are developing Linnea, a code generator for linear algebra problems. As input, Linnea takes a high-level description of a linear algebra problem and produces as output an efficient sequence of calls to high-performance kernels. In 25 application problems, the code generated by Linnea always outperforms Matlab, Julia, Eigen and Armadillo, with speedups up to and exceeding 10x.

💡 Research Summary

The paper presents Linnea, an automatic code generator that bridges the gap between high‑level linear‑algebra programming environments (such as MATLAB, Julia, Eigen, and Armadillo) and low‑level high‑performance libraries (BLAS and LAPACK). The authors observe that while high‑level languages let users express matrix computations in a notation close to the mathematical description, the internal translation to BLAS/LAPACK kernels is usually simplistic. This often leads to sub‑optimal performance because the translators ignore algebraic identities (e.g., distributivity), fail to exploit matrix properties (symmetry, SPD, triangularity), and sometimes even compute explicit inverses instead of solving linear systems.

Linnea addresses these shortcomings by accepting a high‑level description of a linear‑algebra problem together with optional annotations that describe operand properties (upper/lower triangular, diagonal, symmetric, SPD, orthogonal, etc.). The core of Linnea’s approach is an algebraic rewrite engine combined with a graph‑based search for optimal kernel sequences:

1. Expression Rewriting – Using a set of mathematically sound rewrite rules (distributivity, associativity, transpose/inverse manipulations, etc.), the input expression is transformed into multiple equivalent forms. This step reduces arithmetic complexity and reveals opportunities for kernel reuse.

2. Pattern Matching – For each rewritten form, Linnea scans for sub‑expressions that match the signatures of available BLAS/LAPACK kernels (e.g., matrix‑matrix multiply, triangular solve, Cholesky factorization). When a match is found, a new intermediate variable is introduced, the corresponding kernel call is recorded, and the sub‑expression is replaced by the intermediate.

3. Derivation Graph Construction – Nodes represent the remaining symbolic expression at a given stage; edges are annotated with the kernel applied and its estimated cost (flop count, memory traffic, threading). Starting from a single root node, the algorithm expands the frontier (the “active” set) breadth‑first, generating successors via rewriting and matching. Redundant sub‑expressions that appear in different branches are merged, preventing exponential blow‑up.

4. Cost‑Driven Selection – After a termination condition (no active nodes or a sufficient number of leaf nodes) the algorithm searches the graph for the cheapest path from root to a leaf, using a user‑defined cost model. The resulting sequence of kernels constitutes the generated program.

The implementation is written in Python, currently supports only real‑valued arithmetic, and relies on multi‑threaded BLAS for parallelism. Linnea outputs Julia code because Julia provides thin wrappers around BLAS/LAPACK while allowing easy insertion of custom snippets for operations not covered by the libraries.

Experimental Evaluation – The authors evaluate Linnea on 25 real‑world linear‑algebra problems drawn from image restoration, stochastic least‑squares, and other scientific domains. For each benchmark, Linnea’s generated code is compared against hand‑written MATLAB scripts, Julia code using built‑in linear‑algebra functions, Eigen (C++ expression templates), and Armadillo. Results show that Linnea consistently outperforms the competitors, achieving speed‑ups ranging from modest (≈1.5×) to dramatic (≥10×) on large matrix problems. Code‑generation time is on the order of a few minutes, far faster than manual expert tuning.

Contributions and Limitations – The paper contributes (i) a formal grammar for annotated linear‑algebra expressions, (ii) a systematic algebraic rewrite and graph‑search framework that guarantees correctness by construction, (iii) a practical prototype that demonstrates substantial performance gains, and (iv) an empirical study comparing against state‑of‑the‑art high‑level tools. Limitations include lack of complex‑number support, reliance on user‑provided property annotations (no automatic inference), and potential scalability issues for extremely large expression trees where the derivation graph may become very large.

Future Work – The authors plan to extend Linnea with automatic property inference, support for complex arithmetic and GPU kernels, and to explore more powerful search techniques such as Equality Saturation (EQ) to handle a larger set of rewrite axioms without sacrificing scalability.

In summary, Linnea showcases that a disciplined combination of symbolic algebraic rewriting and cost‑aware graph exploration can automatically produce high‑performance linear‑algebra code from high‑level specifications, delivering productivity comparable to MATLAB‑style programming while achieving performance close to hand‑tuned low‑level implementations.