AutoGraph: Optimizing DNN Computation Graph for Parallel GPU Kernel Execution

Deep learning frameworks optimize the computation graphs and intra-operator computations to boost the inference performance on GPUs, while inter-operator parallelism is usually ignored. In this paper, a unified framework, AutoGraph, is proposed to obtain highly optimized computation graphs in favor of parallel executions of GPU kernels. A novel dynamic programming algorithm, combined with backtracking search, is adopted to explore the optimal graph optimization solution, with the fast performance estimation from the mixed critical path cost. Accurate runtime information based on GPU Multi-Stream launched with CUDA Graph is utilized to determine the convergence of the optimization. Experimental results demonstrate that our method achieves up to 3.47x speedup over existing graph optimization methods. Moreover, AutoGraph outperforms state-of-the-art parallel kernel launch frameworks by up to 1.26x.

Authors

• Yuxuan Zhao Wuhan University
• Qi Sun Sakana AI Institute of Science Tokyo
• Zhuolun He The Chinese University of Hong Kong
• Yang Bai
• Bei Yu The Chinese University of Hong Kong

Keywords

• APP: Design
• ML: Scalability of ML Systems