LoopStack: a Lightweight Tensor Algebra Compiler Stack
This work addresses the need for efficient and fast compilation of neural networks, particularly for mobile and embedded devices, though it appears incremental as it builds on existing compiler techniques.
The authors tackled the problem of compiling tensor operations for neural networks by introducing LoopStack, a lightweight compiler stack that generates efficient machine code for various instruction sets, achieving performance that matches or exceeds state-of-the-art frameworks and offering significantly faster compilation times than LLVM.
We present LoopStack, a domain specific compiler stack for tensor operations, composed of a frontend, LoopTool, and an efficient optimizing code generator, LoopNest. This stack enables us to compile entire neural networks and generate code targeting the AVX2, AVX512, NEON, and NEONfp16 instruction sets while incorporating optimizations often missing from other machine learning compiler backends. We evaluate our stack on a collection of full neural networks and commonly used network blocks as well as individual operators, and show that LoopStack generates machine code that matches and frequently exceeds the performance of in state-of-the-art machine learning frameworks in both cases. We also show that for a large collection of schedules LoopNest's compilation is orders of magnitude faster than LLVM, while resulting in equal or improved run time performance. Additionally, LoopStack has a very small memory footprint - a binary size of 245KB, and under 30K lines of effective code makes it ideal for use on mobile and embedded devices.