CUDA 内核 2x2 ZgemmBatched：比 CuBLAS 快 5 倍。能不能快点？

Question

我发现在使用小矩阵时，CuBLAS API 的性能很差，尤其是批量通用矩阵乘法。我的目标是编写一个快速内核来计算 2x2 矩阵大小的批量复数双矩阵乘法。内核应该能够从更大的方阵中提取这些矩阵，即块矩阵乘法。

我编写的内核每块计算 8 个 2x2 双复数矩阵乘法（32 个线程）。我可以在 GTX 970（计算能力 5.2）上实现大约 1.65 双 GFlops（NSight 分析），它应该执行 109 双 GFlops。计算 910 万次矩阵乘法大约需要 300ms。

*a、*b 和 *c 是指向 N 次 nxn 矩阵数组的指针。这些矩阵必须是平方的。

lda、ldb、ldc 是所有输入矩阵的主要维度。使用这些参数，任何方阵数组都可用于提取 2x2 矩阵。

问题：如何提高此内核的性能？

#define THREADMULTIPLIER 8

__global__ void bulkMatrixMul(const cuDoubleComplex *a, int lda, const cuDoubleComplex *b, int ldb, cuDoubleComplex *c, int ldc){
    int pos = (THREADMULTIPLIER * blockIdx.x + threadIdx.z);
    int ia = pos * lda * lda + threadIdx.x;
    int ib = (pos * ldb + threadIdx.y) * ldb;
    c[(pos * ldc + threadIdx.y) * ldc + threadIdx.x] = cuCfma(a[ia], b[ib], cuCmul(a[ia + LD], b[ib + 1]));
};


void bulkBatchedMatrixMul(complex<double> *a, int lda, complex<double> *b, int ldb, complex<double> *c, int ldc, int batchSize, cudaStream_t *stream){
    dim3 threads(2, 2, THREADMULTIPLIER);
    bulkMatrixMul << <batchSize / THREADMULTIPLIER, threads, 0, *stream >> >((cuDoubleComplex*)a, lda, (cuDoubleComplex*)b, ldb, (cuDoubleComplex*)c , ldc);
    if (cudaSuccess != cudaGetLastError())
        printf("Error!\n");
}

通过分析 Christian Sarofeen 的代码，我的内核速度提高了 10 倍：

1. 如果您将总和作为一个整体而不是使用 += 和 -= 写入会更快。
2. 声明输出，尤其是当您拆分输出的实部和虚部计算时会提高速度。

下面的内核是一个优化的内核，每个块计算 16 个矩阵乘法。它可以在 31 毫秒内计算 910 万次矩阵乘法：我使用 10000 的 batchSize 并在 13 个流中每个流调用这个内核 70 次。根据 NSight 的说法，它实现了大约 15 双 GFlops。使用了 GTX970。

__global__ void bulkMatrixMul(const cuDoubleComplex *a, int lda, const cuDoubleComplex *b, int ldb, cuDoubleComplex *c, int ldc){
    int pos = (blockDim.z * blockIdx.x + threadIdx.z);
    int ia = pos * lda * lda + threadIdx.x;
    int ib = (pos * ldb + threadIdx.y) * ldb;
    cuDoubleComplex cR;
    cR.x = a[ia].x * b[ib].x - a[ia].y * b[ib].y - a[ia + lda].y * b[ib + 1].y + a[ia + lda].x * b[ib + 1].x;
    cR.y = a[ia].x * b[ib].y + a[ia].y * b[ib].x + a[ia + lda].x * b[ib + 1].y + a[ia + lda].y * b[ib + 1].x;
    c[(pos * ldc + threadIdx.y) * ldc + threadIdx.x] = cR;
};


void bulkBatchedMatrixMul(complex<double> *a, int lda, complex<double> *b, int ldb, complex<double> *c, int ldc, int batchSize, cudaStream_t *stream){
    dim3 threads(2, 2, 16);
    bulkMatrixMul <<< batchSize / 16, threads, 0, *stream >>>((cuDoubleComplex*)a, lda, (cuDoubleComplex*)b, ldb, (cuDoubleComplex*)c , ldc);
    if (cudaSuccess != cudaGetLastError())
        printf("Error!\n");
}

Answer 1

我已经使用 Christian Sarofeen 的代码调整了我的内核。他的内核需要 60 毫秒才能在 GTX 970 上计算 910 万次矩阵乘法。

我调整后的内核可以在 31 毫秒内完成。谢谢克里斯蒂安！

__global__ void bulkMatrixMul(const cuDoubleComplex *a, int lda, const cuDoubleComplex *b, int ldb, cuDoubleComplex *c, int ldc){
    int pos = (blockDim.z * blockIdx.x + threadIdx.z);
    int ia = pos * lda * lda + threadIdx.x;
    int ib = (pos * ldb + threadIdx.y) * ldb;
    cuDoubleComplex cR;
    cR.x = a[ia].x * b[ib].x - a[ia].y * b[ib].y - a[ia + lda].y * b[ib + 1].y + a[ia + lda].x * b[ib + 1].x;
    cR.y = a[ia].x * b[ib].y + a[ia].y * b[ib].x + a[ia + lda].x * b[ib + 1].y + a[ia + lda].y * b[ib + 1].x;
    c[(pos * ldc + threadIdx.y) * ldc + threadIdx.x] = cR;
};


void bulkBatchedMatrixMul(complex<double> *a, int lda, complex<double> *b, int ldb, complex<double> *c, int ldc, int batchSize, cudaStream_t *stream){
    dim3 threads(2, 2, 16);
    bulkMatrixMul << <batchSize / 16, threads, 0, *stream >> >((cuDoubleComplex*)a, lda, (cuDoubleComplex*)b, ldb, (cuDoubleComplex*)c , ldc);
    if (cudaSuccess != cudaGetLastError())
        printf("Error!\n");
}

Answer 2

我会为每个矩阵乘法分配一个线程，并将所有内容转储到寄存器中。不要担心占用，因为我怀疑这将受到内存限制。一旦进入寄存器，每个线程将通过矩阵乘法来做一个矩阵并将其存储在 c 中。下面是一个演示。

#include <stdio.h>
#include <iostream>

#include <cuComplex.h>
#include <cuda.h>
#include <cuda_runtime.h>

#define N_MATRIX 10000

using namespace std;

__device__ void D_C_Mult_Add(cuDoubleComplex &a, cuDoubleComplex &b, cuDoubleComplex &c){
        c.x+=a.x*b.x;
        c.y+=a.x*b.y;
        c.x-=a.y*b.y;
        c.y+=a.y*b.x;
}

__global__ void multMatrix(cuDoubleComplex* a, cuDoubleComplex* b, cuDoubleComplex* c) {
    int tid = blockIdx.x*blockDim.x + threadIdx.x;
    int stride=gridDim.x*blockDim.x;
    if(tid>=N_MATRIX) return;
    cuDoubleComplex a_sub[4], b_sub[4], c_sub[4];

    for(int i=0; i<4; i++){
        a_sub[i]=a[tid+stride*i];
        b_sub[i]=a[tid+stride*i];
        c_sub[i].x=0.0;
        c_sub[i].y=0.0;
    }

    D_C_Mult_Add(a_sub[0], b_sub[0], c_sub[0]);
    D_C_Mult_Add(a_sub[0], b_sub[1], c_sub[1]);
    D_C_Mult_Add(a_sub[2], b_sub[0], c_sub[2]);
    D_C_Mult_Add(a_sub[2], b_sub[1], c_sub[3]);

    D_C_Mult_Add(a_sub[1], b_sub[2], c_sub[0]);
    D_C_Mult_Add(a_sub[1], b_sub[3], c_sub[1]);
    D_C_Mult_Add(a_sub[3], b_sub[2], c_sub[2]);
    D_C_Mult_Add(a_sub[3], b_sub[3], c_sub[3]);

    for(int i=0; i<4; i++)
        c[tid+stride*i]=c_sub[i];

}

int main(int argc, char **argv) {



    cuDoubleComplex *a, *b, *c;
    cudaMalloc(&a, sizeof(cuDoubleComplex)*N_MATRIX*4);
    cudaMalloc(&b, sizeof(cuDoubleComplex)*N_MATRIX*4);
    cudaMalloc(&c, sizeof(cuDoubleComplex)*N_MATRIX*4);

    multMatrix<<<N_MATRIX/128+1, 128>>>(a, b, c);
}