利用 GPU 张量核实现快速算术归约

本工作提出一种 GPU 张量核方法，将 n 个数的算术归约编码为一组由 GPU 张量核并行执行的链式 $m \times m$ 矩阵乘累加（MMA）操作。所提链式张量核方法的渐近运行时间为 $T(n)=5 log_{m^2}{n}$，相对于经典 $O(n \log n)$ 并行归约算法的加速比为 $S=\dfrac{4}{5} log_{2}{m^2}$。实验性能结果表明，所提归约方法比常规 GPU 归约实现快约 $3.2 \times$，且保持数值精度，因为每条含 R 个 MMA 的链的中间结果在最终归约为一个 32 位结果前均保留为 32 位浮点值。链式 MMA 设计允许线程块的灵活配置；32 或 128 线程的小线程块仍可通过每块 $R=4,5$ 个 MMA 的链达到最大性能，而大线程块在 $R=1$ 时表现最佳。本工作所得结果表明，张量核确实能为非机器学习应用（如作为研究许多科学现象集成工具的算术归约）带来显著性能提升。