DynaTune: Dynamic Tensor Program Optimization in Deep Neural Network Compilation

1 INTRODUCTION . The enormous computational intensity of Deep Neural Network ( DNN ) models has attracted great interest in optimizing their performance . Popular deep learning ( DL ) frameworks such as PyTorch ( Paszke et al. , 2019 ) and TensorFlow ( Abadi et al. , 2016 ) adopt custom optimized kernels such as Intel MKL-DNN or Nvidia cuDNN ( Chetlur et al. , 2014 ) as back-end . However , given the increasing complexity of tensor operations in DNNs and the volatility of DL algorithms , it calls for developing fast and automated compilation frameworks to handle the unprecedented amount of innovations . To imitate or even exceed the success of hand-optimized libraries , recent research has developed neural network compilers , such as XLA ( Leary & Wang , 2017 ) , Glow ( Rotem et al. , 2018 ) , Tensor Comprehension ( Vasilache et al. , 2018 ) , and TVM ( Chen et al. , 2018a ) . Among them , TVM has shown superior performance improvements using a technique called Learning to Compile ( AutoTVM ) ( Chen et al. , 2018b ) . AutoTVM optimizes the code by generating many versions of a tensor operator and chooses the best through a learned cost model and search over a large space of code transformation choices . While the Learning to Compile approach produces highly optimized code of DNN models , it suffers from excessively long optimization time . As an example , although AutoTVM is able to demonstrate close to 2× performance improvement over TensorFlow on ResNet-18 , the optimization time can take several hours or even tens of hours ( Chen et al. , 2018b ) . The long optimization time hinders the turnaround time and even puts the practical utility of the current compiler-based solutions into question . Recent works strive to reduce the optimization time by improving the search strategy for the code transformation plan and lowering the hardware measurement cost ( Ahn et al. , 2020 ; Adams et al. , 2019 ) . However , these approaches mostly focus on accelerating the convergence speed of optimization at the individual tensor operator level ( e.g. , Conv2D , batched GEMM ) , which do not necessarily solve the issue of slow convergence and long optimization time of the entire model , often containing tens of tensor operators . Different from existing methods , we introduce DynaTune , a DL code optimization algorithm that minimizes the sum of the execution time of all operators in a model as much as possible and as ∗Both authors contributed equally . Order of appearance is random . quickly as possible . Specifically , the contributions of our paper consist of ( 1 ) a preliminary analysis that reveals the challenges and opportunities from existing DL code optimization strategies , ( 2 ) a time-slot-based optimization scheme , which simultaneously explores different operators and learns in an online manner that allows to dynamically switch to optimizing more promising tensors operators . ( 3 ) a Bayesian belief model that predicts future performance gains of operators , which helps make better decisions and expedites the convergence speed . ( 4 ) a detailed evaluation of the proposed algorithm with modern DNNs ( ResNet-18 , VGG , SqueezeNet , Transformer ) on both CPU and GPU . Compared with the leading framework , AutoTVM , DynaTune is 1.2–2.4× times faster to obtain the same levels of optimization . 2 BACKGROUND . DL compilation pipeline . A typical DL compiler contains multiple passes to optimize a model trained by popular DL frameworks such as TensorFlow ( Abadi et al. , 2016 ) , PyTorch ( Paszke et al. , 2019 ) , or MXNET ( Chen et al. , 2015 ) , as shown in Fig . 1 . In the first pass ( box with dotted line ) , the compiler frontend applies target-independent and white-box target-dependent optimizations that do not include a measure of actual execution time . The target-independent passes perform optimizations such as operator fusion and data layout transformation , and the white-box target-dependent optimizations apply heuristic rules for code transformation based on domain knowledge . Recent work such as AutoTVM ( Chen et al. , 2018b ) extends the pipeline with another pass , a black-box target-dependent pass , which uses learning machinery to perform optimizations . Black-box target-dependent pass . In this pass , the compiler converts code transformation decisions as code templates . A template contains knobs that control various aspects of the optimization ( e.g. , memory tiling , loop transformations , vectorization ) and determines whether the code ( 1 ) fully utilizes the internal parallelism within processors , ( 2 ) uses the shared memory wisely , and ( 3 ) maximizes data locality . Due to the large transformation space , the compiler makes use of an auto-tuner ( with an optimization algorithm ) and real hardware measurements to find the best transformation on target hardware ( e.g. , CPU , GPU , ARM , or IoT devices ) ( Chen et al. , 2018b ) . 3 CHALLENGES AND MOTIVATIONS . This section presents several studies that reveal the challenges of existing DL compilation that guided our design in Section 4 . Challenge 1 . Existing DL compilation focuses on accelerating the convergence speed of individual tensor operator instead of the entire model , resulting in slow convergence and long optimization time . Prior work ( Chen et al. , 2018a ; b ; Vasilache et al. , 2018 ; Ahn et al. , 2020 ) optimizes one tensor operator at a time in a predefined order ( e.g. , in declaration order ) . However , such an optimization strategy is not always appropriate in practice . For example , there is often an extreme performance difference ( e.g. , an order of magnitude ) between optimized and unoptimized operators . If we optimize operators sequentially , the overall model inference time stays high as long as there are still unoptimized operators . As a result , practitioners may need to wait until all tensor operators have finished optimization to get the desired latency , which results in long optimization time . With the active research that has been pushing the model size to millions or even billion-scale parameters with a training time of only a few hours or less than one hour ( Yamazaki et al. , 2019 ; Goyal et al. , 2017 ; You et al. , 2017 ; Lin et al. , 2019 ; Shoeybi et al. , 2019 ; You et al. , 2019 ) , it becomes even more prominent to reduce the inference optimization cost of the current solution . Furthermore , since major players in the industry have adopted many of these DL compilers ( Wu et al. , 2019a ; b ; Lattner et al. , 2020 ; Liu et al. , 2019 ) , fast convergence is desirable for many users of these pipelines to have a better control of the optimization cost and good performance . For example , deployment engineers may want to obtain an optimized model sooner or quickly get a latency upper-bound estimate of a model in development . Challenge 2 . Static scheduling has only a limited view of the tensor program and has difficulty taking advantage of the actual optimization behavior . We note that from an execution point of view , the optimization of tensor operators is independent of each other , so that we may optimize them in any order and even non-consecutively . As a result , dynamic optimization has a big advantage for iterative DL compilation : We can intelligently order the optimization sequence of operators ( i.e. , scheduling ) to accelerate the convergence of the optimization significantly . For example , it would be better to switch to optimizing another operator if we convincingly identify that the other operator has a higher potential . That being said , is it realistic to assume that all the information concerning optimizing the operators is available before the optimization even starts so that we can decide the schedule from the very beginning ? Our preliminary analysis indicates that the amount of computation of an operator ( known a priori ) has a very disproportionate impact on the optimization time and latency reduction . Fig . 2 shows that although operator 17 of VGG ( Simonyan & Zisserman , 2015 ) takes the longest time to optimize , it yields the least amount of latency reduction1 . Our further investigation shows that the underlying code transformation space is non-linear , as shown in Fig . 3 2 . As a result , the optimization behavior tends to change over time , which is hard to recognize and predict with static knowledge only . Challenge 3 . Even with dynamic information , it is not clear how to best extrapolate estimated performance . Given the optimization results , there is an incentive to adopt a ” predict-thenoptimize ” paradigm that builds a model to learn the correlation between the optimization cost and the observed optimization performance the model can be used to make predictions for potential performance gains . To identify the characteristics of the optimization behavior , we plot 16 optimization curves of best-found GFLOPS ( Giga Floating Point Operations Per Second ) in Fig . 4 to find patterns that can be used for designing a prediction model . We find that most curves ( 1 ) roughly follow an increasing curve with a diminishing return , ( 2 ) saturate towards an unknown final value , and ( 3 ) occasionally exhibit sudden jumps . The curve saturates to an unknown value because the performance can not exceed the hardware peak GFLOPS , which is 9.7-TFLOPS in our case . The curve has sudden jumps because the code transformation space has change points , as shown in Fig . 3 . By taking into account the curve information , we believe it has more opportunity to dynamically optimize operators that likely lead to greater performance improvements . 1The orange bar shows the amount of computation of each operator measured as the floating-point operations ( FLOPs ) , which can be calculated statically before the optimization starts , as described in Molchanov et al . ( 2017 ) . The “ optimization gain ” is calculated as the reduction of wall-clock time from each operator after optimization , and the “ optimization cost ” is calculated as the wall-clock time spent to obtain the optimized latency , both of which are normalized by the total latency reduction and optimization time . 2The figure shows the code transformation space of a Conv2D operator in ResNet-18 . In this case , the performance of this operator varies based on the tiling size along the input channel and output channel while having other knobs fixed . The knobs control various aspects of the optimization and its performance . A summary of the knobs can be found in Ahn et al . ( 2020 ) . 4 METHOD . In this section , we present our design for DynaTune . We illustrate the difference between the existing DL optimization ( Fig . 5 ) and the high-level design of DynaTune ( Fig . 6 ) . 4.1 DYNAMIC MULTI-TENSOR-OPERATOR OPTIMIZATION PROBLEM . In this paper , we take a view of accelerating the convergence speed of multi-tensor-operator optimization by first considering a Multi-Armed Bandits model for the problem . In particular , we partition time into some fixed-length time slots { 1 , 2 , ... , T } . Similar to MAB , we define an operator scheduler that operates in discrete time steps . At the beginning of any time slot t ∈ { 1 , 2 , ... , T } , the scheduler needs to choose an operator kt ∈ [ K ] for tuning in that slot . The scheduler then obtains a list of observations ( i.e. , best-found performance measured in GFLOPS ) Perfk ( t ·Lk : ( t+1 ) ·Lk ) of k , where Lk is the number of iterations an operator has been optimized . The remaining unchosen operators stay the same unless it is selected ( i.e. , rested arms ) . The latency reduction for k at time slot t is rt ( kt ) = op ( k ) Perfk [ t·Lk ] − op ( k ) Perfk [ ( t+1 ) ·Lk ] , where op ( k ) represents the number of floating point operations of k. We further define an optimal schedule , i.e. , π∗ = { k∗1 , ... , k∗t , ... k∗T } , where k∗t is the operator at step t one could have taken in hindsight ( after seeing all performance realizations ) that would yield the maximum latency reduction . The cumulative regret is then defined as : R ( T ) = ∑T t=1 ( rt ( k ∗ t ) − rt ( kt ) ) , where k∗t is the operator at step t one could have taken in hindsight ( after seeing all performance realizations ) that would yield the maximum latency reduction . The goal is therefore to design the scheduler , such that the cumulative regret is minimized . We call this a dynamic multi-tensor-operator optimization problem .

Summary: In this paper, the authors develop DynaTune which achieves faster convergence speed to optimize a DNN model when compared to the state-of-the-art DL compiler, AutoTVM. The key idea is a time-slot-based scheduling method based on UCB-type multi-armed bandit policy. At each time, the scheduler chooses an action to maximize the latency reduction. In practice, A Bayesian belief model via MCMC is used to capture current knowledge of the optimization results to predict future performance, which helps make better decisions and expedites the convergence speed. The idea of using MAB in DL compiler is very interesting. The numerical experiments also demonstrate clear advantage of the proposed DynaTune. My concerns are as follows.