A Learned Performance Model for Tensor Processing Units

TL;DR

This paper introduces a learned performance model for Tensor Processing Units that outperforms traditional analytical models in key tasks and aids autotuning, especially when TPU access is limited or costly.

Contribution

The paper presents a novel machine learning-based performance model for TPUs, improving over existing analytical models in accuracy and utility for autotuning.

Findings

Learned model outperforms analytical models in tile-size selection.

Model improves operator fusion decisions for TPUs.

Assists autotuners in discovering faster programs under limited TPU access.

Abstract

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.