Ctorch开发日志——矩阵乘法优化及数学原理
随着项目的推进，本作者遇到了目前最棘手的问题，即矩阵乘法的优化但是有句话说得好
“你越棘手，我越兴奋”
那么，如下是本作者如何把\(O(MNK)\)（\(O(n^3)\)）的朴素矩阵乘法一步一步优化到\(O(n^{2.81})\) 的全过程
测试环境
macOS Tahoe 26 Beta 2
M3 Pro 11核
18GB
CLion & Cmake
计时器：ctime
矩阵：1024 * 1024 @ 1024 * 1024
为保证准确，时间均为5次测量取平均值
朴素算法实现 & 测速
朴素实现的数学原理
其实就是把矩阵乘的数学公式重写一遍：
\_{m \times n}} \]
\_{n \times s}} \]
\_{m \times s} \]
\[= \left[ \sum \limits_{k=1}^{n}a_{ik}b_{kj}\right]_{m \times s}\]
注意：矩阵乘法不满足交换律
只有左矩阵的列数与右矩阵的行数相同的两个矩阵才能相乘
乘积矩阵的行数等于左矩阵的行数，列数等于右矩阵的列数
概括一下就是
乘积矩阵第i行第j列处的元素等于左矩阵的第i行与右矩阵的第j列对应元素乘积之和
那么，这个很简单，上代码吧
为了测试，所有的矩阵保证满足乘法条件且为2维
时间复杂度 \(O(n^3)\)
// 原始版本（未优化）
void matrix_mult(float* A, float* B, float* C, int N) {
for (int i = 0; i < N; i++)
for (int j = 0; j < N; j++)
for (int k = 0; k < N; k++)
C += A * B;
}在实际测试中，此算法跑出了1898.46ms的优秀成绩
优化一：循环优化
你可能会疑惑，循环优化是什么
故名思义，就是对原有的ijk的循环重新更换顺序为ikj
一个更大的问题来了，凭什么仅仅改变顺序就快了许多
那么不妨看看访问顺序
算法1（朴素实现）：
在这个顺序中，最内层循环是k，它遍历A的一行和B的一列。
对于A的访问是连续的（因为A在内存中是按行存储的，所以k增加时是连续访问），
但是B的访问是不连续的（因为B在内存中是按行存储，k增加时访问的是不同行的同一列，所以是跳跃访问）。这样对B的访问会导致缓存失效。
so，真正影响到速度的，就是缓存，在顺序读取中，缓存可以加载一整行，不必跳跃元素访问
那么，有没有一种循环顺序，使得对三个数组均为顺序访问呢
有的兄弟，有的，让我们欢迎仍为\(O(n^{3})\)的优化算法出场
——“ikj”循环优化
在这个顺序中，最内层循环是j。对于A的访问，固定i和k，所以每次内层循环A是常数。对于B的访问，是B，由于j是连续的，所以B的访问是连续的（因为同一行连续列）。同时，C的访问也是连续的（C）。这样，所有的内存访问都是连续的，因此性能更好。
可以自行验证，对于三个数组，均为顺序访问
给出如下代码：
时间复杂度 \(O(n^3)\)
// 优化后（行优先访问）
void matrix_mult_opt1(float* A, float* B, float* C, int N) {
for (int i = 0; i < N; i++)
for (int k = 0; k < N; k++)// k循环提到中间
for (int j = 0; j < N; j++)
C += A * B;
}在实际测试中，此算法跑出了1462.89ms的优秀成绩
提升：1898.46ms-1462.89ms = 435.57ms 提高22.9%
进阶提升 优化二：矩阵分块算法
顾名思义，分块算法即是把矩阵分为多个小矩阵，对每个矩阵操作后再组合出结果，类似分块算法
那么，它为什么快呢
在计算机中，共有三种CPU缓存以及普通内存，即L1、L2、L3 Cache和内存
前三种的速度要比内存快很多很多，大概只有一个CPU周期的延迟，而普通内存可以达到上百周期延迟
而比他们更快的就是寄存器，直接接触CPU，0延迟
唯一的问题是，寄存器的大小只够存储单个值
新的问题来了，怎样把一个巨大的矩阵放到只有128k-1m的L1缓存中呢
聪明的你一定想到了把原矩阵划分为多个小矩阵，每一个都能放到L1内进行运算
那么恭喜你，你已经知道了矩阵分块算法的原理
更确切的说，先定义一个常数 \(blocks \in N^{+}\) 作为划分的单位矩阵的行列，
对于原矩阵A，我们把其中的 \({blocks \times blocks}\) 个元素划分为一个新矩阵，记为\(a^{'}_{\cdots}\),我们定义:
\[ a^{'}_{11} = \begin{pmatrix}a_{11} & \cdots & a_{1blocks} \\\vdots & \ddots &\vdots \\a_{blocks1} & \cdots & a_{blocks_{ }blocks}\end{pmatrix}\]
其余以此类推，新矩阵\(A^{'}\)即变为
\
正如同我们可以把 \(f(x)\) 中的 \(x\) 替换为任意多项式（函数），我们同样也可以把矩阵中的每个元素换为一个矩阵，其运算规则仍然成立
so，显而易见的，分块矩阵算法有如下公式：
对于整体而言，
\[\begin{array}{c} C={A^{'}\times B^{'}}=\left[ c_{ij}\right]_{m \times s} = \left[ \sum \limits_{k=1}^{n}a^{'}_{ik}b^{'}_{kj}\right]_{m \times s} \end{array}\]
其中：
\[ a^{'}_{ij} = \begin{pmatrix}a_{[(i-1) \times blocks]{[(j-1)\times blocks]}} & \cdots & a_{[(i-1) \times blocks]{}} \\\vdots & \ddots &\vdots \\a_{{[(j-1)\times blocks]}} & \cdots & a_{{}}\end{pmatrix}\]
其中的每个乘法 \(a'{ik} \times b'{kj}\) 是子矩阵乘法
这里为了方便看，默认原矩阵的行列均为 blocks 的倍数
那么下一个很自然的问题就是
若行列不为blocks的倍数，怎么办
分两种情况：
[*]$ min(m,n) < blocks $
[*]$ \exists m,n \nmid blocks $
对于1，直接执行普通矩阵乘法即可，因为整个矩阵均可放于L1、L2 Cache中
对于2，我们定义分块矩阵的大小为$ p,q $
\
\
其中，
\
\
至此，分块矩阵的全部问题已经解决
给出如下代码：
时间复杂度\(O(n^3)\)
void block_mult(float* A, float* B, float* C, int N, int BLOCK) {
// 清除结果矩阵
memset(C, 0, N*N*sizeof(float));
// 三层分块循环
for (int i0 = 0; i0 < N; i0 += BLOCK) {
int i_end = min(i0 + BLOCK, N);// 计算行边界
for (int k0 = 0; k0 < N; k0 += BLOCK) {
int k_end = min(k0 + BLOCK, N);// 计算中间维度边界
for (int j0 = 0; j0 < N; j0 += BLOCK) {
int j_end = min(j0 + BLOCK, N);// 计算列边界
// 核心计算：只处理完整块内的元素
for (int i = i0; i < i_end; i++) {
for (int k = k0; k < k_end; k++) {
float a_val = A;// 一次加载A元素
// 内层循环：连续访问B和C
for (int j = j0; j < j_end; j++) {
C += a_val * B;
}
}
}
}
}
}
}由于矩阵过小时，分块算法优势不大，且会增加调用开销，因此，这里的测试，\(m,n\)为2048
实测结果：\(blocks = 512\) 时，用时 11515.2ms
而不使用分块仅循环优化的算法 用时 16028.9ms
朴素实现 用时 18053.45ms
提升：16028.9ms-11515.2ms = \(4513.2ms\) 提高：\(28.1\%\)
高手过招 优化三 ：并行与SIMD
何为并行与SIMD？
并行：多线程同时处理多个分块
SIMD：乘加一体，即一条CPU指令同时处理乘与加
我们这里使用Apple的AMX(Apple Matrix协处理器)（也属于CPU的一部分，并非GPU优化）
对于x86架构和其余ARM架构的处理器，可以使用AVX、AVX_512、SSE等SIMP指令集
它的特性有：
[*]Apple Silicon芯片(M1/M2/M3等)内置的专用矩阵运算单元
[*]可并行处理大量16位浮点(FP16)或整数(INT8)运算
每个AMX单元包含：
[*]8个32KB的寄存器文件
[*]可同时执行2048次乘加运算(16x16x8矩阵)
[*]专用数据通路减少内存访问延迟
以及自动分块
如下是使用AMX的SIMP的代码：
时间复杂度：\(O(n^3)\)
// 使用Apple的AMX加速BLAS库
cblas_sgemm(CblasRowMajor, // 行主序存储
CblasNoTrans, // 不转置A
CblasNoTrans, // 不转置B
M, // A的行数
N, // B的列数
K, // 公共维度
1.0f, // alpha系数
a_data, // A数据指针
K, // A的列步幅（lda）
b_data, // B数据指针
N, // B的列步幅（ldb）
0.0f, // beta系数
r_data, // 结果数据指针
N); // 结果的列步幅（ldc）那么，本次优化最吓人、最恐怖的一次数据来了：
实测数据：202.831ms（4096*4096）
而标准分块+循环优化算法，在2048*2048时，就已经11515.2ms
提升：11312.37 ms 提高：98.2%
数学手段 优化4:Strassen算法 & 变种
温馨提示：到这里已经是高等数学内容了（实不相瞒，前面其实也是），有点小烧脑，不过欢迎各位继续跟作者一起尝试，本作者大约花了3小时搞完这一部分的证明
介绍：
Strassen算法是一种通过数学变换减少乘法次数的高效矩阵乘/卷积算法
简要推导：
我们设有如下两个矩阵相乘：
\[\begin{pmatrix}a_{11} & a_{12} \\a_{21} & a_{22}\end{pmatrix} \times \begin{pmatrix}b_{11} & b_{12} \\b_{21} & b_{22}\end{pmatrix} = \begin{pmatrix}c_{11} & c_{12} \\c_{21} &c_{22}\end{pmatrix}\]
传统计算需要8次乘法：
\
\
\
\
而Strassen算法只需7次乘
接下来，是Strassen算法最精妙绝伦的一步：
作者定义了7个矩阵：
\[\begin{align*}M_1 &= (a_{11} + a_{22})(b_{11} + b_{22}) \\M_2 &= (a_{21} + a_{22})b_{11} \\M_3 &= a_{11}(b_{12} - b_{22}) \\M_4 &= a_{22}(b_{21} - b_{11}) \\M_5 &= (a_{11} + a_{12})b_{22} \\M_6 &= (a_{21} - a_{11})(b_{11} + b_{12}) \\M_7 &= (a_{12} - a_{22})(b_{21} + b_{22})\end{align*}\]
真正让人惊讶的是下一步：
作者构建的7个矩阵，可以通过有限次的组合成为结果矩阵的一个元素
什么意思呢，让我们尝试展开其中一项：
\
\[\begin{align*}&=(a_{11} + a_{22})(b_{11} + b_{22}) + a_{22}(b_{21} - b_{11}) - (a_{11} + a_{12})b_{22} + (a_{12} - a_{22})(b_{21} + b_{22}) \\ \end{align*}\]
\[\begin{align*}= & \ \ a_{11}b_{11} + \cancel{a_{11}b_{22}} + \cancel{a_{22}b_{11}} + \cancel{a_{22}b_{22}} \\&+ \ \cancel{a_{22}b_{21}} - \cancel{a_{22}b_{11}} \\&- \ \cancel{a_{11}b_{22}} - a_{12}b_{22} \\&+ \ a_{12}b_{21} + \cancel{a_{12}b_{22}} - \cancel{a_{22}b_{21}} - \cancel{a_{22}b_{22}} \\= & \ \ a_{11}b_{11} + a_{12}b_{21}\end{align*}\]
由此，可以类似的推出\(c_{12}、c_{21}、c_{22}\)均与正常计算一致
最后的结果矩阵为
\
接下来，我们证明其对于n>2时仍然成立：
由于( n=2 )时，我们已经证明其正确
所以，在n>2时，我们采取分块
将原矩阵分为4块，此时我们将其中的子矩阵看为一个元素
那么此时又回归到了标准的2x2的Strassen算法
由此在$n,m \mid 2 $时Strassen算法正确
那么，下一个很自然的问题就是
若\(n,m \nmid {2} ,结论是否成立\)
我们的做法是，将矩阵分块，分为几个\(2^k \times 2^k\)的子矩阵以及几个符合矩阵乘规则的小矩阵
显然由于前文的分块算法的正确性，此时的分块仍然正确，对于几个\(2^k \times 2^k\)的矩阵，我们使用Strassen算法进行计算
现在来计算一下Strassen算法的时间复杂度：
设$ T(n) $ 为计算 $ n \times n $ 矩阵乘法的时间：
\(T(n) = 7T\left(\frac{n}{2}\right) + O(n^2)\)
[*]$ 7T(n/2) $：7 个子问题递归计算
[*]$O(n^2) $：矩阵加减法开销（共 18 次加减法）
根据主定理1，
$ a = 7,b = 2,f(n)=\Theta(n^2) $
\(log _b a = log_27 \approx 2.807\)
由于\(log _b a > 2\),所以\(f(n) = O({n^{log_b a-\epsilon})} = O(n^{log_27}) \approx O(n^{2.807})\)
更精确的，复杂度为\(\Theta(n^{2.807})\)
具体的算法为：
1.先将矩阵AB分块，分成大小为 \({blocks \times blocks}\) 的若干块以及几个任意大小的子块
2.对于 \({blocks \times blocks}\) 的子块，我们使用Strassen算法计算
3.在递归过程中，若方阵大小（因为Strassen算法开始时为方阵）n = 128，则使用循环优化的矩阵乘法
4.否则，继续按照Strassen算法递归计算直至n = 128
5.对于不是 \({blocks \times blocks}\) 的子块，使用循环优化的矩阵乘法计算
给出如下代码：
近似时间复杂度\(O(n^{2.81})\)
const int BLOCK_SIZE = 2048;const int STRASSEN_THRESHOLD = 128;// 标准矩阵乘法 (用于小矩阵和边界处理)void standard_matmul(const float* A, const float* B, float* C, int n, int m, int p, int lda, int ldb, int ldc) { for (int i = 0; i < n; ++i) { for (int k = 0; k < m; ++k) { float a = A; for (int j = 0; j < p; ++j) { C += a * B; } } }}// 循环优化矩阵乘法 (n=128时使用)void optimized_matmul(const float* A, const float* B, float* C, int n, int lda, int ldb, int ldc) { for (int i = 0; i < n; ++i) { for (int k = 0; k < n; ++k) { float a = A; for (int j = 0; j < n; ++j) { C += a * B; } } }}// 矩阵加法void matrix_add(const float* A, const float* B, float* C, int n, int lda, int ldb, int ldc) { for (int i = 0; i < n; ++i) { for (int j = 0; j < n; ++j) { C = A + B; } }}// 矩阵减法void matrix_subtract(const float* A, const float* B, float* C, int n, int lda, int ldb, int ldc) { for (int i = 0; i < n; ++i) { for (int j = 0; j < n; ++j) { C = A - B; } }}// 结果累加到目标矩阵void matrix_add_to_target(float* T, const float* S, int n, int ldt, int lds) { for (int i = 0; i < n; ++i) { for (int j = 0; j < n; ++j) { T += S; } }}// Strassen 矩阵乘法 (递归实现)void strassen_matmul(const float* A, const float* B, float* C, int n, int lda, int ldb, int ldc) { // 递归基: n
页:
[1]