Precise hardware performance models play a crucial role in code
optimizations. They can assist compilers in making heuristic decisions or aid
autotuners in identifying the optimal configuration for a given program. For
example, the autotuner for XLA, a machine learning compiler, discovered 10-20%
speedup on state-of-the-art models serving substantial production traffic at
Google. Although there exist a few datasets for program performance prediction,
they target small sub-programs such as basic blocks or kernels. This paper
introduces TpuGraphs, a performance prediction dataset on full tensor programs,
represented as computational graphs, running on Tensor Processing Units (TPUs).
Each graph in the dataset represents the main computation of a machine learning
workload, e.g., a training epoch or an inference step. Each data sample
contains a computational graph, a compilation configuration, and the execution
time of the graph when compiled with the configuration. The graphs in the
dataset are collected from open-source machine learning programs, featuring
popular model architectures, e.g., ResNet, EfficientNet, Mask R-CNN, and
Transformer. TpuGraphs provides 25x more graphs than the largest graph property
prediction dataset (with comparable graph sizes), and 770x larger graphs on
average compared to existing performance prediction datasets on machine
learning programs. This graph-level prediction task on large graphs introduces
new challenges in learning, ranging from scalability, training efficiency, to
model quality.

该论文介绍了基于 Tensor Processing Units（TPUs）上运行的完整张量程序（表示为计算图）的性能预测数据集 TpuGraphs，该数据集提供了大规模图的图级预测任务，涉及规模可比较的图属性预测数据集的 25 倍的图，以及平均比现有的机器学习程序性能预测数据集大 770 倍的大型图，从而带来了从可伸缩性、训练效率到模型质量等方面的新挑战。