[CUDA] 行列の積

  • 2022年8月7日
  • 2022年9月17日
  • CUDA
  • 21View
  • 0件

はじめに

CUDA の練習として、行列の掛け算を実装してみました。
実用上は cuBLAS を利用するのが正解だと思いますが、練習ということでスクラッチで実装しています。

次の環境で動作確認しています。

  • Windows 11 WSL (Ubuntu)
  • GeForce RTX 2080
  • nvcc 10.1

CUDA 実装例 – 行列の積

実装例です。
行列 lhs [NY x NZ] と 行列 rhs [NZ x NX] の積を計算するためのメソッドです。
アウトプットの行列は dest [NY x NX] です。

__global__ が付与されている mul_matrix_dev が GPU 側で実行されるメソッドです。

__global__ void mul_matrix_dev(float *dest, float *lhs, float *rhs, size_t nx, size_t ny, size_t nz) {
    size_t x = blockIdx.x * blockDim.x + threadIdx.x;
    size_t y = blockIdx.y;
    size_t idx = y * nx + x;

    if (x < nx) {
        float acc = 0.0;
        for (size_t i = 0; i < nz; ++i) {
            size_t idx_lhs = y * nz + i;
            size_t idx_rhs = x + i * nx;
            acc += lhs[idx_lhs] * rhs[idx_rhs];
        }
        dest[idx] += acc;
    }
}

void mul_matrix(float *dest, float *lhs, float *rhs, size_t nx, size_t ny, size_t nz) {
    // dest : (ny [rows], nx [columns])
    // lhs  : (ny [rows], nz [columns])
    // rhs  : (nz [rows], nx [columns])

    float *dev_dest, *dev_lhs, *dev_rhs;
    size_t bytes_dest = ny * nx * sizeof(float);
    size_t bytes_lhs = ny * nz * sizeof(float);
    size_t bytes_rhs = nz * nx * sizeof(float);

    cudaMalloc((void **)&dev_dest, bytes_dest);
    cudaMalloc((void **)&dev_lhs, bytes_lhs);
    cudaMalloc((void **)&dev_rhs, bytes_rhs);

    cudaMemcpy(dev_lhs, lhs, bytes_lhs, cudaMemcpyHostToDevice);
    cudaMemcpy(dev_rhs, rhs, bytes_rhs, cudaMemcpyHostToDevice);

    dim3 block(32);
    dim3 grid((nx + block.x - 1) / block.x, ny);

    mul_matrix_dev<<<grid, block>>>(dev_dest, dev_lhs, dev_rhs, nx, ny, nz);

    cudaMemcpy(dest, dev_dest, bytes_dest, cudaMemcpyDeviceToHost);

    cudaFree(dev_dest);
    cudaFree(dev_lhs);
    cudaFree(dev_rhs);
}

メソッド利用例

計算結果を標準出力

mul_matrix メソッドの利用例です。
行列 [3 x 2] と 行列 [2 x 4] の積を計算しています。

#include <stdio.h>

__global__ void mul_matrix_dev(...) {...}
void mul_matrix(...) {...}

void initialize_list(float *dest, size_t n) {
    for (size_t i = 0; i < n; ++i) {
        dest[i] = (float)rand() / RAND_MAX;
    }
}

void print_matrix(float *matrix, size_t nx, size_t ny) {
    for (size_t y = 0; y < ny; ++y) {
        for (size_t x = 0; x < nx; ++x) {
            size_t idx = y * nx + x;
            printf("%.6f, ", matrix[idx]);
        }
        printf("\n");
    }
    printf("\n");
}

int main() {
    float *dest, *lhs, *rhs;
    size_t nx = 4;
    size_t ny = 3;
    size_t nz = 2;

    size_t bytes_dest = ny * nx * sizeof(float);
    size_t bytes_lhs = ny * nz * sizeof(float);
    size_t bytes_rhs = nz * nx * sizeof(float);

    dest = (float *)malloc(bytes_dest);
    lhs = (float *)malloc(bytes_lhs);
    rhs = (float *)malloc(bytes_rhs);

    srand((unsigned int)time(NULL));
    initialize_list(lhs, ny * nz);
    initialize_list(rhs, nz * nx);

    printf("matrix lhs\n");
    print_matrix(lhs, nz, ny);
    printf("matrix rhs\n");
    print_matrix(rhs, nx, nz);

    mul_matrix(dest, lhs, rhs, nx, ny, nz);

    printf("matrix dest = (lhs x rhs)\n");
    print_matrix(dest, nx, ny);

    free(dest);
    free(lhs);
    free(rhs);
}

WSL ターミナルより以下のコマンドで実行します。

$ nvcc main.cu -o main && ./main
matrix lhs
0.041569, 0.849196, 
0.992869, 0.290114, 
0.321827, 0.388423, 

matrix rhs
0.819833, 0.432450, 0.239414, 0.614277, 
0.452702, 0.457569, 0.826649, 0.822087, 

matrix dest = (lhs x rhs)
0.418512, 0.406543, 0.711939, 0.723648, 
0.945322, 0.562113, 0.477529, 0.848395, 
0.439685, 0.316905, 0.398140, 0.517009, 

CPU 計算との結果比較

結果の妥当性を確認するため、CPU での計算結果と突合しました。

ここではサイズを大きくして、行列 [300 x 200] と 行列 [200 x 400] の積を計算しています。
計算誤差が大きいので、判定の閾値は 1e-4 としています。

#include <assert.h>

__global__ void mul_matrix_dev(...) {...}
void mul_matrix(...) {...}
void initialize_list(...) {...}

void mul_matrix_on_host(float *dest, float *lhs, float *rhs, int nx, int ny, int nz) {
    for (size_t x = 0; x < nx; ++x) {
        for (size_t y = 0; y < ny; ++y) {
            size_t idx = y * nx + x;
            for (size_t z = 0; z < nz; ++z) {
                float acc = 0.0;
                for (size_t i = 0; i < nz; ++i) {
                    size_t idx_lhs = y * nz + i;
                    size_t idx_rhs = x + i * nx;
                    acc += lhs[idx_lhs] * rhs[idx_rhs];
                }
                dest[idx] = acc;
            }
        }
    }
}

void diff_list(float *lhs, float *rhs, size_t n) {
    for (size_t i = 0; i < n; ++i) {
        assert(fabs(lhs[i] - rhs[i]) < 1e-4);
    }
}

int main() {
    float *dest, *lhs, *rhs;
    size_t nx = 400;
    size_t ny = 300;
    size_t nz = 200;

    size_t bytes_dest = ny * nx * sizeof(float);
    size_t bytes_lhs = ny * nz * sizeof(float);
    size_t bytes_rhs = nz * nx * sizeof(float);

    dest = (float *)malloc(bytes_dest);
    lhs = (float *)malloc(bytes_lhs);
    rhs = (float *)malloc(bytes_rhs);

    srand((unsigned int)time(NULL));
    initialize_list(lhs, ny * nz);
    initialize_list(rhs, nz * nx);

    mul_matrix(dest, lhs, rhs, nx, ny, nz);

    float *dest_on_host;
    dest_on_host = (float *)malloc(bytes_dest);
    mul_matrix_on_host(dest_on_host, lhs, rhs, nx, ny, nz);
    diff_list(dest, dest_on_host, ny * nx);
    free(dest_on_host);

    free(dest);
    free(lhs);
    free(rhs);
}

assert に引っかかるものがなければ、何事もなく正常終了します。

$ nvcc main.cu -o main && ./main

ダメな場合は次のように異常終了します。

$ nvcc main.cu -o main && ./main
main: main.cu:86: void diff_list(float*, float*, size_t): Assertion `fabs(lhs[i] - rhs[i]) < 1e-6' failed.
Aborted

CPU 計算との実行時間比較

実行時間の比較も行いました。

__global__ void mul_matrix_dev(...) {...}
void mul_matrix(...) {...}
void initialize_list(...) {...}
void mul_matrix_on_host(...) {...}

int main() {
    float *dest, *lhs, *rhs;
    size_t nx = 400;
    size_t ny = 300;
    size_t nz = 200;

    size_t bytes_dest = ny * nx * sizeof(float);
    size_t bytes_lhs = ny * nz * sizeof(float);
    size_t bytes_rhs = nz * nx * sizeof(float);

    dest = (float *)malloc(bytes_dest);
    lhs = (float *)malloc(bytes_lhs);
    rhs = (float *)malloc(bytes_rhs);

    srand((unsigned int)time(NULL));
    initialize_list(lhs, ny * nz);
    initialize_list(rhs, nz * nx);

    long start = clock();
    mul_matrix(dest, lhs, rhs, nx, ny, nz);
    long end = clock();
    printf("Calc on GPU elapsed : %6.3f [sec]\n", (float)(end - start) / CLOCKS_PER_SEC);

    float *dest_on_host;
    dest_on_host = (float *)malloc(bytes_dest);

    start = clock();
    mul_matrix_on_host(dest_on_host, lhs, rhs, nx, ny, nz);
    end = clock();
    printf("Calc on CPU elapsed : %6.3f [sec]\n", (float)(end - start) / CLOCKS_PER_SEC);

    free(dest_on_host);

    free(dest);
    free(lhs);
    free(rhs);
}

GPU の方が圧倒的に爆速です。

$ nvcc main.cu -o main && ./main
Calc on GPU elapsed :  0.123 [sec]
Calc on CPU elapsed : 12.867 [sec]
参考書籍