Joint Optimization of Deep Neural Networks with Dynamic Programming and Spatial Compression

Abstract

Implementing high-performance deep neural network inference requires deciding how to fuse operations, how to tile loops, where to place data and when to move it, how to pack data in buffers, and which instructions and microkernels to select, among other decisions. These choices interact; for instance, increasing a tile size might amortize the cost of packing data, or fusing operators might adjust a trade-off between recomputation and storage. Because these choices interact, finding an optimal implementation requires selecting from the Cartesian product of possible decisions, a space that is typically too large to enumerate exhaustively. Compilers and auto-schedulers manage this combinatorial explosion by exploring only a small fraction of the search space, using some combination of fixed transformation sequences, heuristic search, and computation graph partitioning where sub-graphs are optimized separately. In contrast, this dissertation develops a novel method of exhaustively searching spaces of implementations. The approach has two parts. First, it formulates the optimization problem as one with optimal substructure by developing a compositional intermediate representation and corresponding cost model. Optimization proceeds by recursively rewriting program specifications into partial programs whose leaves may be smaller specifications. Optimal implementations of those sub-specifications compose into an optimal implementation of the parent specification. Second, the approach improve the space efficiency of memoizing optimal implementations with a spatial database that maps specifications to integer coordinates and compresses adjacent, identical optima. We evaluate this approach by developing the Morello compiler. Using Morello, we synthesize and benchmark implementations of matrix multiplication and softmax for two x86 CPUs. The generated code is competitive with, and sometimes faster than, mature baselines. We also evaluate Morello's spatial memoization database, showing that it reduces storage requirements by several orders of magnitude relative to storing each specification's optimum independently.