cuda内核函数--全连接层实现

最近在学习cuda内核函数，将以全连接层为例子来展示神经网络中某一层的完整cuda内核函数实现。

cuda内核函数

首先，介绍一下cuda内核函数的特点：

1. 以__global__开头，用于声明一个内核函数；
2. 无返回值（void），即需要传递的数据（包括输入和输出数据）通常是作为参数传递给内核函数：在GPU上并行计算模型的设计使得内核函数更倾向于通过修改传入的内存（如全局内存）来传递计算结果。
3. 当将一个数组传递给 CUDA 内核函数时，实际上是传递了该数组在 GPU 内存中的起始地址（即指针）。内核函数无法直接感知你传递的数组是几维的，它只看到一个连续的内存块。故实际上内核函数只能接收一维指针，需要通过手动计算索引的方式，在内核函数中实现多维数组的访问逻辑。

并行运算的最小单位是线程。一个内核函数可以启动成千上万个线程，这些线程在GPU上并行运行。
每个内核函数对应一个网格。在GPU中，内核函数会启动大量的线程，这些线程被组织成线程块（thread blocks）。多个线程块组成这个网格（grid），而每个线程块会在一个称为Streaming Multiprocessor (SM) 的硬件单元上执行。每个SM有多个执行单元，每个执行单元负责执行若干线程。
每个线程在执行时通过其索引访问数据：

• threadIdx.x：线程在线程块中的索引。
• blockIdx.x：线程块在网格中的索引。
• blockDim.x：每个线程块中的线程总数。
• gridDim.x：每个网格中线程块的数量。
  使用全局线程索引可以唯一标识一个线程：id_x = blockDim.x * blockIdx.x + threadIdx.x。

创建线程的数量： 在启动内核时会根据任务的大小和线程块的大小来计算所需的线程数（预设值）。在主机代码（CPU）中使用kernel_func_name<<<gridDim, blockDim>>>(...)调用内核函数，默认为该内核函数创建了$\text{gridDim}\times \text{blockDim}$个线程。
合理的线程分配： 为了使用合理的线程数量，节省计算资源和避免不必要的GPU资源消耗，有两种方式。

内核函数的编写都基本一致：

__global__ void kernel_func_name(const int num_of_threads, typename var_name ... ){
	int id_x = blockDim.x * blockIdx.x + threadIdx.x;
	// 只有前num_of_threads个线程执行函数体（只有num_of_threads个并行任务需要执行）
	if (id_x < num_of_threads) {
	    // 函数体
    }
}

两种方式调用内核函数的主机代码编写略有区别：

1. 开足够多的线程，但并不100%使用该内核管理的所有线程，而是通过num_of_threads变量来分配合适的线程来执行并行任务：

void cpu_func_name(const int num_of_threads, xxx){
	int gridDim = 4;
	int blockDim = 512;
	kernel_func_name<<<gridDim, blockDim>>>(num_of_threads, xxx);
}

2.在启动内核时预计算好合适数量的线程，以全连接层计算为例：

const int CUDA_NUM_THREADS = 1024; // 指定每个线程块中的线程数
// 输入期望的线程总数，返回线程块的数量
inline int GET_BLOCKS(const int N)
{
  return (N + CUDA_NUM_THREADS - 1) / CUDA_NUM_THREADS;
}

void cpu_func_name(const int batch_size, const int num_outputs, xxx){
	int bactch_size = 4;
	int num_outputs = 1024;
	int num_of_threads = batch_size * num_outputs;
	kernel_func_name<<<GET_BLOCKS(num_of_threads), CUDA_NUM_THREADS>>>(num_of_threads, xxx);

以全连接层为例

神经网络中的全连接层（fully connected layer, FC layer）的数学表达式是：$y=Wx+b$

函数参数一览

• 必须：

  • 输入：float * x, shape is [batch_size, num_inputs]
  • 输出：float * y, shape is [batch_size, num_outputs]
  • 权重：float * w, shape is [num_outputs, num_inputs]
  • 偏置：float * b, shape is [num_outputs]
  • 输入的神经元数：int num_inputs
  • 输出的神经元数：int num_outputs
  • 训练过程中一次性输入给模型的样本数量：int batch_size

• 可选（用于检验反向传播代码求解的梯度是否正确）：

  • 某权重1D索引：int w_t
  • 某偏置索引：int b_t
  • 一个极小值，用于近似求导数：double delta

代码逻辑

在并行化代码之前，明确可并行化的子任务及其数量是至关重要的。
已知在全连接层的计算中，每个输出的每个神经元的计算过程都是相互独立的：
${(y_k)}_i={\sum }_j{w}_{ij}{(x_k)}_j+b_i$

如果使用独立的线程来处理每个输出每个神经元的计算，即共需$\text{num\_outputs}\times \text{batch\_size}$个线程。

除了前向传播之外，还希望有一个梯度计算校验功能：
想要知道网络$f(\cdot )$的输出对于某参数（微小扰动）的敏感性及网络反向传播中的自动求导是否正确。需使用有限差分法来近似计算梯度，需要对每个参数分别进行扰动，并计算相应的输出变化。
$$\frac{\partial f_{\theta }(x)}{\partial \theta }=\frac{f_{\theta +\delta }(x)-f_{\theta }(x)}{\delta }$$
这里的$\theta$指的是网络中的单个可学习参数$w_{i,j}$或者$b_i$。

1. 运用正常的前向传播代码计算出原始输出$f_{\theta }(x_k{)}_i$。
2. 计算出$f_{\theta +\delta }(x_k{)}_i=f_{\theta }(x_k{)}_i+\delta \frac{\partial f_{\theta }(x_k{)}_i}{\partial \theta }$
3. 对应的grad_wij和grad_bi可以用${\mathbb{E}}_k[\frac{f_{\theta +\delta }(x_k{)}_i-f_{\theta }(x_k{)}_i}{\delta }]$计算得。

当然，该功能可以进化成输出整个grad_w和grad_b张量。

cuda代码编写

__global__ void Linear_forward(const int num_of_threads, int num_inputs, int num_outputs, int batch_size, 
double * x, double * y, double * w, double  * b, int w_t, int b_t, double delta){
	int idx = blockDim.x * blockIdx.x + threadIdx.x;
	// 不同线程的i依次为0, 1, ..., num_outputs-1
	if (idx < num_of_threads){
		int batch_idx = idx / num_outputs;
		int output_idx = idx % num_outputs;
		for (int j=0; j<num_inputs; j++){
			y[batch_idx*num_outputs+output_idx] += w[output_idx*num_inputs+j] * x[batch_idx*num_inputs+j];
		}
		y[batch_idx*num_outputs+output_idx] += b[output_idx];		
		
		// 为了验证反向传播自动求导出来的梯度值是否正确，提供输出某神经元的导数值计算
		if (w_t >= 0 && (w_t / num_inputs) == output_idx){
			int j = w_t % num_inputs;
			y[batch_idx*num_outputs+output_idx] += delta * x[batch_idx*num_inputs+j];
		}
		if (b_t >=0 && b_t == output_idx){
			y[batch_idx*num_outputs+output_idx] += delta * b[output_idx];
	}
}

主机代码编写

const int CUDA_NUM_THREADS = 256;
// 输入期望的线程总数，返回线程块的数量
inline int GET_BLOCKS(const int N)
{
  return (N + CUDA_NUM_THREADS - 1) / CUDA_NUM_THREADS;
}

int main(){
	int num_inputs = 128;
	int num_outputs = 1024;
	int batch_size = 4;
	int num_of_threads = batch_size * num_outputs;
	double delta = 1e-8;
	// x, y, w, b的传递此处略
	// 仅执行前向传播
	// void Linear_forward(int num_inputs, int num_outputs, int batch_size, 
// double * x, double * y, double * w, double  * b, int w_t, int b_t, double delta)
	Linear_forward<<<GET_BLOCKS(num_of_threads), CUDA_NUM_THREADS>>>(num_of_threads, num_inputs, 
		num_outputs, batch_size, x, y, w, b, -1, -1, delta);

	// 求解某1d索引为266的权重参数的梯度
	int w_t = 266;
	int i_idx = w_t / num_inputs;
	int j_idx = w_t % num_inputs;
	Linear_forward<<<gridDim, blockDim>>>(num_inputs,num_outputs, batch_size, x, y, w, b, -1, -1, delta);
	double y_i_ori[batch_size];
	for (int k=0; k<batch_size; k++)y_i_ori[k] = y[k*num_outputs+i_idx];
	Linear_forward<<<gridDim, blockDim>>>(num_inputs,num_outputs, batch_size, x, y, w, b, 266, -1, delta);
	double y_i_plus_delta[batch_size];
	double grad_w_ij = 0.0f;
	for (int k=0; k<batch_size; k++){
		y_i_plus_delta[k] = y[k*num_outputs+i_idx];
		grad_w_ij += (y_i_plus_delta[k] - y_i_ori[k]) / delta; 
	}
	// 对比grad_w_ij和self.weights.grad[i,j]的值是否一致即可
	grad_w_ij /= (double)batch_size;
	return 0;
}