GPU的矩阵转置优化(transpose)

背景

在图像处理，深度学习领域，有很多矩阵运算的工作，而伴随矩阵运算就存在大量的矩阵转置，转置不涉及计算，主要的工作都在数据的读取写入方面，所以如何加快数据搬移是一种很重要的优化点。

CPU代码

假设都是按行存储。

int row = 1024
int col = 512
void transpose_CPU(vector<int>& in, vector<int>& out)
{
  for (int i = 0; i < row; i++)
  {
      for (int j = 0; j < col; j++)
      {
            out[i + j * row] = in[j + i * col];
      } 
  } 
}

GPU代码

//版本1
//假如是block = row; grid = col 
__global__
void transpose_cuda_kernel(int* in, int* out)
{
	int ix = threadIdx.x;
	int iy = blockIdx.x;
	in_pos = iy * blockDim.x + ix;
	out_pos = ix * gridDim.x + iy;
	out[out_pos] = in[in_pos];
}
//版本2
//假如grid和block的设置都为二维（x,y）,矩阵大小为(nrows,ncols)
#define INDEX(ROW, COL, INTER) ((ROW) * (INTER) + (COL))
__global__ 
void naive_cuda_kernel(float *out, float *in, const int nrows, const int ncols)
{
    // matrix coordinate (ix,iy)
    unsigned int row = blockIdx.y * blockDim.y + threadIdx.y;
    unsigned int col = blockIdx.x * blockDim.x + threadIdx.x;

    // transpose with boundary test
    if (row < nrows && col < ncols)
    {
        out[INDEX(col, row, nrows)] = in[INDEX(row, col, ncols)];
    }
}

问题分析

首先来看一下上面的图，这个绿的block中有(3,4)个线程，在gpu中block的thread是以warp为单位运行的，每个warp为32个线程，这32线程执行同样的指令，如果说这32个线程访问的内存的连续，那么就可以合并为一个数据处理事务（具体看数据大小），首先线程去L1/L2缓存中去查找，如果不存在就去全局内存中取数放到缓存中。如果说访问的全局内存地址不是连续的，那么数据处理事务就会分为多次处理，所以最好能形成数据合并，在读取中可以看到，线程访问的内存是连续，但是在写数据的时候是按列去写的，但是内存是按行布局的，所以数据访存无法合并，那么如何解决呢？

解决方案

我们的目的就是让内存能够连续写入，这里引入share memory, 这里逻辑有点绕，为了搞清楚需要好好捋一下，首先看看我们的目的，我们是为了线程在写入内存的时候，连续的线程能够访问连续的内存地址，如下图中的set num 的效果

先看看原始是什么情况, 由于内存是按行保存的，所以在set num的时候，连续的线程拿到的输出地址不是连续的，这样本来三个线程在取数的时候可以合并一个指令取出，结果在赋值的时候三个线程却不能合并指令，因为三个地址不连续，示意图如下：
所以如何解决在赋值的时候连续线程可以对应连续地址就好了，在计算机中有个经典的思路，所有的问题都可以通过一个中间件去解决问题，在这里就是利用share memory.

• 第一步：我们先把一个block的内存放到share memory中
  
• 第二步：将share memory搬移到全局内存中
  这里有思考题，可以看到share memory在读取的时候内存地址也不连续，bank conflict问题严重，难道不影响性能吗？
  
  目前来看解决了连续线程不能访问连续全局内存的问题。

代码
第一步：

    //假设矩阵大小为（M,N）
	int in_x = blockIdx.x * blockDim.x + threadIdx.x; 
    int in_y = blockIdx.y * blockDim.y + threadIdx.y;
     tmp[threadIdx.y][threadIdx.x] = in[in_y * N + in_x];// 这里可以看到当threadIdx连续时候in[pos]中pos也连续

第二步：

   int bid = threadIdx.y*blockDim.x + threadIdx.x;// block拉成一维后
   x = bid/blockDim.y;//对应的share memory的位置
   y = bid%blockDim.y;
   int to_x = blockIdx.y * blockDim.y + y;// blockIdx.y * blockDim.y这是block的新位置
   int to_y = blockIdx.x * blockDim.x + x; 
    // transpose with boundary test
   out[to_y * M + to_x] = tmp[y][x];// 这里可以看到当threadIdx连续的时候out[pos]中pos也连续

总的代码如下

#define INDEX(ROW, COL, INTER) ((ROW) * (INTER) + (COL))
__global__ 
void naive_cuda_kernel(float *out, float *in, const int nrows, const int ncols)
{
    __share__ float tmp[M][N];
    // matrix coordinate (ix,iy)
    int in_x = blockIdx.x * blockDim.x + threadIdx.x; 
    int in_y = blockIdx.y * blockDim.y + threadIdx.y;
   // linear global memory
   int tid = in_y * ncols  + in_x;
   //根据线程求share_memory的位置
   int bid = threadIdx.y*blockDim.x + threadIdx.x;// block拉成一维后
   x = bid/blockDim.y;
   y = bid%blockDim.y;
   /
   int to_x = blockIdx.y * blockDim.y + y;
   int to_y = blockIdx.x * blockDim.x + x; 
   int to = to_y * nrows + to_x;
    // transpose with boundary test
    if (row < nrows && col < ncols)
    {
        tmp[threadIdx.y][threadIdx.x] = in[INDEX(row, col, ncols)];
        __syncthreads();
        out[to] = tmp[y][x];
        
    }
}

结论

从实验结果来看，带宽利用率提升明显