Nova: An End-to-End MLIR Compiler for Deep Learning

Abstract:The performance of deep learning models at scale relies heavily on how effectively high-level mathematical operations are mapped to underlying physical hardware. While high-level tensor frameworks provide flexible abstractions for model design, their eager execution models inherently lack the whole-graph visibility and granular control over hardware and memory required to maximize physical hardware utilization natively. To bridge this gap, we designed Nova, an automated end-to-end JIT compiler whose defining purpose is to achieve absolute control over this hardware mapping: fusing operations across operation boundaries, optimizing complex memory hierarchies, and tuning execution down to the register level. By capturing eager executions and unifying forward and backward passes into a single value-semantic dialect, Nova unlocks aggressive whole-graph optimizations. It then utilizes an Analytic Configurator to deterministically derive optimal execution schedules based on arithmetic intensity, dropping search time to zero. Backed by a structural hashing runtime, Nova synthesizes fine-grained kernels directly from the computation's structure. In our evaluations on an RTX 3060, Nova matches or modestly exceeds cuBLAS and XLA on TF32 matmuls on most shapes, maintaining a stringent < 5e-4 relative error. At the model level, Nova achieves up to 10.6% greater throughput than PyTorch and 4.4% greater than XLA on a 42-million parameter model, without compromising on numerical fidelity. Crucially, by reducing the memory footprint by up to 29% relative to PyTorch, Nova successfully trains a 144-million parameter model at 17,900 tokens/s where PyTorch encounters Out-Of-Memory (OOM) failures on the same 12 GB consumer GPU.

Artificial Intelligence (cs.AI); Hardware Architecture (cs.AR); Machine Learning (cs.LG); Programming Languages (cs.PL)

arXiv:2608.00029 [cs.AI]

(or arXiv:2608.00029v1 [cs.AI] for this version)

https://doi.org/10.48550/arXiv.2608.00029 Focus to learn more arXiv-issued DOI via DataCite

arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.

Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.

Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.