Polyhedral Optimization of TensorFlow Computation Graphs.

We present \({\textsf {R}}\text {-}{\textsf {Stream}}{\cdot }{\textsf {TF}}\), a polyhedral optimization tool for neural network computations. \({\textsf {R}}\text {-}{\textsf {Stream}}{\cdot }{\textsf {TF}}\) transforms computations performed in a neural network graph into C programs suited to the polyhedral representation and uses R-Stream, a polyhedral compiler, to parallelize and optimize the computations performed in the graph. \({\textsf {R}}\text {-}{\textsf {Stream}}{\cdot }{\textsf {TF}}\) can exploit the optimizations available with R-Stream to generate a highly optimized version of the computation graph, specifically mapped to the targeted architecture. During our experiments, \({\textsf {R}}\text {-}{\textsf {Stream}}{\cdot }{\textsf {TF}}\) was able to automatically reach performance levels close to the hand-optimized implementations, demonstrating its utility in porting neural network computations to parallel architectures.