cuda实现任意尺寸的矩阵乘法

行、列数小于1024的cuda矩阵乘法

Nvidia GPU常见的块内线程数最大为1024，当矩阵的行数和列数均小于1024时，我们可以简单的采用行和列点到点依次相乘构建核函数，即块内的每个线程负责一对元素的乘积计算，然后将所有块内线程相乘的结果累加求和，得到结果矩阵对应行和列的元素值。

>>Code：参照CUDA11指导手册，给出核函数代码如下：

// Matrix multiplication kernel called by MatMul()
__global__ void MatMulKernel(Matrix A, Matrix B, Matrix C)
{
	// Each thread computes one element of C
	// by accumulating results into Cvalue
	float Cvalue = 0;
	int row = blockIdx.y * blockDim.y + threadIdx.y;	//结果矩阵C的行索引
	int col = blockIdx.x * blockDim.x + threadIdx.x;	//结果矩阵C的列索引
	for (int e = 0; e < A.width; ++e)
	{
		Cvalue += A.elements[row * A.width + e]			//所有点到点的元素乘积求和
				  * B.elements[e * B.width + col];
		C.elements[row * C.width + col] = Cvalue;
	}
}

这里Matrix是一个结构体，如下：

// Matrices are stored in row-major order:
// M(row, col) = *(M.elements + row * M.width + col)
typedef struct {
	int width;
	int height;
	float* elements;
} Matrix;

任意尺寸的矩阵乘法

上述核函数简单、易理解，且未使用块的共享内存，下面将介绍使用子矩阵划分和共享内存的方法实现任意尺寸的矩阵乘法。
思想如下：
设置块为宽高相等的二维形状，尺寸大小BLOCK_SIZE为32（32*32=1024），从而块内包含的线程个数为1024个，源矩阵A和B可以被nA和nB个块划分，如下图所示：

上图中，每个蓝色的块代表矩阵的划分得到的子矩阵，每个子矩阵对应一个线程块，它们的大小均为BLOCK_SIZE * BLOCK_SIZE, 为了计算C(1, 0)处的值，必须将每个块内第一行和B里面对应块（按矩阵乘法规则相对应）的第0列相乘得到结果，然后求和， 块内则是线程求对应元素相乘的结果求和。

值得注意的是，最后一行和最后一列的子矩阵行或列可能小于BLOCK_SIZE（图二中绿框内的子矩阵）， 因此在核函数计算时，需要作判断，约束仅计算在行列范围内的值。
比如说：矩阵A的shape为（33， 32）， 矩阵B的shape为（32， 35），则矩阵A的第二行第一个块的行数为1， 矩阵B的第二列第一个块的列数为3，核函数内在该线程块下仅计算1行和3列的值，详见代码及注释。

>>Code： 代码如下：

template<int BLOCK_SIZE> __global__ void MatrixMulCUDA(float* C, float* A, float* B,
	int wA, int wB, int hA, int hB)
{
	//Block index
	int bx = blockIdx.x;
	int by = blockIdx.y;

	//Thread index
	int tx = threadIdx.x;
	int ty = threadIdx.y;

	//将矩阵划分为子矩阵，对子矩阵的乘法应用块内线程并行计算，最后将它们的值相加得到C的一个元素值
	int aBegin = by * BLOCK_SIZE * wA;	//A的子矩阵的行坐标
	int aStep = BLOCK_SIZE;				//A的子矩阵列坐标的移动步长
	int aEnd = aBegin + wA - 1;			//限定一个终点

	int bBegin = bx * BLOCK_SIZE;
	int bStep = BLOCK_SIZE * wB;

	float Csub = 0;		//定义在block(x,. y)块中（ty, tx）对应位置的C的元素值

	int subAw = BLOCK_SIZE;
	int subAh = BLOCK_SIZE;
	int subBh = BLOCK_SIZE;
	int subBw = BLOCK_SIZE;
	for (int a = aBegin, b = bBegin; a <= aEnd; a += aStep, b += bStep)
	{
		
		if (a + aStep - 1 > aEnd)		//A矩阵最后一列的块的列数少于BLOCK_SIZE
		{			
			subAw = aEnd - a + 1;
		}
		else
		{
			subAw = BLOCK_SIZE;
		}
		subBh = subAw;

		if ((by + 1) * BLOCK_SIZE > hA)		//A矩阵最后一行的块的行数少于BLOCK_SIZE
		{
			subAh = hA - by * BLOCK_SIZE;
		}
		else
		{
			subAh = BLOCK_SIZE;
		}

		if ((bx + 1) * BLOCK_SIZE > wB)		//B矩阵最后一列的块的列数少于BLOCK_SIZE
		{
			subBw = wB - bx * BLOCK_SIZE;
		}
		else
		{
			subBw = BLOCK_SIZE;
		}
		
		/* 开辟块内共享内存 */
		__shared__ float As[BLOCK_SIZE][BLOCK_SIZE];
		__shared__ float Bs[BLOCK_SIZE][BLOCK_SIZE];

		/* 为行和列范围内的子矩阵对应元素赋值 */
		if (ty < subAh && tx < subAw)
		{
			As[ty][tx] = A[a + ty * wA + tx];
		}
		if (ty < subBh && tx < subBw)
		{
			Bs[ty][tx] = B[b + ty * wB + tx];
		}
		__syncthreads();

		//展开循环来 编译以加速		
#pragma unroll
		//内循环计算每个子矩阵内对应行和列的向量乘积，累加到之前得到的值上
		for (int k = 0; k < subAw; k++)
		{
			if (ty < subAh && tx < subBw)	//满足行和列约束内的元素计算乘积并求和
			{
				Csub += As[ty][k] * Bs[k][tx];
			}			
		}
		__syncthreads();
	}

	//满足行和列约束内的元素计算乘积并求和
	if (ty < subAh && tx < subBw)
	{
		C[by * BLOCK_SIZE * wB + bx * BLOCK_SIZE + ty * wB + tx] = Csub;
	}	
}

此时函数的调用举例为：

/* 参数设置 */
dim3 dimsA(1055, 2137, 1);		//矩阵的宽、高和未使用参数1
dim3 dimsB(108, 1055, 1);		//矩阵的宽、高和未使用参数1
const int block_size = 32;

/* 矩阵初始化、内存传递等常规步骤
....
*/

/* 调用核函数计算 */
dim3 threads(block_size, block_size);
dim3 grid((dimsB.x -1) / threads.x + 1, (dimsA.y - 1) / threads.y + 1);

MatrixMulCUDA<block_size> <<<grid, threads >>> (d_C, d_A, d_B,
			dimsA.x, dimsB.x, dimsA.y, dimsB.y);

当然，也可以对原始矩阵padding使其成为BLOCK_SIZE的整数倍来实现，理论上能有更快的计算速度。