# A Learned Performance Model for Tensor Processing Units

arXiv: https://arxiv.org/abs/2008.01040 · HTML: https://mlsanity.com/paper/2008.01040/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 3 Aug 2020.

## Built on by (most cited first, 2 of 2)

- [A Full-Stack Search Technique for Domain Optimized Deep Learning Accelerators](https://mlsanity.com/paper/2105.12842/) (2021)
- [MetaTune: Meta-Learning Based Cost Model for Fast and Efficient Auto-tuning Frameworks](https://mlsanity.com/paper/2102.04199/) (2021)

## Abstract (arXiv)

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.