TensorIR: An Abstraction for Automatic Tensorized Program Optimization
Deploying deep learning models on various devices has become an important topic. The wave of hardware specialization brings a diverse set of acceleration primitives for multi-dimensional tensor computations. These new acceleration primitives, along with the emerging machine learning models, bring tremendous engineering challenges. In this paper, we present TensorIR, a compiler abstraction for optimizing programs with these tensor computation primitives. TensorIR generalizes the loop nest representation used in existing machine learning compilers to bring tensor computation as the first-class citizen. Finally, we build an end-to-end framework on top of our abstraction to automatically optimize deep learning models for given tensor computation primitives. Experimental results show that TensorIR compilation automatically uses the tensor computation primitives for given hardware backends and delivers performance that is competitive to state-of-art hand-optimized systems across platforms.
Code (2)
Tasks
BIG-bench Machine LearningDeep LearningMethods 이 논문이 사용한 방법론
Similar Papers 제목 키워드 기반
UNIT: Unifying Tensorized Instruction Compilation
Because of the increasing demand for computation in DNN, researchers develope both hardware and software mechanisms to reduce the compute and memory burden. A widely adopted approach is to use mixed precision data types.…
CPUTensor Program Optimization with Probabilistic Programs
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 e…
Deep LearningProbabilistic Programmingminimax: Efficient Baselines for Autocurricula in JAX
Unsupervised environment design (UED) is a form of automatic curriculum learning for training robust decision-making agents to zero-shot transfer into unseen environments. Such autocurricula have received much interest f…
CPUDecision MakingGPUBayesian Tensorized Neural Networks with Automatic Rank Selection
Tensor decomposition is an effective approach to compress over-parameterized neural networks and to enable their deployment on resource-constrained hardware platforms. However, directly applying tensor compression in the…
Model CompressionTensor DecompositionTorchOpt: An Efficient Library for Differentiable Optimization
Recent years have witnessed the booming of various differentiable optimization algorithms. These algorithms exhibit different execution patterns, and their execution needs massive computational resources that go beyond a…
CPUGPU