CUDA矩阵向量乘的多种优化

29 Nov 2019 11101字 38分次高性能计算
如果这篇博客帮助到你，可以请我喝一杯咖啡~
CC BY 4.0 （除特别声明或转载文章外）

使用下面一种或多种优化方法完成 CUDA 的矩阵向量乘法$y=A\times x$,其中$A$是$2^{14}\times 2^{14}$的方阵，$x$为$2^{14}$维向量。假设矩阵$A$的元素为$a_{i,j}=i-0.1\times j+1$，向量$x$的元素为$b_i=\log\sqrt{i\times i-i+2}$。

使用 global memory
使用合并访存
使用 constant memory 存放向量
使用 shared memory 存放向量和矩阵
使用 warp 直接访问寄存器
使用 cublasSgemv

实验环境

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

$ 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                                                 |
+-----------------------------------------------------------------------------+

实验原理

优化 CUDA 架构上的程序，一般从以下几个方面考虑：

选择好的并行算法，发掘更多的数据并行性
保持 SM 尽可能忙碌，尽量利用所有的 SM 参与计算
- 加大数据量
- 减小线程块大小
优化存储器的使用
- 全局存储器合并访问
- 使用更快的 constant memory 或 shared memory
使用一些已有的库，如<cuBlas_v2.h>

实验过程

由于都是 CUDA 架构上的核函数对比性能，下面的计时都只测了用于核函数计算的时间，而不包含数据拷贝的部分（否则运行时间都在 300ms 左右，基本上都是拷贝的时间而没有参考价值了）。当然，由于没有计入拷贝等预处理的时间，那些需要计算转置（列优先）或者预读取的算法在这里会有优势一些。

使用 global memory

这是最基础的矩阵向量乘法。这里假设线程块都是一维组织的，每个 CUDA 线程计算矩阵的一行与向量乘积，这样各线程之间没有读写冲突，不需要使用原子操作。

void __global__ wkSgemvGlobalMemory(
	const float *Ar, //行优先形式，下同
	const float *x,
	float *y,
	const size_t m, //A的行数
	const size_t n) //A的列数
{
	const size_t i = blockDim.x * blockIdx.x + threadIdx.x;
	if (i < m)
	{
		float res = 0; //将结果先存在寄存器里，减少对向量y的访存
		for (size_t j = 0; j < n; ++j)
			res += Ar[i * n + j] * x[j];
		y[i] = res;
	}
}

运行时间为4.694240ms。

使用合并访存

所谓合并访存，指的是相邻的线程访问段对齐的地址。比如在之前的代码中，j == 0时线程 0 访问Ar[0]，线程 1 访问Ar[nCol]，线程 2 访问Ar[2 * nCol]…它们并不相邻，因此不满足合并访问的要求。在这里我们把原来的行优先矩阵$A$转换成列优先表示形式（即行优先下的转置$A^T$），此时j == 0时线程 0 访问Ac[0]，线程 1 访问Ac[1]，线程 2 访问Ac[2]…此时满足了合并访问的要求。

void __global__ wkSgemvGlobalMemoryAlign(
	const float *Ac, //列优先形式，下同
	const float *x,
	float *y,
	const size_t m,
	const size_t n)
{
	const size_t i = blockDim.x * blockIdx.x + threadIdx.x;
	if (i < m)
	{
		float res = 0;
		for (size_t j = 0; j < n; ++j)
			res += Ac[j * m + i] * x[j];
		y[i] = res;
	}
}

运行时间为1.551584ms，性能提高了将近三倍，充分说明了合并访存的重要性。

使用 constant memory 存放向量

注意到向量在计算过程中不会改变，且每个线程访问相同地址，因此考虑把它放在 constant memory 中。

NVIDIA 硬件提供了 64KB 的常量内存，并且常量内存采用了不同于标准全局内存的处理方式。在这里我们大小为$2^{14}$的单精度浮点数向量$x$大小恰好为 64KB，正好可以完整保存。如果向量超过了 constant memory 的 64KB 上限，那就需要分批进行，多次传输和启动内核。

float __constant__ d_cx[(1 << 16) / sizeof(float)]; //64KB
void __global__ wkSgemvConstantMemory(
	const float *Ac,
	const float *x,
	float *y,
	const size_t m,
	const size_t n)
{
	const size_t i = blockDim.x * blockIdx.x + threadIdx.x;
	if (i < m)
	{
		float res = 0;
		for (size_t j = 0; j < n; ++j)
			res += Ac[j * m + i] * d_cx[j];
		y[i] = res;
	}
}

运行时间为1.516992ms，在上一步的基础上略微提高。使用常量内存可以提升运算性能的原因主要有两个：

对常量内存的单次读操作可以广播到同个半线程束的其他$15$个线程，这种方式产生的内存流量只是使用全局内存时的$\frac{1}{16}$。
硬件将主动把常量数据缓存在 GPU 上。在第一次从常量内存的某个地址上读取后，当其他半线程束请求同一个地址时，那么将命中缓存，这同样减少了额外的内存流量。

使用 shared memory 存放向量和矩阵

对于 block 内内存来说，向量都是共享的，因此我们可以使用比 constant memory 更快的 shared memory 来存储，此时相比较使用常量内存，我们免掉了向量比较大的时候多次数据拷贝和启动核函数的开销，也没有使用全局变量，增加了代码的可扩展性。当然，shared memory 更小（48K），因此需要对向量进行分块处理。

另外需要更正的一个问题是，并不需要使用 shared memory 去存矩阵，因为在这个矩阵向量乘的过程中，每个矩阵元素只被访问了一次。此外，shared memory 的大小也并不足以存下完整的矩阵（甚至是向量）。

template <size_t reduce_size>
void __global__ wkSgemvSharedMemory(
	const float *Ac,
	const float *x,
	float *y,
	const size_t m,
	const size_t n)
{
	extern float __shared__ sx[];
	const size_t
		i = blockDim.x * blockIdx.x + threadIdx.x,
		jBegLast = n / reduce_size * reduce_size;
	float res = 0;
	for (size_t jBeg = 0; jBeg < jBegLast; jBeg += reduce_size)
	{
		__syncthreads(); //防止有的进程还在读sx
		sx[threadIdx.x] = x[jBeg + threadIdx.x];
		__syncthreads();
		if (i < m)
			for (size_t j = 0; j < reduce_size; ++j) //能够自动展开
				res += Ac[(j + jBeg) * m + i] * sx[j];
	}
	{
		__syncthreads(); //防止有的进程还在读sx
		if (jBegLast + threadIdx.x < n)
			sx[threadIdx.x] = x[jBegLast + threadIdx.x];
		__syncthreads();
		if (i < m)
			for (size_t j = 0; j < n - jBegLast; ++j) //不能自动展开
				res += Ac[(j + jBegLast) * m + i] * sx[j];
	}
	if (i < m)
		y[i] = res;
}

运行时间为1.400672ms。注意这里我们将循环展开了，好处是减少了核函数运行时的分支。如果不展开的话，其运行时间将退化到比之前的还慢。

使用 shuffle

template <size_t warp_size>
void __global__ wkSgemvWarp(
	const float *Ac,
	const float *x,
	float *y,
	const size_t m,
	const size_t n)
{
	const size_t
		i = blockDim.x * blockIdx.x + threadIdx.x,
		lane_id = i % warp_size,
		jBegLast = n / warp_size * warp_size;
	float res = 0;
	for (size_t jBeg = 0; jBeg < jBegLast; jBeg += warp_size)
	{
		const float val = x[jBeg + lane_id];
		for (size_t j = 0; j < warp_size; ++j) //能够自动展开
			res += Ac[(j + jBeg) * m + i] * __shfl_sync(0xffffffff, val, j, warp_size);
	}
	{
		const float val = jBegLast + lane_id < n ? x[jBegLast + lane_id] : 0;
		for (size_t j = 0; j < n - jBegLast; ++j) //不能自动展开
			res += Ac[(j + jBegLast) * m + i] * __shfl_sync(0xffffffff, val, j, warp_size);
	}
	if (i < m)
		y[i] = res;
}

运行时间1.594368ms，反而有一定的下降。分析一下原因：老师集群上的显卡性能过于强悍（在今年十一月 SC 超算大会刚发布 Tesla V100S 前，Tesla V100 一直都是市面能买到的最强算力），内存读写性能比以往的显卡都要强很多，因此对本来已经很快的 shared memory 的优化效果没有那么明显了，而由于warp_size小于reduce_size导致循环分支次数却比上一步多，效果变差。

使用 `<cublas_v2.h>`

该函数的官方文档见https://docs.nvidia.com/cuda/cublas/index.html#cublas-lt-t-gt-gemv。下面是传入行优先矩阵的调用。要尤其注意的是，cublas库为了与FORTRAN中的接口保持一致，默认的稠密矩阵是按照列优先方法存储的，因此对于行优先存储的形式反而要标记转置。

cublasSgemv(
	wk_cublas_handle,
	CUBLAS_OP_T,
	m,
	n,
	&alpha,
	d_Ar,
	m,
	d_x,
	1,
	&beta,
	d_y,
	1);

运行时间1.863520ms。下面是传入列优先矩阵的调用。

cublasSgemv(
	wk_cublas_handle,
	CUBLAS_OP_N,
	m,
	n,
	&alpha,
	d_Ac,
	m,
	d_x,
	1,
	&beta,
	d_y,
	1);

运行时间1.381888ms。（艹，我做了这半天的优化结果还没它快）

`Sgemv.o12727`

分别是上面几种方法函数的运行时间。

4.694240ms
1.551584ms
1.516992ms
1.400672ms
1.594368ms
1.863520ms
1.381888ms

可以看到，由于现在的硬件性能已经大大强于数年前，做存储器的优化效果已经比较小（并不是说没有，只是甚至已经小于多次循环跳转的开销了）。因此，对这个问题来说，最主要的是要选一个优秀的并行算法，再对程序代码做好访存分析和优化。

当然也不是说存储器结构就不再重要，还是要具体问题具体分析。上面很多算法都是要对矩阵或者向量进行预处理的，而并没有把对应的代价（时间、内存空间、可扩展性等）计入在内，实际上在运用到生产环境的时候这些仍然是必须要考虑的。最后，虽然前面的代码还没有调库跑得快，实际上还是有优化余地的，比如：

使用#pragma unroll进行循环展开。
- 我在测试的时候用这种方法一度将wkSgemvConstantMemory优化至1.35ms，但是这种方法实际上是面向硬件和数据本身的（需要提前知道有多少数据被优化，才能确定展开多少次是安全的），在真正的生产环境中可移植性不够高，我不是非常喜欢。
多路启动，每个线程负责多个向量元素，可以减少调度时开销。
- 不喜欢的理由同上。

总之结论就是，在实际生产力环境中还是直接调库来的省事，库中的黑科技代码在绝大多数情况下都是要优秀于自己写的代码的。当然，自己手动做过这个实验也可以帮助我们对库的代码做进了一步的理解，比如列优先矩阵的向量乘法为什么优秀于行优先。

`Sgemv.pbs`

调度脚本。

#PBS -N Sgemv
#PBS -l nodes=1:ppn=32:gpus=1
#PBS -j oe
#PBS -q gpu
source /public/software/profile.d/cuda10.0.sh
cd $PBS_O_WORKDIR
nvcc Sgemv.cu -run -lcublas

`Sgemv.cu`

完整代码。

#include <stdio.h>
#include <math.h>
#include <cuda_runtime.h>
#include <cublas_v2.h>
void __global__ wkSgemvGlobalMemory(
	const float *Ar, //行优先形式，下同
	const float *x,
	float *y,
	const size_t m, //A的行数
	const size_t n) //A的列数
{
	const size_t i = blockDim.x * blockIdx.x + threadIdx.x;
	if (i < m)
	{
		float res = 0; //将结果先存在寄存器里，减少对向量y的访存
		for (size_t j = 0; j < n; ++j)
			res += Ar[i * n + j] * x[j];
		y[i] = res;
	}
}
void __global__ wkSgemvGlobalMemoryAlign(
	const float *Ac, //列优先形式，下同
	const float *x,
	float *y,
	const size_t m,
	const size_t n)
{
	const size_t i = blockDim.x * blockIdx.x + threadIdx.x;
	if (i < m)
	{
		float res = 0;
		for (size_t j = 0; j < n; ++j)
			res += Ac[j * m + i] * x[j];
		y[i] = res;
	}
}
float __constant__ d_cx[(1 << 16) / sizeof(float)]; //64KB
void __global__ wkSgemvConstantMemory(
	const float *Ac,
	const float *x,
	float *y,
	const size_t m,
	const size_t n)
{
	const size_t i = blockDim.x * blockIdx.x + threadIdx.x;
	if (i < m)
	{
		float res = 0;
		for (size_t j = 0; j < n; ++j)
			res += Ac[j * m + i] * d_cx[j];
		y[i] = res;
	}
}
template <size_t reduce_size>
void __global__ wkSgemvSharedMemory(
	const float *Ac,
	const float *x,
	float *y,
	const size_t m,
	const size_t n)
{
	extern float __shared__ sx[];
	const size_t
		i = blockDim.x * blockIdx.x + threadIdx.x,
		jBegLast = n / reduce_size * reduce_size;
	float res = 0;
	for (size_t jBeg = 0; jBeg < jBegLast; jBeg += reduce_size)
	{
		__syncthreads(); //防止有的进程还在读sx
		sx[threadIdx.x] = x[jBeg + threadIdx.x];
		__syncthreads();
		if (i < m)
			for (size_t j = 0; j < reduce_size; ++j) //能够自动展开
				res += Ac[(j + jBeg) * m + i] * sx[j];
	}
	{
		__syncthreads(); //防止有的进程还在读sx
		if (jBegLast + threadIdx.x < n)
			sx[threadIdx.x] = x[jBegLast + threadIdx.x];
		__syncthreads();
		if (i < m)
			for (size_t j = 0; j < n - jBegLast; ++j) //不能自动展开
				res += Ac[(j + jBegLast) * m + i] * sx[j];
	}
	if (i < m)
		y[i] = res;
}
template <size_t warp_size>
void __global__ wkSgemvWarp(
	const float *Ac,
	const float *x,
	float *y,
	const size_t m,
	const size_t n)
{
	const size_t
		i = blockDim.x * blockIdx.x + threadIdx.x,
		lane_id = i % warp_size,
		jBegLast = n / warp_size * warp_size;
	float res = 0;
	for (size_t jBeg = 0; jBeg < jBegLast; jBeg += warp_size)
	{
		const float val = x[jBeg + lane_id];
		for (size_t j = 0; j < warp_size; ++j) //能够自动展开
			res += Ac[(j + jBeg) * m + i] * __shfl_sync(0xffffffff, val, j, warp_size);
	}
	{
		const float val = jBegLast + lane_id < n ? x[jBegLast + lane_id] : 0;
		for (size_t j = 0; j < n - jBegLast; ++j) //不能自动展开
			res += Ac[(j + jBegLast) * m + i] * __shfl_sync(0xffffffff, val, j, warp_size);
	}
	if (i < m)
		y[i] = res;
}
int main()
{
	const size_t
		m = 1 << 14,
		n = 1 << 14;
	float
		*h_Ar,
		*h_Ac,
		*h_x,
		*d_Ar,
		*d_Ac,
		*d_x,
		*d_y;
	cudaHostAlloc(
		(void **)&h_Ar,
		sizeof(float) * m * n,
		cudaHostAllocWriteCombined);
	cudaHostAlloc(
		(void **)&h_Ac,
		sizeof(float) * m * n,
		cudaHostAllocWriteCombined);
	cudaHostAlloc(
		(void **)&h_x,
		sizeof(float) * n,
		cudaHostAllocWriteCombined);
	for (size_t j = 0; j < n; ++j)
	{
		h_x[j] = log(sqrt(j * j - j + 2));
		for (size_t i = 0; i < m; ++i)
			h_Ac[j * m + i] = h_Ar[i * n + j] = i - 0.1 * j + 1;
	}
	cudaMalloc(
		(float **)&d_Ar,
		sizeof(float) * m * n);
	cudaMalloc(
		(float **)&d_Ac,
		sizeof(float) * m * n);
	cudaMalloc(
		(float **)&d_x,
		sizeof(float) * n);
	cudaMalloc(
		(float **)&d_y,
		sizeof(float) * m);
	cudaMemcpy(
		d_Ar,
		h_Ar,
		sizeof(float) * m * n,
		cudaMemcpyHostToDevice);
	cudaMemcpy(
		d_Ac,
		h_Ac,
		sizeof(float) * m * n,
		cudaMemcpyHostToDevice);
	cudaMemcpy(
		d_x,
		h_x,
		sizeof(float) * n,
		cudaMemcpyHostToDevice);
	cudaMemcpyToSymbol(
		d_cx,
		h_x,
		sizeof(float) * n,
		cudaMemcpyHostToDevice);
	cudaFreeHost(h_Ar);
	cudaFreeHost(h_Ac);
	cudaFreeHost(h_x);
	cublasHandle_t wk_cublas_handle;
	cublasCreate(&wk_cublas_handle);
	float
		alpha = 1,
		beta = 0;
	for (int i = 0; i < 7; ++i)
	{
		cudaEvent_t beg, end;
		cudaEventCreate(&beg);
		cudaEventCreate(&end);
		cudaEventRecord(beg, 0);
		const size_t
			blocks = 1 << 7,
			grids = (m + blocks - 1) / blocks;
		if (i == 0)
			wkSgemvGlobalMemory<<<grids, blocks>>>(
				d_Ar,
				d_x,
				d_y,
				m,
				n);
		else if (i == 1)
			wkSgemvGlobalMemoryAlign<<<grids, blocks>>>(
				d_Ac,
				d_x,
				d_y,
				m,
				n);
		else if (i == 2)
			wkSgemvConstantMemory<<<grids, blocks>>>(
				d_Ac,
				d_x,
				d_y,
				m,
				n);
		else if (i == 3)
			wkSgemvSharedMemory<blocks><<<grids, blocks, sizeof(float) * blocks>>>(
				d_Ac,
				d_x,
				d_y,
				m,
				n);
		else if (i == 4)
			wkSgemvWarp<32><<<grids, blocks>>>(
				d_Ac,
				d_x,
				d_y,
				m,
				n);
		else if (i == 5)
			cublasSgemv(
				wk_cublas_handle,
				CUBLAS_OP_T,
				m,
				n,
				&alpha,
				d_Ar,
				m,
				d_x,
				1,
				&beta,
				d_y,
				1);
		else if (i == 6)
			cublasSgemv(
				wk_cublas_handle,
				CUBLAS_OP_N,
				m,
				n,
				&alpha,
				d_Ac,
				m,
				d_x,
				1,
				&beta,
				d_y,
				1);
		cudaDeviceSynchronize();
		cudaEventRecord(end, 0);
		cudaEventSynchronize(beg);
		cudaEventSynchronize(end);
		float elapsed_time;
		cudaEventElapsedTime(&elapsed_time, beg, end);
		printf("%fms\n", elapsed_time);
	}
	cublasDestroy(wk_cublas_handle);
	cudaFree(d_Ar);
	cudaFree(d_Ac);
	cudaFree(d_x);
	cudaFree(d_y);
}