# A Deep Learning Based Cost Model for Automatic Code Optimization arXiv: https://arxiv.org/abs/2104.04955 ยท HTML: https://mlsanity.com/paper/2104.04955/ ## Key facts - Cited by 33 papers in the map (data through 14 Oct 2026). - Builds on 10 of its 36 references that are in the map. - First posted on arXiv 11 Apr 2021. ## Builds on (most cited first, 10 of 10) - [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019) - [XLNet: Generalized Autoregressive Pretraining for Language Understanding](https://mlsanity.com/paper/1906.08237/) (2019) - [Decoupled Weight Decay Regularization](https://mlsanity.com/paper/1711.05101/) (2017) - [Super-Convergence: Very Fast Training of Neural Networks Using Large Learning Rates](https://mlsanity.com/paper/1708.07120/) (2017) - [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) - [Tiramisu: A Polyhedral Compiler for Expressing Fast and Portable Code](https://mlsanity.com/paper/1804.10694/) (2018) - [Ithemal: Accurate, Portable and Fast Basic Block Throughput Estimation using Deep Neural Networks](https://mlsanity.com/paper/1808.07412/) (2018) - [TIRAMISU: A Polyhedral Compiler for Dense and Sparse Deep Learning](https://mlsanity.com/paper/2005.04091/) (2020) - [PENCIL: Towards a Platform-Neutral Compute Intermediate Language for DSLs](https://mlsanity.com/paper/1302.5586/) (2013) ## Built on by (most cited first, 30 of 33) - [Tensor Program Optimization with Probabilistic Programs](https://mlsanity.com/paper/2205.13603/) (2022) - [TpuGraphs: A Performance Prediction Dataset on Large Tensor Computational Graphs](https://mlsanity.com/paper/2308.13490/) (2023) - [Supersonic: Learning to Generate Source Code Optimizations in C/C++](https://mlsanity.com/paper/2309.14846/) (2023) - [TLP: A Deep Learning-based Cost Model for Tensor Program Tuning](https://mlsanity.com/paper/2211.03578/) (2022) - [LayerDAG: A Layerwise Autoregressive Diffusion Model for Directed Acyclic Graph Generation](https://mlsanity.com/paper/2411.02322/) (2024) - [LOOPer: A Learned Automatic Code Optimizer For Polyhedral Compilers](https://mlsanity.com/paper/2403.11522/) (2024) - [A Reinforcement Learning Environment for Polyhedral Optimizations](https://mlsanity.com/paper/2104.13732/) (2021) - [Transfer-Tuning: Reusing Auto-Schedules for Efficient Tensor Program Code Generation](https://mlsanity.com/paper/2201.05587/) (2022) - [Agentic Auto-Scheduling: An Experimental Study of LLM-Guided Loop Optimization](https://mlsanity.com/paper/2511.00592/) (2025) - [Moses: Efficient Exploitation of Cross-device Transferable Features for Tensor Program Optimization](https://mlsanity.com/paper/2201.05752/) (2022) - [Application Performance Modeling via Tensor Completion](https://mlsanity.com/paper/2210.10184/) (2022) - [CDMPP: A Device-Model Agnostic Framework for Latency Prediction of Tensor Programs](https://mlsanity.com/paper/2311.09690/) (2023) - [PyGim: An Efficient Graph Neural Network Library for Real Processing-In-Memory Architectures](https://mlsanity.com/paper/2402.16731/) (2024) - [A Reinforcement Learning Environment for Automatic Code Optimization in the MLIR Compiler](https://mlsanity.com/paper/2409.11068/) (2024) - [Using Graph Neural Networks to model the performance of Deep Neural Networks](https://mlsanity.com/paper/2108.12489/) (2021) - [Pruner: A Draft-then-Verify Exploration Mechanism to Accelerate Tensor Program Tuning](https://mlsanity.com/paper/2402.02361/) (2024) - [AGO: Boosting Mobile AI Inference Performance by Removing Constraints on Graph Optimization](https://mlsanity.com/paper/2212.01005/) (2022) - [COMET: Neural Cost Model Explanation Framework](https://mlsanity.com/paper/2302.06836/) (2023) - [Combining Neural Architecture Search and Automatic Code Optimization: A Survey](https://mlsanity.com/paper/2408.04116/) (2024) - [Tadashi: Enabling AI-Based Automated Code Generation With Guaranteed Correctness](https://mlsanity.com/paper/2410.03210/) (2024) - [LOOPerSet: A Large-Scale Dataset for Data-Driven Polyhedral Compiler Optimization](https://mlsanity.com/paper/2510.10209/) (2025) - [LOOPRAG: Enhancing Loop Transformation Optimization with Retrieval-Augmented Large Language Models](https://mlsanity.com/paper/2512.15766/) (2025) - [CHEHAB RL: Learning to Optimize Fully Homomorphic Encryption Computations](https://mlsanity.com/paper/2601.19367/) (2026) - [FamilySeer: Towards Optimized Tensor Codes by Exploiting Computation Subgraph Similarity](https://mlsanity.com/paper/2201.00194/) (2022) - [SONAR: Joint Architecture and System Optimization Search](https://mlsanity.com/paper/2208.12218/) (2022) - [ALT: Boosting Deep Learning Performance by Breaking the Wall between Graph and Operator Level Optimizations](https://mlsanity.com/paper/2210.12415/) (2022) - [ML-driven Hardware Cost Model for MLIR](https://mlsanity.com/paper/2302.11405/) (2023) - [Performance Embeddings: A Similarity-based Approach to Automatic Performance Optimization](https://mlsanity.com/paper/2303.08142/) (2023) - [PolyTOPS: Reconfigurable and Flexible Polyhedral Scheduler](https://mlsanity.com/paper/2401.06665/) (2024) - [Data-efficient Performance Modeling via Pre-training](https://mlsanity.com/paper/2501.14438/) (2025) ## Abstract (arXiv) Enabling compilers to automatically optimize code has been a longstanding goal for the compiler community. Efficiently solving this problem requires using precise cost models. These models predict whether applying a sequence of code transformations reduces the execution time of the program. Building an analytical cost model to do so is hard in modern x86 architectures due to the complexity of the microarchitecture. In this paper, we present a novel deep learning based cost model for automatic code optimization. This model was integrated in a search method and implemented in the Tiramisu compiler to select the best code transformations. The input of the proposed model is a set of simple features representing the unoptimized code and a sequence of code transformations. The model predicts the speedup expected when the code transformations are applied. Unlike previous models, the proposed one works on full programs and does not rely on any heavy feature engineering. The proposed model has only 16% of mean absolute percentage error in predicting speedups on full programs. The proposed model enables Tiramisu to automatically find code transformations that match or are better than state-of-the-art compilers without requiring the same level of heavy feature engineering required by those compilers.