提问者:小点点

AVX2稀疏矩阵乘法


我试图利用新的AVX2收集指令来加速稀疏矩阵向量乘法。矩阵采用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 intrinsics来加速这一点。以下代码适用于最新的Intel或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) ) );
}

欢迎提出任何建议。

多谢,克里斯


共1个答案

匿名用户

在哈斯韦尔集合很慢。我用几种不同的方法实现了一个16位值的8位索引LUT查找(GF16乘以par2),以找出最快的方法。在Haswell上,版本所用时间是版本的1.7倍。(在Gathers之外只需要几条/shift指令。)这里的代码,如果有人想要运行基准。

正如第一次引入指令时所常见的那样,他们不会投入大量的硅来使其超速。它只是为了得到硬件支持,所以可以编写代码来使用它。为了在所有CPU上实现理想的性能,您需要做x264为所做的事情:为Core2这样的CPU设置标志,并将其考虑到您的最佳例程查找函数-指针设置中,而不是只考虑CPU支持的insns。

对于那些不太热衷于为每一个CPU调整asm版本的项目来说,引入一个无加速比的指令版本会让人们更快地使用它,这样当下一个设计出现时,更快,更多的代码会加速。发布像Haswell这样的设计,而gather实际上是一种放慢速度的做法是有点冒险的。也许他们想看看人们会如何使用它?它确实增加了代码密度,当收集不是在一个紧密的循环中时,这会有所帮助。

Broadwell应该有一个更快的gather实现,但是我没有访问的权限。Intel手册列出了延时/吞吐量的说明,它说Broadwell's gather的速度大约是1.6倍,因此它仍然比手工构建的循环要慢一些,后者将GP regs中的索引移位/解包,并将它们用于向量。

如果能够利用多个元素具有相同索引的情况,甚至是指向相同32B提取块的索引,那么根据输入数据,可能会有一些大的加速。

希望Skylake能进一步改善。我以为我读到什么东西说它会,但一查,我什么也没发现。

RE:稀疏矩阵:难道不存在一种重复数据的格式,所以你可以对行或列进行连续读取吗?这不是我必须为之编写代码的东西,但我想我已经在一些答案中看到提到过它。