如何为AVX-VNNI与Arm Neon MMLA指令高效加载存储数据？

核心问题分析 #

这类矩阵乘累加指令（比如Arm的SMMLA、x86的_mm256_dpbusd_epi32）对输入数据布局有明确要求：例如SMMLA需要2×8和8×2的矩阵输入，生成2×2的32位结果；_mm256_dpbusd_epi32要求4组相邻无符号8位对与对应有符号8位数据匹配。常规行/列主序存储无法直接适配指令的向量加载需求，这是效率瓶颈的核心原因。

三种方案的优劣与最优选择 #

你提到的三种方案中，重排输入数据是指令设计预期的高效用法，另外两种方案的问题如下：

• 拆分与合并：频繁的拆分、合并指令会占用大量CPU周期，抵消向量化的性能收益，同时代码冗余、维护成本高。
• 更小向量加载：多次小数据块加载会增加内存访问次数，破坏缓存局部性，导致内存带宽利用率下降，效率极低。

重排数据的具体思路 #

重排的核心是让输入数据的布局匹配指令的向量加载要求，将原本分散在不同行/列的数据整合到连续的向量加载范围内，避免运行时的拆分重组：

1. 预处理阶段重排
   • 对于行主序矩阵A，将每2行的同位置8个元素打包成连续内存块，后续可直接用vld1q_u8这类指令一次性加载，无需运行时合并。
   • 对于列主序矩阵B，将每2列的同位置8个元素打包成连续内存块，同样支持直接向量加载。
2. 内循环简化
   重排后，内循环的加载操作会大幅简化，减少冗余指令，同时最大化寄存器的有效使用率。

优化后的示例代码 #

// 假设已完成预处理：A的行数据按2行一组打包，B的列数据按2列一组打包
for (size_t k = 0; k < 64; k += 8) {
    // 直接加载预打包的2行×8列数据
    uint8x16_t row01 = vld1q_u8(a_ptr);
    uint8x16_t row23 = vld1q_u8(a_ptr + 16);
    a_ptr += 32;

    // 直接加载预打包的8行×2列数据
    uint8x16_t col01 = vld1q_u8(b_ptr);
    uint8x16_t col23 = vld1q_u8(b_ptr + 16);
    b_ptr += 32;

    // 直接执行矩阵乘累加
    out01x01 = vmmlaq_u32(out01x01, row01, col01);
    out01x23 = vmmlaq_u32(out01x23, row01, col23);
    out23x01 = vmmlaq_u32(out23x01, row23, col01);
    out23x23 = vmmlaq_u32(out23x23, row23, col23);
}

额外优化建议 #

• 缓存适配分块：结合CPU的L1/L2缓存大小设计分块，确保处理的数据块能完全放入缓存，减少缓存 miss。
• 寄存器复用：利用更多寄存器存储中间结果，减少内存回写次数，进一步提升效率。
• 批量预处理：一次性完成整个矩阵的重排，避免在循环中零散处理。

内容的提问来源于stack exchange，提问作者fabian