Efficiently Emulating High-Bitwidth Computation with Low-Bitwidth Hardware

作者/机构: ZIXUAN MA, Tsinghua University, Beijing, China; HAOJIE WANG, Tsinghua University, Beijing, China; GUANYU FENG, Tsinghua University, Beijing, China; CHEN ZHANG, Tsinghua University, Beijing, China; LEI XIE, Tsinghua University, Beijing, China; JIAAO HE, Tsinghua University, Beijing, China; Shengqi Chen, Tsinghua University, Beijing, China; Jidong Zhai, Tsinghua University, Beijing, China

A1 主要贡献

本文旨在解决领域特定加速器(DSAs)虽然计算能力强大但原生支持的数据类型有限的问题。以往工作虽然探索了软件模拟某些数据类型,但未能对模拟数据类型的误差进行深入分析,从而错失了在表示能力(范围和精度)和性能方面的潜在优化机会。同时,如何为给定应用高效地设计更多模拟数据类型并选择一个高性能的方案,同时保证正确性和精度,仍然是一个开放问题。

为应对这些挑战,本文提出了Ape系统,其核心目标是:
1. 提供多种使用低位宽数据类型模拟高位宽数据类型的策略,并进行深入的误差分析。
2. 为给定的计算任务,在细粒度级别上动态、自动地选择合适的数据类型,并生成高效代码,以在无需人工干预的情况下同时保证正确性、精度和高性能。

本文的主要贡献如下:
* 提出多种高位宽数据类型模拟算法: 本文提出了多种使用低位宽数据类型模拟高位宽数据类型的算法,并进行了深入分析以优化表示能力和性能。与之前的工作相比,这些算法能够扩展数据范围、提高数据精度并显著提升计算性能。
* 设计细粒度的轻量级数据自适应算法: 该算法能根据特定硬件上给定应用的精度要求,自动选择合适的数据类型。
* 在多种硬件上实现Ape系统: 本文在NVIDIA Tensor Core GPU和华为昇腾处理器上实现了Ape,支持浮点和整数类型的模拟,并提供了用户友好的API以便将遗留代码移植到该库。
* 显著的性能提升: 评估结果显示,在NVIDIA Tesla A100 GPU上,Ape能将通用矩阵乘法(GEMM)和卷积的性能分别提升高达3.12倍和1.86倍,并能将各种应用的性能加速高达1.77倍。

下图展示了NVIDIA A100上原生浮点类型与本文提出的模拟数据类型的精度规格和峰值性能对比。

Table 1: NVIDIA A100上的原生浮点和模拟数据类型及其相应的精度规格(位数)和峰值性能(TFLOPS)。

A3 背景知识与相关工作

2.1 领域特定加速器

多种为加速特定领域应用而设计的DSA。为了加速特定领域的应用,各种领域特定加速器(DSAs)正在被设计出来。与通用芯片相比,DSAs在特定应用上可以实现更高的性能。例如,现有的DSAs,包括NVIDIA Tensor Core【索引33,Tesla NVIDIA. V100 GPU architecture. the world’s most advanced data center GPU. 2017】,Google Tensor Processing Unit (TPU)【索引21,In-datacenter performance analysis of a tensor processing unit. 2017. ISCA】和华为昇腾AI处理器【索引25,Ascend: a Scalable and Unified Architecture for Ubiquitous Deep Neural Network Computing: Industry Track Paper. 2021. HPCA】,都是为AI应用设计的,主要关注使用低位宽数据类型(如半精度浮点或低位宽整数类型)的线性算子。

NVIDIA Tensor Core。NVIDIA Tensor Core最早在NVIDIA Volta【索引33,Tesla NVIDIA. V100 GPU architecture. the world’s most advanced data center GPU. 2017】GPU中引入。Tensor Core可以为矩阵乘加(MMA)操作提供强大的计算能力。例如,Tesla A100 GPU中的Tensor Core浮点计算吞吐量高达312 TFLOPS。图1展示了在Tensor Core上执行的半精度浮点计算:Tensor Cores读取两个半精度输入数据,并执行一次半精度矩阵乘法而无精度损失。中间结果存储在单精度寄存器中,累加也以单精度进行。Top500【索引9,TOP500 supercomputer sites. 1997. Supercomputer】榜单上的第二和第三名超级计算机,即Summit和Sierra【索引37,The design, deployment, and evaluation of the CORAL pre-exascale systems. 2018. SC18】,都配备了NVIDIA Tensor Core GPU。


图 1: Tensor Core上的MMA操作。

华为昇腾AI处理器。华为昇腾AI处理器专为AI应用设计,支持训练和推理。以Ascend 910A为例,其在FP16上的峰值性能为256 TFLOPS,在INT8上为512 TOPS。910A配备32 GB HBM,带宽为1228 GB/s。Ascend 910A不支持FP64操作。鹏城云脑II【索引34,AIPerf: Automated machine learning as an AI-HPC benchmark. 2021. Big Data Mining and Analytics】是一个大型AI计算平台,配备了4096个Ascend 910A,FP16峰值性能达到1 EFLOPS。

2.2 低位宽数据类型

浮点数据类型。FP16(IEEE 754半精度浮点)【索引1,IEEE Standard for Floating-Point Arithmetic. 2008】,BF16(Brain Floating Point)【索引38,BFloat16: the secret to high performance on cloud TPUs. 2019. Google Cloud Blog】,和TF32(Tensor Floating Point)【索引23,TensorFloat-32 in the A100 GPU Accelerates AI Training, HPC up to 20x. 2020. the NVIDIA Blog】是DSAs普遍支持的数据类型。如图2所示,FP16有5个指数位和10个尾数位。与FP32相比,FP16的数据范围更窄,精度更低。NVIDIA Tensor Core和昇腾处理器支持此类型。BF16由Google Brain提出,有8个指数位和7个尾数位,可以表示与FP32范围相似但精度低得多的浮点数。NVIDIA Ampere GPU和TPU支持此类型。TF32由NVIDIA提出,有8个指数位和10个尾数位,表示的浮点数范围与FP32相似,精度与FP16相同,但占用与FP32相同的存储空间(4字节)。目前只有Nvidia Ampere GPU支持此数据类型。

定点数据类型。定点是表示实数的另一种方法,通常用整数实现。低位宽定点数据类型在AI领域普遍使用。不仅在推理【索引22,Reduced-precision strategies for bounded memory in deep neural nets. 2015. arXiv】、【索引24,Deep convolutional neural network inference with floating-point weights and fixed-point activations. 2017. arXiv】,在训练中【索引18,Fixed-point feedforward deep neural network design using weights+ 1, 0, and- 1. 2014. SiPS】、【索引26,Fixed point quantization of deep convolutional networks. 2016. ICML】,低位宽数据类型也能在不损失太多精度的情况下减少内存使用并提高性能。为满足定点计算的需求,现有的DSAs支持对各种整数类型的加速。例如,Edge TPU【索引15,Advanced neural network processing for low-power devices. 2020. https://coral.ai/technology】和昇腾支持INT8 。NVIDIA Ampere GPU上的Tensor Core支持INT8、INT4和INT1。所有这些DSAs都使用INT32来累加这些数据类型的乘积,以避免在特定计算中发生溢出。

2.3 相关工作

数据类型扩展。一些先前的工作【索引7,A floating-point technique for extending the available precision. 1971. Numer. Math.】、【索引16,Compiling KB-Sized Machine Learning Models to Tiny IoT Devices. 2019. PLDI】、【索引27,Software for doubled-precision floating-point computations. 1981. TOMS】、【索引36,Extended-precision floating-point numbers for GPU computation. 2006. SIGGRAPH】在GPU和物联网设备上使用低位宽类型来模拟高位宽数据类型。由于缺乏硬件支持,它们需要发出数十条低位宽指令来模拟高位宽运算符,导致效率低下。例如,Thall等人【索引36,Extended-precision floating-point numbers for GPU computation. 2006. SIGGRAPH】在GPU上使用20条单精度指令来模拟一个双精度加法。


图 2: 不同数据类型的比较。

利用Tensor Core模拟单精度GEMM。Markidis等人【索引29,Nvidia tensor core programmability, performance & precision. 2018. IPDPSW】利用Tensor Core上的FP32累加,基于四个半精度(FP16)GEMM来模拟一个单精度GEMM。然而,如第4.1节所述,他们模拟的类型只支持一个有限的范围,即[0.25, 6.55 × 104]。EGEMM-TC【索引11,EGEMM-TC: accelerating scientific computing on tensor cores with extended precision. 2021. PPoPP】通过使用SASS汇编风格的指令在Turing架构上高度优化了Markidis的方法。然而,EGEMM-TC也受限于Markidis方法的范围限制。同时,基于SASS的实现不能直接用于其他GPU,如NVIDIA Tesla V100和A100,这限制了其使用。我们的系统为FP16模拟的单精度类型提供了更大的数据表示范围和显著的性能提升,同时还提供了其他模拟数据类型和自动自适应方法。

数据类型自适应。在准确性和性能之间的权衡使得选择一个合适的数据类型变得必要。一些近期的工作【索引5,Rigorous Floating-Point MixedPrecision Tuning. 2017. POPL】、【索引6,Scalable yet Rigorous Floating-Point Error Analysis. 2020. SC】、【索引35,Rigorous estimation of floatingpoint round-off errors with symbolic taylor expansions. 2018. TOPLAS】专注于严格的浮点误差分析,以降低精度来获得更好的性能,同时在所有程序输入中保证给定的误差界限【索引5,Rigorous Floating-Point MixedPrecision Tuning. 2017. POPL】。尽管这些工作为任何输入提供了足够的精度,但它们受限于其静态和形式化方法所带来的冗长分析。例如,目前最先进的系统Satire【索引6,Scalable yet Rigorous Floating-Point Error Analysis. 2020. SC】需要大约10分钟来完成对128 × 128矩阵的矩阵乘法分析。与这种代价高昂的分析不同,我们的方法是轻量级的,并专注于在运行时根据浮点数据类型的表示范围选择合适的类型。

A3 Ape概览

3.1 Ape框架

Ape的设计定位。Ape被设计为一个数据类型模拟器和性能助推器,它可以使用原生数据类型模拟高位宽数据类型,并自动选择性能最佳且不损害精度的数据类型。通过对给定程序的输入数据进行动态分区和分析,Ape可以为每个细粒度数据块选择合适的数据类型来优化计算,从而提高计算效率。

Ape的组成模块。图3展示了Ape的概览。Ape由两个主要模块组成:Emulator(模拟器)和Adapter(适配器)。模拟器提供了DSA所有可用的原生和模拟数据类型。适配器负责为给定应用选择一个混合数据类型(包括原生和模拟数据类型),以更好地利用DSA的计算能力。


图 3: Ape概览。

Ape的工作流程。对于一个给定的计算Op1,输入为数据类型为T1A和数据类型为T2B,输出为数据类型为T3C(为简化起见,我们假设任何计算都有两个输入和一个输出),Ape会首先将输入和输出数据划分为小块,分别表示为$A_i$、$B_j$和$C_k$。这保证了输入块是输出块的计算依赖,例如,$A_i$ Op1 $B_j$ 是 $C_k$ 计算的一部分。然后Ape会为每个输入生成一个掩码矩阵,记录每个输入块中元素的最大指数位长度。接着,这个掩码矩阵会被进一步分析,以指导策略决策,从Emulator中选择最合适的数据类型和相应的计算。最后,Ape为每个块上的计算生成优化代码。在这个例子中,Ape生成了一个计算Op2,k,其输入为数据类型为T4的$A_i$,数据类型为T5的$B_j$,输出为数据类型为T6的$C_k$。


图 4: Ape的一个示例。$LenE(B_{ij})$表示给定矩阵$B_{ij}$中所有元素的最大指数位长度。

3.2 示例

以GEMM为例说明Ape工作原理。为了更具体地说明Ape的工作原理,我们以一个GEMM计算为例,如图4所示。给定一个GEMM计算$C = A × B$,我们首先将数据划分为几个小块。在这种情况下,我们假设Ape将$A$划分为6个(3 × 2)块,将$B$划分为6个(2 × 3)块,相应地将$C$划分为9个(3 × 3)块。Ape将分析输入块,并指出在$A_{11}$、$A_{22}$和$B_{13}$中至少有一个值必须用8个指数位来表示,而其他输入块中的所有值都可以用少于5个指数位来表示。因此,我们生成它们对应的掩码矩阵$MA$和$MB$。$MA$和$MB$的每个元素表示该块需要使用的最小指数位长度。考虑到GEMM的计算模式,C中的每个块,即$C_{ik}$,将与A中的特定块(即$A_{i:}$)和B中的特定块(即$B_{:k}$)进行块级GEMM计算。分析掩码矩阵后,Ape决定在总共18个块级GEMM计算中,有8个应该使用长指数位的数据类型进行计算,而其他的则可以使用短指数位的数据类型进行计算。最后,Ape为这种情况生成计算代码,并选择一个合适的策略以实现更高的性能。

A2 方法细节

4 模拟高位宽数据类型

本节阐述了我们使用不同低位宽数据类型模拟高位宽数据类型的方法,并对这些模拟数据类型进行了详细分析。对于每种模拟数据类型,在分析了其表示精度、范围和计算准确性之后,我们提出了两种优化表示方法的技术。一种是通过移位来减小两个元素的指数差距,从而充分利用指数位的能力。另一种是省略那些不影响结果精度的计算。通过这两种优化,我们可以扩大数据表示范围并实现更高的性能。

为简要展示我们的方法,我们选择三种广泛支持的半精度数据类型,即FP16、TF32和BF16,来模拟常用的单精度FP32数据类型。我们提出了三种模拟数据类型:FP32-F(4.1节)、FP32-T(4.2节)和FP32-B(4.3节)。表2展示了先前工作和我们工作的表示能力和持续性能。基于BF16和TF32的模拟数据类型是本文首次提出的。尽管基于FP16的模拟数据类型在先前的工作中已有介绍【索引11,EGEMM-TC: accelerating scientific computing on tensor cores with extended precision. 2021. PPoPP】、【索引29,Nvidia tensor core programmability, performance & precision. 2018. IPDPSW】,但我们的模拟算法FP32-F具有更大的表示范围,并受益于1.33倍的理论加速和1.19倍的持续加速。

此外,在4.4节中,我们简要讨论了我们的模拟算法如何应用于模拟整数类型。

Table 2: 不同模拟数据类型的表示能力及其在NVIDIA A100 GPU上的性能。

pair <FP16 , FP16 > toFP32_F ( FP32 A ) { 
    FP16 Hi_A = FP16 ( A ); 
    FP16 Lo_A = FP16 ( A - FP32 ( Hi_A )); 
    return { Hi_A , Lo_A }; 
}

代码清单 1: 使用两个FP16表示一个FP32。

4.1 使用FP16模拟FP32 (FP32-F)

FP32拆分为两个FP16。代码清单1中的算法描述了将一个FP32拆分为两个FP16的过程,这两个FP16分别表示为Hi(高位域)和Lo(低位域)。由于一个FP16只有5位用于指数,而FP32有8个指数位,因此FP32指数的最高3位被丢弃,剩下的5位存储在Hi的指数位中,记为HiE。这个过程与先前的工作【索引29,Nvidia tensor core programmability, performance & precision. 2018. IPDPSW】类似。接下来,通过分析数值误差,我们可以优化表示和计算,以实现更大的表示范围和更高的性能。

表示精度。图5展示了FP32到FP32-F的映射。对于FP32的23个尾数位,它们被拆分为以下几个部分:
- 前10位(HiS)存储在Hi的尾数位中。
- 第11位称为舍入位,记为R,并使用“四舍六入五成双”的原则舍入到第10位。
- 舍入位后第一个等于1的位被视为编码位,记为E,它可以被隐式编码到Lo的尾数位中。
- 编码位后的接下来10位存储在Lo的尾数位中,记为LoS。

精度损失分析。LoS可以保留编码位之后的所有尾随位,只要它们的长度不超过10。否则,即如果有11个尾随位,FP32尾数位的最后一位将被丢弃。为简单起见,在图5、图6以及本节余下部分,我们可以假设编码位始终是FP32的第12个尾数位。在这种情况下,当一个FP32的第12个尾数位为1时,其最后一位被丢弃。因此,这两个FP16只能表示FP32数23个尾数位中的前22位。结果是,FP32-F的表示精度为$2^{-22}$,比标准FP32的精度少一位。


图 5: 使用FP32-F的模拟方法。

表示范围。指数误差的发生是因为FP16只有5位用于指数,所以FP32中的最高3位指数位被截断了。因此,如果一个FP32的指数大于15或小于-14,它就不能用两个FP16来模拟。为简化起见,本文中的表示范围指的是给定类型的绝对值范围。

通过缩放扩大表示范围。在FP32-F中,低位域Lo表示原始FP32数的较低部分。因此,Lo的指数位(记为LoE)等于HiE通过编码位E的索引进行移位,即在图5的例子中$LoE = HiE - 12$。通过这种表示,舍入位R等于Lo的符号位。当使用FP32-F时,一个FP32数A被分成两部分,$Hi_A$和$Lo_A$,每个部分的范围都是$[6.10 \times 10^{-5}, 6.55 \times 10^{4}]$。如果我们像之前的工作那样简单地累加乘法计算的结果,为了保持23位的精度,HiE和LoE都需要在这个范围内。这将原始FP32数的范围限制在$[0.25, 6.55 \times 10^{4}]$。结果是,表示范围的下界太大了。

解决下界问题。为了解决这个问题,$Lo_A$被放大了4096倍,即将指数移动了12位。因此,$Lo_A$可以同时保持A的精度和与$Hi_A$等效的索引。这个补丁将FP32-F的范围扩大到了$[6.10 \times 10^{-5}, 6.55 \times 10^{4}]$。

计算准确性。为了展示Ape如何进行计算模拟,本节以融合乘加(FMA)计算为例。

FP32-F的FMA计算。代码清单2展示了如何用FP32-F进行乘法运算,并正确处理尾数位。我们将乘数和被乘数表示为A和B。当用FP32-F模拟时,它们的HiS位和LoS位分别表示为$HiS_A$、$LoS_A$、$HiS_B$和$LoS_B$。

FP32-F乘法细节。如图6所示,对于每次乘法$A \times B$,需要执行四次FP32-F乘法:$HiS_A \times HiS_B$, $HiS_A \times LoS_B$, $LoS_A \times HiS_B$, 和 $LoS_A \times LoS_B$。然后,这四个FP16乘法的结果根据$Hi_A$和$Hi_B$的指示进行移位相加。结果是原始FP32结果的尾数位。注意,由于FP32尾数位的长度限制,只保留23位。

指数计算和优化。FP32的指数等于$Hi_A$的指数和$Hi_B$的指数之和,这与浮点乘法类似,但只保留了A和B指数位的最低5位。值得注意的是,由于$Lo_A \times Lo_B$被移位了12+12位,它对最终的FP32结果没有影响。因此,可以跳过这一步以节省计算,并理论上相比之前的工作带来1.33倍的加速。

精度分析。由于两个FP16相乘的结果存储在一个FP32中,结果的指数位不会发生上溢或下溢。由于累加是直接在FP32片段上进行的,模拟计算不会引入额外的误差。标准FP32结果的最后一个尾数位在FP32-F乘法的结果中可能不精确,因为两个乘数的最后一个尾数位可能被丢弃了。如4.1节所述,其余位是精确的,FP32-F的计算精度为$2^{-22}$。


图 6: FP32-F乘法和累加。

FP32 FMA ( FP32 C , FP32 A , FP32 B ) {
    auto [ Hi_A , Lo_A ] = toFP32_F ( A );
    auto [ Hi_B , Lo_B ] = toFP32_F ( B );
    C += Hi_A * Hi_B ;
    C += Hi_A * Lo_B + Lo_A * Hi_B ;
    // C += Lo_A * Lo_B ; ( no effect on C )
    return C ;
}

代码清单 2: 使用FP32-F计算FMA。

4.2 使用TF32模拟FP32 (FP32-T)

首次提出使用TF32模拟FP32。本文首次提出使用两个TF32模拟一个FP32。这种模拟与使用两个FP16类似,唯一的区别是TF32有8位用于指数(与FP32相同)和10位用于尾数(与FP16相同)。10位的尾数使得FP32-T的尾数模拟与FP32-F相同。同样,其精度为$2^{-22}$。由于TF32的指数与FP32相同,Lo不需要缩放。因此,FP32-T的范围是$[4.81 \times 10^{-35}, 3.40 \times 10^{38}]$。

4.3 使用BF16模拟FP32 (FP32-B)

首次提出使用BF16模拟FP32。我们首次提出了一种使用BF16模拟FP32的方法。代码清单3展示了这种模拟的算法。第一步仍然是将FP32拆分为三个BF16,分别表示为Hi、Mi和Lo。

指数和尾数处理。与FP32-F类似,FP32的指数位存储在Hi的指数位中。由于BF16有8位用于指数,FP32的指数位可以被完全保留。对于尾数位,由于BF16只有7位用于尾数,需要三个BF16来模拟一个FP32。

tuple <BF16 , BF16 , BF16 > toFP32_B ( FP32 A ) {
    BF16 Hi_A = BF16 ( A );
    BF16 Mi_A = BF16 ( A - FP32 ( Hi_A ));
    BF16 Lo_A = BF16 ( A - FP32 ( Hi_A ) - FP32 ( Mi_A ))
    return { Hi_A , Mi_A , Lo_A };
}

代码清单 3: 使用三个BF16表示一个FP32。


图 7: 使用FP32-B的模拟方法。

表示精度。图7展示了FP32到FP32-B的映射。基本概念与FP32-F类似,列举如下:
- FP32的前7个尾数位存储在Hi中,记为HiS。
- 第8和第9位是舍入位和编码位,分别记为$R_1$和$E_1$。(同样,我们假设第9位为1。)
- 接下来的7位存储在Mi中,记为MiS。
- 再接下来的2位是另一个舍入位和另一个编码位,分别记为$R_2$和$E_2$。
- 最后的5位存储在Lo中,记为LoS。

精度保持。由于FP32的所有尾数位都被完整地存储在FP32-B中,FP32-B具有与FP32相同的精度。

表示范围。由于Mi和Lo没有被缩放,Mi和Lo的表示会影响FP32-B的范围。因此,FP32-B的表示范围是$[7.70 \times 10^{-34}, 3.40 \times 10^{38}]$。

FP32 FMA ( FP32 C , FP32 A , FP32 B ) {
    auto [ HiA , MiA , LoA ] = toFP32_B ( A );
    auto [ HiB , MiA , LoB ] = toFP32_B ( B );
    C += Hi_A * Hi_B ;
    C += Hi_A * Mi_B + Mi_A * Hi_B ;
    C += Hi_A * Lo_B + Mi_A * Mi_B + Lo_A * Hi_B ;
    // C += Mi_A * Lo_B + Lo_A * Mi_B ;
    // C += Lo_A * Lo_B ;
    return C ;
}

代码清单 4: 使用FP32-B计算FMA。

计算准确性。FP32-B的FMA计算与FP32-F类似。我们仍然只介绍Ape如何模拟乘法。需要进行六次乘法运算,$Hi_A \times Hi_B$、$Hi_A \times Mi_B$、$Mi_A \times Hi_B$、$Hi_A \times Lo_B$、$Mi_A \times Mi_B$和$Lo_A \times Hi_B$,如代码清单4所示。$Mi_A \times Lo_B$、$Lo_A \times Mi_B$和$Lo_A \times Lo_B$的计算被跳过,因为它们不影响结果的精度,类似于FP32-F中的$Lo_A \times Lo_B$。

4.4 模拟整数数据类型

利用现有硬件支持模拟高位宽整数。由于当前的DSAs支持低位宽整数类型,包括带有INT32累加器的INT8、INT4和INT1,因此使用低位宽数据类型模拟高位宽整数数据类型是可行的。例如,我们在NVIDIA Tesla A100和华为昇腾910A上用两个INT8实现了模拟的INT16。其他整数类型的模拟是类似的。

整数模拟方法。数据表示和计算的模拟方法与浮点数据类型相似。一个INT16通过字节序直接拆分为两个INT8来模拟。但整数中有一些独特的技术。为了保证模拟的整数数据类型能与原生数据类型完全一样地执行,模拟INT16的计算需要四次INT8计算,我们将在第7节进行评估。

5 将APE应用于真实硬件

尽管第4节提供了用低位宽数据类型模拟高位宽数据类型的方法,但在将它们应用于特定硬件时仍然存在挑战。

硬件约束带来的挑战。第一个挑战来自硬件约束。第4节的理论保证了不同模拟方法在元素级计算中的数学误差。然而,对于一个真实的DSA,计算资源是有限的,计算粒度并不总是元素级的。例如,在NVIDIA Tensor Core上,输入数据和输出数据有不同的存储位宽,并且计算是通过一个融合的MMA而不是FMA进行的。因此,在将模拟方法应用于特定硬件时,我们需要进行进一步的分析,这将在第5.1节中详述。

性能与正确性的平衡挑战。第二个挑战是如何在保证计算正确性的同时实现更高的性能。如表2所示,不同的模拟方法具有不同的性能和表示能力。虽然FP32-F性能最好,但由于其表示范围较窄,可能会引入上溢或下溢。考虑到现实世界的矩阵数据通常包含数千个元素,并具有只能在运行时确定的不同数据模式,选择正确且高效的数据类型同时防止大的运行时开销是具有挑战性的。为了解决这个问题,我们提出了一种细粒度的分块方法,该方法支持对给定计算使用混合模拟方法(5.2节),并设计了一个基于轻量级性能模型的策略决策器来最小化运行时开销(5.3节)。

5.1 硬件约束下的数据正确性

特定算子的数值正确性。对于一个特定的算子,其数值正确性与计算模式有关。在一个真实的DSA中,计算是一个包含许多步骤的融合计算。以GEMM为例,会执行多次乘法,其结果相加作为输出。在这个过程中,每一步都可能发生溢出,从而影响正确性。值得注意的是,一个寄存器(正式名称为累加器)用于存储乘法结果的部分和,它可能与输入数据类型不同。因此,在适配器的设计中,需要仔细检查和考虑目标DSA的架构。

以Tensor Core为例的分析。以Tensor Core为例,通过检查MMA指令的精度和范围,我们发现了以下有用的特性:
- 当在FP16或BF16矩阵乘法中使用FP32作为累加器时,结果具有与FP32计算相同的精度。换句话说,它与原始输入的结果完全相同。
- Tensor Core的累加器会累加乘法的结果,并在每个MMA指令中进行舍入。舍入方法是向零舍入(round-to-zero)。

结论与适配器设计。我们可以发现,假设GEMM计算不会在FP32上溢出或下溢,那么在模拟计算期间也不会发生溢出或下溢。MMA指令的舍入方法与CUDA Core上的FMA指令不同。这意味着Tensor Core上的结果与CUDA Core上的不完全相同,但两者在数值上都是正确的。因此,我们可以得出结论,Ape的模拟方法可以实现在那些累加器支持至少与高位宽数据类型相同位宽的硬件上。

基于数据动态性选择数据类型。基于这些事实,我们可以通过检查输入矩阵的范围来为GEMM计算选择合适的指数。因此,我们可以潜在地挖掘输入数据的动态性。我们设计了一个适配器,它会扫描每个GEMM计算的每个输入矩阵,并选择合适的数据类型。

5.2 分块混合计算

不同类型性能差距的动机。如表2所示,精度相似的不同类型之间存在巨大的性能差距。以NVIDIA A100上的Tensor Core进行GEMM计算为例,对于1024×1024大小的矩阵,FP32-F达到27.24 TFLOPS,而FP32-B只能达到17.20 TFLOPS,仅为前者的63.1%。直观上,为了实现更高的吞吐量,应尽可能多地使用FP32-F。然而,当存在需要更多指数位的计算时,使用FP32-B是不可避免的。

挖掘数据模式的分块算法。此外,数据模式在数据的不同区域通常是变化的。为了挖掘这种模式,我们提出了一种分块算法。我们

type_t getMask ( Partition P ) {
    type_t mask = FP32_F ;
    for ( auto v : P ) {
        if ( v < MIN_FP32_B || v > MAX_FP32_B )
            return FP32 ;
        if ( v < MIN_FP32_F || v > MAX_FP32_F )
            mask = FP32_B ;
    }
    return mask ;
}

代码清单 5: 为一个分区生成掩码。

分块与数据类型确定。我们首先根据硬件的计算粒度将数据划分为块。然后,通过分析数据和计算依赖关系,我们可以为每个输入块确定性能最高的可能数据类型。通过这种方式,我们在一次计算中对不同的块使用不同的数据类型,以获得更高的性能。

输入数据分析过程。分区后,分析输入数据的过程如代码清单5所示。通过检查分区中的每个元素,分析相关的计算过程,以找到能够满足分区中指数位长度需求的最合适类型。例如,在我们的Tensor Core上的GEMM实现中,分区后,一个GPU核函数用于扫描每个分区中的每个元素,并选择性能最佳且能力足够的表示方法。之后,对整个分区的统计数据进行汇总,以确定该分区最合适的数据类型,这被称为掩码(mask)。

混合计算决策。在这个例子中,FP32-F和FP32-B是候选类型。对于一个给定的二元运算符$C = op(A, B)$,其中A、B和C都是块,性能更高的FP32-F只有在A和B都可以用FP32-F表示时才能使用。否则,必须使用FP32-B。

计算拆分与执行。因此,计算被分为两部分。必须使用FP32-B的块和其余部分分别由两个不同的核函数计算。这两个核函数的决定只涉及查找两个标签。因此,它可以在开始实际计算之前在运行时做出,而不会引入显著的开销。

混合计算的优势。结果是,Ape通过同时使用不同类型来平衡能力和性能,从而获益。

5.3 轻量级策略决策器

混合计算的性能权衡。如5.1节所述,通过分析数据掩码,可以得到一个计算的候选类型列表。性能最高的情况是所有数据都可以用FP32-F表示,此时直接使用FP32-F进行计算。相反,如果输入数据只能用原生的FP32表示,我们不得不回退到CUDA Core,导致性能不佳。在更常见的情况下,部分输入数据可以用FP32-F表示,其余的用FP32-B表示。此时,使用5.2节中提到的分块混合计算方法。但混合方法并不总是最优选择。为了清楚地展示这种情况,我们以FP32-F和FP32-B的混合为例,展示FP32-F、FP32-B和混合方法的性能。结果如图8所示。我们可以观察到,随着FP32-B块比例的增长,最佳数据类型会发生变化。因此,应该设计一个在线算法来为每次计算确定合适的数据类型。


图 8: 不同FP32-B块比例下的GEMM性能。

轻量级决策算法。为了降低在线决策所需的延迟,开发了一种轻量级的决策算法。建立了一个性能模型来预测使用专用FP32-B或分块混合类型计算方法的GEMM计算的运行时间。

基于掩码的性能模型。该模型基于数据掩码。为简单起见,我们假设计算过程中的调度方法和负载均衡是理想的。对于每个块$C_{i,j,k}$,计算中使用的数据类型根据$A_{i,k}$和$B_{k,j}$来确定。当任何一个输入块是FP32-B时,两个计算块都必须使用FP32-B。否则,两个输入块都是FP32-F,计算块可以使用FP32-F。因此,使用以下算法计算对应于每种类型的块的数量。首先,在数据掩码中将FP32-F和FP32-B分别编码为0和1。然后,执行一个小型的GEMM $Mask_{C} = Mask_{A} \times Mask_{B}$。这里,每个元素$Mask_{C_{ij}}$表示输出矩阵中可以使用FP32-F的块的数量。最后,将$Mask_{C}$中的所有元素相加,得到计算过程中FP32-F计算块的数量。

最终决策。通过对每种数据类型的块性能进行采样,我们可以计算出给定计算的时间。因此,我们预测使用专用FP32-F、FP32-B或混合计算的运行时间,并选择最佳方案。

6 实现

多平台实现。我们在NVIDIA Tensor Core GPU上实现了功能齐全的Ape,包括Tesla V100、T4和A100。由于我们没有在华为昇腾处理器上进行编程的开放权限,因此我们没有在昇腾上实现分块混合计算。相反,我们使用厂商提供的封装好的核函数,如GEMM和卷积,来验证我们在昇腾上的方法。尽管现有的库如cuBLAS【索引30,cuBLAS. https://developer.nvidia.com/cublas】提供了高度优化的核函数,但这不足以支持Ape的所有功能。例如,我们需要为不同类型实现具有细粒度调度的核函数以支持我们的混合方法。此外,由于Ape支持不同平台,我们提出了一个代码生成器来生成具有各种配置的不同核函数,并为每个平台选择最佳方案 。

核函数设计与代码生成。例如,在Tensor Core上,核函数应该重叠计算和内存访问以充分利用计算吞吐量。流水线需要精心设计。此外,数据转换可以放在核函数内部或外部。更好的方法取决于模拟的类型和平台。我们实现了一些代码模板,描述了几种实现方式,包括使用封装好的核函数或使用具有不同流水线策略的核函数。代码生成器将根据模板生成各种核函数。Ape将为每个平台上的每种情况选择性能最佳的核函数。

用户接口。Ape为用户提供了与其他BLAS【索引3,An updated set of basic linear algebra subprograms (BLAS). 2002. ACM Trans. Math. Software】库类似的接口,目前支持两种线性操作:GEMM和卷积。用户可以通过修改代码或库链接,将原始的GEMM和卷积库调用替换为Ape的库调用,从而轻松使用Ape,而无需修改应用程序的其余部分。如果用户对数据范围和精度有专业知识,他们也可以手动选择特定数据类型的核函数以进行进一步优化。

A4 实验环境

我们的实验在NVIDIA GPU和华为昇腾AI处理器上进行。
* 硬件配置:
* NVIDIA GPUs: NVIDIA Tesla V100, T4, 和 A100。GPU的内存和应用时钟频率固定在最大值以避免性能波动。
* Huawei Ascend: Ascend 910A,其FP16峰值性能为256 TFLOPS,INT8峰值性能为512 TOPS。

  • 软件配置:

    • CUDA版本: 统一使用11.0.2。
    • 依赖库: cuBLAS, cuDNN。
    • 代码实现: Ape通过替换库调用(如cuBLAS中的GEMM)集成到应用中。
  • 数据集/应用:

    • 量子电路模拟 (QS): 使用GEMM和张量转置实现。
    • 主成分分析 (PCA): 基于cuBLAS和cuSolver实现。
    • K最近邻 (kNN): 基于一个开源的GPU实现。
    • K均值聚类 (kMeans): 采用NVIDIA的开源GPU实现。
    • BERT: 采用开源的cuBERT实现。
    • 卷积: 使用ResNet-18中的卷积层配置。
    • 通用矩阵乘法 (GEMM): 测试不同尺寸的方阵。
  • 说明: 除非另有说明,所有评估都包含了模拟和自适应的开销。

A4 实验结果

7.2 端到端性能

实验内容: 我们在NVIDIA GPU上使用五个应用来评估Ape的端到端性能。基线是使用CUDA Core上的cuBLAS GEMM实现。我们仅将GEMM库从cuBLAS替换为Ape。
实验结果与分析: 如图9所示,Ape在所有应用中都取得了显著的性能提升。
* 量子电路模拟 (QS): 最大矩阵尺寸为512,Ape的平均加速比为1.65倍。
* 主成分分析 (PCA): 在原始版本中,GEMM占总执行时间的39.25%。Ape可将PCA加速1.19倍。
* K-最近邻 (kNN): 在基线中,GEMM占执行时间的76.9%。Ape为kNN提供了1.78倍的加速。
* K-均值聚类 (kMeans): GEMM占运行时间的77.2%,Ape的加速比为1.75倍。
* BERT: 在单精度评估中,FP32-F可实现高达1.77倍的加速。


图 9: 应用相对于CUDA Core FP32的加速比。

7.3 GEMM性能

浮点GEMM性能
* 实验内容: 我们在NVIDIA Tesla V100, T4, 和 A100上评估了不同浮点模拟的性能。基线包括cuBLAS中的CUDA Core GEMM,以及先前的工作FP32-M【索引29,Nvidia tensor core programmability, performance & precision. 2018. IPDPSW】和EGEMM-TC【索引11,EGEMM-TC: accelerating scientific computing on tensor cores with extended precision. 2021. PPoPP】。
* 实验结果与分析 (图10):
* A100: FP32-F和FP32-B相比cuBLAS FP32分别实现了3.12倍和1.85倍的加速。与之前的方法相比,FP32-F不仅提供了更好的数据范围,性能也比FP32-M高出1.22倍。
* V100和T4: 在这些不支持BF16或TF32的平台上,Ape只启用FP32-F。FP32-F相比cuBLAS FP32分别实现了1.81倍和1.49倍的加速,并且持续优于FP32-M(分别高出1.29倍和1.23倍)。与T4上的EGEMM-TC相比,Ape在小矩阵上性能更好,但在大矩阵上稍差,这归因于T4较低的内存带宽和EGEMM-TC使用了不可移植的SASS汇编优化。


图 10: NVIDIA GPU上浮点数据类型的模拟性能。

定点GEMM性能

  • 实验内容: 我们在NVIDIA Tesla A100的Tensor Core和华为Ascend 910A处理器上评估了定点模拟。我们使用两个INT8定点数来模拟一个INT16,这两种硬件本身都不支持INT16。
  • 实验结果与分析 (图11): 实验表明,通过Ape的模拟,加速器可以支持非原生数据类型的计算,并展示了模拟INT16定点GEMM的性能。


图 11: Tesla A100和Ascend 910A上定点数据类型的模拟性能。

7.4 卷积性能

  • 实验内容: 为了验证Ape在其他线性计算中的正确性,我们在NVIDIA Tesla A100上评估了使用FP32-F(Tensor Core)和FP32(CUDA Core)的卷积性能,两者均基于cuDNN【索引4,cudnn: Efficient primitives for deep learning. 2014. arXiv】实现。我们选择了ResNet-18中的几个卷积层。
  • 实验结果与分析 (图12): 在某些情况下,FP32-F能达到比FP32更高的性能,最高加速比达到1.86倍。这表明Ape不仅能支持Tensor Core上的卷积,还可能优化cuDNN的实现。这个实验验证了Ape的方法可以应用于卷积并扩展到卷积DSA。


图 12: Tesla A100上的卷积性能。

7.5 精度

应用精度
* 实验内容: 我们以量子电路模拟(QS)为例,评估Ape在真实应用中的精度。我们将不同数据类型配置下的结果与FP64结果进行比较。
* 实验结果与分析 (图13): 与FP64的结果相比,FP32-F和FP32-B的误差与FP32相当。然而,直接使用FP16在大多数情况下会导致明显错误。例如,在电路qv_25(#4)中,FP32-F、FP32-B和FP32的误差都在9%左右,而FP16的误差高达172%。这表明Ape提供的模拟方法能满足量子模拟的要求,而FP16不能。


图 13: 量子模拟的最大误差。x轴上的ID代表不同的电路。

模拟精度
* 实验内容: 我们生成了多组随机数据来验证Ape模拟类型的精度损失。
* 实验结果与分析 (图14):
* FP32-F和FP32-M在其表示范围内有稳定的误差,约为$1.2 \times 10^{-7}$,与理论分析一致。但FP32-F的数据范围比FP32-M更大。
* FP32-B和FP32-T的数据范围接近FP32。FP32-B在其范围内与FP32相比几乎没有误差。FP32-T的相对误差约为$1.2 \times 10^{-7}$。这表明FP32-B和FP32-T在精度和范围上与FP32相似。


图 14: 不同模拟类型的最大误差。

计算精度

  • 实验内容: 我们通过在不同大小的随机矩阵上执行GEMM操作来检验计算中的精度。
  • 实验结果与分析 (图15): 所有模拟方法的误差都随着矩阵大小的增加而增加。这是因为Tensor Core的截断方法是“向零舍入”而非“四舍五入”,导致截断误差在归约过程中累积。这是Tensor Core的硬件设计选择,并非由模拟引入。此外,不同矩阵大小下,Ape会选择不同核函数以获得更好性能,导致算法带来的误差随之变化。


图 15: GEMM算子的最大误差。

7.6 分解分析

  • 实验内容: 我们对混合模拟方法的性能进行了分解分析。实验场景是执行GEMM,其中60%的块由FP32-F模拟,其余由FP32-B模拟。
  • 实验结果与分析 (图16): 结果显示,与实际GEMM计算所花费的时间相比,为每个块在模拟方法之间做出细粒度决策所花费的时间可以忽略不计。通过决策策略利用FP32-F节省了大量时间。总之,适配器引入的开销很小,并有效地提高了性能。


图 16: 混合模拟方法的计算时间分解分析。

A7 补充细节

8 讨论

Ape方法的普适性。Ape的方法可以支持不同的基于融合乘加(FMA)的线性计算,如GEMM和卷积,这些计算在AI和HPC等不同领域都扮演着至关重要的角色。本文主要在NVIDIA Tensor Core GPU和昇腾加速器上评估了Ape。但如第5节所述,Ape可以实现在所有满足特定要求的硬件上。Ape不仅限于本文评估的计算和硬件,还有可能支持为其他基于FMA的线性计算设计的新硬件。

对硬件设计的启示。同时,Ape表明我们可以使用低位宽数据类型来模拟高位宽数据类型,而不会损失太多精度或损害正确性,这可以为硬件设计提供一些启示。例如,如果为一个FP16计算单元提供一个FP64累加器,Ape也可以用来模拟FP64计算。因此,我们可以使用同一个FP16计算单元来模拟FP32和FP64数据类型,从而节省硬件片上面积,这是芯片设计中的一个重要瓶颈。

A5 结论

本文提出了Ape,这是第一个在DSAs上提供多种模拟数据类型并自动选择合适数据类型的框架。Ape能以细粒度方式加速给定的计算,以实现更高的性能,同时保持正确性和精度。Ape扩展了DSAs的计算场景,实现了更好的利用率,且无需任何人工干预。我们在NVIDIA Tensor Core和华为昇腾上都实现了Ape。评估结果显示,与CUDA Core相比,Ape在Tensor Core上可将FP32通用矩阵乘法和卷积的性能分别提升高达3.12倍和1.86倍,并能将各种应用的性能优化高达1.78倍(平均1.61倍)。