稀疏矩阵乘法 wu-kan

分别使用 MPI 和 CUDA 实现稀疏矩阵乘法,输入输出矩阵都用压缩行格式存储;并与串行程序比较。

实验环境

实验在老师提供的计算集群的一个节点上进行。单节点的显卡配置如下:

$ nvdia-smi
Mon Dec  2 08:38:49 2019
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 410.48                 Driver Version: 410.48                    |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  Tesla V100-PCIE...  On   | 00000000:3B:00.0 Off |                    0 |
| N/A   30C    P0    24W / 250W |      0MiB / 16130MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

实验原理

稀疏矩阵乘法在工程上有着重要的应用,然而稀疏矩阵本身的存储方式决定了它不容易像稠密矩阵一样划分到多线程并行计算。下面介绍我完成本次实验的一些思考和用到的一些记号。

稀疏矩阵

本例中,用到的稀疏矩阵在内存中的存储形式有如下三种:

  • 坐标格式(Coordinate Format, COO)
    • 用三元组(RowInd, ColInd, Val)定义矩阵中的元素。
    • 分别表示元素的行坐标、列坐标、取值。
  • 行压缩格式(Compressed Sparse Row Format, CSR)
    • 用三元组(RowPtr, ColInd, Val)定义矩阵中的元素。
    • 将矩阵内的元素按照先行坐标后列坐标的顺序排序。
    • 此时同一行的元素在内存上的排列是相邻的。
    • 左闭合区间[ RowPtr[i] , RowPtr[i + 1] )内可以唯一确定同一行的所有元素。
    • 同一行的元素之间使用列坐标ColInd区分,元素值为Val
  • 列压缩格式(Compressed Sparse Column Format, CSC)
    • 用三元组(ColPtr, RowInd, Val)定义矩阵中的元素。
    • 类似于行压缩格式转置后的结果。

由于稀疏矩阵的操作比较复杂,此处我将它封装成了一个简单的稀疏矩阵处理库<WuKSPARSE.h>。下面是一个稀疏矩阵的存储方法,同时支持上述三种存储格式。

typedef struct Scoo
{
	int
		m,   //行
		n,   //列
		nnz; //非零元素数量
	IntVector
		RowPtr, //大小m+1
		ColPtr, //大小n+1
		RowInd, //大小nnz,行坐标
		ColInd; //大小nnz,列坐标
	FloatVector
		Val; //大小为nnz
} Scoo;

命名方式类似于Nvidia<cusparse.h>。由于操作的过程中可能会丢失行优先的性质但是坐标对的性质始终得到保留,因此命名Scoo表明这是一个单精度(S)的坐标格式(coo)稀疏矩阵。后续实验的时候会确保乘法之前是使用行优先的存储。

由于稀疏矩阵中存在着大量内存管理的过程,这里我封装了IntVectorFloatVector,功能类似于 C++ STL 中的vector<int>vector<float>。没有直接使用 C++是因为想保持语法的纯粹性,而不想让 C++中很多语法糖影响自己对代码结构的判断。这样的好处是后续实验中对内存的分配管理都是通过这两个类型进行,更加稳定。

我实现了如下的接口供调用。

void ScooFree(Scoo *a); //释放矩阵申请的内存空间
void ScooMalloc(Scoo *a); //根据矩阵的m和n和nnz为矩阵申请内存空间
void ScooPushBack(Scoo *a, int r, int c, float v); //向矩阵插入一个新的元素,可能会破坏原先的索引
void ScooRead(Scoo *a, FILE *f); //读取MatrixMarket格式的矩阵
void ScooWrite(const Scoo *a, FILE *f); //将矩阵转化为MatrixMarket格式输出
void ScsrInit(Scoo *a); //假设矩阵三元组的存储已经符合列压缩的要求,并为其生成行索引
void ScooToScsr(const Scoo *a, Scoo *b); //将矩阵元素排序并转换成行压缩格式,并为其生成行索引
void ScscInit(Scoo *a);  //假设矩阵三元组的存储已经符合列压缩的要求,并为其生成列索引
void ScooToScsc(const Scoo *a, Scoo *b); //将矩阵元素排序并转换成列压缩格式,并为其生成列索引

假设要求的矩阵乘法$A\times B=C$,这里老师给的数据生成器使用默认参数会生成如下两个 MatrixMarket 格式的矩阵。注意到这里生成的矩阵中有很多不合理的下标(如-1),我封装的函数会将它们从输入中筛去,以保证程序正确运行。

第一个矩阵$A$对应的matrix12.mat的前十行:

%%MatrixMarket matrix coordinate real general
8191 8191 784379
0 -1 1.0
0 0 1.0
1 -1 1.0
1 0 1.0
2 -1 1.0
2 0 1.0
0 1 1.0
0 2 1.0
...

第二个矩阵$B$对应的matrix.mat的前十行:

%%MatrixMarket matrix coordinate real general
8191 8191 46796
0 -1 1.0
0 0 1.0
1 -1 1.0
1 0 1.0
2 -1 1.0
2 0 1.0
0 1 1.0
0 2 1.0
...

动态向量

这里给出<WuKfloatVector.h>FloatVector的相关定义和接口。<WuKintVector.h>IntVector的定义和接口类似。

typedef struct FloatVector
{
    int size, cap; //大小和实际容量
    float *data; //指向元素的指针
} FloatVector;
void FloatVectorFree(FloatVector *v); //释放v申请的内存资源
void FloatVectorInit(FloatVector *v, int size); //初始化一个大小为size的向量
void FloatVectorPushBack(FloatVector *v, float val); //向向量末尾插入新的元素,如果容量不够会重新分配两倍的空间,从而保证动态插入的均摊时间复杂度是常数

串行稀疏矩阵乘法

矩阵乘法的本质是对矩阵$A$的所有行向量和矩阵的$B$的所有列向量做点积操作,因此要获取两个矩阵的行向量和列向量。

由于初始矩阵是按照行优先的顺序存储的,因此可以直接获得矩阵的行向量。要获得矩阵的列向量,需要将矩阵内的元素重新排序转换成列优先的顺序,所需要的时间复杂度为$O(nnz\log nnz)$。当然,也可以通过定位行向量的索引,用$m$路归并排序实现更优的线性复杂度。

随后就可以对行优先矩阵和列优先矩阵做稀疏向量点积运算。注意到这里算点积的顺序天然就是行优先的,因此生成的矩阵$C$也是行优先矩阵,无需再做转换。时间复杂度为$O(m\times n\times l)$,三个变量分别是矩阵的宽、高、稀疏向量点积的平均时间。容易看出,当稀疏向量点积的平均时间小于稠密向量点积的时间时,稀疏矩阵乘法是要远远优于稠密矩阵乘法的。

此外还有一个 trick,在第二层循环开始之前先判断了 A 这一行是否为空行,可以根据矩阵 A 的稀疏情况大大减少循环次数,实现性能上的调优。

void ScsrMulScsc(const Scoo *a, const Scoo *b, Scoo *c)
{
	c->m = a->m;
	c->n = b->n;
	c->nnz = 0;
	ScooMalloc(c);
	for (int i = 0; i < c->m; ++i)
		if (a->RowPtr.data[i] < a->RowPtr.data[i + 1])
			for (int j = 0; j < c->n; ++j)
			{
				float res = 0;
				for (int
						 kb = b->ColPtr.data[j],
						 ka = a->RowPtr.data[i];
					 kb < b->ColPtr.data[j + 1] &&
					 ka < a->RowPtr.data[i + 1];
					 ++kb)
					if (b->Val.data[kb])
					{
						while (a->ColInd.data[ka] < b->RowInd.data[kb])
							++ka;
						while (a->ColInd.data[ka] == b->RowInd.data[kb])
							res += a->Val.data[ka] * b->Val.data[kb], ++ka;
					}
				if (res)
					ScooPushBack(c, i, j, res);
			}
	ScsrInit(c);
}
void ScsrMulScsr(const Scoo *a_csr, const Scoo *b_csr, Scoo *c_csr)
{
	Scoo B_csc;
	ScooToScsc(b_csr, &B_csc);
	ScsrMulScsc(a_csr, &B_csc, c_csr);
	ScooFree(&B_csc);
}

串行稀疏矩阵乘法的用时为6.199s,其输出结果sparseScsrmm_out.txt的前十行如下。

%MatrixMarket matrix coordinate real general
8191 8191 4185601
0 0 511.000000
0 1 511.000000
0 2 511.000000
0 3 448.000000
0 4 511.000000
0 5 511.000000
0 6 511.000000
0 7 282.000000

MPI 稀疏矩阵乘法

重新分析一下串行算法的过程,可以分成如下两各阶段:

  1. 将第二个矩阵转换成列优先的格式
  2. 对第一个矩阵的行向量和第二个矩阵的列向量分别做稀疏向量点积

第一步的本质就是对三元组按照列优先的顺序对矩阵进行重排序,因此可以使用常见的并行排序方法对其进行优化,优化后的期望时间复杂度是$O(\frac{nnz}{p}\log\frac{nnz}{p})$,其中$p$是并行的核数。或者也可以在$m$路归并的基础上将其优化。

我对第二步的优化方法是,将第一个矩阵$A$按顺序拆分到各个进程上分别执行子矩阵的稀疏矩阵乘法,最后使用MPI_Gatherv汇集成完整的答案$C$。

由于使用 MPI 写出来的代码有很多 API 调用过于繁琐,这里先描述一下 MPI 算法的流程,详细细节可以看后面的代码sparseScsrmm_mpi.c

  1. 使用并行排序方法将矩阵$B$转化成列优先存储格式,并广播到每个进程
  2. 使用MPI_Scatter将行优先矩阵$A$按照顺序均分到各个进程。由于$A$是行优先的,拆分后各个进程获得的$A_{local}$也是行优先的。
  3. 各个进程做串行稀疏矩阵乘法$A_{local}\times B=C_{local}$
  4. 各个进程将$C_{local}$汇总到根进程得到答案$C$。由于$C_{local}$的行与$A_{local}$的行是高度对应的,因此$C_{local}$也是行优先矩阵,汇总得到的$C$也是行优先矩阵,无需特殊处理。

MPI 稀疏矩阵乘法的用时为23.416s,相对于串行算法增加了。这是因为算法在最开始的时候有一个将矩阵$A$散射到各个进程,而结束的时候将矩阵$C$收集到主进程输出,通信开销非常大。

但是实际生产情况中,通常要求各个进程持有待处理矩阵的一部分即可,便于进行多轮矩阵乘法(类似 Cannon 并行稠密矩阵乘法的优点,多次执行A = A * B时不用汇总到主进程)。在这种情况下,MPI 并行化的稠密矩阵乘法还是很有优势的;并且,矩阵$B$转成列优先也只用转一次(但是多轮矩阵乘法之后稀疏矩阵逐渐会变得稠密…)。

输出结果sparseScsrmm_mpi_out.txt的前十行如下。

%MatrixMarket matrix coordinate real general
8191 8191 4187926
0 0 511.000000
0 1 511.000000
0 2 511.000000
0 3 448.000000
0 4 511.000000
0 5 511.000000
0 6 511.000000
0 7 282.000000
...

注意到这里输出的答案中非零元素的数量4187926大于之前的4185601,是因为我按照均分的元素拆分方式导致某些行向量从中间「断开」了,导致结果出现同一个坐标上的多个值,在稀疏矩阵的坐标存储方式是允许的。如果不希望同一个坐标多次出现可以用线性的去重unique算法,这里就没有做了。

CUDA 稀疏矩阵乘法

CUDA 上的稀疏矩阵乘法不能完全按照 MPI 版本中的并行算法来运行,这是由显卡的如下两个特性决定的:

  • 没有全局的同步函数(当然可以通过多次启动核函数解决,但是效率不高)
  • 核函数内部不能动态对一个向量进行扩张

关于后一点,我尝试使用了英伟达提供的动态向量类型thrust::device_vector,结果发现其push_back操作不能由核函数调用(那我要你有何用)。因此需要为 CUDA 版本设计新的并行算法。

我使用二维划分的线程块,每个线程对应答案矩阵的一个元素,或者一次稀疏向量点积运算。这种算法没有用于之前的两个算法,是因为它需要预先分配$m\times n$的空间用于存储每个位置的值,使得之前的算法退化成稠密矩阵算法。但是,显卡的另一个特性就是显存大,线程极多,新开线程开销极小。因此,可以新开一个矩阵来收集每个线程计算的结果。当然,这样计算得到的结果是稠密矩阵,需要使用线性算法扫描整个矩阵,重新转成符合要求的行优先矩阵。

下面是 CUDA 稀疏矩阵乘法的核函数,出人意料地简洁。

void __global__ cuScsrMulScs(
	const int *a_csr_RowPtr,
	const int *a_csr_CowInd,
	const float *a_csr_Val,
	const int *b_csc_ColPtr,
	const int *b_csc_RowInd,
	const float *b_csc_Val,
	float *c,
	int m,
	int n)
{
	int
		i = blockIdx.x * blockDim.x + threadIdx.x,
		j = blockIdx.y * blockDim.y + threadIdx.y;
	float res = 0;
	for (int
			 kb = b_csc_ColPtr[j],
			 ka = a_csr_RowPtr[i];
		 kb < b_csc_ColPtr[j + 1] &&
		 ka < a_csr_RowPtr[i + 1];
		 ++kb)
	{
		while (a_csr_CowInd[ka] < b_csc_RowInd[kb])
			++ka;
		while (a_csr_CowInd[ka] == b_csc_RowInd[kb])
			res += a_csr_Val[ka] * b_csc_Val[kb], ++ka;
	}
	c[i * n + j] = res;
}

分析一下核函数的运行情况:

  • 沿 x 方向对b_csc_ColPtr访存连续
  • 沿 y 方向对a_csr_RowPtr访存连续
  • c[i * n + j]的访存连续
  • a_csr_Val[ka]b_csc_Val[kb]访存不连续
  • 多个while存在分支

后两者应该是这个算法的瓶颈,也是稀疏矩阵乘法难以被并行化的主要原因。

在主机上用下面的代码调用核函数。

dim3
	block(16, 16),
	grid(a_csr->m / block.x, B_csc.n / block.y);

cuScsrMulScs<<<grid, block>>>(
	a_RowPtr,
	a_ColInd,
	a_Val,
	b_ColPtr,
	b_RowInd,
	b_Val,
	d_c,
	a_csr->m,
	B_csc.n);

CUDA 稀疏矩阵乘法的用时为3.942s,相比于串行稀疏矩阵乘法提速了 57%。其输出结果sparseScsrmm_mpi_out.txt的前十行如下。

%MatrixMarket matrix coordinate real general
8191 8191 4177936
0 0 511.000000
0 1 511.000000
0 2 511.000000
0 3 448.000000
0 4 511.000000
0 5 511.000000
0 6 511.000000
0 7 282.000000
...

注意到这里输出的答案中非零元素的数量4177936小于之前的4185601,是因为原来的输入中含有多个元素的位置相同,而这个 CUDA 版本的并行算法将其去重了。

源代码

调度脚本sparseScsrmm.pbs

这里由于临近期末,集群资源比较紧张,所有的调度都只占用集群上的一个节点进行。MPI 使用一台主机上的 32 个核进行并行;CUDA 使用节点上的单张 v100 显卡。为方便起见,我将所有的测试代码写进了一个测试脚本sparseScsrmm.pbs中,这样直接对这个脚本进行调度运行即可自动运行完整实验并进行计时。

#PBS -N sparseScsrmm
#PBS -l nodes=1:ppn=32:gpus=1
#PBS -j oe
#PBS -q gpu
source /public/software/profile.d/mpi_openmpi-intel-2.1.2.sh
source /public/software/profile.d/cuda10.0.sh
cd $PBS_O_WORKDIR

gcc sparseScsrmm.c -o sparseScsrmm -std=c99
time ./sparseScsrmm matrix12.mat matrix.mat > sparseScsrmm_out.txt

mpicc sparseScsrmm_mpi.c -o sparseScsrmm_mpi -std=c99
time mpiexec -machinefile $PBS_NODEFILE ./sparseScsrmm_mpi matrix12.mat matrix.mat > sparseScsrmm_mpi_out.txt

nvcc sparseScsrmm_cuda.cu -o sparseScsrmm_cuda
time ./sparseScsrmm_cuda matrix12.mat matrix.mat > sparseScsrmm_cuda_out.txt

运行结果sparseScsrmm.o17997

自上而下分别是串行算法的时间(6.199s),MPI 算法的时间(23.416s),CUDA 算法的时间(3.942s)。可以看到,这里实现的两个并行算法都在不过分增加并行开销的情况下增加了串行算法的可扩展性。其中 CUDA 算法在原来的算法上得到了 57%的提速,效果还是非常不错的。


real	0m6.199s
user	0m5.969s
sys	0m0.195s

real	0m23.416s
user	0m39.292s
sys	0m56.808s

real	0m3.942s
user	0m2.988s
sys	0m0.848s

串行算法sparseScsrmm.c

终端执行下述指令,可以计时执行并将结果写入sparseScsrmm_out.txt

gcc sparseScsrmm.c -o sparseScsrmm -std=c99
time ./sparseScsrmm matrix12.mat matrix.mat > sparseScsrmm_out.txt
#include "WuKSPARSE.h"
void ScsrMulScsc(const Scoo *a, const Scoo *b, Scoo *c)
{
	c->m = a->m;
	c->n = b->n;
	c->nnz = 0;
	ScooMalloc(c);
	for (int i = 0; i < c->m; ++i)
		if (a->RowPtr.data[i] < a->RowPtr.data[i + 1])
			for (int j = 0; j < c->n; ++j)
			{
				float res = 0;
				for (int
						 kb = b->ColPtr.data[j],
						 ka = a->RowPtr.data[i];
					 kb < b->ColPtr.data[j + 1] &&
					 ka < a->RowPtr.data[i + 1];
					 ++kb)
					if (b->Val.data[kb])
					{
						while (a->ColInd.data[ka] < b->RowInd.data[kb])
							++ka;
						while (a->ColInd.data[ka] == b->RowInd.data[kb])
							res += a->Val.data[ka] * b->Val.data[kb], ++ka;
					}
				if (res)
					ScooPushBack(c, i, j, res);
			}
	ScsrInit(c);
}
void ScsrMulScsr(const Scoo *a_csr, const Scoo *b_csr, Scoo *c_csr)
{
	Scoo B_csc;
	ScooToScsc(b_csr, &B_csc);
	ScsrMulScsc(a_csr, &B_csc, c_csr);
	ScooFree(&B_csc);
}
int main(int argc, char **argv)
{
	Scoo a_coo, b_coo, a_csr, b_csr, c_csr;

	ScooRead(&a_coo, fopen(argv[1], "r"));
	ScooRead(&b_coo, fopen(argv[2], "r"));

	ScooToScsr(&a_coo, &a_csr);
	ScooToScsr(&b_coo, &b_csr);

	ScooFree(&a_coo);
	ScooFree(&b_coo);

	ScsrMulScsr(&a_csr, &b_csr, &c_csr);

	ScooWrite(&c_csr, stdout);

	ScooFree(&a_csr);
	ScooFree(&b_csr);
	ScooFree(&c_csr);
}

MPI 算法sparseScsrmm_mpi.c

终端执行下述指令,可以计时执行并将结果写入sparseScsrmm_mpi_out.txt

mpicc sparseScsrmm_mpi.c -o sparseScsrmm_mpi -std=c99
time mpiexec -machinefile $PBS_NODEFILE ./sparseScsrmm_mpi matrix12.mat matrix.mat > sparseScsrmm_mpi_out.txt
#include <mpi.h>
#include "WuKSPARSE.h"
void ScsrMulScsc(const Scoo *a, const Scoo *b, Scoo *c)
{
	c->m = a->m;
	c->n = b->n;
	c->nnz = 0;
	ScooMalloc(c);
	for (int i = 0; i < c->m; ++i)
		if (a->RowPtr.data[i] < a->RowPtr.data[i + 1])
			for (int j = 0; j < c->n; ++j)
			{
				float res = 0;
				for (int
						 kb = b->ColPtr.data[j],
						 ka = a->RowPtr.data[i];
					 kb < b->ColPtr.data[j + 1] &&
					 ka < a->RowPtr.data[i + 1];
					 ++kb)
					if (b->Val.data[kb])
					{
						while (a->ColInd.data[ka] < b->RowInd.data[kb])
							++ka;
						while (a->ColInd.data[ka] == b->RowInd.data[kb])
							res += a->Val.data[ka] * b->Val.data[kb], ++ka;
					}
				if (res)
					ScooPushBack(c, i, j, res);
			}
	ScsrInit(c);
}
void ScsrMulScsr(const Scoo *a_csr, const Scoo *b_csr, Scoo *c_csr)
{
	Scoo B_csc;
	ScooToScsc(b_csr, &B_csc);
	ScsrMulScsc(a_csr, &B_csc, c_csr);
	ScooFree(&B_csc);
}
int main(int argc, char **argv)
{
	int comSize, comRank;
	MPI_Init(&argc, &argv);
	MPI_Comm_size(MPI_COMM_WORLD, &comSize);
	MPI_Comm_rank(MPI_COMM_WORLD, &comRank);
	Scoo a_coo, b_coo, a_csr, a_csr_local, b_csr, c_csr, c_csr_local;

	if (!comRank)
	{
		ScooRead(&a_coo, fopen(argv[1], "r"));
		ScooRead(&b_coo, fopen(argv[2], "r"));
		ScooToScsr(&a_coo, &a_csr);
		ScooToScsr(&b_coo, &b_csr);
		ScooFree(&a_coo);
		ScooFree(&b_coo);

		for (; a_csr.nnz % comSize; ++a_csr.RowPtr.data[a_csr.m])
			ScooPushBack(&a_csr, a_csr.m - 1, a_csr.n - 1, 0);
		a_csr_local.m = a_csr.m;
		a_csr_local.n = a_csr.n;
		a_csr_local.nnz = a_csr.nnz / comSize;
	}

	MPI_Bcast(
		&a_csr_local.m,
		1,
		MPI_INT,
		0,
		MPI_COMM_WORLD);

	MPI_Bcast(
		&a_csr_local.n,
		1,
		MPI_INT,
		0,
		MPI_COMM_WORLD);

	MPI_Bcast(
		&a_csr_local.nnz,
		1,
		MPI_INT,
		0,
		MPI_COMM_WORLD);

	ScooMalloc(&a_csr_local);

	MPI_Scatter(
		a_csr.RowInd.data,
		a_csr_local.nnz,
		MPI_INT,
		a_csr_local.RowInd.data,
		a_csr_local.nnz,
		MPI_INT,
		0,
		MPI_COMM_WORLD);

	MPI_Scatter(
		a_csr.ColInd.data,
		a_csr_local.nnz,
		MPI_INT,
		a_csr_local.ColInd.data,
		a_csr_local.nnz,
		MPI_INT,
		0,
		MPI_COMM_WORLD);

	MPI_Scatter(
		a_csr.Val.data,
		a_csr_local.nnz,
		MPI_FLOAT,
		a_csr_local.Val.data,
		a_csr_local.nnz,
		MPI_FLOAT,
		0,
		MPI_COMM_WORLD);
	ScsrInit(&a_csr_local);

	MPI_Bcast(
		&b_csr.m,
		1,
		MPI_INT,
		0,
		MPI_COMM_WORLD);

	MPI_Bcast(
		&b_csr.n,
		1,
		MPI_INT,
		0,
		MPI_COMM_WORLD);

	MPI_Bcast(
		&b_csr.nnz,
		1,
		MPI_INT,
		0,
		MPI_COMM_WORLD);
	if (comRank)
		ScooMalloc(&b_csr);

	MPI_Bcast(
		b_csr.RowInd.data,
		b_csr.nnz,
		MPI_INT,
		0,
		MPI_COMM_WORLD);

	MPI_Bcast(
		b_csr.ColInd.data,
		b_csr.nnz,
		MPI_INT,
		0,
		MPI_COMM_WORLD);

	MPI_Bcast(
		b_csr.Val.data,
		b_csr.nnz,
		MPI_FLOAT,
		0,
		MPI_COMM_WORLD);

	ScsrInit(&b_csr);
	ScsrMulScsr(&a_csr_local, &b_csr, &c_csr_local);

	ScooFree(&a_csr_local);
	ScooFree(&b_csr);

	c_csr.m = c_csr_local.m;
	c_csr.n = c_csr_local.n;

	IntVector count, disp;
	IntVectorInit(&count, comSize);
	IntVectorInit(&disp, comSize);

	MPI_Gather(
		&c_csr_local.nnz,
		1,
		MPI_INT,
		count.data,
		1,
		MPI_INT,
		0,
		MPI_COMM_WORLD);

	if (!comRank)
	{
		for (int i = disp.data[0] = 0; i < comSize - 1; ++i)
			disp.data[i + 1] = disp.data[i] + count.data[i];
		c_csr.nnz = disp.data[comSize - 1] + count.data[comSize - 1];
		ScooMalloc(&c_csr);
	}

	MPI_Gatherv(
		c_csr_local.RowInd.data,
		c_csr_local.nnz,
		MPI_INT,
		c_csr.RowInd.data,
		count.data,
		disp.data,
		MPI_INT,
		0,
		MPI_COMM_WORLD);

	MPI_Gatherv(
		c_csr_local.ColInd.data,
		c_csr_local.nnz,
		MPI_INT,
		c_csr.ColInd.data,
		count.data,
		disp.data,
		MPI_INT,
		0,
		MPI_COMM_WORLD);

	MPI_Gatherv(
		c_csr_local.Val.data,
		c_csr_local.nnz,
		MPI_FLOAT,
		c_csr.ColInd.data,
		count.data,
		disp.data,
		MPI_FLOAT,
		0,
		MPI_COMM_WORLD);

	ScooFree(&c_csr_local);
	IntVectorFree(&count);
	IntVectorFree(&disp);

	if (!comRank)
	{
		ScooFree(&a_csr);
		ScsrInit(&c_csr);
		ScooWrite(&c_csr, stdout);
		ScooFree(&c_csr);
	}

	MPI_Finalize();
}

CUDA 算法sparseScsrmm_cuda.cu

终端执行下述指令,可以计时执行并将结果写入sparseScsrmm_cuda_out.txt

nvcc sparseScsrmm_cuda.cu -o sparseScsrmm_cuda
time ./sparseScsrmm_cuda matrix12.mat matrix.mat > sparseScsrmm_cuda_out.txt
#include "WuKSPARSE.h"
void __global__ cuScsrMulScs(
	const int *a_csr_RowPtr,
	const int *a_csr_CowInd,
	const float *a_csr_Val,
	const int *b_csc_ColPtr,
	const int *b_csc_RowInd,
	const float *b_csc_Val,
	float *c,
	int m,
	int n)
{
	int
		i = blockIdx.x * blockDim.x + threadIdx.x,
		j = blockIdx.y * blockDim.y + threadIdx.y;
	float res = 0;
	for (int
			 kb = b_csc_ColPtr[j],
			 ka = a_csr_RowPtr[i];
		 kb < b_csc_ColPtr[j + 1] &&
		 ka < a_csr_RowPtr[i + 1];
		 ++kb)
	{
		while (a_csr_CowInd[ka] < b_csc_RowInd[kb])
			++ka;
		while (a_csr_CowInd[ka] == b_csc_RowInd[kb])
			res += a_csr_Val[ka] * b_csc_Val[kb], ++ka;
	}
	c[i * n + j] = res;
}
void ScsrMulScsr(const Scoo *a_csr, const Scoo *b_csr, Scoo *c_csr)
{
	Scoo B_csc;
	ScooToScsc(b_csr, &B_csc);

	int
		*a_RowPtr,
		*a_ColInd,
		*b_ColPtr,
		*b_RowInd;
	float
		*a_Val,
		*b_Val,
		*d_c;

	cudaMalloc(&a_RowPtr, sizeof(int) * (a_csr->m + 1));
	cudaMemcpy(a_RowPtr, a_csr->RowPtr.data, sizeof(int) * (a_csr->m + 1), cudaMemcpyHostToDevice);

	cudaMalloc(&a_ColInd, sizeof(int) * a_csr->nnz);
	cudaMemcpy(a_ColInd, a_csr->ColInd.data, sizeof(int) * a_csr->nnz, cudaMemcpyHostToDevice);

	cudaMalloc(&a_Val, sizeof(float) * a_csr->nnz);
	cudaMemcpy(a_Val, a_csr->Val.data, sizeof(float) * a_csr->nnz, cudaMemcpyHostToDevice);

	cudaMalloc(&b_ColPtr, sizeof(int) * (B_csc.n + 1));
	cudaMemcpy(b_ColPtr, B_csc.ColPtr.data, sizeof(int) * (B_csc.n + 1), cudaMemcpyHostToDevice);

	cudaMalloc(&b_RowInd, sizeof(int) * B_csc.nnz);
	cudaMemcpy(b_RowInd, B_csc.RowInd.data, sizeof(int) * B_csc.nnz, cudaMemcpyHostToDevice);

	cudaMalloc(&b_Val, sizeof(float) * B_csc.nnz);
	cudaMemcpy(b_Val, B_csc.Val.data, sizeof(float) * B_csc.nnz, cudaMemcpyHostToDevice);

	cudaMalloc(&d_c, sizeof(float) * a_csr->m * B_csc.n);

	dim3
		block(16, 16),
		grid(a_csr->m / block.x, B_csc.n / block.y);

	cuScsrMulScs<<<grid, block>>>(
		a_RowPtr,
		a_ColInd,
		a_Val,
		b_ColPtr,
		b_RowInd,
		b_Val,
		d_c,
		a_csr->m,
		B_csc.n);

	FloatVector c;
	FloatVectorInit(&c, a_csr->m * B_csc.n);
	cudaMemcpy(c.data, d_c, sizeof(float) * a_csr->m * B_csc.n, cudaMemcpyDeviceToHost);

	cudaFree(a_RowPtr);
	cudaFree(a_ColInd);
	cudaFree(b_ColPtr);
	cudaFree(b_RowInd);
	cudaFree(a_Val);
	cudaFree(b_Val);
	cudaFree(d_c);

	c_csr->m = a_csr->m;
	c_csr->n = B_csc.n;
	c_csr->nnz = 0;
	ScooMalloc(c_csr);

	for (int i = 0; i < c_csr->m; ++i)
		for (int j = 0; j < c_csr->n; ++j)
			if (c.data[i * c_csr->n + j])
				ScooPushBack(c_csr, i, j, c.data[i * c_csr->n + j]);
	FloatVectorFree(&c);
	ScooFree(&B_csc);
	ScsrInit(c_csr);
}
int main(int argc, char **argv)
{
	Scoo a_coo, b_coo, a_csr, b_csr, c_csr;

	ScooRead(&a_coo, fopen(argv[1], "r"));
	ScooRead(&b_coo, fopen(argv[2], "r"));

	ScooToScsr(&a_coo, &a_csr);
	ScooToScsr(&b_coo, &b_csr);

	ScooFree(&a_coo);
	ScooFree(&b_coo);

	ScsrMulScsr(&a_csr, &b_csr, &c_csr);

	ScooWrite(&c_csr, stdout);

	ScooFree(&a_csr);
	ScooFree(&b_csr);
	ScooFree(&c_csr);
}

稀疏矩阵处理库WuKSPARSE.h

#ifndef __WuKSPARSE_hpp__
#define __WuKSPARSE_hpp__
#include <stdio.h>
#include "WuKfloatVector.h"
#include "WuKintVector.h"
typedef struct Scoo
{
	int
		m,   //行
		n,   //列
		nnz; //非零元素数量
	IntVector
		RowPtr, //大小m+1
		ColPtr, //大小n+1
		RowInd, //大小nnz,行坐标
		ColInd; //大小nnz,列坐标
	FloatVector
		Val; //大小为nnz
} Scoo;
void ScooFree(Scoo *a)
{
	IntVectorFree(&a->RowPtr);
	IntVectorFree(&a->ColPtr);
	IntVectorFree(&a->RowInd);
	IntVectorFree(&a->ColInd);
	FloatVectorFree(&a->Val);
}
void ScooMalloc(Scoo *a)
{
	IntVectorInit(&a->RowPtr, a->m + 1);
	IntVectorInit(&a->ColPtr, a->n + 1);
	IntVectorInit(&a->RowInd, a->nnz);
	IntVectorInit(&a->ColInd, a->nnz);
	FloatVectorInit(&a->Val, a->nnz);
}
void ScooPushBack(Scoo *a, int r, int c, float v)
{
	IntVectorPushBack(&a->RowInd, r);
	IntVectorPushBack(&a->ColInd, c);
	FloatVectorPushBack(&a->Val, v);
	++a->nnz;
}
void ScooRead(Scoo *a, FILE *f)
{
#define MAXLEN 511
	for (char s[MAXLEN]; fgets(s, MAXLEN, f) && !feof(f);)
		if (s[0] != '%')
		{
			int nnz;
			sscanf(s, "%d%d%d", &a->m, &a->n, &nnz);
			a->nnz = 0;
			ScooMalloc(a);
			for (int line = 0; line < nnz; ++line)
			{
				int r, c;
				float v;
				fscanf(f, "%d%d%f", &r, &c, &v);
				if (0 <= r && r < a->m &&
					0 <= c && c < a->n)
					ScooPushBack(a, r, c, v);
			}
			break;
		}
#undef MAXLEN
}
void ScooWrite(const Scoo *a, FILE *f)
{
	fprintf(
		f,
		"%%MatrixMarket matrix coordinate real general\n%d %d %d\n",
		a->m,
		a->n,
		a->nnz);
	for (int j = 0; j < a->nnz; ++j)
		fprintf(
			f,
			"%d %d %f\n",
			a->RowInd.data[j],
			a->ColInd.data[j],
			a->Val.data[j]);
}
void ScsrInit(Scoo *a)
{
	int cur = 0;
	for (int i = a->RowPtr.data[0] = 0; i < a->nnz; ++i)
		for (; cur < a->RowInd.data[i]; ++cur)
			a->RowPtr.data[cur + 1] = i;
	for (; cur < a->m; ++cur)
		a->RowPtr.data[cur + 1] = a->nnz;
}
void ScooToScsr(const Scoo *a, Scoo *b)
{
	IntVector beg, s;
	IntVectorInit(&beg, a->nnz);
	IntVectorInit(&s, 2);
	for (int i = 0; i < a->nnz; ++i)
		beg.data[i] = i;
	s.data[0] = 0, s.data[1] = a->nnz - 1;
	while (s.size)
	{
		int right = s.data[--s.size], left = s.data[--s.size], l = left, r = right, pivot = beg.data[l];
		while (l < r)
		{
#define LESS(x, y) (a->RowInd.data[x] < a->RowInd.data[y] || a->RowInd.data[x] == a->RowInd.data[y] && a->ColInd.data[x] < a->ColInd.data[y])
			while (l < r && !LESS(beg.data[r], pivot))
				--r;
			beg.data[l] = beg.data[r];
			while (l < r && !LESS(pivot, beg.data[l]))
				++l;
			beg.data[r] = beg.data[l];
#undef LESS
		}
		beg.data[l] = pivot;
		if (l - 1 > left)
		{
			IntVectorPushBack(&s, left);
			IntVectorPushBack(&s, l - 1);
		}
		if (l + 1 < right)
		{
			IntVectorPushBack(&s, l + 1);
			IntVectorPushBack(&s, right);
		}
	}
	IntVectorFree(&s);
	*b = *a;
	ScooMalloc(b);
	for (int i = 0; i < b->nnz; ++i)
	{
		b->RowInd.data[i] = a->RowInd.data[beg.data[i]];
		b->ColInd.data[i] = a->ColInd.data[beg.data[i]];
		b->Val.data[i] = a->Val.data[beg.data[i]];
	}
	ScsrInit(b);
}
void ScscInit(Scoo *a)
{
	int cur = 0;
	for (int i = a->ColPtr.data[0] = 0; i < a->nnz; ++i)
		for (; cur < a->ColInd.data[i]; ++cur)
			a->ColPtr.data[cur + 1] = i;
	for (; cur < a->m; ++cur)
		a->ColPtr.data[cur + 1] = a->nnz;
}
void ScooToScsc(const Scoo *a, Scoo *b)
{
	IntVector beg, s;
	IntVectorInit(&beg, a->nnz);
	IntVectorInit(&s, 2);
	for (int i = 0; i < a->nnz; ++i)
		beg.data[i] = i;
	s.data[0] = 0, s.data[1] = a->nnz - 1;
	while (s.size)
	{
		int right = s.data[--s.size], left = s.data[--s.size], l = left, r = right, pivot = beg.data[l];
		while (l < r)
		{
#define LESS(x, y) (a->ColInd.data[x] < a->ColInd.data[y] || a->ColInd.data[x] == a->ColInd.data[y] && a->RowInd.data[x] < a->RowInd.data[y])
			while (l < r && !LESS(beg.data[r], pivot))
				--r;
			beg.data[l] = beg.data[r];
			while (l < r && !LESS(pivot, beg.data[l]))
				++l;
			beg.data[r] = beg.data[l];
#undef LESS
		}
		beg.data[l] = pivot;
		if (l - 1 > left)
		{
			IntVectorPushBack(&s, left);
			IntVectorPushBack(&s, l - 1);
		}
		if (l + 1 < right)
		{
			IntVectorPushBack(&s, l + 1);
			IntVectorPushBack(&s, right);
		}
	}
	IntVectorFree(&s);
	*b = *a;
	ScooMalloc(b);
	for (int i = 0; i < b->nnz; ++i)
	{
		b->RowInd.data[i] = a->RowInd.data[beg.data[i]];
		b->ColInd.data[i] = a->ColInd.data[beg.data[i]];
		b->Val.data[i] = a->Val.data[beg.data[i]];
	}
	ScscInit(b);
	IntVectorFree(&beg);
}
#endif

浮点动态向量库WuKfloatVector.h

#ifndef __WuKfloatVector_h__
#define __WuKfloatVector_h__
#include <malloc.h>
#include <string.h>
typedef struct FloatVector
{
    int size, cap;
    float *data;
} FloatVector;
void FloatVectorFree(FloatVector *v) { free(v->data); }
void FloatVectorInit(FloatVector *v, int size)
{
    v->size = v->cap = size;
    if (v->cap < 1)
        v->cap = 1;
    v->data = (float *)malloc(sizeof(float) * v->cap);
}
void FloatVectorPushBack(FloatVector *v, float val)
{
    while (v->cap <= v->size)
    {
        FloatVector v_new;
        FloatVectorInit(&v_new, v->cap << 1);
        memcpy(v_new.data, v->data, sizeof(float) * v->size);
        FloatVectorFree(v);
        *v = v_new;
        v->size >>= 1;
    }
    v->data[v->size++] = val;
}
#endif

整型动态向量库WuKintVector.h

#ifndef __WuKintVector_h__
#define __WuKintVector_h__
#include <malloc.h>
#include <string.h>
typedef struct IntVector
{
    int size, cap;
    int *data;
} IntVector;
void IntVectorFree(IntVector *v) { free(v->data); }
void IntVectorInit(IntVector *v, int size)
{
    v->size = v->cap = size;
    if (v->cap < 1)
        v->cap = 1;
    v->data = (int *)malloc(sizeof(int) * v->cap);
}
void IntVectorPushBack(IntVector *v, int val)
{
    while (v->cap <= v->size)
    {
        IntVector v_new;
        IntVectorInit(&v_new, v->cap << 1);
        memcpy(v_new.data, v->data, sizeof(int) * v->size);
        IntVectorFree(v);
        *v = v_new;
        v->size >>= 1;
    }
    v->data[v->size++] = val;
}
#endif