Morello: Compiling Fast Neural Networks with Dynamic Programming and Spatial Compression
This work addresses the challenge of high-throughput neural network inference for developers and researchers, though it is incremental as it builds on existing optimization techniques with a novel method.
The authors tackled the problem of optimizing neural network inference by introducing a dynamic programming approach to explore a larger search space of program optimizations, resulting in Morello, a compiler that synthesized high-throughput programs, such as a bfloat16-to-float32 vector-matrix multiply integrated into Google's gemma.cpp.
High-throughput neural network inference requires coordinating many optimization decisions, including parallel tiling, microkernel selection, and data layout. The product of these decisions forms a search space of programs which is typically intractably large. Existing approaches (e.g., auto-schedulers) often address this problem by sampling this space heuristically. In contrast, we introduce a dynamic-programming-based approach to explore more of the search space by iteratively decomposing large program specifications into smaller specifications reachable from a set of rewrites, then composing a final program from each rewrite that minimizes an affine cost model. To reduce memory requirements, we employ a novel memoization table representation, which indexes specifications by coordinates in $Z_{\geq 0}$ and compresses identical, adjacent solutions. This approach can visit a much larger set of programs than prior work. To evaluate the approach, we developed Morello, a compiler which lowers specifications roughly equivalent to a few-node XLA computation graph to x86. Notably, we found that an affine cost model is sufficient to surface high-throughput programs. For example, Morello synthesized a collection of matrix multiplication benchmarks targeting a Zen 1 CPU, including a 1x2048x16384, bfloat16-to-float32 vector-matrix multiply, which was integrated into Google's gemma.cpp.