通用矩阵乘法GEMM优化与实现（1)

执行长老

已于 2023-06-01 22:56:57 修改

阅读量2.3k

点赞数 4

分类专栏：算法手册文章标签：矩阵线性代数

于 2022-11-11 20:14:26 首次发布

本文链接：https://blog.csdn.net/CQRuler/article/details/127812444

版权

算法手册专栏收录该内容

2 篇文章 1 订阅

订阅专栏

1 矩阵乘法

矩阵乘法（Matrix Multiplication）的一般形式是C=AB ，只有在矩阵A 的列数和矩阵B 的行数相同是才有定义。若A 为m×p 的矩阵，B 为p×n 的矩阵，则C 是一个m×n 的矩阵，其元素为：

$c_{ij}=\sum_{k=1}^{p}a_{ik}b_{kj} \left ( i=1,2,...,m;j=1,2,...,n \right )$

矩阵乘法的伪代码：

// Algorithm 1: matrix multiplication
for i = 1 to m do
    for j = 1 to n do
        C[i][j] = 0
        for k = 1 to p do
            C[i][j] += A[i][k] * B[k][j]
        end
    end
end

计算操作次数为2MNP，其中 M 、N 、P 分别为三层循环执行次数，内层循环执行一次乘法和一次加法。内存访问次数为4MNP ，其中内层循环分别访问A 和B 各一次，访问C 两次，先读取，累加后再存储。

1.1 矩阵分块乘法

分块矩阵是将矩阵划分成小块（Block），极端情况下，每个小块只有一个元素,分块矩阵将退化成普通矩阵。两个分块矩阵相乘，可以把每个块当一个元素，得到与一般矩阵乘法相同的公式。

矩阵分块乘法的伪代码：

// Algorithm 2：matrix block multiplication
for i = 1 to m step 4 do
    for j = 1 to n step 4 do
        C[i + 0][j + 0...3] = 0
        C[i + 1][j + 0...3] = 0
        C[i + 2][j + 0...3] = 0
        C[i + 3][j + 0...3] = 0
        for k = 1 to p step 4 do
            C[i + 0][j + 0...3] += A[i + 0][k + 0...3] * B[k + 0...3][j + 0...3]
            C[i + 1][j + 0...3] += A[i + 1][k + 0...3] * B[k + 0...3][j + 0...3]
            C[i + 2][j + 0...3] += A[i + 2][k + 0...3] * B[k + 0...3][j + 0...3]
            C[i + 3][j + 0...3] += A[i + 3][k + 0...3] * B[k + 0...3][j + 0...3]
        end
    end
end

分块矩阵乘法不会改变计算操作次数，但可以优化内存访问次数。在对M 和N 展开时，可以复用A 和B 的数据，复用4次；在对P 展开时，可以在寄存器中完成累加，内层循环结束后一次性写入内存。内存访问次数减少为MN+MNP/2。

1.2 缓存优化

程序是运行在内存（RAM）之中，称之为主存（Main Memory），主存的访问速度大约是120个时钟周期（CPU Cycle）,相对于CPU的速度差了两个量级。为了缓解CPU与主存之间速度不匹配的问题，增加了速度快但容量小的缓存（Cache）。缓存通常由静态存储器（Static Random Access Memory）组成，用于存储近期可能需要运行的指令和数据。

缓存可分为一级缓存（L1 cache）、二级缓存（L2 cache）和三级缓存（L3 cache）。通常，L1 cache和L2 cache集成在CPU里面，也称为On-chip cache；L3 cache位于CPU外面，称为Off-chip cache。其中，L1 cache访问速度约1ns；L2 cache访问速度约3ns；L3 cache访问速度约12ns。缓存等级越高，速度越慢，容量越大。但是相对于主存约65ns的访问速度而言，依然很快。

当CPU试图从某地址读取数据时，首先从L1 cache中查询是否命中（Hit），如果命中则把数据返回给CPU。如果L1 cache缺失（Miss），则继续从L2 cache中查找。当L2 cache命中时，数据会返回给L1 cache以及CPU。如果L2 cache也缺失，则继续从L3 cache中查找。当L3 cache命中时，数据会返回给L2 cache, L1 cache以及CPU。如果L3 cache也缺失，则需要从主存中读取数据，将数据返回给L3 cache、L2 cache、L1 cache以及CPU。

程序的缓存优化，主要是通过提高缓存的命中率，减少CPU对主存的访问次数，来获取性能的提升。但由于不同CPU存在硬件结构上的差异，各级缓存大小，映射方式都不尽相同，因此，缓存优化是一项复杂而繁琐的工作。缓存优化尽管繁复，但这确是非常有效的代码优化手段。

1.2.1 矩阵存储顺序

矩阵在计算机中有两种存储方式：行优先顺序存储和列优先顺序存储，简称行主序（row-major order）和列主序（column-major order）。行主序中，同一行的元素在内存中是相邻的；列主序中，同一列的元素在内存中是相邻的。

由于在C/C++语言中，矩阵采用的是行主序，因此下面的讨论均基于行主序进行。对于列主序，可以通过转置转换为行主序。

1.2.2 循环顺序优化

分析矩阵乘法的循环顺序对内存访问的影响，暂不考虑内存访问次数。先按照M 、N 、P 的顺序进行循环，伪代码：

// Algorithm 3：matrix multiplication
for i = 1 to m do
    for j = 1 to n do
        for k = 1 to p do
            C[i][j] += A[i][k] * B[k][j]
        end
    end
end

在最内层循环的每次迭代中，k 值都在变化，意味着对矩阵B 是按列读取的。由于矩阵采用的是行主序，因此，按列读取时，每次都需要跳过整行数据。如果矩阵B 的规模足够大，那么每次迭代都可能出现缓存缺失的情况。

上述算法对矩阵A 、B 、C 的内存访问顺序如下图。其中，横坐标是循环迭代次数，纵坐标是内存访问位置的偏移量百分比（相对于矩阵数据总量）。从内存访问顺序图中可以看出，对矩阵B 的内存访问位置存在大幅度跳变，很容易出现缓存缺失；对矩阵C 的内存访问位置变化非常平稳，缓存命中率高。该算法的综合缓存命中率低。

交换第二层和第三次循环的顺序，即按照M 、P 、N 的顺序进行循环，伪代码：

// Algorithm 4：matrix multiplication
for i = 1 to m do
    for k = 1 to p do
        for j = 1 to n do
            C[i][j] += A[i][k] * B[k][j]
        end
    end
end

在最内层循环的每次迭代中，一直改变的是j 值，意味着对矩阵B 是按行读取的，出现缓存缺失的概率会大幅度降低。改进后，对矩阵A 、B 、C 的内存访问顺序如下图。

对矩阵B 的内存访问位置的连续性得到明显改善；对矩阵A 和矩阵C 的内存访问位置顺序虽然有所改变，但几乎没带来任何改善。改进后的算法，综合缓存命中率得到了提升，优化是有效的。

1.2.3 之字形重排

之字形顺序（Z-Order）可以将多维数据映射到一维。矩阵作为二维数据，可以利用Z-Order映射到一维内存空间。对于二维空间，Z-Order有两种不同的排列顺序：行主序之字形顺序（Row-major z-order）和列主序之字形顺序（Column-major z-order）。

Z-Order重排的递归实现流程：对于任何一个大小为m×n 的矩阵，先拆分成2x2排列的4个子矩阵，然后再按照下图顺序（行主序与列主序的拆分顺序不同）拆分每个子矩阵，重复上述过程，直到每个子矩阵足够小，不需要再拆分为止。

Z-Order重排的关键在于如何计算重排后的每个子矩阵的内存地址。每个子矩阵的内存地址等于上一个子矩阵的内存地址偏移上一个子矩阵的面积： $p_{i}=p_{i-1}+s_{i-1}$ ，其中 $2\leqslant i\leqslant 4$ 。

行主序Z-Order重排的伪代码：

// Algorithm 5：row-major Z-Order pack
fuction z_order_rm(B, A, lda, m, n)
    block_m = 4
    block_n = 4
    if m > block_m then
        m1 = 1 << (int)log2(m - 1)
        m2 = m - m1
        if n > block_n then
            n1 = 1 << (int)log2(n - 1)
            n2 = n - n1
            z_order_rm(B, A, lda, m1, n1)
            z_order_rm(B + m1 * n1, A + n1, lda, m1, n2)
            z_order_rm(B + m1 * n, A + m1 * lda, lda, m2, n1)
            z_order_rm(B + m1 * n + m2 * n1, A + m1 * lda + n1, lda, m2, n2)
        else
            z_order_rm(B, A, lda, m1, n)
            z_order_rm(B + m1 * n, A + m1 * lda, lda, m2, n)
        end if
    else
        if n > block_n then
            n1 = 1 << (int)log2(n - 1)
            n2 = n - n1
            z_order_rm(B, A, lda, m, n1)
            z_order_rm(B + m * n1, A + n1, lda, m, n2)
        else
            impl_copy_block(B, A, lda, m, n)
        end if
    end if

列主序Z-Order重排的伪代码：

// Algorithm 6：column-major Z-Order pack
fuction z_order_cm(B, A, lda, m, n)
    block_m = 4
    block_n = 4
    if m > block_m then
        m1 = 1 << (int)log2(m - 1)
        m2 = m - m1
        if n > block_n then
            n1 = 1 << (int)log2(n - 1)
            n2 = n - n1
            z_order_cm(B, A, lda, m1, n1)
            z_order_cm(B + m1 * n1, A + m1 * lda, lda, m2, n1)
            z_order_cm(B + m * n1, A + n1, lda, m1, n2)
            z_order_cm(B + m * n1 + m1 * n2, A + m1 * lda + n1, lda, m2, n2)
        else
            z_order_cm(B, A, lda, m1, n)
            z_order_cm(B + m1 * n, A + m1 * lda, lda, m2, n)
        end if
    else
        if n > block_n then
            n1 = 1 << (int)log2(n - 1)
            n2 = n - n1
            z_order_cm(B, A, lda, m, n1)
            z_order_cm(B + m * n1, A + n1, lda, m, n2)
        else
            impl_copy_block(B, A, lda, m, n)
        end if
    end if

1.2.4 递归矩阵乘法

递归矩阵乘法（Recursive matrix multiplication）是将矩阵拆分成2x2排列的4个子矩阵，根据矩阵分块乘法规则，计算8次子矩阵乘法。每次子矩阵乘法，又可以继续拆分，重复此过程，最终完成矩阵乘法运算。矩阵分块乘法公式：

$\begin{bmatrix} C_{00} & C_{01}\\ C_{10} & C_{11} \end{bmatrix}=\begin{bmatrix} A_{00} & A_{01}\\ A_{10} & A_{11} \end{bmatrix}\times \begin{bmatrix} B_{00} & B_{01}\\ B_{10} & B_{11} \end{bmatrix}=\begin{bmatrix} A_{00}B_{00}+A_{01}B_{10} & A_{00}B_{01}+A_{01}B_{11}\\ A_{10}B_{00}+A_{11}B_{10} & A_{10}B_{01}+A_{11}B_{11} \end{bmatrix}$

矩阵C 和矩阵A 采用行主序Z-Order重排, 矩阵B 采用列主序Z-Order重排。

上述公式可写成：

$\begin{bmatrix} C_{0} & C_{1}\\ C_{2} & C_{3} \end{bmatrix} = \begin{bmatrix} A_{0} & A_{1}\\ A_{2} & A_{3} \end{bmatrix} \times \begin{bmatrix} B_{0} & B_{2}\\ B_{1} & B_{3} \end{bmatrix}=\begin{bmatrix} A_{0}B_{0}+A_{1}B_{1} & A_{0}B_{2}+A_{1}B_{3}\\ A_{2}B_{0}+A_{3}B_{1} & A_{2}B_{2}+A_{3}B_{3} \end{bmatrix}$

递归矩阵乘法计算过程中，对矩阵A 的内存访问位置顺序是： $A_{0}$ 、 $A_{1}$ 、 $A_{0}$ 、 $A_{1}$ 、 $A_{2}$ 、 $A_{3}$ 、 $A_{2}$ 、 $A_{3}$ ；对矩阵B 的内存访问位置顺序是： $B_{0}$ 、 $B_{1}$ 、 $B_{2}$ 、 $B_{3}$ 、 $B_{0}$ 、 $B_{1}$ 、 $B_{2}$ 、 $B_{3}$ ；对矩阵C 的内存访问位置顺序是： $C_{0}$ 、 $C_{0}$ 、 $C_{1}$ 、 $C_{1}$ 、 $C_{2}$ 、 $C_{2}$ 、、。

递归矩阵乘法的伪代码：

// Algorithm 7：recursive matrix multiplication
function rmm(C, A, B, m, p, n)
    block_m = 4
    block_p = 4
    block_n = 4
    if m > block_m then
        m1 = 1 << (int)log2(m - 1)
        m2 = m - m1
        if p > block_p then
            p1 = 1 << (int)log2(p - 1)
            p2 = p - p1
            if n > block_n then
                n1 = 1 << (int)log2(n - 1)
                n2 = n - n1
                rmm(C, A, B, m1, p1, n1)
                rmm(C, A + m1 * p1, B + p1 * n1, m1, p2, n1)
                rmm(C + m1 * n1, A, B + p * n1, m1, p1, n2)
                rmm(C + m1 * n1, A + m1 * p1, B + p * n1 + p1 * n2, m1, p2, n2)
                rmm(C + m1 * n, A + m1 * p, B, m2, p1, n1)
                rmm(C + m1 * n, A + m1 * p + m2 * p1, B + p1 * n1, m2, p2, n1)
                rmm(C + m1 * n + m2 * n1, A + m1 * p, B + p * n1, m2, p1, n2)
                rmm(C + m1 * n + m2 * n1, A + m1 * p + m2 * p1, B + p * n1 + p1 * p2, m2, p2, n2)
            else
                rmm(C, A, B, m1, p1, n)
                rmm(C, A + m1 * p1, B + p1 * n, m1, p2, n)
                rmm(C + m1 * n, A + m1 * p, B, m2, p1, n)
                rmm(C + m1 * n, A + m1 * p + m2 * p1, B + p1 * n, m2, p2, n)
            end if
        else
            if n > block_n then
                n1 = 1 << (int)log2(n - 1)
                n2 = n - n1
                rmm(C, A, B, m1, p, n1)
                rmm(C + m1 * n1, A, B + p * n1, m1, p, n2)
                rmm(C + m1 * n, A + m1 * p, B, m2, p, n1)
                rmm(C + m1 * n + m2 * n1, A + m1 * p, B + p * n1, m2, p, n2)
            else
                rmm(C, A, B, m1, p, n)
                rmm(C + m1 * n, A + m1 * p, B, m2, p, n)
            end if
        end if
    else
        if p > block_p then
            p1 = 1 << (int)log2(p - 1)
            p2 = p - p1
            if n > block_n then
                n1 = 1 << (int)log2(n - 1)
                n2 = n - n1
                rmm(C, A, B, m, p1, n1)
                rmm(C, A + m * p1, B + p1 * n1, m, p2, n1)
                rmm(C + m * n1, A, B + p * n1, m, p1, n2)
                rmm(C + m * n1, A + m * p1, B + p * n1 + p1 * n2, m, p2, n2)
            else
                rmm(C, A, B, m, p1, n)
                rmm(C, A + m * p1, B + p1 * n, m, p2, n)
            end if
        else
            if n > block_n then
                n1 = 1 << (int)log2(n - 1)
                n2 = n - n1
                rmm(C, A, B, m, p, n1)
                rmm(C + m * n1, A, B + p * n1, m, p, n2)
            else
                impl_mat_block_mul(C, A, B, m, p, n)
            end if
        end if
    end if

由于在矩阵乘法中，对矩阵B 的内存访问是最糟糕的，因此矩阵乘法的缓存优化主要是优化矩阵B 的内存访问。采用递归矩阵乘法，矩阵B 的内存连续性可以得到很大改善，而矩阵A 的内存连续性反而有所下降。对矩阵A 、B 、C 的内存访问顺序如下图。

递归矩阵乘法的内存访问顺序图具有自相似性，即整体与局部相似。自相似性是一个重要的特性，这就意味着，无论矩阵规模多大，其内存访问顺序图在整体上是相同的，区别仅仅在于细节的丰富程度。从图中可以看出，矩阵C 和矩阵A 的斜率基本相同，前者跳变量小于后者；矩阵B 的斜率更大，跳变量也更大。因此，在递归矩阵乘法运算中，内存连续性最好的矩阵C ，其次是矩阵A ，稍差的是矩阵B 。

递归矩阵乘法与普通矩阵分块乘法相比，前者无论矩阵B 规模有多大，只有1次跳变量为100%，二级跳变量仅为20%；后者随着矩阵B 规模的增加，跳变量为100%的跳变次数也将线性增加。矩阵规模越大，递归矩阵乘法的缓存优势就越明显。

1.3 矩阵数据重排

递归矩阵乘法能够有效优化缓存，分块矩阵乘法可以减少内存访问次数，综合两者的优势，对矩阵数据进行重排。

分块矩阵乘法的内存访问次数为MN+MNP/2，显然，增大P，能够更有效地减少内存访问次数。在递归矩阵乘法的中，调整子矩阵的形状，使得block_p≫block_m、block_p≫block_n，可以减少内存访问次数。当然，block_p也不是越大越好，否则子矩阵过大，会导致cache命中率降低。因此，在减少内存访问次数与提高cache命中率之间需要进行折中，block_p可以作为参数，根据实际情况进行调整。

对矩阵A进行行主序Z-Order重排，拆分成block_m×block_p的子矩阵与拆分成block_p×block_p的子矩阵是完全等价的，而后者的递归次数更少。

对矩阵B 进行列主序Z-Order重排，拆分成block_p×block_n 的子矩阵与拆分成block_p×block_p 的子矩阵几乎等价，但需要对子矩阵块进行分块转置。

针对矩阵B 的每个block_p×block_p 的子矩阵进行分块转置，如下图所示：

由于在矩阵进行转置时，输入矩阵和输出矩阵总有一个是按列访问的，需要对cache进行优化。矩阵转置的cache优化与Z-Order重排类似，这里就不再赘述。

子矩阵进行分块转置的伪代码：

// Algorithm 8：submatrix block transpose
function block_trp(B, ldb, A, lda ,m, n)
    block_m = 4 
    block_n = 4 
    if m > block_m then
        m1 = 1 << (int)log2(m - 1)
        m2 = m - m1
        if n > block_n then
            n1 = 1 << (int)log2(n - 1)
            n2 = n - n1
            block_trp(B, ldb, A, lda, m1, n1)
            block_trp(B + n1 * ldb, ldb, A + n1, lda, m1, n2)
            block_trp(B + m1 * block_m, ldb, A + m1 * lda, lda, m2, n1)
            block_trp(B + n1 * ldb + m1 * block_m, ldb, A + m1 * lda + n1, lda, m2, n2)
        else
            block_trp(B, ldb, A, lda, m1, n)
            block_trp(B + m1 * block_m, ldb, A + m1 * lda, lda, m2, n)
        end if
    else
        if n > block_n then
            n1 = 1 << (int)log2(n - 1)
            n2 = n - n1
            block_trp(B, ldb, A, lda, m, n1)
            block_trp(B + n1 * ldb, ldb, A + n1, lda, m, n2)
        else
            impl_copy_block(B, A, lda, m, n)
    end if

执行长老

关注

4
点赞
踩
21

收藏

觉得还不错? 一键收藏
2
评论
通用矩阵乘法GEMM优化与实现（1)

通用矩阵乘法（General Matrix Multiply）是最常见底层算法的之一，广泛用于线性代数、机器学习、统计学等各个领域。矩阵乘法的时间复杂度是O(n^3)，虽然有Strassen等优化算法可以降低算法复杂度，但实际效果仍然有限。本文不讨论矩阵乘法的算法优化，仅讨论矩阵乘法的代码优化技术。内容并非介绍和照搬目前流行的BLAS库的做法，而是基于完备的代码优化理论体系，讨论并给出矩阵乘法的最优实现。
复制链接

扫一扫

专栏目录