6

新しい 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 ops が期待どおりに表示されます。キャッシュに収まる小さな行列であっても、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) ) );
}

どんな提案でも大歓迎です。

どうもありがとう、クリス

4

1 に答える 1