Posted inプログラミング

RustにおけるSIMDとベクトル化の極意：モダンCPUハードウェアアクセラレーションの実践

Programming tutorial - IT technology blog
Programming tutorial - IT technology blog

インシデント：4K動画の取り込み処理がクラスタを飽和させた日

昨年の秋、私たちのチームはある画像処理サービスをRustへ移行しました。このパイプラインは、連続する4Kビデオフィード（3840×2160、1フレームあたり約830万画素）を正規化するものです。32ビット浮動小数点数の配列全体に対し、ピクセルごとのコントラストスケーリングとバウンディングボックスのアルファブレンディングを適用します。実行時オーバーヘッドのない、ネイティブC言語に匹敵する速度を期待してRustを採用しました。

初期実装は、極めてクリーンで慣用的なRustコードでした：

pub fn normalize_pixel_luma(samples: &mut [f32], gain: f32) {
    for sample in samples.iter_mut() {
        *sample = (*sample * gain).clamp(-1.0, 1.0);
    }
}

ローカルでのベンチマーク結果は良好でした。合成テストも何の問題もなくパスしました。しかし、32系統のRTSPカメラストリーム全体で本番トラフィックが120 FPSに達した途端、フレーム遅延が3.2ミリ秒から18.7ミリ秒へと跳ね上がったのです。16個のvCPUすべてが使用率100%に張り付きました。Rayonによってスライスはワーカースレッドへ適切に分散されていたものの、すべてのコアがスカラー命令の実行にサイクルを浪費していました。CPUは一度に1つの32ビット浮動小数点数しか処理していなかったのです。コア数は十分にあったにもかかわらず、コア単体のスループットが完全に枯渇していました。

根本原因：なぜスカラー演算ループは最新CPUリソースを無駄にするのか

最新のサーバープロセッサは、スカラー演算パイプラインと並んで専用のベクトルユニットを搭載しています。AVX2対応のx86_64コアは256ビットのベクトルレジスタを備え、AVX-512では512ビットに拡張されます。GravitonコアのARM NEONは128ビットのベクトルレーンを提供します。これらのベクトルユニットは、SISD（単一命令・単一データ）を実行する代わりに、SIMD（単一命令・複数データ）を実行します。

単一の256ビットAVX2レジスタには、8個の32ビット浮動小数点数が収まります（8 × 32 = 256ビット）。理論上、実行ユニットは1クロックサイクルでこれら8つの値を同時に乗算可能です。それにもかかわらず、私たちのバイナリはスカラーアセンブリを実行していました。何が起きていたのかを確かめるため、コンパイル後のアセンブリを検証しました：

# コンパイルされたアセンブリの確認
cargo rustc --release -- --emit asm

頻繁に実行されるホットループ内は、パックドベクトル命令であるvmulpsではなく、スカラー命令のvmulssで埋め尽くされていました。LLVMはこのループの自動ベクトル化を拒否していたのです。浮動小数点数の乗算とclampを組み合わせたことで、境界チェックとNaN処理に伴う潜在的な分岐が発生していました。IEEE-754のセマンティクスでは厳密なNaNの伝播が規定されているため、コンパイラは準拠性の保証を損なうことなく安全に計算順序を並べ替えることができなかったのです。AMD EPYC 7763コアの演算性能の8分の7が眠ったままになっていました。

本番環境における解決策の比較検証

その後の2四半期にわたり、実際の本番ワークロードを用いてこのループをベクトル化する3つの手法をベンチマークしました。

1. コンパイラフラグと自動ベクトル化への依存

最も手軽なアプローチは、アグレッシブなフラグを指定してコンパイラに自動ベクトル化を委ねる方法です。

# .cargo/config.toml
[build]
rustflags = ["-C", "target-cpu=native", "-C", "llvm-args=-ffast-math"]
  • メリット: コードの変更が一切不要。ソースコードの完全なポータビリティを維持できる。
  • デメリット: 極めて壊れやすい。軽微なリファクタリング、境界チェックの追加、またはインデックス指定ロジックの些細な変更だけで、LLVMのベクトル化ヒューリスティクスが警告なしに無効化される。コンパイラは何の通知もなくスカラー演算ループへとフォールバックする。

2. アーキテクチャ固有の組み込み関数（std::arch::x86_64）

Rustはcore::archを介してハードウェアの組み込み関数（Intrinsics）を直接公開しています。これを利用することで、アセンブリと同等のベクトル命令を手作業で記述できます。

#[cfg(target_arch = "x86_64")]
use std::arch::x86_64::*;

pub unsafe fn normalize_avx2(samples: &mut [f32], gain: f32) {
    let gain_vec = _mm256_set1_ps(gain);
    let min_vec = _mm256_set1_ps(-1.0);
    let max_vec = _mm256_set1_ps(1.0);
    
    let chunks = samples.chunks_exact_mut(8);
    let remainder = chunks.into_remainder();
    
    for chunk in chunks {
        let mut data = _mm256_loadu_ps(chunk.as_ptr());
        data = _mm256_mul_ps(data, gain_vec);
        data = _mm256_max_ps(min_vec, _mm256_min_ps(max_vec, data));
        _mm256_storeu_ps(chunk.as_mut_ptr(), data);
    }
    
    for item in remainder {
        *item = (*item * gain).clamp(-1.0, 1.0);
    }
}
  • メリット: 予測可能な機械語コードの生成により、ハードウェア性能を最大限に引き出せる。
  • デメリット: unsafeブロックの使用と、割り切れない端数スライスに対する手動のループピーリングが必須。また、実行時CPU機能検出（is_x86_feature_detected!("avx2")）による呼び出しの保護も必要。さらに、別途NEON向けの実装を用意しない限り、ARM64 Gravitonインスタンス上ではコンパイルエラーとなる。

3. ポータブルで安全なベクトル化（wideクレートまたはNightly std::simd）

ポータブルSIMDは、f32x8やf32x4などのハードウェアに依存しないベクトル型を提供します。x86ホスト上ではAVX2命令へとコンパイルされ、ARMインスタンス上では自動的にNEON命令が生成されるため、unsafeブロックを一切書く必要がありません。

use wide::f32x8;

pub fn normalize_portable_wide(samples: &mut [f32], gain: f32) {
    let gain_vec = f32x8::splat(gain);
    let min_vec = f32x8::splat(-1.0);
    let max_vec = f32x8::splat(1.0);
    
    let (chunks, remainder) = samples.as_chunks_mut::<8>();
    
    for chunk in chunks {
        let mut v = f32x8::from(*chunk);
        v = (v * gain_vec).fast_max(min_vec).fast_min(max_vec);
        *chunk = v.to_array();
    }
    
    for item in remainder {
        *item = (*item * gain).clamp(-1.0, 1.0);
    }
}
  • メリット: 100%安全なRustコード。簡潔な構文で、x86とARMの両プラットフォーム向けに問題なくコンパイルできる。
  • デメリット: 安定版Rustでは外部依存関係（wideやpulpなど）が必要となり、標準のstd::simdを利用するにはNightlyツールチェーンが必須となる。

本番環境で採用したパターン：マルチスレッドSIMD

自動ベクトル化はCI/CDビルドにおいて挙動が不安定すぎることが判明し、アーキテクチャごとに個別の組み込み関数実装を維持することはテスト工数の倍増を招きました。そこで私たちは、コア間の並行処理にはRayonを採用し、演算部分にはwideクレートを用いた安全なチャンク分割ベクトル配列の手法へと標準化しました。

ステップごとの実装手順

Cargo.tomlに以下の依存関係を追加します：

[dependencies]
wide = "0.7"
rayon = "1.8"

以下が本番環境で稼働しているルーチンです。Rayonがバッファをチャンクに分割して複数コアに作業を分散し、各コアのL1キャッシュ内でwide::f32x8が演算をベクトル化します：

use wide::f32x8;
use rayon::prelude::*;

pub fn parallel_simd_normalize(data: &mut [f32], factor: f32) {
    // 4096個のf32 = 16 KB（L1データキャッシュに余裕で収まるサイズ）
    data.par_chunks_mut(4096).for_each(|slice| {
        let factor_vec = f32x8::splat(factor);
        let min_vec = f32x8::splat(-1.0);
        let max_vec = f32x8::splat(1.0);
        
        let chunks = slice.chunks_exact_mut(8);
        let remainder = chunks.into_remainder();
        
        for chunk in chunks {
            let mut vec = f32x8::new([
                chunk[0], chunk[1], chunk[2], chunk[3],
                chunk[4], chunk[5], chunk[6], chunk[7],
            ]);
            
            vec = (vec * factor_vec).fast_max(min_vec).fast_min(max_vec);
            
            let arr = vec.to_array();
            chunk.copy_from_slice(&arr);
        }
        
        // 8で割り切れない端数要素の処理
        for item in remainder {
            *item = (*item * factor).clamp(-1.0, 1.0);
        }
    });
}

運用開始6か月後の実績

この実装は本番環境ですでに6か月間稼働しています。Datadogのメトリクスが示した成果は以下の通りです：

  • レイテンシ: 4Kフレームの処理時間が18.7ミリ秒から2.4ミリ秒へと大幅短縮され、スループットは7.8倍向上しました。
  • CPU使用率: クラスタ全体の平均CPU負荷が100%から32%へと激減。EC2フリートを20台のc5.4xlargeインスタンスから10台へと半減でき、月額約2,400ドルのコスト削減を達成しました。
  • クロスアーキテクチャのポータビリティ: ワーカーノードをAWS Graviton3（c7g.4xlarge、ARM NEON）へ移行した際も、全く同一のコードが一切の修正なしでコンパイル・実行できました。パフォーマンスもAVX-512 x86環境のベンチマークと比較して5%以内の差に収まりました。

近年のコンパイラは非常に優秀ですが、厳密なIEEE浮動小数点ルールに縛られている限り、開発者のアーキテクチャ上の意図を汲み取ることはできません。明示的なベクトル化を行うことで、コンパイラの推測への依存を排除し、遊休状態だったハードウェアレーンを本来のスループットへと変換して、高負荷パイプラインを予測可能かつ軽快に稼働させ続けることができます。

Share:
Tags: