SmartMem: Layout Transformation Elimination and Adaptation for Efficient DNN Execution on Mobile
This work addresses efficiency bottlenecks for mobile DNN execution, offering significant performance improvements for applications like ChatGPT and Stable Diffusion, but it is incremental as it builds on existing optimization techniques.
This paper tackles the problem of layout transformations causing slowdowns in DNN inference on mobile devices, particularly for transformers and large models, by presenting SmartMem, a framework that eliminates most layout transformations through careful layout choices and operator grouping, achieving average speedups of up to 7.9× compared to state-of-the-art frameworks.
This work is motivated by recent developments in Deep Neural Networks, particularly the Transformer architectures underlying applications such as ChatGPT, and the need for performing inference on mobile devices. Focusing on emerging transformers (specifically the ones with computationally efficient Swin-like architectures) and large models (e.g., Stable Diffusion and LLMs) based on transformers, we observe that layout transformations between the computational operators cause a significant slowdown in these applications. This paper presents SmartMem, a comprehensive framework for eliminating most layout transformations, with the idea that multiple operators can use the same tensor layout through careful choice of layout and implementation of operations. Our approach is based on classifying the operators into four groups, and considering combinations of producer-consumer edges between the operators. We develop a set of methods for searching such layouts. Another component of our work is developing efficient memory layouts for 2.5 dimensional memory commonly seen in mobile devices. Our experimental results show that SmartMem outperforms 5 state-of-the-art DNN execution frameworks on mobile devices across 18 varied neural networks, including CNNs, Transformers with both local and global attention, as well as LLMs. In particular, compared to DNNFusion, SmartMem achieves an average speedup of 2.8$\times$, and outperforms TVM and MNN with speedups of 6.9$\times$ and 7.9$\times$, respectively, on average.