# Tensor Program Optimization with Probabilistic Programs

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

## Key facts

- Cited by 30 papers in the map (data through 14 Oct 2026).
- Builds on 26 of its 63 references that are in the map.
- First posted on arXiv 26 May 2022.

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

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [MobileNetV2: Inverted Residuals and Linear Bottlenecks](https://mlsanity.com/paper/1801.04381/) (2018)
- [Competition-Level Code Generation with AlphaCode](https://mlsanity.com/paper/2203.07814/) (2022)
- [Video Swin Transformer](https://mlsanity.com/paper/2106.13230/) (2021)
- [Pyro: Deep Universal Probabilistic Programming](https://mlsanity.com/paper/1810.09538/) (2018)
- [Deep Learning Recommendation Model for Personalization and Recommendation Systems](https://mlsanity.com/paper/1906.00091/) (2019)
- [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)
- [Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyro](https://mlsanity.com/paper/1912.11554/) (2019)
- [Learning to Optimize Tensor Programs](https://mlsanity.com/paper/1805.08166/) (2018)
- [Fast Algorithms for Convolutional Neural Networks](https://mlsanity.com/paper/1509.09308/) (2015)
- [Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions](https://mlsanity.com/paper/1802.04730/) (2018)
- [TVM: An Automated End-to-End Optimizing Compiler for Deep Learning](https://mlsanity.com/paper/1802.04799/) (2018)
- [Tiramisu: A Polyhedral Compiler for Expressing Fast and Portable Code](https://mlsanity.com/paper/1804.10694/) (2018)
- [Autonomous Driving with Deep Learning: A Survey of State-of-Art Technologies](https://mlsanity.com/paper/2006.06091/) (2020)
- [Synthesize, Execute and Debug: Learning to Repair for Neural Program Synthesis](https://mlsanity.com/paper/2007.08095/) (2020)
- [Latent Execution for Neural Program Synthesis](https://mlsanity.com/paper/2107.00101/) (2021)
- [A Deep Learning Based Cost Model for Automatic Code Optimization](https://mlsanity.com/paper/2104.04955/) (2021)
- [Optimizing CNN Model Inference on CPUs](https://mlsanity.com/paper/1809.02697/) (2018)
- [Chameleon: Adaptive Code Optimization for Expedited Deep Neural Network Compilation](https://mlsanity.com/paper/2001.08743/) (2020)
- [Simple, Distributed, and Accelerated Probabilistic Programming](https://mlsanity.com/paper/1811.02091/) (2018)
- [Divide, Conquer, and Combine: a New Inference Strategy for Probabilistic Programs with Stochastic Support](https://mlsanity.com/paper/1910.13324/) (2019)
- [ProTuner: Tuning Programs with Monte Carlo Tree Search](https://mlsanity.com/paper/2005.13685/) (2020)
- [MetaTune: Meta-Learning Based Cost Model for Fast and Efficient Auto-tuning Frameworks](https://mlsanity.com/paper/2102.04199/) (2021)
- [C3: Lightweight Incrementalized MCMC for Probabilistic Programs using Continuations and Callsite Caching](https://mlsanity.com/paper/1509.02151/) (2015)

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

- [Efficient Large Language Models: A Survey](https://mlsanity.com/paper/2312.03863/) (2023)
- [SparseTIR: Composable Abstractions for Sparse Compilation in Deep Learning](https://mlsanity.com/paper/2207.04606/) (2022)
- [Hidet: Task-Mapping Programming Paradigm for Deep Learning Tensor Programs](https://mlsanity.com/paper/2210.09603/) (2022)
- [Allo: A Programming Model for Composable Accelerator Design](https://mlsanity.com/paper/2404.04815/) (2024)
- [Relax: Composable Abstractions for End-to-End Dynamic Machine Learning](https://mlsanity.com/paper/2311.02103/) (2023)
- [Optimal Kernel Orchestration for Tensor Programs with Korch](https://mlsanity.com/paper/2406.09465/) (2024)
- [SuperCoder: Assembly Program Superoptimization with Large Language Models](https://mlsanity.com/paper/2505.11480/) (2025)
- [Tawa: Automatic Warp Specialization for Modern GPUs with Asynchronous References](https://mlsanity.com/paper/2510.14719/) (2025)
- [QiMeng-TensorOp: Automatically Generating High-Performance Tensor Operators with Hardware Primitives](https://mlsanity.com/paper/2505.06302/) (2025)
- [Improving Efficiency of GPU Kernel Optimization Agents using a Domain-Specific Language and Speed-of-Light Guidance](https://mlsanity.com/paper/2603.29010/) (2026)
- [CAKE: Compiler-Agent Co-Design for Frontier Kernel Evolution](https://mlsanity.com/paper/2608.12629/) (2026)
- [Neptune: Advanced ML Operator Fusion for Locality and Parallelism on GPUs](https://mlsanity.com/paper/2510.08726/) (2025)
- [PyGraph: Robust Compiler Support for CUDA Graphs in PyTorch](https://mlsanity.com/paper/2503.19779/) (2025)
- [Tilus: A Tile-Level GPGPU Programming Language for Low-Precision Computation](https://mlsanity.com/paper/2504.12984/) (2025)
- [REASONING COMPILER: LLM-Guided Optimizations for Efficient Model Serving](https://mlsanity.com/paper/2506.01374/) (2025)
- [Pruner: A Draft-then-Verify Exploration Mechanism to Accelerate Tensor Program Tuning](https://mlsanity.com/paper/2402.02361/) (2024)
- [Tensor Program Optimization for the RISC-V Vector Extension Using Probabilistic Programs](https://mlsanity.com/paper/2507.01457/) (2025)
- [Evolution of Kernels: Automated RISC-V Kernel Optimization with Large Language Models](https://mlsanity.com/paper/2509.14265/) (2025)
- [GraphPerf-RT: A Graph-Driven Performance Model for Hardware-Aware Scheduling of OpenMP Codes](https://mlsanity.com/paper/2512.12091/) (2025)
- [Record-Remix-Replay: Hierarchical GPU Kernel Optimization using Evolutionary Search](https://mlsanity.com/paper/2604.11109/) (2026)
- [PassNet: Scaling Large Language Models for Graph Compiler Pass Generation](https://mlsanity.com/paper/2605.29357/) (2026)
- [DLAS: An Exploration and Assessment of the Deep Learning Acceleration Stack](https://mlsanity.com/paper/2311.08909/) (2023)
- [FTuner: A Fast Dynamic Shape Tensors Program Auto-Tuner for Deep Learning Compilers](https://mlsanity.com/paper/2407.21418/) (2024)
- [GraphNet: A Large-Scale Computational Graph Dataset for Tensor Compiler Research](https://mlsanity.com/paper/2510.24035/) (2025)
- [PerfDojo: Automated ML Library Generation for Heterogeneous Architectures](https://mlsanity.com/paper/2511.03586/) (2025)
- [DCC: Data-Centric Compilation of Machine Learning Kernels for Processing-In-Memory Architectures](https://mlsanity.com/paper/2511.15503/) (2025)
- [Accelerated Execution of Bayesian Neural Networks using a Single Probabilistic Forward Pass and Code Generation](https://mlsanity.com/paper/2511.23440/) (2025)
- [nncase: An End-to-End Compiler for Efficient LLM Deployment on Heterogeneous Storage Architectures](https://mlsanity.com/paper/2512.21571/) (2025)
- [Machine Learning for Energy-Performance-aware Scheduling](https://mlsanity.com/paper/2601.23134/) (2026)
- [AsmEvo: Agentic Assembly-Level Optimization of AMD GPU Kernels with Functional Equivalence Verification](https://mlsanity.com/paper/2608.20711/) (2026)

## Abstract (arXiv)

Automatic optimization for tensor programs becomes increasingly important as we deploy deep learning in various environments, and efficient optimization relies on a rich search space and effective search. Most existing efforts adopt a search space which lacks the ability to efficiently enable domain experts to grow the search space. This paper introduces MetaSchedule, a domain-specific probabilistic programming language abstraction to construct a rich search space of tensor programs. Our abstraction allows domain experts to analyze the program, and easily propose stochastic choices in a modular way to compose program transformation accordingly. We also build an end-to-end learning-driven framework to find an optimized program for a given search space. Experimental results show that MetaSchedule can cover the search space used in the state-of-the-art tensor program optimization frameworks in a modular way. Additionally, it empowers domain experts to conveniently grow the search space and modularly enhance the system, which brings 48% speedup on end-to-end deep learning workloads.