[论文解读] Automatic Kernel Generation for Volta Tensor Cores

Somashekaracharya G. Bhaskaracharya, Julien Demouth+1 more|arXiv (Cornell University)|Jun 22, 2020

Parallel Computing and Optimization Techniques参考文献 19被引用 13

一句话总结

本文提出一种多面体编译方法，可自动为NVIDIA Volta张量核心生成高度优化的CUDA内核，针对融合了ReLU和偏置加法等点对点操作的矩阵乘法。通过利用mma.sync.m8n8k4 PTX指令及宏-MMA组合，该方法在手动调优的库实现上实现了最高2.55倍的性能提升，证明了自动化内核生成在张量核心工作负载中的有效性。

ABSTRACT

A commonly occurring computation idiom in neural networks is to perform some pointwise operations on the result of a matrix multiplication. Such a sequence of operations is typically represented as a computation graph in deep learning compilers. When compiling to a GPU target, these computations can be individually mapped to manually tuned implementations provided by libraries such as cuBLAS and cuDNN. These libraries also provide off-the-shelf support for targeting tensor cores in NVIDIA GPUs, which can lead to huge performance boosts through their specialized support for mixed-precision matrix math. Alternatively, tensor cores can be programmed directly using CUDA APIs or inline assembly instructions, which opens up the possibility of generating efficient CUDA kernels automatically for such computations. Automatic kernel generation is particularly crucial when it is beneficial to generate efficient code for an entire computation graph by fusing several operations into a single device function instead of invoking a separate kernel for each of them. Polyhedral compilation techniques provide a systematic approach for the analysis and transformation of a sequence of affine loop-nests. In this paper, we describe a polyhedral approach to generate efficient CUDA kernels for matrix multiplication using inline assembly instructions for programming tensor cores on NVIDIA Volta GPUs. Furthermore, we build on this approach to generate fused kernels for computation sequences involving matrix multiplication and pointwise operations such as bias addition, ReLU activation etc. Experimental evaluation of these techniques show that automatically generated kernels can provide significantly better performance than manually tuned library implementations, with speedups ranging up to 2.55X.

研究动机与目标

• 解决在执行深度学习计算图时，因每个操作使用独立内核而导致的重复全局内存访问性能瓶颈。
• 实现高效CUDA内核的自动生成，以利用Volta张量核心进行混合精度矩阵乘法及融合的点对点操作。
• 将多面体编译技术扩展至直接针对mma.sync.m8n8k4 PTX指令，实现对张量核心使用的细粒度控制。
• 证明自动生成的内核在融合计算场景下可超越手动调优的库实现（如cuBLAS、cuDNN）。
• 提供一种系统化、可复用的框架，用于为张量核心加速的深度学习工作负载生成高性能内核。

提出的方法

• 该方法使用多面体框架分析并转换表示矩阵乘法和点对点操作的仿射循环嵌套。
• 生成直接使用mma.sync.m8n8k4 PTX指令编程Volta张量核心的内核，该指令在线程组上操作8×8×8的矩阵片段。
• 通过将低层级的mma.sync.m8n8k4操作组合为16×16×8和32×32×8大小的高阶宏-MMA操作，提升内存带宽利用率。
• 支持将矩阵乘法与上游（前导）和下游（尾随）点对点操作（如ReLU、Sigmoid和偏置加法）进行融合。
• 该框架与Diesel DSL编译器集成，并使用自定义调度策略优化张量核心内核的分块大小和寄存器使用。
• 采用系统化的代码生成流水线，将高层计算图映射为优化后的CUDA内核，几乎无需手动调优。

实验结果

研究问题

• RQ1 多面体编译技术能否有效应用于为Volta张量核心生成高性能CUDA内核？
• RQ2 在融合的矩阵乘法与点对点操作场景下，自动内核生成在多大程度上能超越手动调优的库（如cuBLAS和cuDNN）？
• RQ3 使用mma.sync.m8n8k4 PTX指令在融合内核中实现高算术强度和内存带宽效率的效果如何？
• RQ4 在不同工作负载下，张量核心内核中不同分块大小和融合策略的性能权衡是什么？
• RQ5 能否通过多面体分析自动生成并优化16×16×8和32×32×8大小的宏-MMA组合，以超越原子mma操作的性能？

主要发现

• 所提出的方法在融合的矩阵乘法与点对点操作上，相较于cuBLAS和cuDNN最高实现了2.55倍的性能提升，尤其在较大问题规模下表现显著。
• 对于单个矩阵乘法内核，自动生成方法相较cuBLAS实现了1.46倍的性能提升，相较cuDNN实现了1.35倍，即使在无融合场景下也表现出有效性。
• 包含多个矩阵乘法的融合内核表现出更高的寄存器压力，需采用更小的分块大小，虽降低了性能增益但未完全消除。
• 与Halide和TVM等现有框架相比，该方法在张量核心内核生成方面表现更优，主要得益于对mma.sync.m8n8k4指令的直接使用和宏-MMA组合。
• 该方法在广泛的问题规模范围内实现了显著的性能提升，最大增益出现在涉及矩阵乘法与ReLU或偏置加法的融合计算图中。
• 本工作证明，多面体编译可有效扩展至低层级张量核心指令，实现高度优化内核的自动生成。