使用cuda完成 矩阵乘法

2023-11-29  本文已影响0人  leon0514

cuda矩阵乘法

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
上一篇 下一篇

猜你喜欢

热点阅读