Polyhedral Optimization of TensorFlow Computation Graphs Benoˆıt Pradelle, Benoˆıt Meister, Muthu Baskaran, Jonathan Springer and Richard Lethin Reservoir Labs [email protected] Abstract—We present R-Stream·TF, a polyhedral optimization currently being developed. The most representative approach tool for neural network computations. R-Stream·TF transforms is XLA, a just-in-time compiler for TensorFlow computation computations performed in a neural network graph into C graphs. XLA has shown its ability to significantly speed up the programs suited to the polyhedral representation and uses R-Stream, a polyhedral compiler, to parallelize and optimize computations performed in TensorFlow, but it seems limited the computations performed in the graph. R-Stream·TF can to basic pattern-matching optimizations. Only simple cases of exploit the optimizations available with R-Stream to generate fusion and array contraction can be realistically achieved with a highly optimized version of the computation graph, specifically this method. mapped to the targeted architecture. During our experiments, Our contribution is to extend and generalize the approach of R-Stream·TF was able to automatically reach performance levels close to the hand-optimized implementations, demonstrating graph optimizers through polyhedral optimization techniques. its utility in porting neural network computations to parallel Compilers and optimizers based on the polyhedral model can architectures. apply powerful code transformations on programs, using a precise mathematical representation of the code. I. INTRODUCTION Polyhedral optimizations encompass fusion and array con- Deep Convolutional Neural Networks (DCNN) [1], and traction, but they also subsume any combination of loop more generally deep learning, recently reached maturity. Im- fusion/fission, interchange, skewing, and reversals. Data de- pressive results achieved in recent years demonstrated the tech- pendencies are exact in the polyhedral model, enabling au- nology was ripe for general, practical use. New applications tomatic parallelization and other common memory-oriented are developed every day, and deep learning is already ubiqui- optimizations such as loop tiling [6] and data layout transfor- tous in our lives. This considerable activity around machine mations [7]. The polyhedral model is most precise on regions learning is becoming increasingly structured around a few with affine constructs [8], which include most of the classical common tools. For instance, Caffe [2], Torch [3], CNTK [4], neural network operators. and TensorFlow [5] are popular frameworks commonly used In this paper, we present R-Stream·TF, a new optimizer to develop and exploit neural networks. These frameworks for TensorFlow computation graphs. R-Stream·TF emits high- are based on a similar concept: high-level operations such as level sequential C code implementing the exact computations convolutions and pooling are exposed to the user, who can performed in the input TensorFlow graph. The generated C design networks simply by composing them as operators. The code is specific to the graph: it is specialized to the exact frameworks also empower users by facilitating data prepro- tensor shapes and types used as the input and output of cessing and streamlining back-propagation for training. every operation. The generated C code is then automatically Applications based on neural networks often require strict parallelized and optimized by R-Stream [9], a polyhedral com- performance constraints to be enforced, such as when per- piler developed by Reservoir Labs. R-Stream optimizes the forming interactive tasks. They also require high throughput computation specifically for the target architecture. R-Stream (bandwidth) such as when performing many queries simulta- supports numerous targets including common x86 CPUs and neously. To minimize the latency and bandwidth required to GPUs and can generate the code to parallel programming process an input sample, neural networks frameworks rely on models including OpenMP, CUDA, OpenCL, POSIX threads highly optimized implementations. A common approach for and task-based runtimes APIs [10], [11]. speeding up neural network processing consists in building a R-Stream·TF extracts and merges TensorFlow operator sub- hand-optimized library of DCNN operators. graphs, and lets R-Stream apply a full range of polyhedral While this method significantly improves the performance optimizations to the underlying computations. Such transfor- of the computations, the hand optimization effort is tedious and mations are specific to the target platform, which can have error-prone. It is also inherently unable to exploit optimization several levels of parallelism and any combination of caches opportunities available by combining successive operations. and scratchpads [12]. The result is a set of highly optimized For instance, an element-wise operation and a convolution can parallel TensorFlow operators, which are reintegrated into the be computed more efficiently if both operations are fused. original TensorFlow computation graph. In order to benefit from these optimization opportunities, The main benefit of our approach is the ability to use several graph-based optimizers have been proposed and are the full set of polyhedral optimizations within computation Fig. 3. Three specializations of the element-wise addition. The generated functions have specific data types, data sizes, and broadcasting even though the TensorFlow operator is the same. number can be quite large, the number of nodes in current DC- NNs is typically larger. In order to ensure that the optimization remains tractable, R-Stream·TF pre-processes the input graph, extracting subgraphs that are optimized independently from each other. Fig. 1. TensorFlow graphs are converted into simple sequential C code, optimized using the R-Stream compiler, wrapped in a custom TensorFlow Partitioning operator graphs in order to expose better opti- operators, and finally stiched back in the computation graph. mization opportunities and maintain scalability (i.e., tractable optimization times) has been studied many times over, from instruction set synthesis [13] to streaming graphs (with e.g., subgraphs. This ability is superior to current approaches based [14], [15]). Optimality of subgraphs is typically defined by the on domain-specific optimizations, since it enables several amount of computation and data reuse within the subgraph, additional optimizations to be performed automatically on and the amount (and weight) of resulting in- and out-edges. computation graphs. Because the optimizations applied to the In parallel computing frameworks, grain of parallelism and graph are both specialized to the graph itself and to the load balancing are also important optimality criteria. target architecture, R-Stream·TF generates highly optimized The most impactful constraints are similar in our case. code, specifically tailored for the target platform. This makes Additionally, code generators may not be available for some R-Stream·TF an adequate automatic porting tool, providing operators, which should then not be included in a subgraph an optimized support for TensorFlow computations on new to optimize. While we plan to implement more sophisticated architectures. subgraph selection algorithms in the future, we meet the proof- The rest of the paper is organized as follows. The design of-concept objective of this paper using a simple two-step of R-Stream·TF is presented in details in Section II. The tool approach. First, we identify connected subgraphs in the overall has been implemented and evaluated on popular DCNNs. The computation graph that are exclusively made of operations for evaluation results are presented in Section III. We compare which a code generator is available. Second, the connected R-Stream·TF to existing systems in Section V, before con- subgraphs are partitioned when they are estimated to be too cluding in Section VI. large for the optimizer. These steps are illustrated in Figure 2. The second step is expressed as a balanced k-way partitioning II. DESIGN problem, where the objective is to minimize the number of A. Overview graph edges between two partitions. Edges in the computation The overall flow of R-Stream·TF, described in Figure 1, graph represent tensors on which the operations are performed. starts with a TensorFlow computation graph and results in an While R-Stream is free to change the data layout in the optimized graph with the same semantics. The optimization tensors within partitions, transforming the layout of tensors process is performed in several successive steps. First, opti- used across several partitions is illegal. Thus, by minimizing mizable subgraphs of the overall operator graph are identified. the number of edges between the partitions, R-Stream·TF Restricting the process to well-structured subgraphs ensures increases data layout optimization opportunities for R-Stream, that optimization is possible and tractable. Simple sequential including array contraction, expansion, and spatial locality C code is then generated for every identified operator. The optimizations. sequential source code is sent as-is to R-Stream to be paral- R-Stream works from C code, which is generated as the lelized and optimized. The resulting parallel code is wrapped next phase of the R-Stream·TF() optimization process. in a custom C++ TensorFlow operator automatically generated by R-Stream·TF. The operator itself implements the operation C. Operator Code Generators