A Learned Performance Model for Tensor Processing Units

Accurate hardware performance models are critical to efficient code generation. They can be used by compilers to make heuristic decisions, by superoptimizers as a minimization objective, or by autotuners to find an optimal configuration for a specific program. However, they are difficult to develop...

Full description

Saved in:
Bibliographic Details
Published in:arXiv.org 2021-03
Main Authors: Kaufman, Samuel J, Phothilimthana, Phitchaya Mangpo, Zhou, Yanqi, Mendis, Charith, Roy, Sudip, Sabne, Amit, Burrows, Mike
Format: Article
Language:English
Subjects:
Accelerators
Compilers
Cost analysis
Machine learning
Mathematical analysis
Neural networks
Optimization
Tensors
Online Access:Get full text
Tags: Add Tag
No Tags, Be the first to tag this record!
Description
Summary:Accurate hardware performance models are critical to efficient code generation. They can be used by compilers to make heuristic decisions, by superoptimizers as a minimization objective, or by autotuners to find an optimal configuration for a specific program. However, they are difficult to develop because contemporary processors are complex, and the recent proliferation of deep learning accelerators has increased the development burden. We demonstrate a method of learning performance models from a corpus of tensor computation graph programs for Tensor Processing Unit (TPU) instances. We show that our learned model outperforms a heavily-optimized analytical performance model on two tasks -- tile-size selection and operator fusion -- and that it helps an autotuner discover faster programs in a setting where access to TPUs is limited or expensive.
ISSN:2331-8422