Convolution is one of the fundamental operations of deep neural networks with
demanding matrix computation. In a graphic processing unit (GPU), Tensor Core
is a specialized matrix processing hardware equipped with reduced-precision
matrix-multiply-accumulate (MMA) instructions to increase throughput. However,
it is challenging to achieve optimal performance since the best scheduling of
MMA instructions varies for different convolution sizes. In particular,
reduced-precision MMA requires many elements grouped as a matrix operand,
seriously limiting data reuse and imposing packing and layout overhead on the
schedule. This work proposes an automatic scheduling method of
reduced-precision MMA for convolution operation. In this method, we devise a
search space that explores the thread tile and warp sizes to increase the data
reuse despite a large matrix operand of reduced-precision MMA. The search space
also includes options of register-level packing and layout optimization to
lesson overhead of handling reduced-precision data. Finally, we propose a
search algorithm to find the best schedule by learning from the distinctive
candidates. This reduced-precision MMA optimization method is evaluated on
convolution operations of popular neural networks to demonstrate substantial
speedup on Tensor Core compared to the state of the arts with shortened search
time.

这篇论文提出了一种基于 Tensor Core 和 MMA 指令的卷积操作优化方法，该方法通过自动调度探索线程块和 warp 尺寸的搜索空间，包括寄存器级 Packing 和布局优化等选项，使用学习算法找到最佳调度，从而比现有技术在 Tensor Core 上实现了大幅度加速。