基于 SIMD 架构的三种张量布局的高性能 Im2win 与直接卷积
机器学习
2024-08-02 v1 人工智能
神经与进化计算
摘要
卷积是深度神经网络的核心组件,计算密集且耗时。张量数据布局对卷积操作的内存访问和计算效率有显著影响。然而,目前仍缺乏对 SIMD 架构上数据布局在卷积方法方面的全面性能表征。本文为 im2win 卷积提出了三种新型数据布局:NHWC、CHWN 和 CHWN8,并引入了一套针对直接卷积和 im2win 卷积的通用优化技术。我们将优化后的 im2win 卷积与直接卷积以及 PyTorch 的 im2col 卷积在上述布局上进行比较。实验表明,采用新 NHWC 布局的 im2win 卷积相比 NCHW 布局实现了高达 355% 的性能加速。我们的优化还显著提升了 im2win 和直接卷积的性能。优化后的 im2win 和直接卷积分别达到了机器理论峰值的 95% 和 94%。
引用
@article{arxiv.2408.00278,
title = {High Performance Im2win and Direct Convolutions using Three Tensor Layouts on SIMD Architectures},
author = {Xiang Fu and Xinpeng Zhang and Jixiang Ma and Peng Zhao and Shuai Lu and Xu T. Liu},
journal= {arXiv preprint arXiv:2408.00278},
year = {2024}
}