Learning to Optimize Tensor Programs

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.

arXiv commentNeurIPS 2018

References in corpus (8)

• Sequence to Sequence Learning with Neural Networks
• MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications
• Unsupervised Representation Learning with Deep Convolutional Generative Adversarial Networks
• Practical Bayesian Optimization of Machine Learning Algorithms
• MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems
• Theano: new features and speed improvements
• Scalable Bayesian Optimization Using Deep Neural Networks
• Device Placement Optimization with Reinforcement Learning

Cited by in corpus (38)

• MCUNet: Tiny Deep Learning on IoT Devices
• Tiny Machine Learning: Progress and Futures
• Learning to Perform Local Rewriting for Combinatorial Optimization
• Ansor: Generating High-Performance Tensor Programs for Deep Learning
• Optimizing CNN Model Inference on CPUs
• A Learned Performance Model for Tensor Processing Units
• Reinforced Genetic Algorithm Learning for Optimizing Computation Graphs
• MLGO: a Machine Learning Guided Compiler Optimizations Framework
• UNIT: Unifying Tensorized Instruction Compilation
• Pure Tensor Program Rewriting via Access Patterns (Representation Pearl)
• Nimble: Efficiently Compiling Dynamic Neural Networks for Model Inference
• NeurObfuscator: A Full-stack Obfuscation Tool to Mitigate Neural Architecture Stealing
• Supersonic: Learning to Generate Source Code Optimizations in C/C++
• DNNFusion: Accelerating Deep Neural Networks Execution with Advanced Operator Fusion
• Reinforcement Learning and Adaptive Sampling for Optimized DNN Compilation
• Relax: Composable Abstractions for End-to-End Dynamic Machine Learning
• Optimal Kernel Orchestration for Tensor Programs with Korch
• Cortex: A Compiler for Recursive Deep Learning Models
• A Multi-level Compiler Backend for Accelerated Micro-kernels Targeting RISC-V ISA Extensions
• Integration of a systolic array based hardware accelerator into a DNN operator auto-tuning framework
• Bring Your Own Codegen to Deep Learning Compiler
• Automating Generation of Low Precision Deep Learning Operators
• A Highly Configurable Hardware/Software Stack for DNN Inference Acceleration
• Compiler-Level Matrix Multiplication Optimization for Deep Learning
• The CoRa Tensor Compiler: Compilation for Ragged Tensors with Minimal Padding
• A Runtime-Based Computational Performance Predictor for Deep Neural Network Training
• Simulating Execution Time of Tensor Programs using Graph Neural Networks
• HASCO: Towards Agile HArdware and Software CO-design for Tensor Computation
• PolyDL: Polyhedral Optimizations for Creation of High Performance DL primitives
• AI Powered Compiler Techniques for DL Code Optimization
• Search for Optimal Systolic Arrays: A Comprehensive Automated Exploration Framework and Lessons Learned
• Optimizing Memory Placement using Evolutionary Graph Reinforcement Learning
• Learning to Accelerate Decomposition for Multi-Directional 3D Printing
• OpEvo: An Evolutionary Method for Tensor Operator Optimization
• Transferable Graph Optimizers for ML Compilers
• ISA Mapper: A Compute and Hardware Agnostic Deep Learning Compiler
• Self-Adaptive Reconfigurable Arrays (SARA): Using ML to Assist Scaling GEMM Acceleration
• Learning to Superoptimize Real-world Programs