CUDA矩阵转置

Udacity的CUDA编程课程中介绍了CUDA实现矩阵转置的六种方式，本文介绍其中的一种方式

如果矩阵为N*N的方阵。该方式让每个线程处理一个矩阵元素，总共需要N*N个线程。首先，声明两个常量并配置blocks，threads：

const int N=1024;
const int K=32;
dim3 blocks(N/K,N/K); 
dim3 threads(K,K);

内核函数：

__global__ void 
transpose_parallel_per_element_tiled(float in[], float out[])
{
	// (i,j) locations of the tile corners for input & output matrices:
	int in_corner_i  = blockIdx.x * blockDim.x, in_corner_j  = blockIdx.y * blockDim.y;
	int out_corner_i = in_corner_j, out_corner_j = in_corner_i;

	int x = threadIdx.x, y = threadIdx.y;

	__shared__ float tile[K][K];

	// coalesced read from global mem, TRANSPOSED write into shared mem:
	tile[y][x] = in[(in_corner_i + x) + (in_corner_j + y)*N];
	__syncthreads();
	// read from shared mem, coalesced write to global mem:
	out[(out_corner_i + x) + (out_corner_j + y)*N] = tile[x][y];
}

内核涉及两个输入参数，in代表输入矩阵，out代表in的转置矩阵。为了更好地理解这段代码，我们将矩阵规模缩小并画图展示。假设N=4，K=2。blocks(2,2)，threads(2,2)。

如图所示，矩阵被分成了4个区域，每个区域包含2*2=4个数据，分别由4个线程块进行处理。

矩阵 $in$ 左下那个方形区域 $S$ 由block(0,1)处理。该线程块corner $（ici,icj）$ （threads(0,0)）所处理数据在in矩阵中的坐标为 $(ici,icj)$ 即 $(0,2)$ 。S区域的其他数据坐标可由[(icj+y)*N+ici+x]得到，x，y是线程索引。blocks，threads的二维索引值x，y分别代表了矩阵的水平方向和竖直方向。

内核声明了一个和数据块大小相同的共享内存tile，用于存储S区域的数据。block读取S区域数据并按原样复制到tile当中，一个线程复制一个数据。随后block将tile中的数据转置，再把经过转置的tile中的数据写入到out矩阵当中 $S'$ 区域。根据转置行列互换的性质可知，block(0,1) corner所处理数据在out矩阵中的坐标为(oci=icj,ocj=ici)即(2,0)， $S'$ 区域的其他数据索引可由[(ocj+y)*N+oci+x]得到。转置后，in矩阵 $S$ 区域的数据会被block(0,1)转置到out矩阵 $S'$ 。tile转置和写入数据到out是由最后一行代码完成的。

该例中，数据读取和写入都利用了全局内存合并访问，这也说明，对于二维矩阵块(x,y)，CUDA按x优先顺序处理线程。

出于某种目的，我需要线程块的尺寸为非方形。假设blocks(4,2)，threads(1,2)。内核函数__shared__ float tile[K][K];需要修改为__shared__ float tile[2][1]; 这样修改之后，问题在于无法在tile中对数据进行转置。所以我在将数据写入到out的时候，将位置索引[(oci + x) + (ocj + y)*N]转置为(oci + y) + (ocj + x)*N，实现非方矩阵块的矩阵转置。

这段解释不要看了。内核函数中的in_corner_i，in_corner_j分别代表线程块左上角线程（也就是threads(0,0)）在矩阵中的x索引（水平方向）和y索引（竖直方向），作为线程块内部线程索引数据的参照点（线程块(0,0)线程所处理数据在矩阵中的位置）。对于block(0,1)，in_corner_i=0，in_corner_j=2。转置后，S区域内的元素会被转置到右上方区域S’。所以block(0,1)的功能就是把in矩阵S区域的数据转置并复制到out矩阵S'区域。这就需要block(0,1)找到out矩阵右上方区域的参照点索引。根据转置行列互换的性质，得到out_corner_i(x索引)=in_corner_j=2，out_corner_j(y索引)=in_corner_i=0。x，y变量分别是线程在线程块中的二维索引。

共享内存tile的尺寸和block一样，block读取in数据并按原样复制到tile当中，一个线程复制一个数据。同步后，再把tile中的数据按转置顺序写入到out矩阵当中的对应位置。转置的过程相当于在共享内存中完成，所以线程读取和写入都用[(in_corner_i + x) + (in_corner_j + y)*N]来索引数据。

猜你喜欢