TpuGraphs: A Performance Prediction Dataset on Large Tensor Computational Graphs
Phitchaya Mangpo Phothilimthana, Sami Abu-El-Haija, Kaidi Cao, Bahare, Fatemi, Mike Burrows, Charith Mendis, Bryan Perozzi

TL;DR
TpuGraphs is a large, detailed dataset of tensor program computational graphs on TPUs, designed to improve performance prediction models for optimizing machine learning workloads.
Contribution
The paper introduces TpuGraphs, a comprehensive dataset of full tensor program graphs on TPUs, enabling advanced performance prediction research for large-scale machine learning workloads.
Findings
Provides 25x more graphs than existing datasets.
Includes 770x larger graphs on average.
Highlights new challenges in learning from large graphs.
Abstract
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…
Peer Reviews
No public reviews on file for this paper yet. If you reviewed it on a platform where reviews are public (OpenReview, ICLR, NeurIPS, ICML), you can paste yours below so the community can read it here.
Code & Models
Videos
Taxonomy
TopicsParallel Computing and Optimization Techniques · Software System Performance and Reliability · Software Engineering Research
Methods*Communicated@Fast*How Do I Communicate to Expedia? · Multi-Head Attention · Attention Is All You Need · Depthwise Convolution · Pointwise Convolution · Depthwise Separable Convolution · RoIAlign · Region Proposal Network · Batch Normalization · Absolute Position Encodings
