全文 part1 - DGEMM Using Tensor Cores, and Its Accurate and Reproducible Versions

摘要

本文提出了一种在 NVIDIA 图形处理器（GPU）的张量核心（Tensor Cores，仅含 FP16、INT8 等 GEMM 计算功能）上实现 FP64（双精度，DGEMM）和 FP32（单精度，SGEMM）稠密矩阵乘法的方法。Tensor Cores 是一种特殊的处理单元，可对 FP16 输入执行内联图形矩阵乘法并以 FP32 精度运算，最终返回 FP32 精度的结果。该方法采用 Ozaki 方案 —— 一种基于无误差变换的精确矩阵乘法算法。所提方法具有三大显著优势：其一，可基于 cuBLAS 库的 cublasGemmEx 例程，借助 Tensor Cores 操作构建；其二，能实现比标准 DGEMM 更高的精度，包括正确舍入结果；其三，即便核心与线程数量不同，也能确保位级可复现性。

该方法的可达性能取决于输入矩阵各元素的绝对值范围。例如，当矩阵以动态范围为 1E+9 的随机数初始化时，在 Titan RTX GPU（Tensor Cores 算力为 130 TFlops）上，我们实现的等效 DGEMM 操作可达约 980 GFlops 的 FP64 运算性能，而 cuBLAS 库的 cublasDgemm 在 FP64 浮点单元上仅能达到 539 GFlops。研究结果表明，可利用 FP32/FP64 资源有限但低精度处理单元（如面向 AI 的处理器）运算速度快的硬件，来处理通用工作负载。

关键词：Tensor Cores ；FP16；半精度；低精度；矩阵乘法；GEMM；线性代数；精度；可复现性

1. 引言

近年来，深度学习应用的日益增多推动了特殊处理单元的发展，例如 NVIDIA GPU 上的 Tensor Cores 以及谷歌的张量处理单元（TPUs）。此类任务的核心是矩阵乘法，其并不需要 IEEE 754-2008 binary32（即单精度，FP32，含 8 位指数和 23 位尾数）和 binary64（即双精度，FP64，含 11 位指数和 52 位尾数）这样的高精度。相应地，硬件更支持快速的低精度运算，如 binary16（即半精度，FP16，含 5 位指数和 10 位尾数）以及 INT8/INT16 位整数运算。

一个广为熟知的例子是 Volta 架构中引入的 Tensor Cores ，其每时钟周期可通过融合乘加操作完成一次 4x4 矩阵乘法。尽管 Tensor Cores 支持多种数据格式和计算精度，但本文聚焦于 FP32 精度模式下的 FP16 计算 —— 该模式以 FP32 精度执行 $d=a \times b +c$ 运算（图 1）。其中， $a$ 和 $b$ 为 FP16 值， $c$ 和 $d$ 为 FP32 值。Tensor Cores 的运算速度最高可达 CUDA 核心上标准 FP32 浮点单元（FPUs）的 8 倍。已有诸多研究在通用任务中充分利用了 Tensor Cores 的这一卓越性能。

本文提出一种在 Tensor Cores 上计算 Level-3 基本线性代数子程序（BLAS）中通用矩阵乘法例程（GEMM）的方法，可支持 FP64（DGEMM）和 FP32（SGEMM）精度。GEMM 是众多科学计算工作负载以及高性能 Linpack 测试中的核心操作之一。该方法基于 Ozaki 等人提出的一种基于无误差变换的精确矩阵乘法算法（亦称 Ozaki 方案）[13]。该方法的优势如下：

高效实用：基于 NVIDIA 提供的 cuBLAS 库中的 cublasGemmEx 例程构建，开发成本低；
精度更高：即便采用正确舍入，也能实现比标准 SGEMM 和 DGEMM 更高的精度；
可复现性强：对于相同输入，即便每次执行时核心与线程数量不同，也能得到相同（位级一致）的结果；
适应性好：其理念可适配其他精度场景。

部分研究通过利用低精度硬件来加速无需高精度的计算，而本文则尝试借助低精度硬件实现更高精度的计算。我们的 DGEMM 实现仅在 FP64 支持有限的处理器上性能优于 cuBLAS 的 DGEMM。但相比之下，比 FP64 浮点单元的性能提升并非关键，更重要的是旨在提升低精度硬件（如面向人工智能的处理器）的潜力。此外，该方法为兼顾 AI 与传统高性能计算（HPC）工作负载的高效硬件设计提供了新视角。例如，可减少 FP64 资源数量，转而提供大规模低精度支持。

本文其余部分结构如下：第 2 节介绍相关工作；第 3 节阐述基于 Ozaki 方案的方法；第 4 节介绍该方法的实现；第 5 节在 Titan RTX 和 Tesla V100 GPU 上进行精度与性能评估；第 6 节探讨基于该方法的未来硬件设计方向；第 7 节总结全文。

2. 相关工作

已有多项研究尝试将为 AI 工作负载设计的低精度硬件用于其他场景。例如，Haidar 等人 [7] 在稠密与稀疏线性系统中，结合迭代精化技术利用了标准 FP16 和 FP32 精度模式下的张量核心操作。还有研究针对能效提升展开探讨 [6]，Carson 与 Higham 则对此进行了误差分析 [1]。Yang 等人 [16] 在谷歌 TPUs 上采用 bfloat16（BF16，含 8 位指数和 7 位尾数）实现了伊辛模型的蒙特卡洛模拟。这些研究将低精度运算应用于代码中无需高精度的部分 —— 此类部分可在该精度水平下完成计算，因此其适用性取决于具体算法或问题。

与本文类似，部分研究尝试借助低精度硬件实现比硬件原生精度更高的运算。例如，Markidis 等人 [11] 提出了一种提升张量核心矩阵乘法计算精度的方法。尽管其方法在概念上与本文相似，但仅能对 FP16 支持的动态范围内的矩阵进行计算，且精度与 SGEMM 相当。Henry 等人 [8] 探讨了在 BF16 浮点单元上采用双双精度算术 [2]（一种经典的二重精度算术技术）执行高精度运算的性能。Sorna 等人 [15] 提出了一种提升张量核心上二维快速傅里叶变换精度的方法。需注意的是，这些研究中，获得超过 FP32 或 FP64 浮点单元的性能并非关键，重点同样在于提升低精度硬件的潜力。因此，可能需要重新设计硬件，以平衡浮点单元所支持的精度。本文的探讨方向与之相近。

本文所提方法的核心 ——Ozaki 方案，最初是为通过标准浮点运算实现精确矩阵乘法而提出的。OzBLAS [12] 基于 Ozaki 方案在 CPU 和 GPU 上实现了精确且可复现的 BLAS 例程。OzBLAS 基于 FP64 浮点单元上的 DGEMM 构建【注，结果的精度高于原始的 FP64 gemm】，而本文中的 Ozaki 方案则借助 Tensor Cores 上的 GEMM 执行 DGEMM/SGEMM 操作。Ichimura 等人 [10] 还报道了在多核 CPU 上基于 FP64 运算的 Ozaki 方案高性能实现。

3. 方法

本节首先介绍通过改进的 Ozaki 方案在 Tensor Cores 上计算 DGEMM 的基本方案，随后介绍用于提升计算速度的附加技术。本文中， $\text{fl}_\text{FP64}(\cdots)$ 和 $\text{fl}_\text{FP32}(\cdots)$ 分别表示 FP64 和 FP32 算术下的计算； $\mathbf{u}_\text{FP64}$ 和 $\mathbf{u}_\text{FP32}$ 分别表示 FP64（ $\mathbf{u}_\text{FP64}=2^{-53}$ ）和 FP32（ $\mathbf{u}_\text{FP32}=2^{-24}$ ）的单位舍入误差【注, 单位舍入误差解释】； $\mathbb{F}_\text{FP64}$ 和 $\mathbb{F}_\text{FP16}$ 分别表示 FP64 和 FP16 浮点数的集合； $\mathbb{N}$ 表示包含零的自然数集合。

Fig. 2. Schematic of matrix multiplication ( $C=AB$ ) by Ozaki scheme (in this figure, scaling is omitted).

3.1 适用于张量核心的 Ozaki 方案

Ozaki 方案对矩阵乘法进行无误差变换，具体而言，将矩阵乘法转化为多个矩阵乘法的求和 —— 这些矩阵乘法可通过浮点运算执行且无舍入误差。图 2 为整个 Ozaki 方案的示意图，该方法包含三个主要步骤：

拆分：对输入矩阵按元素拆分得到多个子矩阵；
计算：计算所有子矩阵间的两两矩阵乘积；
求和：对上述所有子矩阵间的两两矩阵乘积按元素求和。

下面详细描述各步骤。为简洁起见，我们以两个向量 $x,y \in {\mathbb{F}_{\text{FP64}}}^n$ 的内积为例进行说明，但该方法可自然扩展到矩阵乘法（矩阵乘法本质上由多个内积构成）。此外，尽管此处仅描述 DGEMM 的情况，但其理念同样适用于 SGEMM。

步骤 1：拆分。算法 1 将 FP64 输入向量拆分为多个 FP16 向量，过程如下：

首先得到 FP64 格式的子向量，再将其转换（降尺度）为 FP16 格式。该转换仅调整指数，不会导致有效数字的 bit 丢失。其中， $2^{c^{(p)}}$ 和 $2^{c^{(q)}}$ 分别为向量 $x^{(p)}$ 和 $y^{(q)}$ 的指数从 FP64 到 FP16 的缩小因子。在算法 1 的第 7 行中，当 $\mu =\text{DBL\_MAX}$ 时， $\tau$ 达到 1024，这意味着 $c^{(p)}$ 和 $c^{(q)}$ 可存储为 2 字节 short 类型的整数。拆分算法需满足以下特性：

1. 若 ${x^{(p)}}_i$ 和 ${y^{(q)}}_j$ 为非零元素，则

2. $(x^{(p)})^Ty^{(q)}$ 在 FP32 计算中必须无误差：

拆分可理解为从浮点表示到定点表示的转换。上述第一个特性意味着，通过省略最低阶的部分子向量，可控制最终结果的精度（降低精度）。在一定数量的子向量下可获得的最终结果精度，取决于内积的长度以及输入向量各元素的绝对值范围。需注意的是，若要将 Tensor Cores 替换为其他精度的浮点单元，需修改算法 1 中的参数 $\rho$ ，且被分割的向量（ $x^{(p)}$ 和 $y^{(q)}$ ）的 bit 数取决于浮点单元的精度。

步骤 2：计算。接下来，内积 $x^Ty$ 按如下方式计算：

此处需计算子向量间所有两两内积，共需计算 $s_xs_y$ 个内积。 $2^{c^{(x_p)}+c^{(y_q)}}$ 为放大因子，用于补偿拆分过程中的缩小操作。根据算法 1 的第二个特性，由于输入为 FP16 格式，子向量的内积可通过 Tensor Core 操作计算。将该示例扩展到矩阵乘法时，必须基于标准内积算法对子矩阵进行乘法 —— 不允许采用 Strassen 算法等分治方法。

步骤 3：求和。最后，对子向量的内积进行求和。若所需精度为标准 DGEMM 的精度，求和可通过 FP64 算术完成。但由于步骤 2 中的 $\text{fl}_\text{FP32}((x^{(p)})^Ty^{(q)})$ 无舍入误差（无误差），通过例如 NearSum [14] 等正确舍入方法进行求和，可得到 $x^Ty$ 的正确舍入结果。即便在 FP64 算术中，若采用可复现的求和方法，结果也具有可复现性。由于求和是按元素进行的，计算顺序易于固定。

矩阵乘法的完整流程。算法 2 在 Tensor Cores 上实现了矩阵乘法的完整 Ozaki 方案。其中，SplitA 和 SplitB 分别沿矩阵 $A$ 和 $B$ 的内积方向（k 维），采用算法 1 进行拆分。需注意的是，由于 $A_\text{split}$ 和 $B_\text{split}$ 可存储为 FP16 格式， $\text{GEMM}_\text{FP32}$ 可通过 cuBLAS 库的 cublasGemmEx 例程，在 Tensor Core 上以 FP32 精度执行 FP16 计算。

3.2 快速计算技术

为进一步提升性能，可通过以下方法修改算法 1 或算法 2。第 4 节将讨论不改变算法的基于实现的加速技术。

快速模式。与 OzBLAS 的实现类似，我们定义参数 $d\in \mathbb{N},d\le \max{(s_x,s_y)}$ 来确定计算中子矩阵的数量。指定后，可通过省略 $(x^{(p)})^Ty^{(q)}$ 中 $p+q>d+1$ 的计算来减少计算量，但会导致精度略有损失。若所需精度为 FP64（与标准 DGEMM 相当，通过下述确定子矩阵数量的方法实现），则精度损失可忽略不计。该技术最多可将矩阵乘法的数量从 $d^2$ 减少到 $d(d+1)/2$ 。

估算实现 FP64 等效精度所需的子矩阵数量。算法 1 在 $\mu=0$ 时会自动停止拆分，即此时最终结果的精度达到最大。但如果所需精度为 FP64 算术下标准 DGEMM 的精度（FP64 等效精度），可基于概率误差界 [9] 通过算法 3 估算所需的最小拆分次数：

其中， $e=(1,\cdots,1)^T$ 的引入是为了避免在估算中进行矩阵乘法（需注意，在算法 3 的第 7 行中 $2^{c_A[d]_i} A_{\text{split}} [d]_{ij}$ ，为第 d 个未缩小的 FP64 子矩阵。因此，第 2 行的 SplitA到 $s_A$ 无需执行到，且可将 SplitA 与该算法集成）。该算法设计用于快速模式。若拆分次数按算法 1 执行到的方式确定，精度可能低于标准 DGEMM。此时，必须禁用快速模式。需注意的是，由于子矩阵数量仅基于概率误差界估算，且会受向量 $e$ 影响，该方法中期望精度（标准 DGEMM 所达精度）与实际获得精度之间可能存在一定差异。

内积分块。本研究暂未实现此步骤。由于算法 1 中的 $\rho$ 包含内积维度 $n$ ，实现一定精度所需的拆分次数取决于矩阵的内积维度。通过对内积操作采用分块策略，可避免拆分次数的增加。分块大小可设置为实现最佳性能的最小尺寸。但该策略会增加求和成本，且除非执行正确舍入计算，否则改变分块大小可能会影响可复现性。