AVX2 稀疏矩阵乘法

Question

我正在尝试利用新的 AVX2 GATHER 指令来加速稀疏矩阵 - 向量乘法。该矩阵采用 CSR（或 Yale）格式，带有指向列索引数组的行指针，而列索引数组又包含列。这种 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 内在函数来加速它。以下代码适用于基于 https://blog.fox-toolkit.org/?p=174 的最新 Intel 或 GCC。我在这里删除了其余部分，因为无论如何我的行都对齐了 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

Gather on Haswell 运行缓慢。我以几种不同的方式实现了 16 位值的 8 位索引 LUT 查找（对于 GF16 乘以 par2），以找出最快的方法。在 Haswell 上，VPGATHERDD 版本花费的时间是 movd / pinsrw 版本的 1.7 倍。（除了收集之外只需要几个 VPUNPCK / shift 指令。）code here, if anyone wants to run the benchmark.

在首次引入指令时很常见，他们不会投入大量硅来使其超快。它只是为了获得硬件支持，所以可以编写代码来使用它。为了在所有 CPU 上获得理想的性能，那么您需要执行 x264 为 pshufb 所做的事情：为 CPU 设置一个 SLOW_SHUFFLE 标志，例如 Core2，并将其考虑到您的寻找最佳例程的函数指针设置，而不仅仅是 CPU 支持的 insns。

对于不那么热衷于为每个 CPU 调整 asm 版本的项目，他们可以运行引入指令的无加速版本将使人们更快地使用它，所以当下一个设计出现了，它的速度更快，更多的代码加速了。发布像 Haswell 这样的设计，其中收集实际上是一种减速，这有点冒险。也许他们想看看人们会如何使用它？它确实增加了代码密度，这在收集不处于紧密循环中时很有帮助。

Broadwell 应该有一个更快的收集实现，但我无法访问它。列出 latency/throughput 说明的英特尔手册说 Broadwell 的收集速度大约快 1.6 倍，因此它仍然比 shifts/unpacks GP regs 中的索引并将它们用于PINSRW 转化为向量。

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

希望Skylake能更上一层楼。我以为我读过一些说它会的东西，但在检查时，我没有找到任何东西。

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

AVX2 稀疏矩阵乘法

AVX2 sparse matrix multiplication

c

intrinsics

sparse-matrix

matrix-multiplication

avx2