稀疏矩阵向量乘法 (SpMV) 是几乎所有隐式 稀疏线性代数求解器 的核心计算内核。从简单的 Krylov 算法到多重网格方法，算法的性能在很大程度上取决于 SpMV 实现的速度。由于 SpMV 的算术强度（定义为每内存访问的浮点运算次数）非常低，因此实现的性能受内存带宽的限制。最大化内存带宽的实现将比简单方法获得更高的性能。或者，利用稀疏矩阵固有结构的实现，从而最小化所需内存访问次数的实现，也将获得更高的性能。在这篇博文中，我们将开发广义 SpMV 操作 的 HIP 实现，包括几种标准的 SpMV 实现：
- 标量压缩稀疏行 (CSR) 格式，
- 向量 CSR，
- Ellpack，和
- 块状 Ellpack。
许多常见的 SpMV API，例如 rocSparse，都使用这种通用接口。我们希望在我们的实现与 ROCm 中提供的实现之间进行公平的比较。在 SpMV 公式中， 和 是标量值， 是一个大小为 的稀疏矩阵， 是一个大小为 的稠密向量，而 是一个大小为 的稠密向量。在本篇博文中将贯穿使用的一个关键矩阵统计量是每行的平均非零元个数，。GPU 计算资源的分配通常基于此度量。
首先，我们将回顾广泛使用的 CSR 和 Ellpack 格式的稀疏矩阵。在描述完实现后，我们将比较它们在 AMD MI250X 架构上的各种稀疏矩阵的相对性能。一种矩阵存储格式转换为另一种格式的转换算法可能是一项成本高昂的工作。我们的代码 示例 将提供 CSR 和 Ellpack 之间转换算法的未优化设备实现，但我们将在本文中不详细讨论它们。
存储格式
我们对不同稀疏矩阵存储类型的回顾，最好用一个简单的例子来说明。在下图中，我们展示了一个包含 12 个非零元的 矩阵。该稀疏矩阵的坐标 (COO) 格式由三个长度相同的数组组成，分别存储行、列和值。我们使用 0 基索引来表示行和列。数据按行主序排序，以便先列出所有行 0 的数据。然后是行 1、行 2 等。
压缩稀疏行 (CSR) 格式
CSR 格式存储稀疏矩阵，源于上面显示的行数据结构的简单压缩。只需计算每行非零元的数量。然后，对行计数执行累积求和算法即可完成计算。压缩行数据结构（以下称为行偏移量）的大小为 。第一个值是 0，其余值是累积求和的结果。此数据结构非常有用，因为它通过相邻读取提供了稀疏矩阵行乘以稠密向量（即稀疏点积）计算的起始和结束索引。
Ellpack 格式
当每行非零元数量的变化很小时，Ellpack 格式是一种有效的 SpMV 数据结构。与 CSR 相比，该格式将存储顺序反转为列主序，并使用零填充来实现数据结构的统一。每行的条目数由所有行中的最大值确定。任何非零元少于最大值的行都用 0 填充，并且列索引也使用“合理”的值。例如，可以用行中最后一个有效列索引进行填充，以确保我们重新读取先前加载的索引和值对。
可以将此数据结构可视化为将所有值（和列索引）向左移动到数组中。
通过将 Ellpack 数据结构分成多行块，可以进一步优化 Ellpack 格式。这样可以最小化零填充量，同时实现跨数据的统一性。我们将此称为块状 Ellpack。
GPU 内核实现
标量 CSR 内核
最简单的 GPU 加速 SpMV 实现之一是标量内核方法。标量内核分配一个线程来处理 SpMV 中的每个稀疏点积。每个线程以顺序方式处理稀疏点积，因此无需更高级的技术，如共享内存和/或 warp 级规约。下面代码片段显示了标量 CSR 内核的一个直接实现。
__global__ void scalar_csr_kernel(const int m, const int *__restrict__ row_offsets, const int *__restrict__ cols, const double *__restrict__ vals, const double *__restrict__ x, double *__restrict__ y, const double alpha, const double beta){ const int row = threadIdx.x + blockDim.x * blockIdx.x; if (row < m) { // determine the start and ends of each row int p = row_offsets[row]; int q = row_offsets[row+1];
// execute the full sparse row * vector dot product operation double sum = 0; for (int i = p; i < q; i++) { sum += vals[i] * x[cols[i]]; }
// write to memory if (beta == 0) { y[row] = alpha * sum; } else { y[row] = alpha * sum + beta * y[row]; } }}
void scalar_csr(int m, int threads_per_block, int * row_offsets, int * cols, double * vals, double * x, double * y, double alpha, double beta){ int num_blocks = (m + threads_per_block - 1) / threads_per_block; dim3 grid(num_blocks, 1, 1); dim3 block(threads_per_block, 1, 1); scalar_csr_kernel<<<grid, block>>>(m, row_offsets, cols, vals, x, y, alpha, beta);}如上所述，计算仅在行之间并行化。令 为每个块中的线程数，则完成计算所需的独立块数为： 。在此上下文中， 是一个可调参数，我们可以使用它来最大化此实现的性能。在内核定义中，我们然后：
- 通过简单的线程和块映射来计算我们的行，
- 确保我们不会超出行偏移量数组的任何数组边界，
- 我们通过相邻读取行偏移量数组来确定要加载的列（因此是向量）索引和矩阵值，
- 我们进行稀疏点积的标量计算，最后
- 我们将结果写入内存。
在最后一步，我们可以通过测试 向量是覆盖还是更新来节省一些内存带宽。
标量 CSR 实现将在最广泛的稀疏矩阵范围内获得中等性能。然而，在某些情况下，该内核可能效果适中，即当每行的平均非零元数量很小时。事实上，当每行的平均非零元数量很高时，将需要多次读取来处理内存事务，并且合并效果会很小。相比之下，当每行的平均非零元数量非常低时，高内存带宽可以与此内核的简单性质相结合以获得良好的性能。
有关如何从矩阵市场文件输入运行标量 CSR 内核，请参阅以下 示例。
向量 CSR 内核
GPU 加速 SpMV 最通用的高性能实现是向量内核方法。与标量内核不同，向量内核使用 warp 中的多个线程来规约每个稀疏点积。它需要共享内存或 warp 混洗操作来规约跨线程计算的部分规约值。用于选择分配给每个稀疏行-向量乘积的每个 warp 的线程数的启发式方法基于 。通常，我们发现通过找到小于或等于 的最小 2 的幂来获得最佳性能。此时，研究代码示例以检查实现的细节很有帮助。这里我们重点介绍仅使用 warp 混洗操作来完成稀疏点积的情况。应该注意的是，每个输出结果最多可以由 wavefront 大小的线程计算，在 MI200 架构上是 64。当每行的平均非零元数量远大于 wavefront 大小时，这可能会限制性能。
template <int THREADS_PER_ROW>__global__ void vector_csr_kernel(const int m, const int *__restrict__ row_offsets, const int *__restrict__ cols, const double *__restrict__ vals, const double *__restrict__ x, double *__restrict__ y, const double alpha, const double beta){ const int row = threadIdx.y + blockDim.y * blockIdx.x; if (row < m) { // determine the start and ends of each row int p = row_offsets[row]; int q = row_offsets[row+1];
// start the sparse row * vector dot product operation double sum = 0; for (int i = p + threadIdx.x; i < q; i += THREADS_PER_ROW) { sum += vals[i] * x[cols[i]]; }
// finish the sparse row * vector dot product operation#pragma unroll for (int i = THREADS_PER_ROW >> 1; i > 0; i >>= 1) sum += __shfl_down(sum, i, THREADS_PER_ROW);
// write to memory if (!threadIdx.x) { if (beta == 0) { y[row] = alpha * sum; } else { y[row] = alpha * sum + beta * y[row]; } } }}
void vector_csr(int m, int threads_per_block, int warpSize, int * row_offsets, int * cols, double * vals, double * x, double * y, double alpha, double beta){ int nnz_per_row = nnz / m; int threads_per_row = prevPowerOf2(nnz_per_row); // limit the number of threads per row to be no larger than the wavefront (warp) size threads_per_row = threads_per_row > warpSize ? warpSize : threads_per_row; int rows_per_block = threads_per_block / threads_per_row; int num_blocks = (m + rows_per_block - 1) / rows_per_block;
dim3 grid(num_blocks, 1, 1); dim3 block(threads_per_row, rows_per_block, 1); if (threads_per_row <= 2) vector_csr_kernel<2><<<grid, block>>>(m, row_offsets, cols, vals, x, y, alpha, beta); else if (threads_per_row <= 4) vector_csr_kernel<4><<<grid, block>>>(m, row_offsets, cols, vals, x, y, alpha, beta); ... else if (threads_per_row <= 32) vector_csr_kernel<32><<<grid, block>>>(m, row_offsets, cols, vals, x, y, alpha, beta); else vector_csr_kernel<64><<<grid, block>>>(m, row_offsets, cols, vals, x, y, alpha, beta);}此实现与标量实现有几个关键区别。我们使用 来构建二维线程块。
-
线程的 x 维度表示分配给每行的线程数。这由 `prevPowerOf2` 函数计算得出，该函数计算小于或等于输入变量的最小 2 的幂。prevPowerOf2。
-
y 维度表示每个线程块的行数。
-
总线程块数由每个块处理的行数决定。
模板用于以编译时常量形式启动内核，以确定每行的线程数。这使得在循环展开规约算法期间可以进行优化。
内核的开始方式与标量内核类似，但一个关键的区别在于如何从块/网格启动启发式中计算行索引。实际上，请注意我们使用线程索引的 y 分量和块大小来计算此值。然后，稀疏点积被分解为两个步骤。
-
每个线程以模板参数的倍数循环遍历行，计算稀疏点积的部分结果。
-
完成后，子 warp 将使用 shuffle down 操作将累加值汇总到每个 wavefront 的线程 0。
在此，我们注意到任一稀疏点积步骤中都缺少同步。这是因为我们将自己限制在至少等于 wavefront 大小的规约。
有关如何从矩阵市场文件输入运行向量 CSR 内核，请参阅以下 示例。
Ellpack 内核
Ellpack SpMV 实现沿着行并行化计算。由于数据已重新排序为按列主序存储，因此沿着 Ellpack 数据连续行的内存访问是合并的。在下面的实现中，我们假设输入的 `cols` 和 `vals` 数组已转换为 ellpack 格式。此格式的一个关键部分是元数据参数，即每行的最大非零元数量，该参数也作为参数传入。
__global__ void ellpack_kernel(const int m, const int max_nnz_per_row, const int *__restrict__ cols, const double *__restrict__ vals, const double *__restrict__ x, double *__restrict__ y, const double alpha, const double beta){ const int row = threadIdx.x + blockDim.x * blockIdx.x; if (row < m) { double sum = 0; for (int i=row; i < max_nnz_per_row*m; i+=m) { sum += vals[i] * x[cols[i]]; }
// write to memory if (beta == 0) { y[row] = alpha * sum; } else { y[row] = alpha * sum + beta * y[row]; } }}
void ellpack(int m, int threads_per_block, int max_nnz_per_row, int * cols, double * vals, double * x, double * y, double alpha, double beta){ int num_blocks = (m + threads_per_block - 1) / threads_per_block; dim3 grid(num_blocks, 1, 1); dim3 block(threads_per_block, 1, 1); ellpack_kernel<<<grid, block>>>(m, max_nnz_per_row, cols, vals, x, y, alpha, beta);}线程块之间的工作分解与标量 csr 格式相同。此外，稀疏点积也以标量方式执行。此内核与标量 CSR 内核之间的关键区别在于矩阵以合并的方式访问。数据结构中的填充使我们无需昂贵的条件语句即可编写稀疏点积实现。
有关如何从矩阵市场文件输入运行 Ellpack 内核，请参阅以下 示例。此示例包括从 CSR 格式到 Ellpack 的设备转换。
块状 Ellpack 内核
块状 Ellpack 数据结构旨在最小化 Ellpack 结构中引入的额外内存读取。每行的最大条目数在一个等于 wavefront 大小倍数的行集中计算。然后，在每个块中计算一个局部 Ellpack 结构。然后将每个块连接起来形成全局数据结构。这需要一个额外的短数组，类似于 CSR 行偏移量，该数组指示每个块在每行最大列数方面的开始和结束。
下面的代码片段展示了一个实现。首先要注意的是，代码是针对以 2 为底的 wavefront 大小进行模板化的，以避免用于计算全局 wave 索引的整数除法。全局 wave 索引用于立即从内核返回那些不参与计算的尾随 wave。
template<int LOG2WFSIZE>__global__ void blocked_ellpack_kernel(const int m, const int tw, const int *__restrict__ max_nnz_per_row_per_warp, const int *__restrict__ cols, const double *__restrict__ vals, const double *__restrict__ x, double *__restrict__ y, const double alpha, const double beta){ const int row = threadIdx.x + blockDim.x * blockIdx.x; const int warp = row>>LOG2WFSIZE; if (warp>=tw) return;
const int lane = row&(warpSize-1);
int p = max_nnz_per_row_per_warp[warp]; int q = max_nnz_per_row_per_warp[warp+1];
p<<=LOG2WFSIZE; q<<=LOG2WFSIZE;
if (row < m) { // sparse dot product implementation double sum = 0; for (p+=lane; p<q; p+=warpSize) { sum += vals[p] * x[cols[p]]; }
// write to memory if (beta == 0) { y[row] = alpha * sum; } else { y[row] = alpha * sum + beta * y[row]; } }}有关如何从矩阵市场文件输入运行块状 Ellpack 内核，请参阅以下 示例。此示例包括从 CSR 格式到块状 Ellpack 的设备转换。
性能
数据
为了衡量我们各种实现的性能，我们选择来自各种来源的矩阵。首先，我们从 SuiteSparse 矩阵集合中收集了一些经典示例。这些矩阵在下面的图中用黑色标签表示。一些博文的作者与几个 Exascale Computing Projects (ECP) 团队密切合作。他们对加速这些代码相关的某些矩阵计算特别感兴趣。其中包括来自 ExaWind 和 Pele ECP 项目的矩阵。所有数据都已公开，并托管在 SuiteSparse 集合中。
我们从 ExaWind 项目收集了几种类型的矩阵。ExaWind 项目旨在模拟大气边界层中整个风电场的复杂流体运动。求解器由“背景”（AMR-Wind）和“叶片解析”（Nalu-Wind）求解器组成，分别尝试捕获大尺度和小尺度行为。背景求解器有几个“类泊松”求解器，它们会产生来自 7 点和 27 点样条的结构化矩阵。这些矩阵在下面的图中用红色标签表示。叶片解析求解器使用非结构化网格来模拟风力涡轮机叶片和塔架周围的流体运动。我们使用来自背景求解器压力泊松方程（另请参阅 SC21）的各种大小的网格的矩阵。Nalu-Wind 利用 Hypre Boomer AMG（代数多重网格）算法来求解压力泊松方程。AMG 求解器层次结构产生了具有值得在此博文中探讨的有趣结构的稀疏矩阵。Nalu-Wind 和 Hypre 的非结构化矩阵在下面的图中用蓝色标签表示。
Pele ECP 项目旨在以超高分辨率模拟 复杂几何结构内的燃烧过程。Pele 利用 自适应网格细化技术 结合嵌入边界算法来实现前所未有的计算规模。与 AMR-Wind 类似，该应用程序需要求解类泊松矩阵才能推进底层物理方程求解器。我们从这些类泊松问题中提取了几个大小问题不同的矩阵。这些矩阵具有值得进一步研究的有趣结构，并在下面的图中用绿色标记。
结果
性能结果如下图所示，针对单个 MI250X GCD[1]。除了上面提供的代码示例外，我们还提供了一个 示例，说明如何从矩阵市场文件输入运行 Rocsparse CSR 内核。
对于每个内核，我们使用每块 256 个线程。还生成了 128 和 512 个线程/块的结果，这些结果通常比 256 个线程/块略差。在左图中，我们显示了平均每次 SpMV 执行时间，并按默认 RocSparse 性能进行缩放。在右图中，我们显示了设置成本，按需要从 CSR 格式转换（Ellpack 和块状 Ellpack）或需要分析 CSR 矩阵（RocSparseWithAnalysis）的三种格式的平均 SpMV 时间进行缩放。
数据显示在考虑的矩阵集合中存在显著差异。使用带分析的 Rocsparse 通常是最快的，但它具有最显著的设置成本。该算法对于具有每行 nnz 分布长尾的矩阵特别有效。块状 Ellpack 实现通常相当不错，尤其是在每行 nnz 方差很小时。块状 Ellpack 的设置成本通常远小于 RocsparseWithAnalysis 的成本，而且收益通常相当可观。
矩阵 *nalu_large_level1*、*nalu_large_level2* 和 *nalu_large_level3* 的每行 nnz 分布具有长尾。对于 nalu 矩阵的中型和小型版本也是如此。*nalu_large_level2* 的分布和稀疏模式如下图所示。由于稀疏模式的分散性以及每行 nnz 分布的长尾，这些矩阵通常会在 Ellpack 存储格式中产生大量的零填充。最终，Ellpack 内核的简单性将无法克服由于零填充值导致的内存读取量的巨大增加。这解释了为什么这些实现的性能如此显著下降。相比之下，带分析的 RocSparse 内核旨在很好地处理这些类型的矩阵。有点令人惊讶的是，简单的向量 CSR 内核也能很好地处理这些情况。
结论
这标志着使用 HIP 开发和优化一些基本 SpMV 内核的第一部分。我们已经展示了各种矩阵上几种不同存储格式的性能。我们还展示了从 CSR 格式转换为 Ellpack 和内部 Rocsparse 格式的性能成本。总的来说，Rocsparse 配合 Analysis 算法在所有方面都提供了最佳性能，但与其他考虑的格式相比，它具有更高的设置成本。由于 SpMV 算法通常嵌入在更大的算法中，例如 AMG，因此在为特定问题选择最佳格式时，考虑完整的应用程序配置文件非常重要。
如果您有任何问题或评论，请在 GitHub 讨论区 与我们联系
[1]
测试使用 ROCm 版本 5.4.22804 进行。基准测试结果并非经过验证的性能数据，仅用于演示代码修改的相对性能改进。实际性能结果取决于包括系统配置和环境设置在内的多个因素，不保证结果的可重现性。
