# LoopTune: Optimizing Tensor Computations with Reinforcement Learning

arXiv: https://arxiv.org/abs/2309.01825 · HTML: https://mlsanity.com/paper/2309.01825/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 27 of its 71 references that are in the map.
- First posted on arXiv 4 Sep 2023.

## Builds on (most cited first, 27 of 27)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [Asynchronous Methods for Deep Reinforcement Learning](https://mlsanity.com/paper/1602.01783/) (2016)
- [IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures](https://mlsanity.com/paper/1802.01561/) (2018)
- [In-Datacenter Performance Analysis of a Tensor Processing Unit](https://mlsanity.com/paper/1704.04760/) (2017)
- [Distributed Prioritized Experience Replay](https://mlsanity.com/paper/1803.00933/) (2018)
- [cuDNN: Efficient Primitives for Deep Learning](https://mlsanity.com/paper/1410.0759/) (2014)
- [Ansor: Generating High-Performance Tensor Programs for Deep Learning](https://mlsanity.com/paper/2006.06762/) (2020)
- [Learning to Optimize Tensor Programs](https://mlsanity.com/paper/1805.08166/) (2018)
- [Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions](https://mlsanity.com/paper/1802.04730/) (2018)
- [NVIDIA Tensor Core Programmability, Performance & Precision](https://mlsanity.com/paper/1803.04014/) (2018)
- [Glow: Graph Lowering Compiler Techniques for Neural Networks](https://mlsanity.com/paper/1805.00907/) (2018)
- [Dissecting the Graphcore IPU Architecture via Microbenchmarking](https://mlsanity.com/paper/1912.03413/) (2019)
- [Action Space Shaping in Deep Reinforcement Learning](https://mlsanity.com/paper/2004.00980/) (2020)
- [CompilerGym: Robust, Performant Compiler Optimization Environments for AI Research](https://mlsanity.com/paper/2109.08267/) (2021)
- [A Survey on Compiler Autotuning using Machine Learning](https://mlsanity.com/paper/1801.04405/) (2018)
- [MLGO: a Machine Learning Guided Compiler Optimizations Framework](https://mlsanity.com/paper/2101.04808/) (2021)
- [NeuroVectorizer: End-to-End Vectorization with Deep Reinforcement Learning](https://mlsanity.com/paper/1909.13639/) (2019)
- [Chameleon: Adaptive Code Optimization for Expedited Deep Neural Network Compilation](https://mlsanity.com/paper/2001.08743/) (2020)
- [Intel nGraph: An Intermediate Representation, Compiler, and Executor for Deep Learning](https://mlsanity.com/paper/1801.08058/) (2018)
- [Incidence Networks for Geometric Deep Learning](https://mlsanity.com/paper/1905.11460/) (2019)
- [Fast Stencil-Code Computation on a Wafer-Scale Processor](https://mlsanity.com/paper/2010.03660/) (2020)
- [A Reinforcement Learning Environment for Polyhedral Optimizations](https://mlsanity.com/paper/2104.13732/) (2021)
- [High-Performance Tensor Contraction without Transposition](https://mlsanity.com/paper/1607.00291/) (2016)
- [A Tiling Perspective for Register Optimization](https://mlsanity.com/paper/1406.0582/) (2014)
- [LoopStack: a Lightweight Tensor Algebra Compiler Stack](https://mlsanity.com/paper/2205.00618/) (2022)

## Built on by (most cited first, 1 of 1)

- [A Deep Learning Model for Predicting Transformation Legality](https://mlsanity.com/paper/2511.06120/) (2025)

## Abstract (arXiv)

Advanced compiler technology is crucial for enabling machine learning applications to run on novel hardware, but traditional compilers fail to deliver performance, popular auto-tuners have long search times and expert-optimized libraries introduce unsustainable costs. To address this, we developed LoopTune, a deep reinforcement learning compiler that optimizes tensor computations in deep learning models for the CPU. LoopTune optimizes tensor traversal order while using the ultra-fast lightweight code generator LoopNest to perform hardware-specific optimizations. With a novel graph-based representation and action space, LoopTune speeds up LoopNest by 3.2x, generating an order of magnitude faster code than TVM, 2.8x faster than MetaSchedule, and 1.08x faster than AutoTVM, consistently performing at the level of the hand-tuned library Numpy. Moreover, LoopTune tunes code in order of seconds.