硕大的汤姆

硕大的汤姆

The official website of Minhua Chen

24 Aug 2019

Cuda实战入门2: 将矩阵乘法速度提升 5000 倍

本实验采用不同的方法来计算 8192 * 8192 的整型矩阵乘法运算。

C 语言版

C 语言是大家公认的高性能语言,那我们就从 C 语言开始吧。

    // 用一位数组表示二维矩阵
    mat1 = (int*) malloc(m_size * m_size * sizeof(int));
    mat2 = (int*) malloc(m_size * m_size * sizeof(int));
    result = (int*) malloc(m_size * m_size * sizeof(int));

    // initialize
    for (int i = 0; i < m_size * m_size; i++) {
        mat1[i] = rand()/1000000;
        mat2[i] = rand()/1000000;
        result[i] = 0;

    }

    for (int r = 0; r < m_size; r++) {
        for (int c = 0; c < m_size; c++) {
            for (int n = 0; n < m_size; n++) {
                result[r*m_size + c] += mat1[r*m_size+n] * mat2[n*m_size+c];
            }
        }
    }

这代码没什么可说的,值得一提的可能就是用一维数组来存二维矩阵,这样可以让矩阵在内存中的分布更连续。很容易估算出来这段代码需要进行万亿级别的乘法或加法运算,倘若使用 1950 年冯.诺依曼用来预测天气的计算机(每秒 5000 次计算),大概要算个六七年吧。

还好现在是 2019 年,这段代码在我的 i7 cpu 上运行花了 7000 s,差不多两个小时了。(记得开编译器优化,不然可能今天都跑不完)

eigen 版

#include <Eigen/Dense>
using namespace::Eigen;

typedef Eigen::Matrix<int, Dynamic, Dynamic> IntMatrix;

IntMatrix mat1(m_size,m_size);
IntMatrix mat2(m_size,m_size);
IntMatrix result(m_size,m_size);
mat1.setRandom();
mat2.setRandom();

result = mat1 * mat2;

运行此段代码需要下载 eigen 库,并在编译时 link 上去。这段代码在我的机器(i7 单核)上运行花了 120s,从两个小时直接变成了两分钟,代码还简单了很多啊握草。(同样记得开编译器优化)

给 eigen 跪了啊,这比我快了快 60 倍啊,这告诉我们一个道理,轮子该用还是得用啊。( eigen 进行了指令级级别的优化,以后有机会在写篇文章深入讲下。)

CUDA 显然,矩阵乘法是一个天然完美的可并行问题。事实上,大量伟大的理论研究和工程实践都基于矩阵乘法可并行这一前提,比如 google。

CPU 并不擅长并行计算,但是 GPU 擅长啊。GPU 可以一口气拉起一个网格(grid)的线程,像一个计算方针一样同时处理计算的不同部分。关于 CUDA 可以参考前面一篇文章。​ 在矩阵乘法这个问题中,结果矩阵中的每一个位置的数据都不依赖于其他位置数据的计算过程和结果。所以我们完全可以拉一堆进程起来,让他们分别计算某一些位置的结果。

你可以想象为你们学校要算两个 100 * 100 的矩阵的乘法,这个矩阵计算的结果是万级的,而运算量可是百万级的。于是你们校长找到 50 个班主任,给每个班主任分配两行结果计算任务,你们班呢,分到了前两行,也就是要算出结果矩阵中前两行的 200 个数字。班主任再把这两百个数字分给了 50 个同学,每个同学只要算 4 个数字,大约 400 次加法和 400 次乘法。于是一个巨复杂的计算问题,你们学校半天就完成了。

下面是我的代码实现

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

#define BLOCK_NUM 32   //块数量
#define THREAD_NUM 256 // 每个块中的线程数
#define R_SIZE BLOCK_NUM * THREAD_NUM
#define M_SIZE R_SIZE * R_SIZE

__global__ void mat_mul(int *mat1, int *mat2, int *result) {
    const int bid = blockIdx.x;
    const int tid = threadIdx.x;
    // 每个线程计算一行
    const int row = bid * THREAD_NUM + tid;
    for (int c = 0; c < R_SIZE; c++) {
        for (int n = 0; n < R_SIZE; n++) {
            result[row*R_SIZE+c] += mat1[row*R_SIZE+n] * mat2[n*R_SIZE+c];
        }
    }
}

int main(int argc, char *argv[]) {
    int *mat1, *mat2, *result;
    int *g_mat1, *g_mat2, *g_mat_result;

    // 用一位数组表示二维矩阵
    mat1 = (int*) malloc(M_SIZE * sizeof(int));
    mat2 = (int*) malloc(M_SIZE * sizeof(int));
    result = (int*) malloc(M_SIZE * sizeof(int));

    // initialize
    for (int i = 0; i < M_SIZE; i++) {
        mat1[i] = rand()/1000000;
        mat2[i] = rand()/1000000;
        result[i] = 0;

    }

    cudaMalloc((void **)&g_mat1, sizeof(int) * M_SIZE);
    cudaMalloc((void **)&g_mat2, sizeof(int) * M_SIZE);
    cudaMalloc((void **)&g_mat_result, sizeof(int) * M_SIZE);

    cudaMemcpy(g_mat1, mat1, sizeof(int) * M_SIZE, cudaMemcpyHostToDevice);
    cudaMemcpy(g_mat2, mat2, sizeof(int) * M_SIZE, cudaMemcpyHostToDevice);

    mat_mul<<<BLOCK_NUM, THREAD_NUM>>>(g_mat1, g_mat2, g_mat_result);

    cudaMemcpy(result, g_mat_result, sizeof(int) * M_SIZE, cudaMemcpyDeviceToHost);
}

这段代码在我的机器(RTX 2080TI)上运行了多久呢?

22 秒!!!

差不多比 c 写的第一版要快了 300 倍了。

Tensorflow

mat1 = tf.random_uniform([8192, 8192], minval=1, maxval=65536, dtype=tf.int32)
mat2 = tf.random_uniform([8192, 8192], minval=1, maxval=65536, dtype=tf.int32)

sess = tf.Session()

sess.run(mat1)
sess.run(mat2)
sess.run(tf.matmul(mat1, mat2))

sess.close()

这段代码在我的机器上也运行了 22 秒,和上面我自己手写 cuda 程序性能相当。

CUBLAS

22 秒已经到极限了吗?还早得很呢!我们可以使用 CUBLAS 库,不但可以屏蔽掉复杂的底层实现以及不同计算设备带来的参数设计的影响,还有机会把矩阵乘法的效率进一步提升。先贴出 cublas 的一份文档 。 https://developer.nvidia.com/sites/default/files/akamai/cuda/files/Misc/mygpu.pdf​ 下面是我的代码实现

// 用一位数组表示二维矩阵
    mat1 = (float*) malloc(m_size * sizeof(float));
    mat2 = (float*) malloc(m_size * sizeof(float));
    result = (float*) malloc(m_size * sizeof(float));

    // initialize
    for (int i = 0; i < m_size; i++) {
        mat1[i] = rand()/10000000;
        mat2[i] = rand()/10000000;
        result[i] = 0;

    }

    cudaMalloc((void **)&g_mat1, sizeof(*mat1) * m_size);
    cudaMalloc((void **)&g_mat2, sizeof(*mat2) * m_size);
    cudaMalloc((void **)&g_mat_result, sizeof(*result) * 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);

其中的 cublasSgemm 方法就是计算矩阵乘法的函数,具体请参考上面 pdf 的 2.4 小节。

这段代码计算 8192 * 8192 的矩阵乘法要多久呢?

答案是:1.35 秒 !!!我们在第一版 c 语言实现的基础上完成了 5000 倍的提升!!!

总结

本篇文章中,我们将 8192 * 8192 的矩阵运算从 7000 秒一路提升到了 1.35 秒,提升了五千多倍。这是一个成就感与挫败感并存的过程,一方面不断提升计算效率很爽,一方面又感觉这种提升并非来自我的智慧而是他人完成的工作。

这篇文章虽然暂时告一段落,但是我相信,这里面依然有各种优化的空间。另外,如果矩阵更大呢?如果是稀疏矩阵呢?如果想利用多块 GPU 呢?

游戏才刚刚开始呢。

PS: 本文代码我都放到 github 上啦。 github.com/chenminhua/math

PS: 有朋友质疑为什么第一个 C 语言版本的运行时长这么慢,是不是没开编译器优化。事实上我是开了 O3 的,编译器是 gcc49。我的猜测是计算这么慢应该和内存访问有关,因为这两个数组都比较大,而内存是分页管理的,这让 cpu 不停地在不同的页上取数据,cpu 上的页表地址缓存一直不能命中,所以比较慢。网友 Antholiqi 提到,如果把左乘矩阵按行储存,右乘矩阵按列存储,计算效率大幅提升,事实确实如此。并且每个线程计算的区域最好不是一行,而是一块矩形。比如每个线程计算 64*64 的区域,则只需要扫描两个矩阵各 64 行数据即可,而不用读整个右矩阵。总而言之,第一版的 c 语言代码其实还有非常多非常多的优化空间啦。

PS: 还有网友问「Eigen 指令级别优化是啥意思」,大家可以去参考 eigen 的文档 https://github.com/libigl/eigen/blob/1f05f51517ec4fd91eed711e0f89e97a7c028c0e/doc/InsideEigenExample.dox