SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Phitchaya Mangpo Phothilimthana Google
Sami Abu-El-Haija Google
Kaidi Cao* Stanford
Bahare Fatemi Google
Mike Burrows Google
Charith Mendis* UIUC
Bryan Perozzi Google
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, Efficient-Net, Mask R-CNN, and Transformer. TPUGRAPHs provides  $25\mathrm{x}$  more graphs than the largest graph property prediction dataset (with comparable graph sizes), and  $770\mathrm{x}$  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.
Compilers often use performance models to solve optimization problems [28, 48], as collecting performance measurements from real hardware can be expensive, limited, or infeasible. A performance model can also be used by a compiler autotuner to evaluate candidate configurations in a search space [2, 14, 37, 53, 56]. However, developing an accurate analytical model of program performance on a modern processor is challenging and time-consuming because the underlying processor architecture, the compiler, and their interactions are complex and difficult to model analytically.
Many recent methods [2, 14, 41, 51, 66, 65, 5, 79, 45, 3, 24, 37] apply machine learning (ML) to learn performance prediction models. However, there exist only a few datasets for program performance prediction, and they all target small sub-programs. BHive [15] targets small basic blocks of assembly instructions. TenSet [81] targets ML kernels consisting of a small number of tensor operations. The database-query dataset [31] contains larger query programs, but they are still relatively small, most with fewer than 100 nodes.
*Work partially done during internship/visiting researcher term at Google.
Unlike prior datasets, TPUGRAPHs is a performance prediction dataset on full tensor programs, represented as computational graphs. Each graph represents the main computation of an ML program, which is usually one or many training steps or one inference step. The graphs in the dataset are collected from open-source ML programs, featuring popular models (e.g., ResNet, EfficientNet, Mask R-CNN, and a large variety of Transformer) for a wide range of tasks, e.g., vision, NLP, speech, audio, recommendation, and generative AI. Each data sample contains a computational graph, a compilation configuration, and the execution time when executing the graph when compiled with the given configuration on a Tensor Processing Unit (TPU) v3 [39], an accelerator for ML workloads. A compilation configuration controls how the XLA compiler [70] transforms the graph for a specific optimization pass. In particular, the TPUGRAPHs dataset consists of two collections: (i) layout and (ii) tile. Layout configurations control how tensors are laid out in the physical memory, by specifying the dimension order of each input and output of an operation node. A tile configuration controls the tile size of each fused subgraph. We primarily focus on layout and tile configurations because tuning them offers the highest performance gain on average, compared to tuning other compiler optimizations.
The layout collection contains 31 million pairs of graphs and configurations, averaging over 7,700 nodes per graph. The tile collection contains 13 millions pairs of kernels and configurations, averaging 40 nodes per kernel subgraph. The layout collection is unique among existing graph datasets, in that it provides data for graph-level predictions on very large graphs. In contrast, most of the existing graph datasets fall into two categories: graph-level prediction on small graphs [11, 72, 32, 4, 32, 14, 81, 67, 83], and node-level or edge-level prediction on large graphs [29, 12, 34, 77, 86, 9, 49]. TPUGRAPHs provides  $25\mathrm{x}$  more graphs than MalNet [27] — the largest graph property prediction dataset with comparable graph sizes — and  $770\mathrm{x}$  larger graphs on average compared to TenSet [80] — the only existing large-sclae ML program performance dataset — as depicted in Figure 1. The scale of TPUGRAPHs poses several new research challenges:
Figure 1: Scale of TPUGRAPHs compared to other graph property prediction datasets.
- How to train a neural network model that can perform graph-level predictions when the memory required to train the model on a single graph may not fit on a single device?
- How to make a model generalize well to unseen graphs when they are diverse, and the training data may be imbalanced?
- How to improve the efficiency of a training pipeline when multiple data points contain a large amount of redundant data (same core graph but different graph configurations)?
We provide baseline model code $^2$  based on a Graph Neural Network (GNN) [11], following the techniques from the most recent works on TPU learned cost models [41, 10]. The baseline models achieve moderate performance on both layout and tile collections. For competitive baselines, we encourage the reader to follow the Kaggle competition $^3$  [54].
2 Background & Challenges
ML compilers solve multiple optimization problems to translate an ML program, typically represented as a tensor computation graph, to an efficient executable for a hardware target. Recent works have demonstrated that search-based autotuning techniques can be used to generate code with close to optimal performance [13, 82, 2, 3, 71, 36, 65, 56]. However, autotuning requires a relatively large
$^{2}$ https://github.com/google-research-datasets/tpu_graphs
<sup>3</sup>https://www.kaggle.com/competitions/predict-ai-model-routine
Figure 2: Important optimizations in ML compilers include graph-level and kernel-level optimizations. A graph-level optimization requires the context of the entire graph to make optimal decisions and transforms the entire graph accordingly. A kernel-level optimization transforms each kernel (a fused subgraph) at a time, independently of other kernels.
amount of resources to find quality candidates compared to traditional heuristics-based compilers. Therefore, many methods develop a learned cost model to accelerate autotuning [14, 41, 65, 79, 45, 3].
2.1 XLA and Autotuner
XLA [70] is a production-grade heuristics-based compiler for ML programs, capable of generating code for various hardware targets, including CPUs, GPUs, and notably TPUs [38, 39]. Figure 2 depicts important optimizations that are featured in XLA and most ML compilers. Graph-level optimizations require the context of the entire program graph to make good decisions, while kernel-level optimizations can be done independently within each kernel. A tensor computation graph is represented as High Level Operations (HLO) in XLA. Each optimization pass transforms an HLO graph into a functionally-equivalent one. The output of graph-level optimizations are a collection of kernels (represented as fused subgraphs). XLA has an accompanying autotuner [56] that can tune both graph-level and kernel-level configurations for TPUs, unlike most search-based compilers [13, 14, 79, 82, 2, 3, 71, 45, 3], which focus on kernel-level optimizations.
Kernel-Level Optimizations. Each node in a tensor computation graph represents a tensor operation, such as matrix multiplication, convolution, element-wise addition, etc. A kernel, represented as a fused subgraph, is then a fusion of multiple tensor operations. For example, Convolution-BatchNorm is a common fused kernel that appears in Convolutional Neural Networks. The most important optimization at the kernel level is tile size selection: selecting the shape of a tile of the output tensor to maximize compute efficiency of the hardware, while the required regions of input, output, and intermediate data fit in the local cache or scratchpad memory. The XLA tile size autotuner has been deployed in production to optimize the most heavily executed kernels on Google TPU fleet on a daily basis, saving approximately  $2\%$  of the total TPU compute time overall [55]. The learned cost model based on a GNN is used to select the top K most promising tile sizes to execute on real hardware [41], reducing the autotuning search time by approximately 20x.
Graph-Level Optimizations. At the graph level, the XLA autotuner supports tuning layout assignment, fusion, and memory space assignment passes, as well as compiler flags that control multiple optimization passes. The XLA graph-level autotuner has delivered  $10 - 20\%$  speedup state-of-the-art models serving substantial production traffic at Google. However, it often takes at least a few hours for the autotuner to converge when tuning one optimization pass of a single graph, and much longer for larger computation graphs. Therefore, a learned cost model would significantly reduce the search time. This motivates us to release the dataset collected from the autotuning process to advance research in developing learned performance prediction models, by addressing challenges outlined in Section 2.2, and ultimately accelerate the autotuning process for production ML workloads.
We focus on layout tuning because it offers the most speedup in general. The layout assignment pass chooses the physical layouts of the input and output tensors of each node that satisfy constraints, while minimizing program's execution time. A layout determines the order of (minor-to-major) tensor dimensions. Figure 3 displays valid input layouts in blue and the chosen layout in red. If an edge connects an output to an input with a different layout, the compiler inserts a copy (transpose) operator to convert the layout. In Figure 3 (left), layout of  $\{1,0,2\}$  is assigned to the output of add but  $\{0,1,2\}$
Figure 3: A node represents a tensor operator, annotated with its output tensor shape  $[n_0, n_1, \ldots]$ , where  $n_i$  is the size of dimension  $i$ . Layout  $\{d_0, d_1, \ldots\}$  represents minor-to-major ordering in memory. Applied configurations are highlighted in red, and other valid configurations are highlighted in blue. A layout configuration specifies the layouts of inputs and outputs of influential operators (i.e., convolution, dot, and reshape). A copy operator is inserted when there is a layout mismatch.
to the first input of conv, causing a layout mismatch, unless a copy operator is inserted. The compiler must trade off between selecting the best layouts for each specific operator and the overhead of copy operators. The autotuner tunes the input-output layouts of the most layout-performance-critical nodes — i.e., convolution, dot (einsum), and reshape because they are common operations and have the most constrained implementations for TPUs — and propagates layouts from these nodes to others. The autotuner picks one input-output layout combination from the valid options for each configurable node.
2.2 Learning Challenges
TPUGRAPHs is non-trivial because training a neural network model to make an accurate prediction on a large graph comes with multiple challenges, as follows.
Scalability. Existing efforts to scale GNNs have mostly focused on node-level and edge-level prediction using sampled subgraphs [29, 12, 34, 77, 86, 49, 26, 1] or graph transformations followed by local models [8, 9]. However, there is a lack of research on how to train scalable models for property prediction of large graphs. Training on sampled subgraphs alone is insufficient as they may not contain all the necessary information for accurate predictions. Aggregating information from the entire graph is essential for graph property prediction, but it poses challenges due to memory limits on a training device, as the memory required scales at least linearly with the size of the graph [78]. Our layout collection contains graphs of up to 44,000 nodes, so training a GNN model on an entire graph (or a batch of thereof) using a single GPU may run out of memory.
Diversity and Imbalance of Graphs. We want to learn a model that generalizes well to unseen graphs. However, this is non-trivial because the model must be trained on diverse types of graphs with enough samples for each type. TPUGRAPHS consists of graphs for all kinds of important ML workloads, including both inference and training, from past to present. While the dataset may be imbalanced — containing graphs from some types of architectures more than others — each graph has at least 10,000 samples of data from different configurations on average.
Redundancy. Another unique property of our dataset is that many samples share the same graph, which represents a large amount of redundant data. An efficient training pipeline should leverage this knowledge and reduce the redundant computation when possible. Additionally, there is another aspect of redundancy within each graph. A tensor computation graph, representing an ML workload, often consists of repeated blocks of neural network layers. The repeated blocks appear as repeated subgraphs. One may leverage this knowledge to improve the learning algorithm.
The baselines accompanying this dataset attempt to address some of these challenges, but are not close to fully solving them.
3 The TpuGraphs Dataset
The TPUGRAPHs dataset contains execution time data points, where each data point contains an HLO graph, its configuration, and its execution time on a single core of TPU v3. The HLO graph in
each data point is a partially optimized graph before being fed into the corresponding optimization pass. For example, in the layout collection, an HLO graph is the input graph to the layout assignment pass. The layout configuration of a graph is a collection of per-node layout decisions on configurable nodes (i.e., convolution, dot, and reshape). For the tile collection, an HLO graph in each data point is a fused subgraph representing a kernel. The tile configuration of a subgraph is a configuration for the entire subgraph, not specific to any particular node.
Within our dataset, there are multiple collections of data, differing in terms of (1) the compiler optimization (i.e., layout and tile), (2) the source of graphs, and (3) the search strategy.
Graphs Collection. We collect HLO graphs from two sources. The first source, called XLA, is the combination of the XLA regression benchmark — from where we collect all open-source models — and the MLPerf benchmark [50, 35]. The XLA graphs span diverse types of popular ML training and inference models, such as vision, NLP, speech, audio, and recommendation. The second source, called NLP, contains a variety of BERT for training and inference, with varying number of layers, attention heads, and hidden sizes. For each model, we run the program — written in TensorFlow, PyTorch, or JAX — and collect the largest HLO graph compiled by XLA, which represents the model's main computation. Note that a typical way that XLA handles a graph with dynamic shapes is to bucketize the graph into multiple static-shape graphs. During execution, the runtime will pad the input to match the static-shape graph with the larger closet shape. Our dataset includes graphs — for varying sequence length, batch size, model size, etc. — some of which are used for dynamic shape workloads. The TPUGRAPHS dataset is similar to the internal datasets used for prior TPU learned cost models [41, 10], but it exclusively contains graphs from open source-programs, while the internal datasets also include production models that cannot be released publicly.
Configurations Generation. Once we have the graphs, we use the XLA autotuner to generate data samples. The set of configurations being generated depends on how the autotuner explores the search space. For the layout collections, we ran the autotuner in two modes. The first mode explores the search space using a genetic algorithm starting from the default configuration, chosen by the compiler's heuristic. Data collected from this mode is labeled default. The second mode explores the search space by picking random candidates. Data collected from this mode is labeled random. We keep data collected in different modes in separate collections; the default collection tends to contain configurations that are not too different from the default, and have similar execution times, while the random collection includes very different configurations with very different execution times.
For the tile size tuning, the autotuner first invokes the compiler to run the graph-level optimizations and obtain fused subgraphs (kernels). For each subgraph, the autotuner enumerates all possible tile sizes for the kernel in a random order, limited by a timeout. Note that the tile size search space is much smaller than the layout search space, so we can enumerate all possible tile sizes. Therefore, there is one data collection for tile sizes. We use only the XLA source for graphs in this collection.
Appendix A.2 describes how we measure the execution time of a given graph and configuration.
3.2 Dataset Statistics and Related Datasets
Table 1 summarizes the details of the different data collections, where the collection name follows the pattern optimization:source:search. Table 3 in Appendix A.1 compares properties of the TPUGRAPHs dataset (all collections) against existing graph property prediction datasets.
ML Program Performance. The TPUGRAPHs layout collections provide more than 770x larger graphs on average compared to TenSet [80], the only existing large-scale dataset on ML program performance. Our tile collection is similar to TenSet as the configuration controls the optimization at the kernel (fused subgraph) level. However, it compliments TenSet nicely as it provides data points on different hardware. Halide Auto-scheduler [2] releases their evaluation dataset of Halide programs mainly consisting of image processing benchmarks with a few ML benchmarks.
Other Program Performance. Beyond ML programs, the performance prediction dataset with largest graphs is on database queries [31], whose graphs are still more than a few orders of magnitudes
Table 1: Statistics of TPUGRAPHs collections. The collection name follows the pattern optimization:source:search. The search may explore the same configuration multiple times, so the same pair of graph and configuration may appear multiple times with slightly different execution time from multiple measurements. The total number of samples is thus higher than the number of unique pairs.
Collection	Core (Sub) Graphs	Avg. Nodes	Configs per Graph	Total Graphs + Configs	Samples
Layout:XLA:Default	78	14,105 (372–43,615)	10,147 (681–71,574)	771,496	1,272,538
Layout:XLA:Random	11,648 (109–99,783)	908,561	1,115,709
Layout:NLP:Default	244	5,659 (876–21,919)	56,534 (9032–90,985)	13,285,415	15,479,038
Layout:NLP:Random	66,089 (8,843–100,001)	16,125,781	16,135,731
Tile:XLA	6,988	40	1,842	12,870,077	12,870,077
smaller than ours. Another popular performance prediction dataset is BHive [15], consisting of x86 basic blocks sourced from multiple open source programs, with runtime measurements on different Intel hardware platforms. However, the basic blocks are quite small, including four instructions on average. CompilerGym [18] releases a collection of LLVM IR code datasets that can be evaluated in their environment. The largest datasets in their collection include AnghaBench [19] and CSmith [75]. AnghaBench provides a large number of relatively small real-world programs. CSmith programs are large (comparable to ours), but they are randomly generated programs. Additionally, CompilerGym's datasets do not come with performance measurements, so one would have to execute the programs and configurations in the CompilerGym's environment themselves to obtain program execution time.
Program Analysis. Other closely related datasets are on programming tasks. CodeNet [57] is a large dataset to teach AI to code, in which each code sample is a solution to one of the coding problems. OBGB-CODE2 [33] is for code summarization, containing Abstract Syntax Trees obtained from Python functions. TBCNN [52] releases its dataset on program classification from a pedagogical programming open judge system. CuBERT [40] uses Python files extracted from the ETH Py150 dataset [59] for fine-tuning and uses github_repos dataset under BigQuery's public-data project for pre-training. CodeBERT [25] releases its multi-programming-lingual dataset used for pre-training. Works such as inst2vec [7] and ProGraML [17] uses datasets of code in LLVM compiler intermediate representation to learn generic code representation for various program analyses and optimizations.
Other. Apart from code datasets, there are many other graph datasets. Open Graph Benchmark [33] suite presents graphs that are used for machine learning tasks such as GNN inference and training. GAP [6] and Graph Based Benchmark Suite (GBBS) [21] provide large-scale curated sets of graphs, primarily for evaluating traditional graph problems. SuiteSparse [43] consists of a wide variety of sparse matrices, which can be viewed as graphs. Most of these datasets are for node-level or edge-level prediction tasks. TPUGRAPHs is by far one of the largest graph property prediction datasets. TPUGRAPHs' average graph size is comparable to that of MalNet [27] — the largest scale graph property prediction dataset to date — while offering 25x more combinations of graphs and configurations. Other popular graph property prediction datasets include small molecule [58, 61], bioinformatic [22, 33], and social network datasets [62, 74].
3.3 Dataset Split
We split the data using 80-10-10 ratio by graphs in each collection. Splitting data by graphs ensures that graphs in the validation and test sets do not appear in the training set to evaluate the generalization of the model on unseen graphs. The validate and test graphs stay the same across different XLA collections; the same applies to NLP collections. We deliberately holdout the target labels of samples in the test set for competition purposes.
We report test and validation metrics for the tile collection by considering all configurations. For the layout collections, we report final metrics only on 1,000 configurations to reduce the computational demand for the model evaluation. We select these 1,000 configurations by sorting all configurations based on their execution times and extracting the  $[0,m,2m,\dots ,length - 1]$ th configurations. The dataset includes the indices of the selected configurations<sup>4</sup>.
$^{4}$ https://github.com/google-research-datasets(tpu_graphs/tree/main(tpu_graphs/evals
4 Learning a Performance Prediction Model
The goal of a learned cost model is to rank the performance of different configurations of a given graph. This section explains the baseline models we provide and how we train them, primarily based on the TPU learned cost model papers [41, 10].
4.1 Feature Extraction