[论文解读] Implementing Strassen's Algorithm with CUTLASS on NVIDIA Volta GPUs

Jianyu Huang, Chenhan D. Yu+1 more|arXiv (Cornell University)|Aug 24, 2018

Parallel Computing and Optimization Techniques参考文献 21被引用 9

一句话总结

本论文提出了一种新颖的、内存高效的NVIDIA Volta GPU上的Strassen算法实现，采用CUTLASS框架，通过重用GEMM内核中的共享内存和寄存器文件，避免了额外的工作空间。与cublasSgemm相比，其1级Strassen的加速比最高达1.11×，临界点低至1,536；2级Strassen的加速比最高达1.19×，临界点为7,680。该实现利用融合操作、基于任务的并行计算以及预测性性能模型，显著提升了性能。

ABSTRACT

Conventional GPU implementations of Strassen's algorithm (Strassen) typically rely on the existing high-performance matrix multiplication (GEMM), trading space for time. As a result, such approaches can only achieve practical speedup for relatively large, "squarish" matrices due to the extra memory overhead, and their usages are limited due to the considerable workspace. We present novel Strassen primitives for GPUs that can be composed to generate a family of Strassen algorithms. Our algorithms utilize both the memory and thread hierarchies on GPUs, reusing shared memory and register files inherited from GEMM, fusing additional operations, and avoiding extra workspace. We further exploit intra- and inter-kernel parallelism by batching, streaming, and employing atomic operations. We also develop a performance model for NVIDIA Volta GPUs to select the appropriate blocking parameters and predict the performance for GEMM and Strassen. Overall, our 1-level Strassen can achieve up to 1.11x speedup with a crossover point as small as 1,536 compared to cublasSgemm on a NVIDIA Tesla V100 GPU. With additional workspace, our 2-level Strassen can achieve 1.19x speedup with a crossover point at 7,680.

研究动机与目标

• 为克服GPU上Strassen算法的实际限制，如高内存开销以及因工作空间需求导致的问题规模受限。
• 通过消除额外的全局内存和共享内存使用，使Strassen算法在小规模和非方阵矩阵上超越传统GEMM。
• 通过批处理、流式处理和原子操作，在不增加内核启动开销的前提下，有效利用GPU的核内与核间并行性。
• 为V100 GPU开发准确的性能模型，以指导最优分块参数选择并预测运行时性能。
• 证明Strassen算法可在现代GPU上实现高效计算，且内存占用极低、算术强度高，即使在小规模问题上亦表现优异。

提出的方法

• 设计新型GPU Strassen内核，将内存访问与算术操作直接融合进GEMM流水线，重用现有共享内存和寄存器文件，从而消除额外工作空间。
• 通过内核特化和操作融合优化寄存器使用，降低寄存器压力，同时保持高线程占用率。
• 利用线程块、瓦片（warp）和线程之间的核内并行性，以及通过批处理和流式处理实现的核间并行性，实现计算与内存带宽的重叠。
• 使用原子操作管理Strassen计算图中的并发更新，实现可扩展的任务级并行计算。
• 基于算术强度、内存操作和寄存器使用，为V100 GPU开发性能模型，用于预测性能并选择最优分块大小。
• 实现1级和2级Strassen变体，通过精心设计的内存访问模式和操作数预取，最小化延迟并最大化重叠性。

实验结果

研究问题

• RQ1 Strassen算法能否在GPU上实现，而无需额外的全局或共享内存工作空间，仅依赖标准GEMM的内存布局？
• RQ2 在无额外内存使用的前提下，现代V100 GPU上Strassen算法在多小的矩阵规模下可超越cublasSgemm？
• RQ3 如何在Strassen中有效利用核内与核间并行性，以维持高占用率并隐藏内存延迟？
• RQ4 基于算术强度和内存操作的性能模型，能在多大程度上准确预测V100 GPU上Strassen的性能表现？
• RQ5 分块参数的选择如何影响Strassen相对于GEMM的性能表现？如何实现最优参数配置？

主要发现

• 1级Strassen实现的临界点为1,536（在Tesla V100 GPU上），在1,536×1,536矩阵规模下已超越cublasSgemm。
• 2级Strassen实现的临界点为7,680，相较于cublasSgemm的峰值加速比达1.19×，远低于先前最先进方法所需的13,312规模。
• 所提实现相比标准GEMM，未引入任何额外的全局或共享内存，彻底消除了与工作空间相关的内存瓶颈。
• 性能模型能准确预测运行时性能，并有效指导GEMM与Strassen内核的最优分块参数选择。
• 通过操作融合和重叠内存访问与计算阶段，实现高算术强度，显著降低内存带宽限制的影响。
• 该方法证明了Strassen算法在小规模和非方阵矩阵上的GPU实现具有实际可行性，此前因内存与性能开销而难以与传统GEMM竞争。