我正在尝试利用新的 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 内在函数加速这一过程。以下代码适用于最新的 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) ) );
}
欢迎任何建议。
非常感谢,克里斯