一句话：SmartMem 通过让多个算子共享同一张量布局，消除大部分布局变换并适配移动端 2.5D 内存，显著加速 DNN 推理。
该工作针对移动端推理中算子之间的布局变换（layout transformation）造成显著性能损失的问题，关注新兴的 Swin 类高效注意力 Transformer 以及 Stable Diffusion、LLM 等基于 Transformer 的大模型。作者提出 SmartMem 框架，其核心思想是：通过精心选择布局并相应实现算子，使多个算子能够复用同一张量布局，从而消除大部分布局变换。方法上，SmartMem 将算子划分为四类，并系统考虑算子间生产者—消费者边的组合关系，开发了一组布局搜索方法；同时针对移动设备常见的 2.5 维内存设计了高效的内存布局。实验在移动设备上覆盖 18 个不同神经网络（含 CNN、局部与全局注意力 Transformer 以及 LLM），与 5 个最先进的 DNN 执行框架比较：相比 DNNFusion 平均加速 2.8 倍，相比 TVM 和 MNN 分别平均加速 6.9 倍和 7.9 倍。
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×, and outperforms TVM and MNN with speedups of 6.9× and 7.9×, respectively, on average.
DOI 原文 · arXiv · PDF（开放获取） · DBLP
@inproceedings{NiuSSGSYAR24,
title = {SmartMem: Layout Transformation Elimination and Adaptation for Efficient DNN Execution on Mobile},
author = {Wei Niu and Md Musfiqur Rahman Sanim and Zhihao Shu and Jiexiong Guan and Xipeng Shen and Miao Yin and Gagan Agrawal and Bin Ren},
booktitle = {ASPLOS},
year = {2024},
doi = {10.1145/3620666.3651384}
}