And ye shall know the truth, and the truth shall make you free.

引子

本来想在上一篇文章中简单提一下矩阵乘法的优化的，但是我发现，这个事情并没有我想象中那么简单。

矩阵乘法的优化其实是一个非常复杂的事情，需要考虑到访存开销、缓存命中和编码等多方面因素。

简单的实现

在这里，我们不关心非方阵的情况。

根据线性代数的知识，我们知道，两个矩阵相乘，相当于左矩阵的行和右矩阵的列相乘，得到每一个位置的元素，也就是：
$$
\rm A_{ij}=\sum_k B_{ik}C_{kj}
$$
那么，一个简单的实现就是：

void MatrixMatMul(int a[MATRIX_SIZE][MATRIX_SIZE], int b[MATRIX_SIZE][MATRIX_SIZE], int c[MATRIX_SIZE][MATRIX_SIZE]) {
    for(int i = 0; i < MATRIX_SIZE; i++) {
        for(int j = 0; j < MATRIX_SIZE; j++) {
            for(int k = 0; k < MATRIX_SIZE; k++) {
                c[i][j] += a[i][k] * b[k][j];
            }
        }
    }
}

输出为：

1	Time difference = 4204[ms]

完整的代码如下所示：

#include "cuda_runtime.h"
#include "device_launch_parameters.h"

#include <iostream>
#include <chrono>

#define MATRIX_SIZE 1024

void InitMatrix(int** a, int lines, int columns) {
    for (int i = 0; i < lines; i++) {
        for (int j = 0; j < columns; j++) {
            a[i][j] = i * columns + j;
        }
    }
}

void InitMatrixAsZero(int** a, int lines, int columns) {
    for (int i = 0; i < MATRIX_SIZE; i++) {
        for (int j = 0; j < MATRIX_SIZE; j++) {
            a[i][j] = 0;
        }
    }
}

void TrivialMatrixMatMul(int** a, int** b, int** c, int lines, int columns) {
    for (int i = 0; i < lines; i++) {
        for (int j = 0; j < columns; j++) {
            for (int k = 0; k < columns; k++) {
                c[i][j] += a[i][k] * b[k][j];
            }
        }
    }
}


int main(int argc, char* argv[]) {
    int** a, ** b, ** c;
    a = new int* [MATRIX_SIZE];
    b = new int* [MATRIX_SIZE];
    c = new int* [MATRIX_SIZE];
    for (int i = 0; i < MATRIX_SIZE; i++) {
        a[i] = new int[MATRIX_SIZE];
        b[i] = new int[MATRIX_SIZE];
        c[i] = new int[MATRIX_SIZE];
    }
    InitMatrix(a, MATRIX_SIZE, MATRIX_SIZE);
    InitMatrix(b, MATRIX_SIZE, MATRIX_SIZE);
    InitMatrixAsZero(c, MATRIX_SIZE, MATRIX_SIZE);


    std::chrono::steady_clock::time_point begin = std::chrono::steady_clock::now();
    TrivialMatrixMatMul(a, b, c, MATRIX_SIZE, MATRIX_SIZE);
    std::chrono::steady_clock::time_point end = std::chrono::steady_clock::now();
    std::cout << "Time difference = " << std::chrono::duration_cast<std::chrono::milliseconds>(end - begin).count() << "[ms]" << std::endl;


    return 0;
}

当然，我们可以对这个速度进行进一步的优化。

访存的空间局部性：Easy Mode

当一条数据被搬运到缓存中时，这条数据周围的数据也会被加载到缓存中。比如，我们如果访问b[k][0]，那么这个数据后面的数据，即b[k][1]……b[k][15]都会被搬运到缓存中，这些内容构成了一个缓存行（Cache Line，通常是64字节大小）——一个int是四个字节。

正因如此，我们在矩阵乘法中，更想要的是行访问，而非列访问，因为多维数组是按行存储在内存中的。在矩阵乘法中，是左矩阵的行乘以右矩阵的列，我们需要将「右矩阵的列」，转变为一种行的索引，该怎么做呢？其实很简单：

void TrivialMatrixMatMul2(int** a, int** b, int** c, int lines, int columns) {
    for (int i = 0; i < lines; i++) {
        for(int k = 0; k < columns; k++) {
            for (int j = 0; j < columns; j++) {
                c[i][j] += a[i][k] * b[k][j];
            }
        }
    }
}

这个时候，输出是：

1	Time difference = 3120[ms]

我们优化了大概一秒的时间，这一切都拜「访存的空间局部性」所赐。

访存的空间局部性：Lunatic Mode

Easy Mode只要了解过缓存局部性，就基本上都能想出来（或许需要一点点提示），但是Lunatic Mode的版本，却着实让我意外。首先我们来看看Easy Mode的代码有什么可以优化的地方。

假设我们的Cache Line的长度是64，那么在一次最内层循环中，运行64次之后，就要触发一次读取，将新的b[k][j]读取进来，与此同时，新的c[i][j]也要被读取进来。

那么，我们优化的点就确定了——能不能在缓存已有的数据中进行运算？可以的，数学原理就是——分块矩阵的乘法。具体的原理请参照其他文章。

代码如下：

void TiledMatrixMatMul(int** a, int** b, int** c, int lines, int columns) {
    int tile_size = 80;
    for (int i = 0; i < lines; i += tile_size) {
        for (int j = 0; j < columns; j += tile_size) {
            for (int k = 0; k < columns; k += tile_size) {
                for (int ii = i; ii < i + tile_size; ii++) {
                    if (ii >= lines) break;
                    for (int kk = k; kk < k + tile_size; kk++) {
                        if (kk >= columns) break;
                        for (int jj = j; jj < j + tile_size; jj++) {
                            if (jj >= columns) break;
                            c[ii][jj] += a[ii][kk] * b[kk][jj];
                        }
                    }
                }
            }
        }
    }
}

最后的结果是：

1	Time difference = 3129[ms]

看起来并不理想，哈？这是因为，我们只关注了访存局部性，但是，为了访存局部性，又引入了新的运算。这里，我们需要平衡分块大小tile_size，不能过大，否则会导致缓存放不下我们需要的数据；也不能过小，否则会导致缓存放的数据过少。不过，tile_size通常都需要是Cache Line的大小除以元素的大小。在笔者实测的过程中，这个方法不是很稳定，很多时候会比Easy Mode慢，也在有些时候会比Easy Mode快，tile_size的取值通常是64、80和96。

当数据越大，这个方法的优势越大，越容易比Easy Mode快，但是由于时间和粗心原因，笔者没有记录详细的时间，敬请谅解。

在GPU上的优化

矩阵的分配

首先值得一提的是矩阵在GPU上的分配方式，CUDA上的二维数组是一个一维数组加上数据对齐，代码如下：

int *a_gpu, *b_gpu, *c_gpu;
size_t pitch;
cudaMallocPitch(&a_gpu, &pitch, MATRIX_SIZE * sizeof(int), MATRIX_SIZE);
cudaMallocPitch(&b_gpu, &pitch, MATRIX_SIZE * sizeof(int), MATRIX_SIZE);
cudaMallocPitch(&c_gpu, &pitch, MATRIX_SIZE * sizeof(int), MATRIX_SIZE);

这里，分配的二维数组其实就是一个一维数组，索引方式也和一维数组没差。

实现

一个简单的实现依旧是根据矩阵乘法的定义：

__global__ void TrivialMatrixMatMulGPU(int* a, int* b, int* c, int lines, int columns) {
    int curr_idx_x = blockIdx.x * blockDim.x + threadIdx.x;
    int curr_idx_y = blockIdx.y * blockDim.y + threadIdx.y;
    int sum = 0;
    if (curr_idx_x >= columns || curr_idx_y >= lines) {
        return;
    }
    for (int k = 0; k < columns; k++) {
        sum += a[curr_idx_y * columns + k] * b[k * columns + curr_idx_x];
    }
    c[curr_idx_y * columns + curr_idx_x] = sum;
}

在主函数中：

1
2
3

dim3 block(32, 32);
dim3 grid((MATRIX_SIZE+block.x-1) / block.x, (MATRIX_SIZE+block.y-1) / block.y);
TrivialMatrixMatMulGPU<<<grid, block>>>(a_gpu, b_gpu, c_gpu, MATRIX_SIZE, MATRIX_SIZE);

通过cudaEvent_t进行计时，结果如下：

1	Time difference = 23.5479[ms]

到这里，我们能看到在这种「高度并行且任务简单」的场景下，GPU是有极强的优势的。