资讯动态

外积计算重构 GEMM:大幅降低主存加载次数的算法推导与向量化实操

发布时间:2026/10/9 12:36:54 来源:尧图企业网站定制
在通用矩阵乘法GEMM, $C A \times B$的经典教科书算法中几乎所有人都习惯了以“内积Inner Product / Dot Product”的心智模型去理解计算结果矩阵 $C$ 的每一个元素 $C_{i,j}$是矩阵 $A$ 的第 $i$ 行向量与矩阵 $B$ 的第 $j$ 列向量的点积标量和。然而当我们在现代超标量处理器上将内积模型推向极致时很快会遭遇一个难以逾越的微架构瓶颈在计算单点点积时累加操作始终汇聚在单个寄存器中且必须沿着 $K$ 维度频繁在内存中按列跳跃跨步寻址引发大量的非连续内存访问或 Gather 指令更致命的是每一对行与列的点积计算完成后中间累加值就被固化写出导致矩阵 $A$ 和 $B$ 的元素无法在寄存器内部获得最大化的空间复用。真正统治现代工业级 BLAS 库如 GotoBLAS、BLIS 以及各大顶级 AI 推理底座的黄金计算范式是外积Outer Product更新模型。本文我们将从数学公式出发严格推导外积计算如何将矩阵乘法解构为一系列秩-1Rank-1矩阵的累加并运用现代 C 和 AVX2 向量广播指令手写一个大幅压降主存加载次数的高性能外积 GEMM 微内核。一、从内积到外积矩阵乘法的代数视角重塑我们先对比两种完全不同的代数计算视角。假设矩阵 $A \in \mathbb{R}^{M \times K}, B \in \mathbb{R}^{K \times N}, C \in \mathbb{R}^{M \times N}$1. 经典内积视角Dot Product Perspective$$C_{i,j} \sum_{k1}^{K} A_{i,k} \cdot B_{k,j} \text{Row}_i(A) \cdot \text{Col}_j(B)$$视角特征固定结果矩阵的某一个元素坐标 $(i, j)$沿着共享维度 $K$ 进行整条维度的规约求和访存代价为了算出一个 $C_{i,j}$ 标量必须将 $A$ 的一行和 $B$ 的一列完整读取一遍。两个长为 $K$ 的向量只产出了 1 个浮点输出访存与计算之比极差。2. 外积视角Outer Product Perspective / Sum of Rank-1 Matrices我们将矩阵 $A$ 视为由 $K$ 个列向量组成的集合$A [\mathbf{a}_1, \mathbf{a}_2, \dots, \mathbf{a}_K]$其中每个 $\mathbf{a}_k \in \mathbb{R}^{M \times 1}$将矩阵 $B$ 视为由 $K$ 个行向量组成的集合$B [\mathbf{b}_1^T; \mathbf{b}_2^T; \dots; \mathbf{b}_K^T]$其中每个 $\mathbf{b}_k^T \in \mathbb{R}^{1 \times N}$。根据分块矩阵乘法法则矩阵乘积 $C$ 可以等价重写为$K$ 个外积矩阵的直接累加$$C \sum_{k1}^{K} \mathbf{a}_k \otimes \mathbf{b}k^T \sum{k1}^{K} \mathbf{a}_k \mathbf{b}_k^T$$审视这个公式的物理美感对于给定的某一个 $k$$1 \le k \le K$$\mathbf{a}_k$ 是一个长为 $M$ 的列向量$\mathbf{b}_k^T$ 是一个长为 $N$ 的行向量两者的外积 $\mathbf{a}_k \mathbf{b}_k^T$ 瞬间生成一个尺寸为$M \times N$ 的完整二维矩阵最终的矩阵 $C$就是这 $K$ 个 $M \times N$ 矩阵的逐元素累加和。二、为什么外积模型能够降低主存加载次数从计算机微架构的物理视角来看外积模型带来了三大颠覆性的访存优势结果矩阵 $C$ 的累加块长期常驻寄存器考虑一个 $4 \times 16$ 的寄存器微分块Tile。我们可以用 8 个 256 位ymm向量寄存器牢牢锁住 $C$ 矩阵的一个 $4 \times 16$ 局部子块随着 $k$ 从 1 到 $K$ 逐步推进这 8 个寄存器在整整 $K$ 次迭代中永远不需要向内存写回任何中间数据始终保持纯净的原地乘加矩阵 $B$ 的访存变为绝对的连续行加载Stride-1在外积模型中第 $k$ 步需要读取的是行向量 $\mathbf{b}_k^T$在 C 语言默认的行优先Row-Major存储下$\mathbf{b}_k^T$ 的所有元素在物理内存中就是紧挨在一起的整行我们只需要用_mm256_loadu_ps进行极速连续向量加载彻底消除了列优先或 Gather 跨步的噩梦。矩阵 $A$ 的访存转化为寄存器标量广播Broadcast第 $k$ 步中$\mathbf{a}_k$ 的每一个元素只需读取一次然后使用硬件广播指令如_mm256_set1_ps复制到整个向量寄存器的所有通道中1 次标量加载即可与加载的整行向量并行完成 8 次浮点乘加计算访存比Arithmetic Intensity达到了理论上限。三、基于 AVX2 向量广播的外积微内核现代 C 实现我们用现代 C 编写一个纯外积驱动的 $4 \times 16$ 矩阵乘法微内核#include immintrin.h #include span #include cstdint #include iostream // 外积核心微内核计算 C[4x16] A[4xK] * B[Kx16] inline void outer_product_gemm_micro_kernel_4x16( const float* A, // 指向 A 的某个 4 行子块首地址步长为 lda const float* B, // 指向 B 的某个 16 列子块首地址步长为 ldb float* C, // 指向 C 的 4x16 目标子块步长为 ldc size_t lda, size_t ldb, size_t ldc, size_t K ) noexcept { // 1. 初始化 8 个 256 位向量累加器4 行每行 16 个 float占 2 个 ymm 寄存器 // 从 C 内存加载初始值 __m256 c00 _mm256_loadu_ps(C 0 * ldc 0); __m256 c01 _mm256_loadu_ps(C 0 * ldc 8); __m256 c10 _mm256_loadu_ps(C 1 * ldc 0); __m256 c11 _mm256_loadu_ps(C 1 * ldc 8); __m256 c20 _mm256_loadu_ps(C 2 * ldc 0); __m256 c21 _mm256_loadu_ps(C 2 * ldc 8); __m256 c30 _mm256_loadu_ps(C 3 * ldc 0); __m256 c31 _mm256_loadu_ps(C 3 * ldc 8); // 2. 外积推进主循环沿着 K 维度一步一步发射外积更新 for (size_t k 0; k K; k) { // 2.1 连续加载 B 矩阵第 k 行的 16 个元素2 条向量加载指令 __m256 b_row_part0 _mm256_loadu_ps(B k * ldb 0); __m256 b_row_part1 _mm256_loadu_ps(B k * ldb 8); // 2.2 广播加载 A 矩阵第 k 列的 4 个元素并原地执行 FMA 外积乘加 // 处理 A 的第 0 行标量 __m256 a0 _mm256_set1_ps(A[0 * lda k]); c00 _mm256_fmadd_ps(a0, b_row_part0, c00); c01 _mm256_fmadd_ps(a0, b_row_part1, c01); // 处理 A 的第 1 行标量 __m256 a1 _mm256_set1_ps(A[1 * lda k]); c10 _mm256_fmadd_ps(a1, b_row_part0, c10); c11 _mm256_fmadd_ps(a1, b_row_part1, c11); // 处理 A 的第 2 行标量 __m256 a2 _mm256_set1_ps(A[2 * lda k]); c20 _mm256_fmadd_ps(a2, b_row_part0, c20); c21 _mm256_fmadd_ps(a2, b_row_part1, c21); // 处理 A 的第 3 行标量 __m256 a3 _mm256_set1_ps(A[3 * lda k]); c30 _mm256_fmadd_ps(a3, b_row_part0, c30); c31 _mm256_fmadd_ps(a3, b_row_part1, c31); } // 3. 循环彻底结束后一次性将最终结果写回内存 _mm256_storeu_ps(C 0 * ldc 0, c00); _mm256_storeu_ps(C 0 * ldc 8, c01); _mm256_storeu_ps(C 1 * ldc 0, c10); _mm256_storeu_ps(C 1 * ldc 8, c11); _mm256_storeu_ps(C 2 * ldc 0, c20); _mm256_storeu_ps(C 2 * ldc 8, c21); _mm256_storeu_ps(C 3 * ldc 0, c30); _mm256_storeu_ps(C 3 * ldc 8, c31); }四、指令流水线与算力账本核算让我们来精确统计这套外积微内核在每一个时钟周期内的物理动作每次 $k$ 循环迭代内存读取开销仅仅 2 次 256 位向量连续加载来自 $B$ 4 次标量广播加载来自 $A$总计从 L1 缓存读取 $64 16 80$ 字节浮点计算产出执行整整 8 条_mm256_fmadd_ps指令。每条指令完成 8 次乘加16 FLOPs总计完成 $8 \times 16 128$ 次浮点计算计算访存比$128 / 80 1.6$ FLOP/Byte。在 Intel Skylake/Icelake 或 AMD Zen4 架构上CPU 核心配备了 2 个全功能的 FMA 执行端口Port 0 与 Port 1每个周期可以同时发射并退役 2 条 256 位 FMA 指令。这意味着上述 8 条 FMA 指令仅需4 个时钟周期即可被算力核心全速吞吐完毕与此同时处理器的加载单元Load Ports在后台无缝将下一次迭代所需的 $B$ 向量预取就绪计算与访存形成了天衣无缝的双缓冲重叠流水线。五、算子工程落地总结外积计算模型的精髓在于彻底颠覆了“按点算矩阵”的传统思维局限。它把矩阵乘法还原为寄存器空间内的二维张量膨胀让 CPU 的通用向量寄存器真正充当起了超微型、单周期访问的片上 SRAM。搞懂了外积模型你就拿到了通往工业级 BLAS 内核调优的入场券。在接下来的深水区中我们将以此为基石进一步把分块规模扩展到 AVX-512 的 32 个寄存器空间中开启对更高维度算力的极致压榨。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑