Toward Portable GPU Performance: Julia Recursive Implementation of TRMM and TRSM

本文提出了一种高效且可移植的递归实现方法，用于在 Julia 中为 GPU 实现三角矩阵-矩阵乘法 (TRMM) 和三角求解 (TRSM)，这两种核函数是许多线性代数算法的基础。我们对 TRMM 和 TRSM 进行了重构，使得大部分计算可以通过通用矩阵-矩阵乘法 (GEMM) 来执行，从而更好地利用 GPU 的内存层次结构并降低延迟。通过利用 Julia 的多重分派和元编程功能，结合 GPUArrays 和 KernelAbstractions 框架，我们提供了一个与硬件无关的统一 API，能够在 NVIDIA、AMD 和 Apple Silicon 的 GPU 上运行。对于大规模矩阵，递归代码的吞吐量可以与厂商库（如 cuBLAS 和 rocBLAS）相媲美，同时首次在 Apple Silicon 上提供了这些功能。整个实现仅需几百行代码，证明了统一的 Julia 程序能够在异构架构上接近厂商级性能。

该论文试图解决在GPU上高效实现三角矩阵-矩阵乘法（TRMM）和三角求解（TRSM）的问题，这两个核函数是许多线性代数算法的基础。这是一个需要优化性能和跨平台兼容性的问题，尤其是在Apple Silicon等新兴架构上。

论文的关键思路是通过将TRMM和TRSM重构为通用矩阵-矩阵乘法（GEMM）来提高对GPU内存层次结构的利用并减少延迟。同时，利用Julia语言的多重分派和元编程特性，结合GPUArrays和KernelAbstractions框架，提供了一个硬件无关的API，使代码能够在NVIDIA、AMD和Apple Silicon GPU上运行。

论文展示了其递归实现能够达到与厂商库（如cuBLAS和rocBLAS）相当的吞吐量，并首次在Apple Silicon上提供了这些功能。此外，整个实现仅需几百行代码，证明了Julia语言在跨异构架构上的高效性和简洁性。论文未提及具体的数据集，但实验设计针对不同规模的矩阵进行性能测试。代码已开源，为进一步研究提供了基础。

最近的相关研究包括：1) NVIDIA cuBLAS和AMD rocBLAS，它们为特定硬件提供高度优化的BLAS实现；2) Google的XLA（Accelerated Linear Algebra），它通过编译优化提升线性代数运算性能；3) Cutlass库，专注于CUDA上的高性能矩阵乘法。其他相关研究包括使用SYCL或OpenCL实现跨平台线性代数计算的工作，例如clBLAS和oneAPI Math Kernel Library (oneMKL)。