使用int8 AVX512-VNNI指令的有效方法，尤其是将数据加载到zmm寄存器

Question

我想在 int8 数据中使用 AVX512-VNNI 指令优化矩阵乘法运算。

我明白这个

vpdpbusd

是如何工作的，但我不知道如何有效地使用它。

具体来说，我将大矩阵（MxN）划分为更小的矩阵，“较小矩阵”的大小为 4x4，其中一个

vpdpbusd

只需一条指令即可处理。

您还应该注意到，我的矩阵乘法算法是基于“GotoBLAS”中使用的算法。

通过这种方式，我想我可以用多个

4x4

X

4x4

矩阵乘法来完成整个矩阵乘法（某种“分而治之”）。

详细来说，

1. 加载一个

   4x4

   矩阵（比方说

   A

   ）四次到 1

   zmm

   注册
2. 将四个不同的

   4x4

   矩阵（例如

   B

   、

   C

   、

   D

   、

   E

   ）加载到另一个

   zmm

   寄存器。

3. vpdpbusd

   一次进行四次乘法：

   AxB

   、

   AxC

   、

   AxD

   、

   AxE

   。
4. 存储结果。

但问题是将数据加载到

zmm

寄存器。

4x4

矩阵在长

zmm

寄存器上加载非常复杂。

此外，

vpdpbusd

会产生 32 位数据。所以我需要将其转换回8位数据。我认为这可能还会产生另一项开销。

我不知道如何处理这个问题...救命！

Answer 1

如果你的矩阵确实很大（足够大），出于 TLB 的考虑，重新打包图块（也称为“打包”，同样的事情）是值得做的（不重新打包的图块会导致跨越许多不同页面的图块，并且可能会花费大量 TLB错过访问）。无论如何，当您重新包装时，您可以选择对元素进行任何重新排序，而无需为此支付太多额外费用，它并不比重新包装本身贵得多。当然，重新包装是有成本的，但包装好的瓷砖会被重复使用，因此这笔成本可以分摊到多次使用中。这是安排计算的方法：

• 从矩阵 A (4x

  vpbroadcastd

  ) 加载并广播 4 个元素（一小行）的 4 个块。
• 从矩阵 B 的重新打包块中加载 4 个块，其中每个块中的第

  i

  ' 个双字是第

  i

  ' 列中的 4 个元素的小列。 4 个块总共跨越 16 行。听起来很糟糕，但这里没有复杂的洗牌，这是在重新包装过程中发生的。
• 形成所有 16 个乘积，将它们求和为 16 个独立累加器。在每个 DWORD 中，

  vpdpbusds

  计算 A 中的一小行和 B 中的一小列之间的点积。以这种方式计算的 16 个乘积都重用了 A 中的相同小行，但 B 中的不同列。

结果不会形成输出矩阵的 4x4 子矩阵，而是形成 16 宽的行，可以将它们相加然后放入矩阵中（在您的情况下也缩小到 8 位，但正如您所看到的，这发生在内循环，所以影响很小（如果有的话）。

如果您的矩阵太小，以至于重新打包的开销大于其回报，那么我不确定安排计算的好方法是什么。