具有列优先布局的 int8 x uint8 矩阵向量乘积

Question

我希望使用 AVX-1 或更早的指令来加速这个矩阵向量乘积：

    // a is an array N columns of length M each
    // b is length N
    // c is length M
    //
    // M % 32 == N % 32 == 0
    // all memory is nicely aligned and unaliased
    
    void mat_vec_prod(const int8_t** a, const uint8_t* b, int16_t* c) {
        for(int i = 0; i < M; ++i) {
            c[i] = 0;
            for(int j = 0; j < N; ++j)
                c[i] += int16_t(a[j][i]) * int16_t(b[j]);
        }
    }

（我知道交换循环值得考虑）

内在函数 _mm_maddubs_epi16 和 _mm_maddubs_pi16 可以帮助处理 uint8 x int8 点积，但在我的例子中，矩阵的布局很尴尬，它是指向列（而不是行）的指针数组。

一种可能性是加载a 的8x8 块，然后将它们转置并乘以b 的片段。 （我在 8x8 字节矩阵转置上找到了 this thread）。但是，这必须使用_mm_maddubs_pi16，它的吞吐量只有_mm_maddubs_epi16 的一半。

我的问题是：是否值得尝试加载和转置 16x16 补丁，或者我会用完 xmm 寄存器？我的策略应该是什么？

Answer 1

我会采用 chtz 建议的交错方法。

从两行中读取 32 或 64 个字节（也就是一个完整的缓存行），然后交错。

至少 32 字节，因为每行的宽度 % 32 == 0，最好是 64 字节，因为这是一个完整的高速缓存行，它将占用 16 个寄存器中的 8 个累加器。

另外我猜想将输入处理为（8、16 或 32 行）×（32 或 64 列）的块会比处理所有行更好；处理的行越多，将累加器溢出到内存的需求就越少，以非线性顺序处理的行越多，从缓存中驱逐即将需要的行的可能性就越高。 4 行应该是绝对安全的。

交错b 很自然地由

完成

auto b0to7 = _mm_unpacklo_epi16(b,b);
auto b8tof = _mm_unpackhi_epi16(b,b);
auto b01 = _mm_shuffle_epi32(b0to7, 0x00);
auto b23 = _mm_shuffle_epi32(b0to7, 0x55);
...
auto bef = _mm_shuffle_epi32(b8tof, 0xff);

另一种将输入拆分为偶数/奇数序列的可能性是每 16 字节需要 4 条算术指令，或每 32 字节需要 8 条指令：

// common terms
auto b_even = _mm_set1_epi16(b[j] & 0x00ff);
auto b_odd = _mm_set1_epi16(b[j] * 256);
// per 16 input bytes in `a`
auto mul_even = _mm_maddubs_epi16(row_j, b_even);
auto mul_odd = _mm_maddubs_epi16(row_j, b_odd);
sum_even = _mm_add_epi16(sum_even, mul_even);
sum_odd = _mm_add_epi16(mul_odd, mul_even);

这显然没有那么紧

auto prod_lo = _mm_unpacklo_epi8(row_j, row_jplus1);
auto prod_hi = _mm_unpackhi_epi8(row_j, row_jplus1);
prod_lo = _mm_maddubs_epi16(prod_lo, b01);
prod_hi = _mm_maddubs_epi16(prod_hi, b01);
sum_lo = _mm_add_epi16(sum_lo, prod_lo);
sum_hi = _mm_add_epi16(sum_hi, prod_hi);

但是 shuffle 只能在 Port5 上执行，因为 2 mul/adds 可以在每个周期开始。它们的性能可能非常接近。