Simulating Execution Time of Tensor Programs using Graph Neural Networks

Abstract

Optimizing the execution time of tensor program, e.g., a convolution, involves finding its optimal configuration. Searching the configuration space exhaustively is typically infeasible in practice. In line with recent research using TVM, we propose to learn a surrogate model to overcome this issue. The model is trained on an acyclic graph called an abstract syntax tree, and utilizes a graph convolutional network to exploit structure in the graph. We claim that a learnable graph-based data processing is a strong competitor to heuristic-based feature extraction. We present a new dataset of graphs corresponding to configurations and their execution time for various tensor programs. We provide baselines for a runtime prediction task. INTRODUCTIONCurrent deep learning frameworks, such as TensorFlow, PyTorch, allow to optimize a computational graph representation using, e.g., auto differentiation and memory management (Abadi et al., 2016;Paszke et al., 2017). However, they do not tackle optimization of hardware-specific operator-level transformations, but rely on manually tuned and vendor-specific operator libraries. Thus, there is room to further improve a computational graph by optimizing transformations for specific hardware.Recently, this gap has been filled by TVM, a compiler framework that allows both graph-and operator-level optimization in an end-to-end manner (Chen et al., 2018a). TVM specifies a configuration for an operator, e.g., a specific way of performing a convolution, and compiles the resulting tensor program to a target hardware. As a consequence, for each new workload/operator, optimization over a new configuration space must be carried out. This results in a hard optimization problem, e.g., for Nvidia GPU the search space of a single operator consists of more than 10 6 configurations.Recent efforts overcome this issue by learning how to optimize tensor programs from data (Chen et al., 2018b). Instead of running an exhaustive search over an impractically large search space, a surrogate model is trained to predict runtime for a given configuration. This model is in turn used to select the configuration that minimizes the runtime. (Chen et al., 2018b) utilizes XGBoost Guestrin, 2016) and TreeGRU (Tai et al., 2015) as surrogate models.Contribution Similar to (Chen et al., 2018b), we represent a configuration of a tensor operator as an abstract syntax tree (AST) (Allamanis et al., 2017), and extract node features using TVM. We then train a Graph Neural Network (GraphNN) on the resulting graph as the surrogate model. We claim that GraphNNs are a good fit, as, crucially, they preserve the graph structure of the AST and allow propagating information among nodes. The contribution of the paper is threefold: