Menu

再谈CUDA上的Reduce操作

上次我学习了用 Shuffle 加速 CUDA 上的 Reduce 操作,据说这是目前在 CUDA 上最快的区间规约算法。然而运用在实际的情况中却并没有对代码的性能带来多大提升。本文中我再次整理了自己已知的所有 CUDA 上的快速区间规约方法,并以此对写出高性能高可扩展的 CUDA 代码提出一些自己的思考。

  • 多路规约
  • 使用 Shared Memory 和 Warp Shuffle 增加计算带宽
  • 使用 <thrust/reduce.h>
  • 使用 <cublas_v2.h>

实验环境

使用 v100 集群上一个结点的单张 v100 运行。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
$ 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                                                 |
+-----------------------------------------------------------------------------+

实验过程与分析

一些说明

和前一个实验用 Shuffle 加速 CUDA 上的 Reduce 操作不同,这里被规约元素的类型不再是 unsigned 而是 double,更加贴合实际使用场景。单个 unsigned 内存占用 4 字节,而 double 是 8 字节,这就使得 Warp 间互相访问寄存器的流量增加了一倍,很大程度上降低了 Shuffle 的优化效果(推测)。

thrust::reduce

首先是性能标杆 thrust 库,来看一看目前最流行的实现可以达到怎样的性能。

1
sum = thrust::reduce(src.begin(), src.begin() + n);

运行时间 10.561248ms

simpleDasum

先来做最基础的算法优化,几乎无需掌握任何 CUDA 内存分布的知识。

首先,由于这里规约的元素数量高达十亿个,如果按照通常习惯的每个线程对应输入的一个元素,那么显卡的调度开销一定程度上无法忽视。这里我们让一个线程可以对应输入中的多个元素,减少所需要的线程数量,从而减少调度开销。

最后,我们用 template <size_t UNROLL_SIZE> 传入 #pragma unroll(UNROLL_SIZE) 循环展开次数,这样编译器可以在代码生成的时候将循环展开,以减少多次循环跳转的开销。

1
2
3
4
5
6
7
8
9
10
11
12
13
template <size_t UNROLL_SIZE>
void __global__ simpleDasumKernel(
	const size_t n,
	const double *src_d,
	double *tmp_d)
{
	const size_t global_id = threadIdx.x + blockDim.x * blockIdx.x;
	double val = 0;
#pragma unroll(UNROLL_SIZE)
	for (size_t i = global_id; i < n; i += blockDim.x * gridDim.x)
		val += src_d[i];
	tmp_d[global_id] = val;
}

运行时间为10.727680ms,和thrust相比还算可以接受。

naiveDasum

simpleDasum 基础上,优化存储器的使用。

对于同一个 block 内的所有线程,我们可以借助 Shared Memory 再进行一次树形规约,从而减少对内存的写操作和二次规约。

对于同一个 warp 内的所有线程也是同理,warp 间直接访问寄存器的开销比 Shared Memory 更小。此外 warp 同步的开销也要小于__syncthreads(),并且同一个 warp 上执行语句是不需要条件分支的,因为不管怎样都会被执行。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
template <
	size_t UNROLL_SIZE,
	size_t BLOCK_SIZE,
	size_t WARP_SIZE>
void __global__ naiveDasumKernel(
	const size_t n,
	const double *src_d,
	double *tmp_d)
{
	const size_t global_id = threadIdx.x + blockDim.x * blockIdx.x;
	double __shared__ shared[BLOCK_SIZE];
	{
		double val = 0;
#pragma unroll(UNROLL_SIZE)
		for (size_t i = global_id; i < n; i += blockDim.x * gridDim.x)
			val += src_d[i];
		shared[threadIdx.x] = val;
	}
#pragma unroll
	for (size_t offset = BLOCK_SIZE >> 1; offset > (WARP_SIZE >> 1); offset >>= 1)
	{
		__syncthreads();
		if (threadIdx.x < offset)
			shared[threadIdx.x] += shared[threadIdx.x ^ offset];
	}
	if (threadIdx.x < WARP_SIZE)
	{
		double val = shared[threadIdx.x];
#pragma unroll
		for (size_t offset = WARP_SIZE >> 1; offset > 0; offset >>= 1)
			val += __shfl_xor_sync(0xffffffff, val, offset, WARP_SIZE);
		if (threadIdx.x == 0)
			tmp_d[blockIdx.x] = val;
	}
}

运行时间 10.467008ms ,终于比 thrust::reduce 快了一丢丢。

cublasDasum

也可以使用线性代数库 <cublas_v2.h> 中的 asum 系列函数实现。

1
2
3
4
5
6
cublasDasum(
	wk_cublas_handle,
	n,
	thrust::raw_pointer_cast(src.data()),
	1,
	&sum);

运行时间为 11.918208ms,看来cublas库提供的线性代数抽象有一定的开销。

总结

可以看到,随着硬件技术的发展,显卡上的运行速度已经是相当快了,十亿多的数据只用了十毫秒就完成了规约,并且在这种前提下存储优化的效果越来越有限了。

然而,较之略显复杂的访存优化,一些简单的编程习惯反而能有效提高 CUDA 代码的效率。从可扩展性的角度来说,我也更倾向于写 simpleDasum 这样对硬件的依赖程度更低的代码。毕竟未来显卡到底会怎么发展谁也说不准,也许以后一个 warp 或者一个 block 中会有更多的线程。

最后调库大法好,自己做了半天优化最后也只比库快了一丢丢,从开发成本的角度来说还是不要重复造轮子为妙。

源代码

Dasum.pbs

调度脚本。

1
2
3
4
5
6
7
#PBS -N Dasum
#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 Dasum.cu -run -lcublas

Dasum.o18880

运行结果。

1
2
3
4
1073741824.000000 : 10.561248ms elapsed.
1073741824.000000 : 10.727680ms elapsed.
1073741824.000000 : 10.467008ms elapsed.
1073741824.000000 : 11.918208ms elapsed.

Dasum.cu

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
#include <stdio.h>
#include <cuda_runtime.h>
#include <thrust/device_vector.h>
#include <thrust/reduce.h>
#include <cublas_v2.h>
template <size_t UNROLL_SIZE>
void __global__ simpleDasumKernel(
	const size_t n,
	const double *src_d,
	double *tmp_d)
{
	const size_t global_id = threadIdx.x + blockDim.x * blockIdx.x;
	double val = 0;
#pragma unroll(UNROLL_SIZE)
	for (size_t i = global_id; i < n; i += blockDim.x * gridDim.x)
		val += src_d[i];
	tmp_d[global_id] = val;
}
template <
	size_t UNROLL_SIZE,
	size_t BLOCK_SIZE,
	size_t WARP_SIZE>
void __global__ naiveDasumKernel(
	const size_t n,
	const double *src_d,
	double *tmp_d)
{
	const size_t global_id = threadIdx.x + blockDim.x * blockIdx.x;
	double __shared__ shared[BLOCK_SIZE];
	{
		double val = 0;
#pragma unroll(UNROLL_SIZE)
		for (size_t i = global_id; i < n; i += blockDim.x * gridDim.x)
			val += src_d[i];
		shared[threadIdx.x] = val;
	}
#pragma unroll
	for (size_t offset = BLOCK_SIZE >> 1; offset > (WARP_SIZE >> 1); offset >>= 1)
	{
		__syncthreads();
		if (threadIdx.x < offset)
			shared[threadIdx.x] += shared[threadIdx.x ^ offset];
	}
	if (threadIdx.x < WARP_SIZE)
	{
		double val = shared[threadIdx.x];
#pragma unroll
		for (size_t offset = WARP_SIZE >> 1; offset > 0; offset >>= 1)
			val += __shfl_xor_sync(0xffffffff, val, offset, WARP_SIZE);
		if (threadIdx.x == 0)
			tmp_d[blockIdx.x] = val;
	}
}
int main()
{
	const size_t
		n = 1 << 30,
		REDUCE_SIZE = 1 << 22,
		UNROLL_SIZE = n / REDUCE_SIZE,
		BLOCK_SIZE = 1 << 10,
		WARP_SIZE = 1 << 5;
	thrust::device_vector<double> src(n, 1), tmp(REDUCE_SIZE);
	cublasHandle_t wk_cublas_handle;
	cublasCreate(&wk_cublas_handle);
	for (int op = 0; op < 4; ++op)
	{
		double sum;
		cudaEvent_t beg, end;
		cudaEventCreate(&beg);
		cudaEventCreate(&end);
		cudaEventRecord(beg, 0);
		if (op == 0)
			sum = thrust::reduce(src.begin(), src.begin() + n);
		if (op == 1)
		{
			simpleDasumKernel<
				UNROLL_SIZE><<<
				(REDUCE_SIZE + BLOCK_SIZE - 1) / BLOCK_SIZE,
				BLOCK_SIZE>>>(
				n,
				thrust::raw_pointer_cast(src.data()),
				thrust::raw_pointer_cast(tmp.data()));
			sum = thrust::reduce(tmp.begin(), tmp.begin() + REDUCE_SIZE);
		}
		if (op == 2)
		{
			naiveDasumKernel<
				UNROLL_SIZE,
				BLOCK_SIZE,
				WARP_SIZE><<<
				(REDUCE_SIZE + BLOCK_SIZE - 1) / BLOCK_SIZE,
				BLOCK_SIZE>>>(
				n,
				thrust::raw_pointer_cast(src.data()),
				thrust::raw_pointer_cast(tmp.data()));
			sum = thrust::reduce(tmp.begin(), tmp.begin() + (REDUCE_SIZE + BLOCK_SIZE - 1) / BLOCK_SIZE);
		}
		if (op == 3)
			cublasDasum(
				wk_cublas_handle,
				n,
				thrust::raw_pointer_cast(src.data()),
				1,
				&sum);
		cudaEventRecord(end, 0);
		cudaEventSynchronize(beg);
		cudaEventSynchronize(end);
		float elapsed_time;
		cudaEventElapsedTime(
			&elapsed_time,
			beg,
			end);
		printf("%f : %fms elapsed.\n", sum, elapsed_time);
	}
	cublasDestroy(wk_cublas_handle);
}