LP-GEMM 技术总结：将布局传播集成到 GEMM 操作中

1. 研究背景与问题 (Problem)

在现代科学计算和机器学习（ML）工作负载中，连续的依赖通用矩阵乘法（GEMM）序列往往占据了主要的执行时间。尽管现有的顶级 BLAS 库（如 OpenBLAS、Intel MKL）对单个 GEMM 调用进行了极致优化，但它们受限于标准的 BLAS API 规范。

核心痛点：

• 冗余的数据打包与解包（Redundant Packing/Unpacking）： 标准 BLAS 库要求每个 GEMM 调用独立地将输入矩阵打包（Packing）成适合微内核（Micro-kernel）计算的特定内存布局，并在计算完成后将输出解包（Unpacking）回标准的行/列主序布局。
• 流水线效率低下： 在连续的 GEMM 操作中（例如 Transformer 或 MLP 层），前一个操作的输出立即作为下一个操作的输入。标准库会在每个边界处执行“解包 -> 重新打包”的过程，导致大量的内存带宽浪费和计算资源闲置，而这些操作对于连续计算而言是多余的。
• 缺乏跨操作优化： 现有的优化主要集中在单个内核的微架构调优上，忽略了操作序列之间的数据布局连续性。

2. 方法论 (Methodology)

本文提出了 LP-GEMM (Layout Propagation GEMM)，一种对 OpenBLAS GEMM 内核的分解方案，旨在通过**布局传播（Layout Propagation）**消除连续 GEMM 操作间的冗余数据移动。

核心设计思想

LP-GEMM 打破了 BLAS 必须返回标准内存布局的限制，允许数据在连续操作间保持优化的内部打包布局。它将传统的 GEMM 操作分解为三个专门化的内核：

1. 初始 GEMM (Ini-GEMM)：

   • 负责接收原始数据并进行初始打包。
   • 执行计算，并将结果以**传播布局（Propagated Layout）**存储，而不是标准布局。
   • 这是布局传播链的起点。

2. 中间 GEMM (Mid-GEMM)：

   • 接收前一个操作已打包好的数据。
   • 跳过打包步骤，直接利用现有的布局进行计算。
   • 计算结果继续以相同的传播布局存储，供下一个操作使用。
   • 这是消除冗余开销的关键环节。

3. 结束 GEMM (End-GEMM)：

   • 负责传播链的终止。
   • 执行计算，并将最终结果解包回标准的 BLAS 内存布局，以便后续非 GEMM 操作或用户访问。

技术实现细节

• 微内核设计： 定义了两种微内核：
  • Propagate-Layout µkernel： 维持传播布局，计算顺序与存储顺序一致，最大化空间局部性。
  • Default µkernel： 用于结束传播，将数据重新排序为标准布局。
• 内存布局统一： 强制统一了三个要素：(1) 读取输入的打包布局，(2) 计算时的迭代顺序，(3) 最终存储的布局。这使得输出可以直接作为下一个操作的输入，无需中间转换。
• 步长处理 (Strided Loads/Stores)： 为了处理非连续内存访问（如分块矩阵），引入了额外的参数来定义数据块的存储步长，确保在保持布局传播的同时支持复杂的矩阵操作。

3. 关键贡献 (Key Contributions)

1. LP-GEMM 分解架构： 提出了一种将 OpenBLAS GEMM 操作分解为 ini-GEMM、mid-GEMM 和 end-GEMM 的新范式，有效利用了数据布局传播。
2. Llama-3.2 独立实现： 使用纯 C++ 和 BLAS 级 GEMM 调用，独立实现了 Llama-3.2 的推理路径（包括 Attention 层），验证了该方法在实际大模型中的可行性。
3. 其他算子的适配： 在 AVX-512 架构上实现了 Softmax、RoPE（旋转位置编码）和 RMSNorm 等常见 ML 算子，使其能够直接利用 LP-GEMM 的传播布局，减少了中间转换开销。
4. 跨架构验证： 在 x86 (AVX-512) 和 RISC-V (RVV 1.0) 两种截然不同的架构上进行了广泛评估，证明了该方法的通用性。

4. 实验结果 (Results)

研究在 Intel Xeon Gold 6252 (x86) 和 SpacemiT X60 (RISC-V) 平台上进行了评估，对比对象包括 OpenBLAS、Intel MKL、BLIS 和 FlashGEMM。

• 连续 GEMM 性能提升：

  • 在 x86 架构上，针对连续 GEMM 序列，LP-GEMM 相比 OpenBLAS 实现了平均 2.25 倍 的加速。
  • 在 RISC-V 架构上，针对 Attention 类工作负载，加速比高达 5 倍。
  • 对于中间内核（Mid-GEMM），由于消除了打包开销，相比 OpenBLAS 实现了约 1.5x 到 2x 的中位数加速。

• 与商业库对比：

  • 在 x86 上，LP-GEMM 的中间/结束内核性能与高度优化的 Intel MKL 相当，甚至在某些场景下更具优势，尽管 MKL 的初始内核优化更极致。
  • 在 RISC-V 上，LP-GEMM 显著优于其他现有方案，因为 OpenBLAS 在 RISC-V 上的解包开销较大，而 LP-GEMM 完全避免了这一点。

• 实际工作负载 (Llama-3.2)：

  • 在 Llama-3.2 的 Attention 层中，LP-GEMM 在小输入尺寸下实现了高达 2 倍 的加速。
  • 随着序列长度增加，RISC-V 上的加速比呈线性增长趋势，这归因于 LP-GEMM 避免了 OpenBLAS 中随矩阵增大而恶化的非连续内存访问开销。

• 连续 GEMM 基准测试：

  • 在由三个连续 GEMM 组成的基准测试中（模拟 DNN 瓶颈块），LP-GEMM 在大多数测试用例中均优于 OpenBLAS 和 FlashGEMM。

5. 意义与结论 (Significance & Conclusion)

• 范式转变： LP-GEMM 证明了从“优化单个孤立内核”向“优化端到端工作负载流”转变的巨大潜力。它表明，通过理解操作序列间的交互，可以挖掘出比单纯微架构调优更多的性能空间。
• 消除冗余： 研究明确指出，在连续矩阵乘法中，**打包/解包（Packing/Unpacking）**往往是主导成本，而非微内核本身的计算。LP-GEMM 通过消除这些冗余步骤，显著提升了计算效率。
• 实用性与兼容性： 虽然 LP-GEMM 不完全兼容标准 BLAS API（因为输出布局不同），但其对现有代码的侵入性很小。通过引入少量参数（如步长），即可集成到现有的 ML 框架中。
• 未来方向： 该方法不仅适用于 GEMM，还暗示了将布局传播扩展到卷积（Convolutions）等其他算子的可能性，为编译器驱动的优化框架提供了新的思路。

总结： LP-GEMM 是一种高效、通用的优化技术，通过打破 BLAS 的布局限制，利用数据在连续操作间的自然连续性，显著降低了内存带宽压力并提升了计算吞吐量，特别是在资源受限的 RISC-V 边缘设备和大规模 Transformer 模型推理中表现卓越。