TVM对ARM架构GPU上的移动深度学习进行优化

在移动设备上部署深度神经网络的需求正随着深度学习技术的广泛应用而持续上升。与在桌面平台上的部署类似，借助GPU进行推理不仅能够提升处理速度，还能提高能源利用效率。然而，当前大多数深度学习框架对移动GPU的支持尚不完善。这一问题的根源在于移动GPU架构与桌面GPU架构之间存在显著差异，这使得在移动GPU上实施优化工作变得尤为复杂。因此，在移动GPU上实现高效的深度学习推理需要额外的、非平凡的努力，这也正是目前多数框架对移动GPU支持不足的主要原因。TVM通过引入统一的中间表示（IR）堆栈，有效解决了针对不同硬件平台进行部署所面临的难题，从而简化了对各类硬件的优化流程。本文阐述了如何借助TVM/NNVM为ARM Mali GPU生成高效的计算内核，并完成从模型到可执行代码的端到端编译过程。在基于Mali-T860 MP4平台的实际测试中，该方法在VGG-16模型上的推理速度较ARM计算库提升了1.4倍，在MobileNet模型上的推理速度则提升了2.2倍。这些性能提升得益于图形层面和算子层面的多重优化策略共同作用的结果。

图1展示了在ImageNet数据集上不同后端的推理速度