6

CUDA 中的矩阵乘

 2 years ago
source link: https://dingfen.github.io/mpi&openmp/2021/10/20/cuda-with-matmul.html
Go to the source link to view the article. You can view the picture content, updated content and better typesetting reading experience. If the link is broken, please click the button below to view the snapshot at that time.
neoserver,ios ssh client

CUDA 中的矩阵相乘Permalink

矩阵乘中很多计算步骤都十分相似且数据依赖不复杂,所以特别适合使用 GPU 来计算, 利用 GPU 内部的高度并行性,可极大地提高计算速度。

矩阵乘法 CPU 实现Permalink

老规矩,先看看在 CPU 下如何实现矩阵乘法,代码非常简单,直接给出如下:

void matMulCpu(float* c, float* a, float* b, int m, int n, int k) {
    for (int i = 0; i < m; i++) {
        for (int j = 0; j < n; j++) {
            float sum = 0.0f;
            for (int l = 0; l < k; l++) {
                sum += a[i * k + l] * b[l * n + j];
            }
            c[i * n + j] = sum;
        }
    }
}

对于矩阵 $A_{m\times k} \cdot B_{k\times n} = C_{m\times n}$ 而言:

  • 计算次数为$ m \times n \times k $
  • 时间复杂度为$ O(n^3) $

CUDA 的核函数实现Permalink

仍然计算矩阵 $A_{m\times k} \cdot B_{k\times n} = C_{m\times n}$。但这次使用 GPU 计算。不难想到,可以让每个线程读取 A 中的第 m 行和 B 中的第 n 列,进而计算出 C[m, n]。线程数量控制为 $m \times n$ 个,这样就可以计算出矩阵C。更精确地说,控制线程数量为 gridDim.x * blockDim.x == n && girdDim.y * blockDim.y == m

// cuda kernel function
__global__ void matMulNaiveGpu(float* c, float* a, float* b, int m, int n, int k) {
    int y = blockIdx.y * blockDim.y + threadIdx.y;
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    float v = 0.0f;
    for(int i =0; i < k; i++) {
        v += a[y * k + i] * b[i * n + x];
    }
    pfMatrixC[y * n + x] = v;
}

// call in host
dim3 threadsPerBlock(m, n)
matMulNaiveGpu<<<1, threadsPerBlock>>>(c, a, b, m, n, k);
  • 计算次数为$ m \times n \times k $,但有$m \times n$个线程并行
  • 时间复杂度为$ O(k) $

但我发现,仅仅以计算次数为优化依据是不够明智的。在上面的程序中,每个线程需要读取 $2 \times k$ 个数据,然后写入一个数据,一共需要读取 $2 \times n \times m \times k$ 个数据。而上述核函数中所有的数据都放在全局内存中,读取数据的速度很有可能比计算速度慢,成为程序提升性能的瓶颈。

线程束(Warp) :GPU 指令发射和执行的最小单元。由于多种原因,即使用户要求 GPU 只需要用一个线程执行指令,GPU 仍会至少启动一个线程束(大多数为 32 个线程)来发射、执行指令。线程束中的一条指令发射一次, 称为 1 次 “transaction”, 重复发射一次, 称为 1 次 “reply”. 对于 Global Memory 的访问, warp 内的指令需要几次 transaction, 或者说是否会发生 reply, 取决于地址对齐及可合并访问的情况.

核函数的优化——共享内存Permalink

前面提到,对全局内存的访问虽然已经经过优化,尽量实现了访问合并。但我也发现,对矩阵 A、B 仍很多重复读取操作。例如,一个线程束访问矩阵 A 时,其实是合并访问了(32 个线程合并为一次访问)内存的同一地址,无效的重复有 31 次。访问矩阵 B 时,合并访问(8 个线程合并访问一次,取 32 个字节的数据)。

充分利用 GPU 中的共享内存,可以提高读取写入数据的速度。共享内存通常位于芯片内部,访问延时是全局内存的二十到三十分之一,而带宽高 10 倍,具有低延时、高带宽的特性,但比全局内存小得多,因此可以将短时间内经常访问的数据放入到共享内存中,减少重复访问全局内存的频率,从而提高程序运行速度。

那么,如何充分使用 GPU 中的共享内存呢?一般用 __shared__ 指示数组应放在共享内存上。但共享内存比全局内存要小很多,假设矩阵可以完全放入共享内存是不合理的。为此,必须要对矩阵进行切分。我使用了如下切分方式。

matsplit.png

为更好地说明切分和计算过程,以下图为例,GPU 网格中一共有 16 个线程块,将矩阵按照线程块的数量切分,如下图。然后分行列进行计算。在第一次循环中,将红色的数据块放入所有线程块的共享内存中,可以看到第一次循环时,读取了矩阵 A 的一列和矩阵 B 的一行,于是可以计算出矩阵 C 的部分和。第二次循环,将往右移动读取矩阵 A 的一列,往下移动读取矩阵 B 的一行,那么橙色的数据块会被读入共享内存中,再计算矩阵 C 的部分和,不断循环累加,就可以计算出矩阵 C。

matmul_shared.png

下面的程序就是根据上述流程编写出来的。首先声明各个标识变量和共享内存(使用 __shared__)。重点在于下面的循环,先计算矩阵的宽度可以容纳多少个线程块覆盖,然后是分别读取矩阵 A 和 B 的对应列行,对索引计算需要格外仔细。读取完成后,需要加上 __syncthreads() 语句,确保所有线程已经读取完成,再开始下面的计算,同样也要加上 __syncthreads() 来保证所有线程都已经完成。

__global__ void matMulProGpu(float* c, float* a, float* b, int m, int n, int k) {
    int y = blockIdx.y * blockDim.y + threadIdx.y;
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    int tidy = threadIdx.y;
    int tidx = threadIdx.x;
    float v = 0.0f;

    __shared__ float as[BLOCK_SIZE][BLOCK_SIZE];
    __shared__ float bs[BLOCK_SIZE][BLOCK_SIZE];

    int tilesx = (k + BLOCK_SIZE - 1) / BLOCK_SIZE;
    for(int i = 0; i < tilesx; i++) {
        // load data from global memory to shared memory
        as[tidy][tidx] = a[y * k + i * BLOCK_SIZE + tidx];
        bs[tidy][tidx] = b[(i * BLOCK_SIZE + tidy) * n + x];

        // sync to wait for all threads in one block to finish loading datas
        __syncthreads();

        // sub-matrix multiply
        for(int l = 0; l < BLOCK_SIZE; l++) {
            v += as[tidy][l] * bs[l][tidx];
        }

        // sync to wait for all threads in one block to finish compute
        __syncthreads();
    }

    // store results into global memory
    c[y * n + x] = v;
}

核函数的进一步优化Permalink

cuBLAS 版本的实现Permalink

cuBlAS 是 NVidia 对传统 BLAS(Basic Liner Algebra Subprograms,基本线性代数子程序库)的 GPU 实现,因为 BLAS 最初是由 FORTRAN 语言实现,为了最大程度地遵循传统,NVidia 在实现时,保留了大量 BLAS 和 Fortran 语言的特点,比如以列优先存储数组,索引从 1 开始计数等。但这些“落后于时代”的特征,往往给现代程序员制造了不少麻烦。

cuBLAS 最早在 CUDA6.0 中出现,现在(v11.5.0)包含3类 API

  • cuBLAS API
  • cuBLASXt API
  • cuBLASLt API

若用户要使用 cuBLAS API,那么应该:

  1. 分配矩阵或向量所需的 GPU 内存空间,并加载数据到 GPU 上
  2. 调用所需的 cuBLAS 函数
  3. 将计算结果传输至主机,cuBLAS API也提供一些帮助函数来写或者读取数据从 GPU 中
a = (float*) malloc(m_size * sizeof(float));
b = (float*) malloc(m_size * sizeof(float));
c = (float*) malloc(m_size * sizeof(float));

cudaMalloc((void **)&dev_a, sizeof(float) * m_size);
cudaMalloc((void **)&dev_b, sizeof(float) * m_size);
cudaMalloc((void **)&dev_c, sizeof(float) * m_size);

// initialize CUBLAS context
stat = cublasCreate(&handle);

stat = cublasSetMatrix(r_size, r_size, sizeof(*mat1), mat1, r_size, g_mat1, r_size);
stat = cublasSetMatrix(r_size, r_size, sizeof(*mat2), mat2, r_size, g_mat2, r_size);
stat = cublasSetMatrix(r_size, r_size, sizeof(*result), result, r_size, g_mat_result, r_size);

float al = 1.0f;
float bet = 0.0f;
    
stat = cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, 
        r_size, r_size, r_size, &al, g_mat1, 
        r_size, g_mat2, r_size, &bet, g_mat_result, r_size);
stat = cublasGetMatrix(r_size, r_size, sizeof(*result), g_mat_result, r_size, result, r_size);

About Joyk


Aggregate valuable and interesting links.
Joyk means Joy of geeK