中文

swTVM:面向神威众核处理器上深度学习优化的张量代码生成

机器学习 2022-07-12 v3 编程语言 机器学习

摘要

深度学习框架与硬件平台的蓬勃发展要求一种高效的编译器,以屏蔽软硬件的多样性从而提供应用可移植性。在现有深度学习编译器中,TVM 以其在多种硬件设备上代码生成与优化的高效性而著称。同时,神威众核处理器凭借在科学计算与深度学习负载中诱人的算力成为极具竞争力的候选。本文结合了这两方面的趋势。具体地,我们提出 swTVM,对原始 TVM 进行扩展以支持对需要交叉编译的架构(如神威)进行提前编译。此外,我们在编译中利用架构特性,如核组实现大规模并行、DMA 实现高带宽内存传输、本地设备内存实现数据局部性,从而为神威上的深度学习负载生成高效代码。实验结果表明,在六个代表性基准上,swTVM 生成的代码相较神威上最先进的深度学习框架平均达到 1.79 倍加速。本工作是首次从编译器视角弥合深度学习与神威处理器之间鸿沟的尝试,特别兼顾生产力与效率。我们相信该工作将鼓励更多人拥抱深度学习与神威众核处理器的算力。

关键词

引用

@article{arxiv.1904.07404,
  title  = {swTVM: Towards Optimized Tensor Code Generation for Deep Learning on Sunway Many-Core Processor},
  author = {Mingzhen Li and Changxi Liu and Jianjin Liao and Xuegui Zheng and Hailong Yang and Rujun Sun and Jun Xu and Lin Gan and Guangwen Yang and Zhongzhi Luan and Depei Qian},
  journal= {arXiv preprint arXiv:1904.07404},
  year   = {2022}
}