mlsanityData to 14 Oct 2026

Tensor Program Optimization with Probabilistic Programs

arXiv 2205.13603 ↗First posted 26 May 2022cs.LGcited by 30

Not on the map. Why?

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.

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

26 of 63 references.

References · 26

  1. 2015Deep Residual Learning for Image Recognitioncited by 52,096
  2. 2018BERT: Pre-training of Deep Bidirectional Transformers for Language Understandingcited by 40,399
  3. 2018MobileNetV2: Inverted Residuals and Linear Bottleneckscited by 4,795
  4. 2022Competition-Level Code Generation with AlphaCodecited by 1,462
  5. 2021Video Swin Transformercited by 1,091
  6. 2018Pyro: Deep Universal Probabilistic Programmingcited by 380
  7. 2019Deep Learning Recommendation Model for Personalization and Recommendation Systemscited by 327
  8. 2014cuDNN: Efficient Primitives for Deep Learningcited by 245
  9. 2020Ansor: Generating High-Performance Tensor Programs for Deep Learningcited by 201
  10. 2019Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyrocited by 191
  11. 2018Learning to Optimize Tensor Programscited by 140
  12. 2015Fast Algorithms for Convolutional Neural Networkscited by 102
14 more
  1. 2018Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractionscited by 85
  2. 2018TVM: An Automated End-to-End Optimizing Compiler for Deep Learningcited by 62
  3. 2018Tiramisu: A Polyhedral Compiler for Expressing Fast and Portable Codecited by 53
  4. 2020Autonomous Driving with Deep Learning: A Survey of State-of-Art Technologiescited by 41
  5. 2020Synthesize, Execute and Debug: Learning to Repair for Neural Program Synthesiscited by 39
  6. 2021Latent Execution for Neural Program Synthesiscited by 35
  7. 2021A Deep Learning Based Cost Model for Automatic Code Optimizationcited by 33
  8. 2018Optimizing CNN Model Inference on CPUscited by 29
  9. 2020Chameleon: Adaptive Code Optimization for Expedited Deep Neural Network Compilationcited by 29
  10. 2018Simple, Distributed, and Accelerated Probabilistic Programmingcited by 18
  11. 2019Divide, Conquer, and Combine: a New Inference Strategy for Probabilistic Programs with Stochastic Supportcited by 11
  12. 2020ProTuner: Tuning Programs with Monte Carlo Tree Searchcited by 9
  13. 2021MetaTune: Meta-Learning Based Cost Model for Fast and Efficient Auto-tuning Frameworkscited by 9
  14. 2015C3: Lightweight Incrementalized MCMC for Probabilistic Programs using Continuations and Callsite Cachingcited by 3

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2022: 2; 2023: 3; 2024: 4; 2025: 15; 2026: 6.

30 papers.

Sort:

By area:

  1. 2023Efficient Large Language Models: A Surveycited by 186
  2. 2022SparseTIR: Composable Abstractions for Sparse Compilation in Deep Learningcited by 47
  3. 2022Hidet: Task-Mapping Programming Paradigm for Deep Learning Tensor Programscited by 28
  4. 2024Allo: A Programming Model for Composable Accelerator Designcited by 16
  5. 2023Relax: Composable Abstractions for End-to-End Dynamic Machine Learningcited by 12
  6. 2024Optimal Kernel Orchestration for Tensor Programs with Korchcited by 7
  7. 2025SuperCoder: Assembly Program Superoptimization with Large Language Modelscited by 6
  8. 2025Tawa: Automatic Warp Specialization for Modern GPUs with Asynchronous Referencescited by 6
  9. 2025QiMeng-TensorOp: Automatically Generating High-Performance Tensor Operators with Hardware Primitivescited by 5
  10. 2026Improving Efficiency of GPU Kernel Optimization Agents using a Domain-Specific Language and Speed-of-Light Guidancecited by 5
  11. 2026CAKE: Compiler-Agent Co-Design for Frontier Kernel Evolutioncited by 5
  12. 2025Neptune: Advanced ML Operator Fusion for Locality and Parallelism on GPUscited by 4
18 more
  1. 2025PyGraph: Robust Compiler Support for CUDA Graphs in PyTorchcited by 3
  2. 2025Tilus: A Tile-Level GPGPU Programming Language for Low-Precision Computationcited by 3
  3. 2025REASONING COMPILER: LLM-Guided Optimizations for Efficient Model Servingcited by 3
  4. 2024Pruner: A Draft-then-Verify Exploration Mechanism to Accelerate Tensor Program Tuningcited by 2
  5. 2025Tensor Program Optimization for the RISC-V Vector Extension Using Probabilistic Programscited by 1
  6. 2025Evolution of Kernels: Automated RISC-V Kernel Optimization with Large Language Modelscited by 1
  7. 2025GraphPerf-RT: A Graph-Driven Performance Model for Hardware-Aware Scheduling of OpenMP Codescited by 1
  8. 2026Record-Remix-Replay: Hierarchical GPU Kernel Optimization using Evolutionary Searchcited by 1
  9. 2026PassNet: Scaling Large Language Models for Graph Compiler Pass Generationcited by 1
  10. 2023DLAS: An Exploration and Assessment of the Deep Learning Acceleration Stackcited by 0
  11. 2024FTuner: A Fast Dynamic Shape Tensors Program Auto-Tuner for Deep Learning Compilerscited by 0
  12. 2025GraphNet: A Large-Scale Computational Graph Dataset for Tensor Compiler Researchcited by 0
  13. 2025PerfDojo: Automated ML Library Generation for Heterogeneous Architecturescited by 0
  14. 2025DCC: Data-Centric Compilation of Machine Learning Kernels for Processing-In-Memory Architecturescited by 0
  15. 2025Accelerated Execution of Bayesian Neural Networks using a Single Probabilistic Forward Pass and Code Generationcited by 0
  16. 2025nncase: An End-to-End Compiler for Efficient LLM Deployment on Heterogeneous Storage Architecturescited by 0
  17. 2026Machine Learning for Energy-Performance-aware Schedulingcited by 0
  18. 2026AsmEvo: Agentic Assembly-Level Optimization of AMD GPU Kernels with Functional Equivalence Verificationcited by 0