使用cuda完成 矩阵乘法

cuda矩阵乘法

• 数据
  P : m x n m行n列的矩阵
  Q : n x k n行k列的矩阵
  R : m x k m行k列的矩阵
  计算 R = P x Q
  BlockSize设置为 32 x 32
  GridSize 设置为 ((m + block_size - 1) / block_size, (k + block_size - 1) / block_size)。因为最终计算出的矩阵是m x k大小的，设计出来的cuda核函数的所有线程数量 >= m x k。(m + block_size - 1) / block_size表示向上取整。

• 索引计算 (二维举例)
  blockDim 描述block内线程数量，即32 x 32
  blockIdx 描述block线程块在Grid中的id
  threadIdx 描述在block内线程的id

  
  image.png

• 矩阵乘法

  
  计算示意图.png

• 核函数代码
  代码解释均在注释中

template <typename T>
__global__ void cuda_matrix_mul_kernel(T* P, T* Q, int m, int n, int k, T* R)
{

     /**
      *
      * 0 --------------->x
      *  |
      *  |   o (ix, iy)
      *  |
      *  v
      *  y
      */
      // 计算该线程对应R矩阵中的x坐标
     int ix = blockDim.x * blockIdx.x + threadIdx.x;
     // 计算该线程对应R矩阵中的y坐标
     int iy = blockDim.y * blockIdx.y + threadIdx.y;
     // (ix, iy) 为二维矩阵中的坐标，最大的为(k - 1, m - 1)

     T value = 0;
    // 只处理在矩阵范围内的数据。ix，iy如果落在矩阵外则不需要处理
     if (ix < k && iy < m)
     {
          // 矩阵相乘的计算中需要累加一行乘一列的值，所以这里有一个循环
         for (int i = 0; i < n; i++)
         {
             // iy 是 该线程对应R矩阵中的y坐标，即R矩阵的第iy行，同时也是P矩阵的第iy行
             // ix 是 该线程对应R矩阵中的x坐标，即R矩阵的第ix列，同时也是Q矩阵的第ix列
             // iy * n 表示在内存中P矩阵的第iy行的开始， iy * n + i 表示第iy行的第i个元素
             // i * k 表示在内存中R矩阵中每一行的开始（i在循环中变化）i * k + ix 表示第ix列的第i个元素
             value += P[iy * n + i] * Q[i * k + ix];
         }
         // iy * k + ix 表示在内存中R矩阵坐标为 (ix ,iy)的元素
         R[iy * k + ix] = value;
     }
}

• 全部代码

#include <cuda_runtime.h>
#include <stdio.h>

// P : m x n
// Q : n x k
// R : m x k
// blockDim : (BLOCK_SIZE, BLOCK_SIZE)
// gridDim  : ((k + BLOCK_SIZE - 1) / BLOCK_SIZE, (m + BLOCK_SIZE - 1)/BLOCK_SIZE)
template <typename T>
__global__ void cuda_matrix_mul_kernel(T* P, T* Q, int m, int n, int k, T* R)
{

     /**
      *
      * 0 --------------->x
      *  |
      *  |   o (ix, iy)
      *  |
      *  v
      *  y
      */
     int ix = blockDim.x * blockIdx.x + threadIdx.x;
     int iy = blockDim.y * blockIdx.y + threadIdx.y;
     // (ix, iy) 为二维矩阵中的坐标，最大的为(k - 1, m - 1)

     T value = 0;
     if (ix < k && iy < m)
     {
         for (int i = 0; i < n; i++)
         {
             // iy * n  处于第行 i * k 处于第几列
             value += P[iy * n + i] * Q[i * k + ix];
         }
         R[iy * k + ix] = value;
     }

}

void cuda_matrix_mul()
{
    int* P = nullptr;
    int* Q = nullptr;
    int m = 2;
    int n = 3;
    int k = 4;
    cudaMallocManaged(&P, sizeof(int) * m * n);
    cudaMallocManaged(&Q, sizeof(int) * n * k);

    int* R = nullptr;
    cudaMallocManaged(&R, sizeof(int) * m * k);

    for (int i = 0; i < m; i++)
    {
        for (int j = 0; j < n; j++)
        {
            P[i * n + j] = i + j;
            printf("%d ", P[i*n + j]);
        }
        printf("\n");
    }
    printf("\n");
    for (int i = 0; i < n; i++)
    {
        for (int j = 0; j < k; j++)
        {
            Q[i * k + j] = i + j;
            printf("%d ", Q[i*k + j]);
        }
        printf("\n");
    }
    printf("\n");

    int block_size = 16;
    unsigned int grid_rows = (m + block_size - 1) / block_size;
    unsigned int grid_cols = (k + block_size - 1) / block_size;
    dim3 grid(grid_cols, grid_rows);
    dim3 block(block_size, block_size);
    cuda_matrix_mul_kernel<<<grid, block>>>(P, Q, m, n, k, R);
    cudaDeviceSynchronize();
    for (int i = 0; i < m; i++)
    {
        for (int j = 0; j < k; j++)
        {
            printf("%d ", R[i*k + j]);
        }
        printf("\n");
    }
}

int main()
{
    cuda_matrix_mul();
    return 0;
}

执行结果

image.png