SlowGuess/NeurIPS_2023_OCR huggingface dataset:
TPUGRAPHS provides data in two formats: raw protobuf format and numpy arrays similar to the OGBG format [33]. The autotuner produces output results in protobuf format. A data pre-processing script converts data from the protobuf format to the numpy format. The main function of the data pre-processor is feature extraction. Node features describe the node's properties, such as output tensor shape, tensor layout, striding, padding, and operation-specific parameters. Our feature extraction is minimal. To extract a node feature vector, we either copy values from various fields in an HLO instruction (a node in an HLO graph) as they are, or convert categorical values using one-hot encoding. To convert an unbounded list of numbers (e.g., tensor shape) to a fixed-size vector, we truncate the list to six elements and include the summation and/or product of all elements in the list (e.g., the product of dimension sizes represents the volume of the tensor) because the tensors appearing our dataset do not contain more than six dimensions. A per-node layout configuration and tile size can be represented as a nested list with some unbounded dimensions. Similarly, we truncate these unbounded dimensions to six elements. The detailed description of node and configuration features can be found in the GitHub repo.
We provide code for training a variety of models over the numpy format. Nonetheless, the raw format can allow researchers to experiment with different feature extractions and measure impacts on the quality of a learned model.
Figure 4 shows the model architecture we use for our baseline models, which are based on a GNN since the input program is represented as a graph. Node features consist of two parts. The first part is an opcode id, i.e., type of tensor operation (such as convolution). Our baseline models map an opcode id to an opcode embedding via an embedding lookup table. The opcode embedding is then concatenated with the rest of the node features as inputs to a GNN. We combine the node embeddings produced by the GNN to create the embedding of the graph using a simple pooling reduction. The resulting graph embedding is then linearly transformed into the final scalar output by a feedforward layer. Prior work [41] has studied alternative models, including LSTM and Transformer, and shown that GNNs offer the best performance. We provide baseline models with GCN [42] and GraphSAGE [30].
Figure 4: Model architecture.
The primary use case of the model is to rank configurations within a given graph and select top candidates to evaluate on real hardware. Thus, we can train the model using regression losses (e.g., Mean Square Error (MSE)) or ranking losses (e.g., pairwise hinge loss and ListMLE [73]). A ranking loss is computed among sample pairs within the same graph in the same batch, and the losses from different graphs in the batch are reduced to get the total loss. We use Ordered Pair Accuracy (OPA) as a validation metric to select the best model checkpoint.
4.4 Implementation
**Layout model.** Our baseline model is a 3-layer GraphSAGE with residual connections. We concatenate node features and per-node configuration features as inputs to the GNN. If a node is non-configurable (having no layout configuration), we use a zero vector as configuration features. Our baseline code allows both a typical full graph training and a graph segment training [10]. One may improve the compute efficiency further by using historical embeddings of subgraphs and segment dropout, as in the Graph Segment Training paper.
Tile size model. For the tile collection, we implement three baselines: an MLP model and two GNNs (GraphSAGE and GCN with residual connections). The MLP model embeds all opcodes, concatenates with node features, sums across all nodes, and then concatenates with kernel configuration features, feeding into a 3-layer MLP. We experiment with two options to combine the graph-level kernel configuration features with the node-level information (yellow in Figure 4): either late-join or early-join. The first runs the GNN only on node features, reduces the node embeddings, and then concatenates with the graph (configuration) features. The second replicates the graph features onto every node. The early-join GraphSAGE model closely resembles the original TPU learned cost model [41].
For both models, we experiment with two objective functions: MSE and ListMLE. We find that ListMLE gives better empirical performance. Our baseline models are available at https://github.com/google-research-datasets(tpu_graphs. They are implemented using TensorFlow-2 and TF-GNN. The details of hyperparameters can be found in Appendix B.
To evaluate a model, we use two metrics. (1) Kendall's Tau assesses how well a model's ranking of configurations correlates with their corresponding runtimes. (2) Top-K error (or slowdown error@  $K$ ) measures the slowdown of the chosen  $K$  configurations as:
$$
\text {t o p - K e r r o r} = \frac {\text {T h e b e s t r u n t i m e o f t h e t o p - K p r e d i c t i o n s}}{\text {T h e b e s t r u n t i m e o f a l l c o n f i g u r a t i o n s}} - 1 = \frac {\min  _ {i \in K} y _ {i}}{\min  _ {i \in A} y _ {i}} - 1 \tag {1}
$$
where  $K$  is the top-K predictions,  $A$  is all configurations of the given graph from the dataset collection, and  $y$  is the measured execution time.
The choice of metrics is justified as follows. For the tile collection, since the number of configurations per graph is relatively small, one can apply the model to obtain the predictions of all configurations, choose top-K candidates according to the model to measure on real hardware, and finally select the best one according to the real measurements. On the other hand, for the layout collections, the search space is quite large. Therefore, common search strategies, such as Genetic Algorithm and Simulated Annealing, need access to a fitness function (which can be the model). Therefore, it is important that the model can well-preserve the order of the configurations (from fastest to slowest) as reflected by the correlation score.
Layout model. For each model variant, we train the model once with only a few set of hyperparameters, and select the checkpoint with the highest OPA on the validation set to evaluate its ranking correlation and top-K prediction errors. Table 5 in Appendix B reports attempted hyperparameters for modeling layout. We report the performance of the best model based on the validation score.
Tile size model. For each model variant, we perform hyperparameter search on opcode embedding size, hidden size, network depth, and learning rate, considering values specified in Table 5. Unlike in the layout collections where the accuracy of each model is quite stable across multiple training runs, the accuracy of a model in the tile collection fluctuates dramatically. Therefore, we train each model variant three times. For each run, we select the model checkpoint with the highest OPA on the validation set. We report the top-K errors of the run that achieves the median top-1 error on
Table 2: Kendall's Tau correlation and prediction errors (Eq. 1) of our best baseline model on different dataset collections. The values of  $(K_{1}, K_{2}, K_{3})$  are (1, 10, 100) for the layout collections, and (1, 5, 10) for the tile collection.
Collection	Kendall τ	Top-K1 E %	Top-K2 E %	Top-K3 E %
Val	Test	Val	Test	Val	Test	Val	Test
Layout:XLA:Random	0.19	0.34	19.8	10.9	12.3	5.7	9.7	1.6
Layout:XLA:Default	0.12	0.21	3.8	14.1	1.9	0.6	0.3	0.2
Layout:NLP:Random	0.58	0.53	2.1	4.6	2.0	1.0	0.6	0.09
Layout:NLP:Default	0.30	0.28	4.0	4.0	3.7	3.1	3.5	0.13
Tile:XLA	-	-	10.5	9.1	3.0	4.2	1.8	2.8
Figure 5: Prediction errors  $(\%)$  of different model variants on the Tile:XLA collection. Early and Late refer to early-join and late-join options.
the validation set. An aggregated top-K error is an average across all kernels in the graphs in the validation/test set. Note that the number of kernels varies across graphs.
5.3 Results on Layout Collections
Table 2 reports the top-K slowdown errors and Kendall's Tau correlation of the best model across all graphs (programs) in the validation and test sets for each dataset collection. According to the correlation scores, the layout collections on the default search space are more difficult than those on the random search space. This result matches our intuition because the default search space contains many similar configurations near the default, so it is difficult to rank them; whereas, the random search space contains more diverse configurations. The XLA layout collections are also noticeably more difficult than the NLP layout collections. This is also expected because the XLA collections contain more diverse graphs, while the NLP collections contain only graphs with the Transformer architecture. In terms of top-K errors, the model struggles to identify fast candidates on Layout:XLA:Random. If we use the learned cost model to select the top configuration, we will be on average  $10 - 20\%$  slower than the known optimal configuration. Even if we consider the top 10 candidates, we will still be on average  $5 - 13\%$  off. We hypothesize that this is due to the combination of the diversity of both graphs and configurations in this collection. The correlation and top-K errors vary wildly across graphs (programs) as shown in Table 7 in Appendix C.
5.4 Results on Tile Collection
Table 2 also reports the average top-K errors on the tile collection. The average top-1 error of the best model is comparable to the original TPU learned cost model paper's [41]. Figure 5 compares alternative choices. Similar to the original paper, our results show that combining configuration features with node features early (early-join) is superior to combining configuration features with a reduced graph embedding later (late-join). Using a ranking loss (ListMLE) is much more effective than using MSE. Additionally, we compare the choice of a GNN between GraphSAGE and GCN, and find they are comparable. We also provide an MLP baseline without a GNN, and confirm that a GNN is essential to achieve good accuracy.
Appendix C reports additional results including per-graph evaluation metrics, an additional ablation study, and model's prediction overhead. Note that our dataset is not exactly the same as the internal dataset used in the original papers [41, 10], but they share a large number of overlapping graphs.
There are many potential improvements to be made on top of the baseline models we provide. First, we observe that a tensor computation graph typically contains repeated subgraphs, representing repeated blocks of neural network layers. One direction is to leverage this repeated structure to devise a more compact representation that is easier to learn. Second, as mentioned earlier, the dataset may contain some types of graphs, e.g., ResNet, significantly more than others. This skew may make the learned model perform well on common types of graphs, but poorly on uncommon types. One may investigate how to address this data imbalance problem to improve the quality of the learned model. Finally, while we know that developing a purely analytical cost model is extremely difficult, training an accurate learned cost model is not easy either, especially when graphs are large. One idea is to combine the best of both worlds, using analytical modeling when easy to do and letting the learned model make corrections to the analytical estimates.
We plan to continue improving our dataset in multiple aspects. First, we would like to include more diverse graphs. Prior approaches generate random programs for training data [2, 14, 5]. We deliberately avoid randomly generated programs in our dataset because we would like a model trained on the dataset to achieve high performance on realistic programs used in production, instead of achieving moderate performance on both real-word and randomly generated programs. However, we acknowledge that the diversity of graphs in the dataset is extremely important for the generalization of the model. One way to generate more realistic tensor programs is to leverage Neural Architecture Search [84, 68, 69, 85, 47, 46, 60, 64, 76, 23]. We leave this as future work, potentially the next version of the dataset. Second, we would like to include data measured on other hardware platforms beyond TPUs, such as CPUs and GPUs. Nonetheless, we believe that the general techniques of training an accurate learned performance model (e.g., improvements on GNNs, Graph Segment Training method, etc.) are applicable to other hardware targets; therefore, the improvements coming out from experimenting with the current version of the dataset should also benefit other hardware platforms as well. Many compiler optimizations are also common across multiple hardware backends. For example, the tile size selection has shown to be one of the most important optimizations across all widely used hardware (i.e., CPUs, GPUs, and TPUs) and even custom accelerators. Layout optimizations are also applicable on CPUs and GPUs, but the layout options on CPUs and GPUs may be limited if the compiler depends on pre-optimized library kernels.
We hope that TPUGRAPHs will propel advances in compilers. In particular, researchers may be able to extract insights on how to improve code generation for tensor programs. For example, which information in a tensor computation graph is important to make various optimization decisions? How can we build an accurate cost model for an important class of hardware architectures? Can a learned representation for tensor programs guide various tensor compiler optimizations?
We would like to thank Sai Ganesh for a pointer to generate open-source XLA HLO graphs, Mark Daoust for help on the dataset's hosting location, and Amit Sabne and David Majnemer for reviewing the dataset's security information. Charith's contributions were partially supported by ACE, one of the seven centers in JUMP 2.0, a Semiconductor Research Corporation (SRC) program sponsored by DARPA and NSF under grant CCF-2316233.
[1] Sami Abu-El-Haija, Joshua V. Dillon, Bahare Fatemi, Kyriakos Axiotis, Neslihan Bulut, Johannes Gasteiger, Bryan Perozzi, and Mohammadhossein Bateni. Submix: Learning to mix graph sampling heuristics. In Proceedings of the Thirty-Ninth Conference on Uncertainty in Artificial Intelligence, 2023.
[81] Lianmin Zheng, Ruochen Liu, Junru Shao, Tianqi Chen, Joseph E. Gonzalez, Ion Stoica, and Ameer Haj Ali. Tenset: A large-scale program performance dataset for learned tensor compilers. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 1), 2021. URL https://openreview.net/forum?id=aIfp8kLuvc9.
[82] Size Zheng, Yun Liang, Shuo Wang, Renze Chen, and Kaiwen Sheng. FlexTensor: An Automatic Schedule Exploration and Optimization Framework for Tensor Computation on Heterogeneous System. In Proceedings of the Twenty-Fifth International Conference on Architectural Support for Programming Languages and Operating Systems, ASPLOS '20, page 859-873, New York, NY, USA, 2020. Association for Computing Machinery. ISBN 9781450371025. doi: 10.1145/3373376.3378508.
[83] Marinka Zitnik, Rok Sosic, Marcus W Feldman, and Jure Leskovec. Evolution of resilience in protein interactomes across the tree of life. Proceedings of the National Academy of Sciences, 116(10):4426-4433, 2019.
[84] Barret Zoph and Quoc Le. Neural architecture search with reinforcement learning. In International Conference on Learning Representations, 2017. URL https://openreview.net/forum?id=r1Ue8Hcxg.
[85] Barret Zoph, Vijay Vasudevan, Jonathon Shlens, and Quoc Le. Learning transferable architectures for scalable image recognition. In Conference on Computer Vision and Pattern Recognition, CVPR'18, pages 8697-8710, 06 2018. doi: 10.1109/CVPR.2018.00907.
[86] Difan Zou, Ziniu Hu, Yewen Wang, Song Jiang, Yizhou Sun, and Quanquan Gu. Layer-dependent importance sampling for training deep and large graph convolutional networks. Advances in neural information processing systems, 32, 2019.
A Additional Dataset Information
Documentation. The documentation of the dataset can be found at https://github.com/google-research-datasets(tpu_graphs. The github repo contains instructions and code on how to download and use the dataset.
License. The dataset is licensed under the Creative Commons Attribution 4.0 International License (CC-BY). To view a copy of this license, visit http://creativecommons.org/licenses/by/4.0/. All code is licensed under the Apache License, Version 2.0 (Apache 2.0); You may obtain a copy of the Apache 2.0 license at: https://www.apache.org/licenses/License-2.0.
Author Statement. The authors bear all responsibility in case of violation of rights. The authors will monitor the issues and provide necessary maintenance to ensure the access to the data.
A.1 Dataset Comparison
Table 3: Comparison of TPUGRAPHS properties with other large-scale graph property prediction datasets. * provide only programs, but one may use them in CompilerGym [18] environment to obtain performance measurements when compiling with specific configurations. † provides randomly generated programs.
Application	Dataset	Graphs (+ Configs)	Avg. Nodes
ML Program Perf	TPUGRAPHS (Layout)	31,091,253	7,705
TPUGRAPHS (Tile)	12,870,077	40
TenSet [80]	51,577,248	5–10
Other Program Perf	Database [31]	300,000	< 100
BHive [15]	330,018	4
AnghaBench* [19]	1,041,333	62
CSmith*† [75]	530,000	5,845
Program Analysis	CodeNet [25]	13,916,868	200–500
OGBG-CODE2 [33]	452,741	125
TBCNN [52]	52,000	190
Cybersecurity	MalNet [27]	1,262,024	15,378
Molecule	PCBA [58]	437,929	26
MUV [61]	93,087	24
Bioinfomatic	DD [22]	1,178	284
OGBG-PPA [33]	158,100	243
Social Network	Reddit-T [62]	203,088	24
REDDIT-12K [74]	11,929	391
REDDIT-5K [74]	4,999	509
A.2 Execution Time Measurement
We measure the execution of a compiled binary on a single TPU chip using random input data. Note that some of the graphs in the layout collection must be run on multiple TPU chips. However, doing so is not economically viable for autotuning a large number of models and generating the dataset. Therefore, the autotuner modifies the final optimized graph (after all graph-level optimizations) to make it runnable on a single TPU chip in two ways. First, we replace each collective communication operation (e.g., all-reduce and all-gather) with a no-op that simply allocates the right amount of output buffer (with undefined values). This means the measured execution time ignores the time taken by collective operations. We think this is reasonable because layout decisions rarely affect the execution time of collective operations, and we care about the ranking of execution time rather than the absolute time. The second modification we perform is to replace dynamic loop bounds with fixed loop bounds. Without such replacement, a dynamic loop bound may depend on random input data, resulting in an extremely large loop bound, making the program run unrealistically slowly. The use of random or undefined data, however, does not affect the execution time of a compute operation (e.g., convolution) because the timing does not depend on the input data. Because of these modifications, our absolute execution time measurement may be inaccurate in some cases, but this estimation approach has been used in production to tune graph-level optimizations and deliver large speedups on many important
models. Therefore, we believe it is reasonable to use the execution time measured by the approach outlined here as a prediction target.
A.3 Graphs in Dataset
XLA Collection. Graphs in the XLA collection are collected from open-source models from the following sources:
- https://github.com/tensorflow/models
- https://github.com/tensorflow/tensorflow
- https://github.com/tensorflow/tensor2tensor
- https://github.com/tensorflow(tpu
- https://github.com/google/brax
MLPerf [50, 35]
NLP Collection. Graphs in the NLP collection are all collected from TensorFlow Hub. Table 4 reports the architectures and hyperparameters of the models used to generate graphs in this collection.
Table 4: Architectures and hyperparameters of models in the NLP collection.
Model Name	Source	Layers	Hidden Size	Attention Heads
bert_en_uncased	Devlin et al. [20]	12, 24	768, 1024	12, 16
bert_en_wwm_uncased	Devlin et al. [20]	24	1024	16
bert_en_cased	Devlin et al. [20]	12, 24	768, 1024	12, 16
bert_en_wwm_cased	Devlin et al. [20]	24	1024	16
bert-multi_cased	Devlin et al. [20]	12	768	12
small_bert	Devlin et al. [20]	2, 4, 6, 8, 10, 12	128, 256, 512, 768	2, 4, 8, 12
albert_en	Lan et al. [44]	12, 24	768, 1024, 2048, 4096	12, 16, 32, 64
electra	Clark et al. [16]	12, 24	256, 768, 1024	4, 12, 16
experts_pubmed	TensorFlow Hub	12, 24	768, 1024	12, 16
experts_wiki_books	TensorFlow Hub	12, 24	768, 1024	12, 16
talking_heads	Shazeer et al. [63]	12, 24	768, 1024	12, 16
Table 5: Hyperparameters of our baseline models. We report the result of the best version of each baseline variant by tuning opcode embedding size, hidden size, and number of GNN layers from the values specified.
Parameter	Considered Values
Layout Models	Tile Size Models
Opcode embedding size	64, 128	64, 128, 256
Hidden size	100, 200	64, 128
GNN layers	2, 3, 4	2, 3
Batch size	20	100
Learning rate	0.001	0.01, 0.001
Training iterations	200	500
Loss	MSE, ListMLE	MSE, ListMLE
Optimizer	Adam	Adam
C.1 Prediction Accuracy
Layout Collections. Table 6 compares alternative choices in terms of the model architecture and the training method on the Layout:XLA:Random collection. GST implements the Graph Segment Training method [10] without historical embedding and segment dropout. Unlike in the original paper, we partition a graph based on a topological order of nodes. At training time, we randomly pick a segment — covering nodes  $\in [i,i +$  segment length), where  $i$  is drawn at random — for backward
Table 6: Kendall's Tau correlation and prediction errors (%) of different model variants and training methods on the Layout:XLA:Random collection.
Model	Kendall τ	Top-1 E %	Top-10 E %	Top-100 E %
Val	Test	Val	Test	Val	Test	Val	Test
GST	0.13	0.32	29.2	8.1	12.6	6.2	6.0	2.3
Full Graph	0.19	0.37	33.3	12.2	14.5	4.8	6.8	3.3
MSE loss	0.03	0.31	24.5	8.1	17.7	5.4	6.0	1.3
Random	0.002	0.007	24.3	97.1	4.7	10.0	0.1	1.9
Table 7: Per-program Kendall's Tau correlation and prediction errors  $(\%)$  on the validation set of the Layout:XLA collections of a few selected models.
Program	Layout:XLA:Random	Layout:XLA:Default
GST	Full Graph	MSE loss	GST
Kendall	Top-1	Kendall	Top-1	Kendall	Top-1	Kendall	Top-1
bert_pretraining.4x4.fp16	0.23	8.9	0.65	0.0	0.04	3.2	0.25	8.9
inception_v3_batch_128_train	-0.34	70.3	-0.49	63.8	-0.48	63.4	0.27	3.0
mlperf_bert_batch_24_2x2	0.39	0.5	0.24	0.5	-0.01	21.1	0.21	7.2
resnet50.4x4.fp16	0.13	53.6	0.10	69.0	0.12	53.6	0.15	9.6
resnet_v1_50_official_batch_128 bf16	0.10	17.1	0.07	17.1	0.09	66.6	-0.05	0.7
tf2_bert_pretrainDynamic_batch_size	0.52	5.5	0.66	5.9	0.01	19.5	0.18	1.4
unet_3d.4x4.bf16	0.23	1.6	0.19	1.6	0.30	1.6	-0.03	0.0
propagation. Note that all nodes and edges are used in a forward pass. We use a segment length of 5,000. Results in the main paper (Table 2) are from the GST model. Full Graph uses the entire graph for forward and backward passes during training (a typical method). Unlike all other models, MSE loss uses Mean Squared Error loss function: it aims to model the exact runtime of every configuration, rather than the ranking of configurations.
Table 7 further reports per-program ranking correlation and top-K errors of a few selected baseline models on the programs (graphs) in the validation set. As shown in the table, the evaluation scores vary widely across different graphs. We show the scores only on the validation set because we do not want to reveal the details about the test set for competition purposes.
Tile Collection. Table 8 summarizes the overall average prediction errors of all the baseline models. Table 9 further reports per-program average prediction errors across each program's kernel subgraphs of a few selected model variants with their best hyperparameter values.
C.2 Prediction Overhead