# Learning to Optimize Tensor Programs

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

## Key facts

- Cited by 140 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 21 May 2018.

## Built on by (most cited first, 30 of 140)

- [Efficient LLM Scheduling by Learning to Rank](https://mlsanity.com/paper/2408.15792/) (2024)
- [DNNFusion: Accelerating Deep Neural Networks Execution with Advanced Operator Fusion](https://mlsanity.com/paper/2108.13342/) (2021)
- [T-MAC: CPU Renaissance via Table Lookup for Low-Bit LLM Deployment on Edge](https://mlsanity.com/paper/2407.00088/) (2024)
- [SparseTIR: Composable Abstractions for Sparse Compilation in Deep Learning](https://mlsanity.com/paper/2207.04606/) (2022)
- [CudaForge: An Agent Framework with Hardware Feedback for CUDA Kernel Optimization](https://mlsanity.com/paper/2511.01884/) (2025)
- [Mirage: A Multi-Level Superoptimizer for Tensor Programs](https://mlsanity.com/paper/2405.05751/) (2024)
- [MLGO: a Machine Learning Guided Compiler Optimizations Framework](https://mlsanity.com/paper/2101.04808/) (2021)
- [Tiny Machine Learning: Progress and Futures](https://mlsanity.com/paper/2403.19076/) (2024)
- [FLAT: An Optimized Dataflow for Mitigating Attention Bottlenecks](https://mlsanity.com/paper/2107.06419/) (2021)
- [TensorIR: An Abstraction for Automatic Tensorized Program Optimization](https://mlsanity.com/paper/2207.04296/) (2022)
- [A Deep Learning Based Cost Model for Automatic Code Optimization](https://mlsanity.com/paper/2104.04955/) (2021)
- [Tensor Program Optimization with Probabilistic Programs](https://mlsanity.com/paper/2205.13603/) (2022)
- [NNSmith: Generating Diverse and Valid Test Cases for Deep Learning Compilers](https://mlsanity.com/paper/2207.13066/) (2022)
- [A Runtime-Based Computational Performance Predictor for Deep Neural Network Training](https://mlsanity.com/paper/2102.00527/) (2021)
- [Hidet: Task-Mapping Programming Paradigm for Deep Learning Tensor Programs](https://mlsanity.com/paper/2210.09603/) (2022)
- [Flashlight: Enabling Innovation in Tools for Machine Learning](https://mlsanity.com/paper/2201.12465/) (2022)
- [NAAS: Neural Accelerator Architecture Search](https://mlsanity.com/paper/2105.13258/) (2021)
- [Bolt: Bridging the Gap between Auto-tuners and Hardware-native Performance](https://mlsanity.com/paper/2110.15238/) (2021)
- [The CoRa Tensor Compiler: Compilation for Ragged Tensors with Minimal Padding](https://mlsanity.com/paper/2110.10221/) (2021)
- [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)
- [Allo: A Programming Model for Composable Accelerator Design](https://mlsanity.com/paper/2404.04815/) (2024)
- [MPK: A Compiler and Runtime for Mega-Kernelizing Tensor Programs](https://mlsanity.com/paper/2512.22219/) (2025)
- [TLP: A Deep Learning-based Cost Model for Tensor Program Tuning](https://mlsanity.com/paper/2211.03578/) (2022)
- [KernelFoundry: Hardware-aware evolutionary GPU kernel optimization](https://mlsanity.com/paper/2603.12440/) (2026)
- [SECDA: Efficient Hardware/Software Co-Design of FPGA-based DNN Accelerators for Edge Inference](https://mlsanity.com/paper/2110.00478/) (2021)
- [Incremental Randomized Smoothing Certification](https://mlsanity.com/paper/2305.19521/) (2023)
- [Relax: Composable Abstractions for End-to-End Dynamic Machine Learning](https://mlsanity.com/paper/2311.02103/) (2023)
- [Slapo: A Schedule Language for Progressive Optimization of Large Deep Learning Model Training](https://mlsanity.com/paper/2302.08005/) (2023)
- [SwizzlePerf: Hardware-Aware LLMs for GPU Kernel Performance Optimization](https://mlsanity.com/paper/2508.20258/) (2025)

## Abstract (arXiv)

We introduce a learning-based framework to optimize tensor programs for deep learning workloads. Efficient implementations of tensor operators, such as matrix multiplication and high dimensional convolution, are key enablers of effective deep learning systems. However, existing systems rely on manually optimized libraries such as cuDNN where only a narrow range of server class GPUs are well-supported. The reliance on hardware-specific operator libraries limits the applicability of high-level graph optimizations and incurs significant engineering costs when deploying to new hardware targets. We use learning to remove this engineering burden. We learn domain-specific statistical cost models to guide the search of tensor operator implementations over billions of possible program variants. We further accelerate the search by effective model transfer across workloads. Experimental results show that our framework delivers performance competitive with state-of-the-art hand-tuned libraries for low-power CPU, mobile GPU, and server-class GPU.