TpuGraphs: A Performance Prediction Dataset on Large Tensor Computational Graphs
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.