矩阵计算性能优化-初探

山花2024-03-29 14:35

背景

矩阵计算无论是在高性能计算还是深度学习中，都是比较常用的操作，常规的矩阵相乘包括大量的浮点数计算或者乘加运算，是典型的计算密集型算子。矩阵运算的性能直接展示了程序员设计的矩阵算法对CPU的利用能力。

FLOPS和FLOPs

FLOPS 意指每秒浮点运算次数。用来衡量硬件的性能
FLOPs 是浮点运算次数，可以用来衡量算法/模型复杂度

在性能优化的过程中，我们首先需要获取当前硬件的FLOPS，再通过FLOPs / cost_time 获取当前算法每秒的浮点运算次数，通过比较两者之间的差距，衡量自己对CPU的利用能力。

那么，同样是计算，为什么利用能力会有差距呢？现代计算机系统使用多级缓存来减少处理器与主存之间的数据传输延迟。矩阵计算算法设计主要考虑如何使用更好的数据访问模式以减少cache miss。在硬件水平相同，不考虑并行处理与分布式计算带来的通信开销的情况下，降低cache miss，可以显著提高性能。

矩阵计算常规实现

以下是一个矩阵计算的常规实现

C++ 复制代码

void naive_gemm(const float* A, const float* B, float* C, const int M, const int N, const int K) { 
    for (int m = 0; m < M; ++m) {  // 循环1
        for (int n = 0; n < N; ++n) {  // 循环2
            for (int k = 0; k < K; ++k) {  // 循环3
                C[m * N + n] += A[m * K + k] * B[k * N + n];
            }
        } 
    } 
}

为了方便后面讲述，这里为每个循环添加了标记

我们可以看到，每次循环都会获取不同列的值，在矩阵本身按行存储的同时，相当于每次获取同一列的不同数据时，都要读取。这会大大增加cache miss的频率。

矩阵计算优化1

C++ 复制代码

void naive_gemm(const float* A, const float* B, float* C, const int M, const int N, const int K) { 
    for (int m = 0; m < M; ++m) {  // 循环1
        for (int k = 0; k < K; ++k) {  // 循环3
            for (int n = 0; n < N; ++n) {  // 循环2
                C[m * N + n] += A[m * K + k] * B[k * N + n];
            }
        } 
    } 
}

通过对循环进行重排我们发现，

每次循环从获取B的每列不同值变成了每行不同值，这就让我们可以通过一次读取矩阵一行的数据写入cache这个行为可以获得收益，而不是每次循环都会造成cache miss。通过这种方式可以大大增加矩阵计算的性能。