Learning to Optimize Tensor Programs

2018/05/21 by Tianqi Chen, Lianmin Zheng, Chen, Tianqi +14 · 1 voice · 124 citations
Computer Science · Mathematics · Physics and Astronomy · Psychology · #Computational Physics and Python Applications #Computer science #Mathematics #Model Reduction and Neural Networks #Parallel Computing and Optimization Techniques #Psychology #Pure mathematics #Tensor (intrinsic definition) #cs.LG #stat.ML

paper · pdf · doi:10.48550/arxiv.1805.08166

published in arXiv (Cornell University) (Cornell University) · NeurIPS 2018

openalex publication_date 2018/05/21 · arxiv created 2019/01/08 · arxiv updated 2019/01/10 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

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.

Citations

• An analysis of approximations for maximizing submodular set functions—I
• ImageNet classification with deep convolutional neural networks
• Taking the Human Out of the Loop: A Review of Bayesian Optimization
• MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications

Cited by

• Transition-Aware Backend Dispatch for Edge LLM Inference
• LATTICE: Constraint-Directed Scheduling, Memory Planning, and Pipeline Refinement for NPUs
• QiMeng: Fully Automated Hardware and Software Design for Processor Chip
• MPK: A Compiler and Runtime for Mega-Kernelizing Tensor Programs
• Workload-Aware Autotuning of Block Size in Square-Root Decomposition
• Kernel Forge: An Agent Harness for LLM-based Generation and Optimization of CUDA Kernels
• nncase: An End-to-End Compiler for Efficient LLM Deployment on Heterogeneous Storage Architectures
• XTC, A Research Platform for Optimizing AI Workload Operators
• Theoretical Foundations of GPU-Native Compilation for Rapid Code Iteration
• KernelBand: Steering LLM-based Kernel Optimization via Hardware-Aware Multi-Armed Bandits
• DCC: Data-Centric Compilation of Machine Learning Kernels for Processing-In-Memory Architectures
• A Deep Learning Model for Predicting Transformation Legality
• A Data-driven Analysis of Code Optimizations
• Agentic Auto-Scheduling: An Experimental Study of LLM-Guided Loop Optimization
• CONQuER: Hardware-Aware Mixed-Precision Quantisation with Online-Calibrated Surrogates
• CudaForge: An Agent Framework with Hardware Feedback for CUDA Kernel Optimization
• Automatic Generation of High-Performance Convolution Kernels on ARM CPUs for Deep Learning
• LOOPerSet: A Large-Scale Dataset for Data-Driven Polyhedral Compiler Optimization
• AutoMegaKernel: A Statically-Checked Agent Harness for Self-Retargeting Megakernel Synthesis
• Syncopate: Efficient Multi-GPU AI Kernels via Automatic Chunk-Centric Compute-Communication Overlap
• Prompt-Aware Scheduling for Low-Latency LLM Serving
• LLMulator: Generalizable Cost Modeling for Dataflow Accelerators with Input-Adaptive Control Flow
• Evolution of Kernels: Automated RISC-V Kernel Optimization with Large Language Models
• VoltanaLLM: Energy-Efficient and SLO-Aware Disaggregated LLM Serving via Adaptive Frequency Control and State-Space Routing
• SwizzlePerf: Hardware-Aware LLMs for GPU Kernel Performance Optimization
• Meta-Learning for Speeding Up Large Model Inference in Decentralized Environments
• Relax: Composable Abstractions for End-to-End Dynamic Machine Learning
• Tensor Program Optimization for the RISC-V Vector Extension Using Probabilistic Programs
• MCFuser: High-Performance and Rapid Fusion of Memory-Bound Compute-Intensive Operators
• MLGO: a Machine Learning Guided Compiler Optimizations Framework
• QiMeng-TensorOp: Automatically Generating High-Performance Tensor Operators with Hardware Primitives
• Pearl: Automatic Code Optimization Using Deep Reinforcement Learning
• REASONING COMPILER: LLM-Guided Optimizations for Efficient Model Serving
• A Highly Configurable Hardware/Software Stack for DNN Inference Acceleration
• COGNATE: Acceleration of Sparse Tensor Programs on Emerging Hardware using Transfer Learning
• KPerfIR: Towards an Open and Compiler-centric Ecosystem for GPU Kernel Performance Tooling on Modern AI Workloads
• JIR-Arena: The First Benchmark Dataset for Just-in-time Information Recommendation
• Sandwich: Joint Configuration Search and Hot-Switching for Efficient CPU LLM Serving
• SuperCoder: Assembly Program Superoptimization with Large Language Models
• AgentCompile: An LLM-Guided Compiler for Direct CUDA Inference
• Reinforcement Learning and Adaptive Sampling for Optimized DNN Compilation
• Optimizing Memory Placement using Evolutionary Graph Reinforcement Learning
• ParEVO: Synthesizing Code for Irregular Data: High-Performance Parallelism through Agentic Evolution
• Prism: Symbolic Superoptimization of Tensor Programs
• WaveTune: Wave-aware Bilinear Modeling for Efficient GPU Kernel Auto-tuning
• ComFuse: Fusing Complex Memory-Intensive Subgraphs with Compute-Intensive Kernels For Modern GPU Architectures

Discussions

• Learning to Optimize Tensor Programs [hn, 106 points, 13 comments]

Related

• TLP: A Deep Learning-based Cost Model for Tensor Program Tuning
• Ansor: Generating High-Performance Tensor Programs for Deep Learning
• Gensor: A Graph-based Construction Tensor Compilation Method for Deep Learning
• TensorIR: An Abstraction for Automatic Tensorized Program Optimization
• A Learned Performance Model for Tensor Processing Units