如何高效地加载和存储新的 AVX-VNNI 和 Arm Neon MMLA 指令的数据？

Question

最近的 AVX-VNNI 和 Arm Neon MMLA 指令加载数据的正确方法是什么？

例如SMMLA的描述为：

有符号 8 位整数矩阵乘法累加。该指令将第一个源向量中的有符号 8 位整数值的 2x8 矩阵乘以第二个源向量中的有符号 8 位整数值的 8x2 矩阵。生成的 2x2 32 位整数矩阵 [...]

同样，

_mm256_dpbusd_epi32

的描述是：

将 a 中 4 个相邻的无符号 8 位整数对与 b 中相应的有符号 8 位整数相乘，产生 4 个中间有符号 16 位结果。将这 4 个结果与 src 中相应的 32 位整数相加，并将打包的 32 位结果存储在 dst 中。

它们似乎都需要 2[4]x8 和 8x[4]2 形式的输入。并产生 2[4]x[4]2 形式的输出。如何有效地加载和存储这些函数的数据？

我看到使用这些说明的三种广泛可能性，但没有一种有吸引力：

• [拆分和组合]我加载两个连续的128位向量，然后拆分它们。同样，对于 AVX，我会加载 4 128 或 256 个向量，然后拆分它们。存储同样“复杂”，因为我需要在存储之前提取 2[4]x[4]2 矩阵的相关部分。我的代码充满了拆分/合并指令。
• [较小的向量] 或者，我可以加载较小的部分，但这似乎也效率低下。
• [重新排序输入数据] 当然，我可以重新排序输入数据，以便矢量化负载已经跨越多行或多列。这应该是预期用途吗？

小型 4xK 输入矩阵 A（行主）和 Kx4 矩阵 B（列主）的内循环（K 上的约简）示例代码如下：

for (size_t k = 0; k < 64; k += 8) {
    uint8x8_t low = vld1_u8(row0);
    uint8x8_t high = vld1_u8(row1);
    uint8x16_t row01x01234567 = vcombine_u8(low, high);
    row0 += 8;
    row1 += 8;
    low = vld1_u8(row2);
    high = vld1_u8(row3);
    uint8x16_t row23x01234567 = vcombine_u8(low, high);
    row2 += 8;
    row3 += 8;
    low = vld1_u8(col0);
    high = vld1_u8(col1);
    uint8x16_t col01x01234567 = vcombine_u8(low, high);
    col0 += 8;
    col1 += 8;
    low = vld1_u8(col2);
    high = vld1_u8(col3);
    uint8x16_t col23x01234567 = vcombine_u8(low, high);
    col2 += 8;
    col3 += 8;
    out01x01 = vmmlaq_u32(out01x01, row01x01234567, col01x01234567);
    out01x23 = vmmlaq_u32(out01x23, row01x01234567, col23x01234567);

    out23x01 = vmmlaq_u32(out23x01, row23x01234567, col01x01234567);
    out23x23 = vmmlaq_u32(out23x23, row23x01234567, col23x01234567);
}

结果是正确的，但效率似乎非常低。上面的代码只是一个例子。实际上，我会使用更大的图块尺寸来最大化寄存器的使用。

Answer 1

对矩阵A和B进行打包确实是必要的。

有关简短的概述，请考虑 PowerPC 文档（红皮书）。 https://www.redbooks.ibm.com/abstracts/redp5612.html（第 35 页）。 PowerPC 具有与 VNNI 和 Arm Neon 类似的分块矩阵乘法指令。

我已经在矩阵乘法代码中编写了这样的打包函数。打包不会对代码的吞吐量造成任何损害。