AVX命令セットを活用した1次元畳み込み演算の高速化技術

AVXによる1次元畳み込み演算の最適化

基本的なAVX最適化実装

// conv_avx_optimized.cpp

bool Execute1DConvolution_AVX(double* __restrict__ output, const double* __restrict__ input, const double* __restrict__ filter, int64_t data_length) {
    constexpr int64_t filter_width = 5;
    constexpr int64_t half_width = filter_width / 2;

    if (data_length < filter_width) {
        return false;
    }

    // フィルタ係数のAVXレジスタへのロード
    __m256d coeff0 = _mm256_set1_pd(filter[0]);
    __m256d coeff1 = _mm256_set1_pd(filter[1]);
    __m256d coeff2 = _mm256_set1_pd(filter[2]);
    __m256d coeff3 = _mm256_set1_pd(filter[3]);
    __m256d coeff4 = _mm256_set1_pd(filter[4]);

    // AVX命令セットを使用した畳み込み処理
    for (int64_t idx = half_width; idx <= data_length - filter_width; idx += 4) {
        __m256d data0 = _mm256_loadu_pd(&input[idx + 2]);
        __m256d data1 = _mm256_loadu_pd(&input[idx + 1]);
        __m256d data2 = _mm256_loadu_pd(&input[idx]);
        __m256d data3 = _mm256_loadu_pd(&input[idx - 1]);
        __m256d data4 = _mm256_loadu_pd(&input[idx - 2]);

        // 積和演算の実行
        __m256d result_val = _mm256_add_pd(
            _mm256_add_pd(
                _mm256_mul_pd(data0, coeff0),
                _mm256_mul_pd(data1, coeff1)),
            _mm256_add_pd(
                _mm256_mul_pd(data2, coeff2),
                _mm256_add_pd(
                    _mm256_mul_pd(data3, coeff3),
                    _mm256_mul_pd(data4, coeff4))));

        // 結果の保存
        _mm256_storeu_pd(&output[idx], result_val);
    }

    return true;
}

実行結果の考察:

double型データを使用する場合、_mm256_load_pd_mm256_store_pdはメモリアドレスがアライメントされていないためセグメンテーションフォールトを引き起こす可能性があります。そのため、非アライメントアクセスをサポートする_mm256_loadu_pd_mm256_storeu_pdを使用してデータにアクセスする必要があります。また、畳み込み操作ではFMA(Fused Multiply-Add)命令を活用することでさらに最適化できます。

AVX2はAVXに比べてベクトル整数演算のサポートを追加し、gather命令を導入しました。これにより、非連続なメモリ位置からデータをレジスタにロードできます。また、ベクトルビット操作のサポートを追加し、FMA操作を強化しています。

FMA命令による最適化

// conv_avx_fma_optimized.cpp

bool Execute1DConvolution_AVX_FMA(double* __restrict__ output, const double* __restrict__ input, const double* __restrict__ filter, int64_t data_length) {
    constexpr int64_t filter_width = 5;
    constexpr int64_t half_width = filter_width / 2;

    if (data_length < filter_width) {
        return false;
    }

    // フィルタ係数のAVXレジスタへのロード
    __m256d coeff0 = _mm256_set1_pd(filter[0]);
    __m256d coeff1 = _mm256_set1_pd(filter[1]);
    __m256d coeff2 = _mm256_set1_pd(filter[2]);
    __m256d coeff3 = _mm256_set1_pd(filter[3]);
    __m256d coeff4 = _mm256_set1_pd(filter[4]);

    // AVX命令セットを使用した畳み込み処理
    for (int64_t idx = half_width; idx <= data_length - filter_width; idx += 4) {
        __m256d data0 = _mm256_loadu_pd(&input[idx + 2]);
        __m256d data1 = _mm256_loadu_pd(&input[idx + 1]);
        __m256d data2 = _mm256_loadu_pd(&input[idx]);
        __m256d data3 = _mm256_loadu_pd(&input[idx - 1]);
        __m256d data4 = _mm256_loadu_pd(&input[idx - 2]);

        // FMA命令を使用した積和演算
        __m256d result_val = _mm256_setzero_pd();
        result_val = _mm256_fmadd_pd(data0, coeff0, result_val);
        result_val = _mm256_fmadd_pd(data1, coeff1, result_val);
        result_val = _mm256_fmadd_pd(data2, coeff2, result_val);
        result_val = _mm256_fmadd_pd(data3, coeff3, result_val);
        result_val = _mm256_fmadd_pd(data4, coeff4, result_val);

        // 結果の保存
        _mm256_storeu_pd(&output[idx], result_val);
    }

    return true;
}

AVX512による拡張最適化

bool Execute1DConvolution_AVX512(double* __restrict__ output, const double* __restrict__ input, const double* __restrict__ filter, int64_t data_length) {
    constexpr int64_t filter_width = 5;
    constexpr int64_t half_width = filter_width / 2;

    if (data_length < filter_width) {
        return false;
    }

    // フィルタ係数のAVX512レジスタへのロード
    __m512d coeff0 = _mm512_set1_pd(filter[0]);
    __m512d coeff1 = _mm512_set1_pd(filter[1]);
    __m512d coeff2 = _mm512_set1_pd(filter[2]);
    __m512d coeff3 = _mm512_set1_pd(filter[3]);
    __m512d coeff4 = _mm512_set1_pd(filter[4]);

    // AVX512命令セットを使用した畳み込み処理
    for (int64_t idx = half_width; idx <= data_length - filter_width; idx += 8) { // AVX512ではストライドを8に変更
        __m512d data0 = _mm512_loadu_pd(&input[idx + 2]);
        __m512d data1 = _mm512_loadu_pd(&input[idx + 1]);
        __m512d data2 = _mm512_loadu_pd(&input[idx]);
        __m512d data3 = _mm512_loadu_pd(&input[idx - 1]);
        __m512d data4 = _mm512_loadu_pd(&input[idx - 2]);

        // FMA命令を使用した積和演算
        __m512d result_val = _mm512_setzero_pd();
        result_val = _mm512_fmadd_pd(data0, coeff0, result_val);
        result_val = _mm512_fmadd_pd(data1, coeff1, result_val);
        result_val = _mm512_fmadd_pd(data2, coeff2, result_val);
        result_val = _mm512_fmadd_pd(data3, coeff3, result_val);
        result_val = _mm512_fmadd_pd(data4, coeff4, result_val);

        _mm512_storeu_pd(&output[idx], result_val);
    }

    return true;
}

パフォーマンス比較結果

実装方式 基本実装 AVX2 AVX2+FMA AVX512
実行時間 85643μs 45320μs 44870μs 46879μs

考察と分析

1次元畳み込みに対するAVX最適化から、いくつかの重要な現象が観察できます:

  1. 基本的な実装では、手動でのループ展開によってすでに良好な最適化が行われていますが、AVX命令セットによる初回の最適化では約1.89倍(ほぼ2倍)の加速比が達成されます。さらに、畳み込み操作の並列化により、小規模な効率向上が期待できます(空間と時間のトレードオフ)。

  2. FMA命令による最適化では、さらなる効率向上が見込めます。

  3. しかし、AVX512を導入した場合、逆にパフォーマンスが低下する現象が確認されました。調査の結果、これはCPUがAVX512命令を実行する際に動的にクロック周波数を下げる(デクロック)ためであると考えられます。

周波数スケーリングのメカニズムについてIntelの資料を参照すると、CPUは利用可能なヘッドルームに基づいて動的周波数スケーリングを行うことがわかります。これは、異なる負荷に応じてCPUが自ら周波数を上げ下げする仕組みです。特にAVX-512のような複雑で高消費電力の計算を行う場合、CPUの熱的・電気的制限を超えないように、周波数を調整する必要があります。

非連続メモリアクセスの最適化

gather命令はAVX2で導入された、非連続なメモリ位置からデータをレジスタにロードする操作です。gather実装では、インデックス配列を使用して対応するバイト位置から値を取得します。

__m256d _mm256_i64gather_pd(const double* base_addr, __m256i indices, int scale);

permute命令はAVX2で導入された、ベクトル要素を並べ替える(Permute)操作です。permute実装では、imm8値を制御ビットとして使用してベクトルの要素を置換します。

参考資料

タグ: AVX FMA AVX512 畳み込み演算 ベクトル化

5月29日 02:34 投稿