NeoFlow: A Flexible Framework for Enabling Efficient Compilation for High Performance DNN Training

Abstract

Deep neural networks (DNNs) are increasingly deployed in various image recognition and natural language processing applications. The continuous demand for accuracy and high performance has led to innovations in DNN design and a proliferation of new operators. However, existing DNN training frameworks such as PyTorch and TensorFlow only support a limited range of operators and rely on hand-optimized libraries to provide efficient implementations for these operators. To evaluate novel neural networks with new operators, the programmers have to either replace the holistic new operators with existing operators or provide low-level implementations manually. Therefore, a critical requirement for DNN training frameworks is to provide high-performance implementations for the neural networks containing new operators automatically in the absence of efficient library support. In this paper, we introduce NeoFlow, which is a flexible framework for enabling efficient compilation for high-performance DNN training. NeoFlow allows the programmers to directly write customized expressions as new operators to be mapped to graph representation and low-level implementations automatically, providing both high programming productivity and high performance. First, NeoFlow provides expression-based automatic differentiation to support customized model definitions with new operators. Then, NeoFlow proposes an efficient compilation system that partitions the neural network graph into subgraphs, explores optimized schedules, and generates high-performance libraries for subgraphs automatically. Finally, NeoFlow develops an efficient runtime system to combine the compilation and training as a whole by overlapping their execution. In the experiments, we examine the numerical accuracy and performance of NeoFlow. The results show that NeoFlow can achieve similar or even better performance at the operator and whole graph level for DNNs compared to deep learning frameworks. Especially, for novel networks training, the geometric mean speedups of NeoFlow to PyTorch, TensorFlow, and CuDNN are 3.16X, 2.43X, and 1.92X, respectively.

An End-to-End Programming Model for AI Engine Architectures

“…Recently, deep learning compilers (e.g., TVM , MLIR (Lattner et al, 2020b), and Glow ) have demonstrated the ability to reduce dramatically inference latencies , training times (Zheng et al, 2022), and memory usage . These compilers function by extracting intermediate-level representations (IRs) of the DNNs from the representations produced by the frameworks, and performing various optimizations (e.g., kernel fusion (Ashari et al, 2015), vectorization (Maleki et al, 2011), and memory planning ) on those IRs.…”

Venturing into the Age of AI: Insights and Perspectives

“…Therefore, the convolution products for every point are added together, where the result of the computation output is nonetheless but activation output, named output feature map. Moreover, all input feature maps are processed together as a batch (Yang, 2023), resulting in improvement of filter weights. Additionally, there are other optional layers, as observed in the figure above like, nonlinearity (generally it can evaluate the maximum value of two intersecting function), pooling (it makes the network to withstand to any invariance or distortion) and normalization which is nonetheless but, controlling the input distribution through the layers.…”

Amos

“…Halide [49] introduces the concept of compute and schedule to represent software and optimization, while TVM [9] generalizes the concept and allows users to use tensorize primitive to lower part of software to spatial accelerators manually. Automatic schedulers such as Halide Scheduler [2,30,38], FlexTensor [70], Pro-Tuner [19], ALT [63], Rammer [34], NeoFlow [69], and Ansor [68] focus on general-purpose hardware and ignore the mapping problem for spatial accelerators such as Tensor Core. Polyhedral model is widely used in compilers for constrained optimization [5,56,57].…”