从差异化候选中学习以优化张量核上低精度卷积程序
机器学习
2022-02-25 v2 人工智能
摘要
卷积是具有严苛矩阵计算需求的深度神经网络基本操作之一。在图形处理单元（GPU）中，Tensor Core 是一种专用的矩阵处理硬件，配备有低精度矩阵乘加（MMA）指令以提升吞吐量。然而，由于不同卷积尺寸下 MMA 指令的最佳调度各异，实现最优性能颇具挑战。特别是，低精度 MMA 需要将许多元素编组为矩阵操作数，严重限制了数据复用，并为调度带来打包与布局开销。本工作提出一种面向卷积操作的低精度 MMA 自动调度方法。在该方法中，我们设计了一个搜索空间，探索线程块（thread tile）与 warp 大小以提升数据复用，尽管低精度 MMA 的矩阵操作数较大。该搜索空间还包含寄存器级打包与布局优化选项，以减轻处理低精度数据的开销。最后，我们提出一种通过从差异化候选中学习的搜索算法来寻找最佳调度。该低精度 MMA 优化方法在流行神经网络的卷积操作上进行了评估，相较于最先进方法在 Tensor Core 上取得了显著加速且搜索时间更短。
引用
@article{arxiv.2202.06819,
title = {Learning from distinctive candidates to optimize reduced-precision convolution program on tensor cores},
author = {Junkyeong Choi and Hyucksung Kwon and Woongkyu Lee and Jungwook Choi and Jieun Lim},
journal= {arXiv preprint arXiv:2202.06819},
year = {2022}
}
备注
10 pages, 16 figures, preliminary work