Morello: Compiling Fast Neural Networks with Dynamic Programming and Spatial Compression

TL;DR

This paper introduces Morello, a compiler that uses dynamic programming and spatial compression to efficiently explore a large search space of neural network optimization programs, achieving high-throughput inference on x86 CPUs.

Contribution

It presents a novel dynamic programming approach with a memory-efficient memoization technique for neural network program optimization, enabling exploration of larger search spaces than prior methods.

Findings

Morello successfully synthesized high-throughput matrix multiplication programs.

An affine cost model effectively guides program selection for optimization.

The approach enables practical deployment of optimized neural network kernels.

Abstract

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.