AVX2稀疏矩阵乘法

Question

我正在尝试利用新的AVX2 GATHER指令来加速稀疏矩阵 - 向量乘法。 矩阵采用CSR（或耶鲁）格式，行指针指向列索引数组，而列索引数组又保存列。 这样一个mat-vec mul的C代码看起来像这样：

for (int row = 0; row < n_rows - 1; row++) {
    double rowsum = 0;
    for (int col = row_ptr[row]; col < row_ptr[row + 1]; col++) {
        rowsum += values[col] * x[col_indices[col]];
    }
    result[row] = rowsum;
}

现在我的目标是通过AVX2内在函数加速这一点。 以下代码适用于最新的英特尔或GCC，基于https://blog.fox-toolkit.org/?p=174 。 我删除了剩余部分，因为我的行都在4个双打（列％4 == 0）上对齐（幸运的是我）。 如果有人感兴趣，我也有处理余数的代码，但关键是，代码实际上稍慢。 我检查了反汇编，对于上面的版本，只生成FP指令，对于我的AVX2代码，所有AVX2操作都按预期显示。 即使使用适合缓存的小矩阵，AVX2版本也不行。 我在这里很困惑......

double* value_base = &values[0];
double* x_base = &x[0];
int*    index_base = &col_indices[0];


for (int row = 0; row < n_rows - 1; row++) {
    int col_length   = row_ptr[row + 1] - row_ptr[row];

    __m256d rowsum = _mm256_set1_pd(0.);
    for (int col4 = 0; col4 < col_length; col4 += 4) {
        // Load indices for x vector(const __m128i*)
        __m128i idxreg     = _mm_load_si128((const __m128i*)index_base);
        // Load 4 doubles from x indexed by idxreg (AVX2)
        __m256d x_     = _mm256_i32gather_pd(x_base, idxreg, 8);
        // Load 4 doubles linear from memory (value array)
        __m256d v_     = _mm256_load_pd(value_base);
        // FMA: rowsum += x_ * v_
        rowsum = _mm256_fmadd_pd(x_, v_, rowsum);

        index_base += 4;
        value_base += 4;
    }
    __m256d s = _mm256_hadd_pd(rowsum, rowsum);
    result[row] = ((double*)&s)[0] + ((double*)&s)[2];
    // Alternative (not faster):
    // Now we split the upper and lower AVX register, and do a number of horizontal adds
    //__m256d hsum = _mm256_add_pd(rowsum, _mm256_permute2f128_pd(rowsum, rowsum, 0x1));
    //_mm_store_sd(&result[row], _mm_hadd_pd( _mm256_castpd256_pd128(hsum), _mm256_castpd256_pd128(hsum) ) );
}

欢迎任何建议。

非常感谢，克里斯

Answer 1

Haswell的聚会很慢。 我以几种不同的方式实现了16位值的8位索引LUT查找（对于par2的GF16乘法），以找出最快的。 在Haswell上， VPGATHERDD版本的VPGATHERDD是VPGATHERDD movd / pinsrw版本的movd / pinsrw 。 （ 如果有人想要运行基准测试，那么这里只需要几个VPUNPCK /移位指令。） 代码 。

正如在第一次引入指令时常见的那样，它们并没有投入大量的硅来使其超快。 只是为了获得硬件支持，所以可以编写代码来使用它。 为了在所有CPU上实现理想的性能，那么你需要做x264为pshufb做的pshufb ：对像Core2这样的CPU有一个SLOW_SHUFFLE标志，并将其纳入你的最佳例程查找函数 - 指针设置，而不仅仅是确定CPU的内容支持。

对于那些不太热衷于为每个可以运行的CPU调整asm版本的项目，引入一个非加速版本的指令将让人们更快地使用它，所以当下一个设计出现时，它的快速，更多的代码加速。 释放像哈斯威尔这样的设计实际上是一种减速，这有点冒险。 也许他们想看看人们会如何使用它？ 它确实增加了代码密度，这有助于聚集不紧密循环。

Broadwell应该有更快的聚集实现，但我没有访问权限。 列出指令延迟/吞吐量的英特尔手册称，Broadwell的聚集速度大约快1.6倍，因此它仍然比手动制作的循环稍微慢一点，这种循环可以在GP regs中移位/解包索引，并将它们用于PINSRW到矢量。

如果gather可以利用多个元素具有相同索引的情况，或者甚至是指向相同32B获取块的索引，则根据输入数据可能会有一些大的加速。

希望Skylake能够进一步提升。 我以为我会读到一些东西，说它会，但在检查时，我没有找到任何东西。

RE：稀疏矩阵：是不是有重复数据的格式，所以你可以对行或列进行连续读取？ 这不是我必须编写代码的东西，但我想我已经在一些答案中看到过它。