COFFEE

c:cross-layer optimization for fast and efficient executions of the SK algorithm on HPC systems with clusters of compute nodes
在有计算节点集群的HPC系统上快速高效执行SK算法。
算法中:行、列都可以重新缩放列缩放相较行缩放极其缓慢
在多节点上的性能比单节点提升最高7.5倍,平均2倍;与天河一号的MPI Allreduce算法比,最高2.9,平均1.6.
SK算法:a simple but very useful iterative method to approach the double stochastic matrix of Sinkhorn’s theorem by alternately rescaling all rows and all columns of the given matrix.
对矩阵进行缩放列。

现存的SK算法大多用去搞强化学习了(应用层),或者去加速收敛,很少有研究从计算机系统架构的角度考虑改进算法,特别是高性能计算(HPC)系统。HPC有他自己独特的计算、存储以及交流能力,看看是否能发挥全部潜力。SK算法在四个代表性应用的时间占比都超过一半(BALS的卷积也这样)所以就去优化。

这篇论文用MPI,通过多核、多节点集群加速SK算法。
先分析经典算法在天河1上。
列缩放的时间远超行缩放,原因是通过列缩放进行的内存访问是高度非连续的,这导致了较高的缓存未命中率。
解决方法:探险重新设计列缩放以及信息阻塞去减少缓存未命中;设计微核并且重新设计指令去增加并行性

优化思想:通用矩阵乘法、分层
分层的思想在MPI Allreduce算法(MPI_Allreduce 是 MPI(消息传递接口)中的一个函数,用于在所有进程之间进行归约操作并广播结果。)的相关优化中非常常见。

选择SALaR(?)作为基准去研究。发现:实现Allreduce可以与SK算法的其他任务进行overlap(重叠)进一部提高性能

本篇文章的主要contribution

  • 我们分析了 SK 算法在 HPC 集群上的执行行为,并观察到两个主要的性能挑战。首先是其列重新缩放表现出高度非连续的内存访问模式,这导致非常高的缓存未命中率,从而大大降低了整体性能。第二,即使采用 Foster 的方法设计,列重新缩放也会严重限制并行性
  • 我们提出了 COFFEE,这是一种新颖的方法,它实现了多级优化设计,以优化 HPC 系统中大规模 SK 算法的处理(第 IV 节)。我们通过增强 MPI Allreduce 来提高并行效率,采用有效的领导者-工作者机制,尽可能重叠节点间(intra-node)通信、节点内通信和节点内计算
  • 我们在天河一号超级计算机上评估了 COFFEE 的原型实现,证明了其与 SOTA 解决方案相比具有显著的性能优势(第六节)。我们的实验结果表明,COFFEE 分别在单节点和多节点环境中带来了高达 7.5 倍和 2.9 倍的性能提升。

SK算法

双随机矩阵,sk算法就是在行列都归一化后,每行元素相加都为1,每列元素相加也都为1

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
import numpy as np
def sinkhorn(A, max_iter=1000, tol=1e-6):
A = np.array(A, dtype=np.float64)
for _ in range(max_iter):
A_prev = A.copy()
A /= A.sum(axis=1, keepdims=True) # 行归一化
A /= A.sum(axis=0, keepdims=True) # 列归一化
if np.allclose(A, A_prev, atol=tol):
break
return A

# 示例
A = np.array([[4, 5, 6], [3, 4,2],[7,8,5]])
B = sinkhorn(A)
print(B)
print("Row sums:", B.sum(axis=1))
print("Col sums:", B.sum(axis=0))

选SK算法而不是和其他算法有两重原因:

  • 现在主流的线性代数库(如 BLAS、NumPy、PyTorch 等)在进行矩阵乘法时,底层通常使用的是最基础的“三重循环”实现方式,而不是像 Strassen 或 Coppersmith-Winograd 这样的快速乘法算法。就像基本的矩阵乘法实现一样,最原始的 Sinkhorn-Knopp(SK)算法也更容易从计算机系统架构的角度进行优化
  • 现有的 SK 算法研究主要集中于通过减少矩阵缩放迭代次数来加快收敛速度​​,但我们的目标是减少每次迭代的时间

我们不再去限制每一行/每一列的和接近1,而是去最后计算的时候让行/列和为1,在计算结果之前不需要让它接近1。

Intel团队用Python。为了在超算上运行,我们选C实现,是他们SOTApython的重写,有循环展开和数据并行优化,并且性能不逊于他们。

串行处理行列缩放,明显看到列用了十几倍的时间。

介绍典型的并行算法

关键点:如何分割数据和任务让交流少一些,让计算任务更稳定
处理方法:把矩阵按行划分成多个子矩阵,每个处理器只负责其中一些行,这样在行归一化的时候每个处理器只需要处理自己的行,不需要通信;最后让通信发生在列归一化阶段,这部分可以统一优化。这样目的是把本地能算的留在本地,只在必要时跨节点通信
算法执行被拆成四步:

  • 每个进程独立地对自己那一块行子矩阵进行行归一化(和算法1的第1–10行一样),不需要通信;
  • 每个进程计算自己那部分子矩阵的列和;
  • 调用 MPI_Allreduce 汇总所有进程的列和,得到全矩阵每列的总和;
  • 每个进程根据上一步得到的列缩放因子,独立地对自己那块子矩阵做列归一化。
    这样做能最大限度减少通信,仅在列缩放因子计算这一步使用 MPI,有利于并行效率。

图片示例

在计算并行效率的时候,列缩放时处理器数为16的效率也骤降?
原因:行缩放的时候不需要进行通信,所有的处理器被平衡地加载;但到列缩放的时候处理器在等数据。
The reason is that the row rescaling is communication free and all the processors are load balanced. While for the column rescaling, the Allreduce used to find the column sum performs a lot of inter-node and intra-node communication, so that some processors are in the process of waiting for data

Motivation of COFFEE

我们看到了在并行处理行缩放时候节点内的通信花费了太多时间,大大降低效率;因此我们想利用节点的不同通信特点,然后提升效率。

算法设计-微核设计-MPI优化

CPU-ORIENTED OPTIMIZATION

列重排算法设计(Algorithm1和2的对比)

Micro-kernel redesign

采用SIMD:SIMD(Single Instruction, Multiple Data,单指令多数据流)是一种并行计算技术,它让一个指令同时处理多个数据。常用于图像处理、音频处理、科学计算等场景,加速处理速度,提升性能。
采用AVX2指令集:AVX2(Advanced Vector Extensions 2)是Intel推出的SIMD指令集扩展,属于x86架构的一部分。它在AVX的基础上增强了整数运算能力,支持256位宽的YMM寄存器,可以并行处理更多数据,广泛用于图像处理、机器学习等高性能计算中。
修改汇编指令

MPI optimization

节点内Reduce算法优化

二叉树效率低是因为每个处理器开销不同,尤其在根节点,其他处理器都空闲(idle),加载不均
为了解决这种严重的负载不均衡问题,我们重新设计了SK算法的Reduce实现,将本地和数组分成几部分。在节点内Reduce之后,每个worker保留本地最终和的一部分并以非阻塞方式将其发送给leader。我们的节点内Reduce实现基于MPI标准原语MPI_Send和MPI_Recv,与MPICH库中Reduce的实现一致。我们没有使用打包技术,因为要传递的数据几乎是连续的,打包带来的额外开销超过了使用它带来的性能提升。

节点间AllReduce算法优化

我们使用最流行的 Ring 算法实现 Allreduce,以生成列重新缩放的全局最终总和。Ring Allreduce 的一个缺点是它没有考虑节点的层次结构。一般来说,节点之间的带宽远低于节点内的带宽。因此,最近提出了分层 Ring Allreduce。
图 7 显示了我们基于分层环的优化。
主要思想是重叠节点内 Reduce 和节点间 Allreduce 的时间。我们将本地和数组分成几个数据块。如前所述,对数据块进行 allreduce 有三个连续步骤。首先,工作者对本地和的块执行节点内 Reduce,并将本地最终总和发送给领导者(图 7 中时间 1 的红色箭头)。接下来,领导者对全局最终总和执行节点间 Allreduce(图 7 中时间 2 的红色箭头)。最后,领导者将全局最终总和广播给其工作者(图 7 中时间 3 的红色箭头)。发现不同数据块的顺序步骤可以重叠。例如,图 7 中的时间 2 表示第 i 个数据块的节点间 Allreduce(红色箭头)和第 (i + 1) 个数据块的节点内 Reduce(黑色箭头)可以同时处理。因此,我们在为 SK 算法实现 Allreduce 时将管道的思想结合到分层环中。

重叠通信和计算优化

当领导者执行 Allreduce 时,工作者必须停滞。我们利用这段停滞时间让工作者修改节点内的矩阵。在我们对 SK 算法的优化 Allreduce 设计中,矩阵修改的计算被添加到流水线中。
同时完成 Allreduce 的通信任务和修改子矩阵的计算任务

experiment evaluation

Experimental Setup

为了评估 COFFEE 的有效性,我们将其两个版本进行比较,即面向 CPU 的优化(第 IV 节),表示为 COFFEE-CPU,以及面向 MPI 的优化(第 V 节),表示为 COFFEE-MPI,与 SK 算法的两个现有实现进行比较,一个使用 Ring Allreduce 算法(MPICH-Ring),另一个在 MPICH 环境中使用 SALaR(MPICH-SALaR)
高密度矩阵(非零元素占 95%)、中等密度矩阵(非零元素占 50%)、稀疏矩阵(非零元素占 5%)

只去算非零矩阵。

CPU-Oriented Optimization

在AMD平台上使用GCC编译器运行的SK算法通过我们的优化获得了最大的改进。在ARM平台上使用Clang编译器,SK算法的典型实现的性能在所有平台上都是最好的,但我们的优化在M = N = 16,000时仍实现了3.3倍的加速比。

MPI-Oriented Optimization

conclusion and further work

SK算法在机器学习等领域的重要性日益凸显。本文提出并实现了一种针对SK算法实现的计算和通信的跨层优化设计,称为COFFEE。与大多数现有的通过减少缩放迭代次数来加快收敛速度​​的工作不同,COFFEE着重于通过缩短每次缩放迭代来加快收敛速度​​。我们对SK算法实现中影响性能的问题进行了深入研究。发现列缩放会导致较高的缓存未命中率和较低的并行效率。我们使用列缩放重新设计、数据分块和微内核设计等跨层优化来加速列缩放。我们还根据SK算法的特点优化了MPI Reduce和Allreduce,以提高并行效率。最后,我们在天河一号超级计算机上验证了COFFEE 的有效性。未来我们计划进一步探索和利用行缩放和列缩放之间的相关性。此外,我们计划结合GPU,充分利用异构并行计算架构,进一步提高 COFFEE 的性能。最后,我们计划研究 COFFEE 在 SK 算法稀疏矩阵上的性能,其中数据不是以数组格式存储的。

HSMU-SpGEMM

High Shared Memory Utilization for Parallel Sparse General Matrix-Matrix Multiplication on Modern GPUs
在现代 GPU 上实现并行稀疏通用矩阵-矩阵乘法的高共享内存利用率(utilization)

Abstract

传统的基于哈希的方法无法在减少哈希冲突和有效利用快速共享内存之间取得平衡,这严重损害了在 GPU 上执行 SpGEMM 的性能。设计了一种累加器,四个通用库在三种架构上面跑,**HSMU有显著的加速优势。

Introduction

Gustavson算法在GPU上用于实现快速并行稀疏矩阵乘法时有两个流程:符号阶段(symbolic stage)以及数值阶段(numeric stage)

  • 符号阶段的主要任务是去确定矩阵C中非零元素的数量(NNZ),以便在数值计算的时候预先分配内存
  • 数值阶段在已分配的内存上进行实际的乘法和累加,是整个SpGEMM最耗时的部分
    高效的累加器设计对数值阶段的性能至关重要。
    单纯增加哈希表的容量会降低GPU共享内存的利用率。
    主流SpGEMM库的不足
  • Nsparse 虽利用最大 NNZ 设置哈希表长度以提高共享内存利用率,但哈希冲突严重,性能下降
  • spECK 通过分配 1.5× 空间减少冲突但造成约 34% 内存浪费
  • OpSparse 建议设为 2× 最大 NNZ,性能好但共享内存利用率仅达 50%
    HSMU-SpGEMM 通过为每个累加器内核维护一个按列排序的数组(长度为分配行的最大 NNZ)来避免哈希冲突,并针对小规模与大规模矩阵设计不同的符号阶段,从而在优化 GPU 共享内存利用率下的同时保持低冲突率,实现高性能 SpGEMM。

Background

Gustavson 算法具有天然的并行性,因为它可以独立计算矩阵 C 的每一行。在数值阶段,矩阵 A 中每个非零元素 𝑎𝑖𝑗会与矩阵 B 中对应行𝑏𝑗∗的非零元素相乘,生成大量中间结果,这些结果的列索引与 B 中元素的列索引一致。最终,通过累加器将这些中间结果累加到矩阵 C 的相应位置。累加器的具体设计将在下一小节介绍。
分为两种累加器:稠密型稀疏型

  • 稠密累加器使用稠密数组存储中间结果,通常由三个向量组成:一个存储实际数值,一个用于标记列索引是否插入,另一个记录列索引。这种方法在处理稠密行时效率高,但对稀疏行内存需求大、性能较差。

  • 稀疏累加器则按累加方式分为三类:基于合并、ESC 和哈希的累加器。

  • 基于合并的稀疏累加器在 RMerge和 bhSPARSE等库中实现。这些累加器执行多次迭代,每次迭代将一个 NZ(Non-zero) 元素垂直合并到最终的稀疏向量中。由于在合并过程中使用了大小相同的临时数组,因此基于合并的稀疏累加器对于密度变化较大的矩阵表现出较低的内存利用率

  • ESC 方法在处理生成大量中间产品的矩阵时存在不足。由于存储和分类大量中间产品会产生大量的空间开销和时间成本,因此效率低下。

Shared Memory

除了寄存器文件之外,共享内存是 GPU 上最快的内存类型。共享内存的主要优点是它允许多个线程共享数据,这使得共享内存成为高效并行计算的关键 GPU 组件。共享内存可供单个线程块内的所有线程访问,并且可以将其视为可编程缓存或暂存器,在其中放置经常访问的数据。但是,GPU 上的共享内存容量有限。所以高性能共享内存是非常重要的对于现代GPU。

MOTIVATION OF THE WORK

Principles of Hash-based Accumulators

现有的哈希累加器要去平衡共享内存利用率哈希碰撞率

Philosophy of HSMU-SpGEMM Accumulator Design

  • 引入一个预排序列索引数组(sorted column indices array)表示 C 的非零列。使用 findInSorted(colIp, sortedColArray) 函数来定位每个中间乘积,直接查找中间乘积该落在哪个已知列上。完全消除哈希冲突,查找位置准确,无需哈希函数或冲突处理。

HSMU-SPGEMM

HSMU-SpGEMM Accumulator Design

采用二分查找而不是哈希,可以低开销,高效率,并且很稳定
有以下好处

  • 哈希法相对于二分查找法的最大优势在于,在哈希表中添加或删除项目的成本要低得多。然而,在 HSMU-SpGEMM 中,排序数组是在符号阶段预先确定的,因此我们的新累加器不需要在数字阶段更改排序数组。因此,二分查找法中维护排序结构的缺点不存在;
  • 哈希表的一个缺点是,当发生碰撞时,它会影响其他哈希位置,并可能导致链式碰撞,导致哈希性能低下。而对于二分查找,其性能稳定,最坏情况为O(logN)。在这种情况下,二分查找优于哈希查找方法;
  • 二分查找更适合于范围查询等复杂操作。在这种情况下,每个线程通过不断更新变量pos逐渐缩小共享col数组上的搜索范围,从而在一定程度上减少查找次数;尽管如此,对于密集和大数据,二分查找的最坏时间复杂度为O(logN),而哈希表的理想时间为O(1)。时间复杂度的增加可能会在某些情况下降低我们的累加器设计的搜索性能。

Generate the Sorted Ccol Array


先生成maskB,对于每一行B,非零元素位置标记为1,其他为0。
不太懂生成maskC的原理。