缩减共享内存占用以利用 Tensor Core 高吞吐及其灵活 API 扩展库
分布式、并行与集群计算
2023-08-30 v1
摘要
NVIDIA Tensor Core 是一种混合精度矩阵-矩阵乘加计算单元，在 NVIDIA A100 GPU 上理论峰值性能超过 300 TFlop/s。NVIDIA 提供 WMMA API 以在自定义核函数中使用 Tensor Core。使用 Tensor Core 最常见的方式是从共享内存提供输入矩阵，其带宽高于全局内存。然而，由于 Tensor Core 性能很高，共享内存与 Tensor Core 的 Bytes-per-Flops (B/F) 比值较小。因此，减少共享内存占用对于高效使用 Tensor Core 十分重要。本文通过 roofline 模型分析 Tensor Core 上的简单矩阵-矩阵乘法，发现使用 WMMA API 时共享内存带宽可能成为性能瓶颈。为缓解该问题，我们提供一个 WMMA API 扩展库来提升计算吞吐，其包含两个组件。第一个组件允许灵活操纵输入 Tensor Core 的寄存器阵列。我们评估了该库的性能提升，结果表明我们的库减少了共享内存占用并加速了基于 Tensor Core 的计算。第二个组件是在 Tensor Core 上无需额外共享内存使用的 SGEMM 仿真 API。我们展示了使用该库在 Tensor Core 上实现的单精度批量 SGEMM 仿真在 A100 GPU 上达到 54.2 TFlop/s，优于 FP32 SIMT Core 的理论峰值性能，同时达到与 cuBLAS 同等的精度。若不通过我们的库在同等寄存器用量下缩减共享内存占用，则无法达到该吞吐。
引用
@article{arxiv.2308.15152,
title = {Reducing shared memory footprint to leverage high throughput on Tensor Cores and its flexible API extension library},
author = {Hiroyuki Ootomo and Rio Yokota},
journal= {arXiv preprint arXiv:2308.15152},
year = {2023}
}
备注
HPC Asia 2023