AGO: Boosting Mobile AI Inference Performance by Removing Constraints on Graph Optimization
This work addresses performance bottlenecks in mobile AI inference for developers and users by enabling more flexible graph optimization, though it is incremental as it builds on existing compiler techniques.
The paper tackles the problem of constrained subgraph generation in deep learning compilers, which limits graph optimization, and proposes AGO, a framework that removes these constraints to boost mobile AI inference performance, achieving up to 3.3x speedup compared to state-of-the-art compilers.
Traditional deep learning compilers rely on heuristics for subgraph generation, which impose extra constraints on graph optimization, e.g., each subgraph can only contain at most one complex operator. In this paper, we propose AGO, a framework for graph optimization with arbitrary structures to boost the inference performance of deep models by removing such constraints. To create new optimization opportunities for complicated subgraphs, we propose intensive operator fusion, which can effectively stitch multiple complex operators together for better performance. Further, we design a graph partitioning scheme that allows an arbitrary structure for each subgraph while guaranteeing the acyclic property among all generated subgraphs. Additionally, to enable efficient performance tuning on complicated subgraphs, we devise a novel divide-and-conquer tuning mechanism to orchestrate different system components. Through extensive experiments on various neural networks and mobile devices, we show that our system can improve the inference performance by up to 3.3x when compared with state-of-the-art deep compilers.