ARM NEONによる1秒あたり数十ギガバイトのプレフィックス和

1日あたりの売上記録があるとします。各日について、年初からの累計売上を知りたい場合があるでしょう。

このような操作はプレフィックス和（prefix sum）またはスキャンと呼ばれます。

C言語での実装は難しくありません。単純なループです。

  for (size_t i = 1; i < length; i++) {
    data[i] += data[i - 1];
  }

この関数はどれくらい高速にできるでしょうか？ 速度の上限は簡単に導けます。現在の値を計算するには、前の値を計算済みでなければならず、その前も同様です。

data[0] -> data[1] -> data[2] -> ...

最善の場合でも、テーブル内の各エントリにつき1 CPUサイクルが必要です。したがって4 GHzのプロセッサでは、1秒あたり40億個の整数値を処理できる可能性があります。これは上限ですが、多くの最新システムでは実際にこれに近い値に到達できるかもしれません。もちろん、ロード、ストア、分岐などの他の命令も関わりますが、現代のプロセッサは1サイクルで多数の命令を実行でき、分岐予測も効果的です。そのため、今日のほとんどのプロセッサで1秒あたり数十億個の整数を処理できるはずです。

悪くない結果です！ しかし、さらに速くできるでしょうか？

SIMD命令を使えます。SIMD命令は複数の値を同時に処理する特殊な命令です。すべての64ビットARMプロセッサはNEON命令をサポートしています。NEON命令は、1つのSIMDレジスタにパックされた4つの整数を一度に処理できます。

しかし、4値レジスタ上でプレフィックス和を行うにはどうすればよいでしょうか？ 2回のシフトと2回の加算で実現できます。理論上は、ベクトルレジスタ内の要素数Nに対してlog(N)でスケールします。

input   = [A B   C     D]
shift1  = [0 A   B     C]
sum1    = [A A+B B+C   C+D]
shift2  = [0 0   A     B+A]
result  = [A A+B A+B+C A+B+C+D]

最後の値（A+B+C+D）を取り出し、次の4値ブロックに加算できるように全位置にブロードキャストできます。

これはスカラー方式より高速でしょうか？ 4命令が連続し、さらに次の4値ブロックで合計値を使うために少なくとも1命令が必要です。

したがってSIMD方式の方が遅くなる可能性があります。残念な結果です。

解決策として、より多くの整数値にスケールアップする方法があります。

ARM NEONのインターリーブド・ロード／ストア命令を考えてみましょう。一度に16個の値をロードし、最初の値群、2番目の値群などをそれぞれまとめて取得できます。

original data : ABCD EFGH IJKL MNOP
loaded data   : AEIM BFJN CGKO DHLP

その後、4ブロックに対して並列にプレフィックス和を実行できます。3命令で完了します。3命令の終了時点で、ローカル和を含む1つのレジスタが得られます。

A+B+C+D E+F+G+H I+J+K+L M+N+O+P

次に、このレジスタに対してプレフィックス和の手法を適用します（4命令）。結果として、16値ブロックあたり8命令程度の逐次命令になる可能性があります。

理論上はスカラー方式の2倍高速です。

C言語でintrinsicsを使うと、次のように記述できます。

void neon_prefixsum_fast(uint32_t *data, size_t length) {
  uint32x4_t zero = {0, 0, 0, 0};
  uint32x4_t prev = {0, 0, 0, 0};

  for (size_t i = 0; i < length / 16; i++) {
    uint32x4x4_t vals = vld4q_u32(data + 16 * i);

    // Prefix sum inside each transposed ("vertical") lane
    vals.val[1] = vaddq_u32(vals.val[1], vals.val[0]);
    vals.val[2] = vaddq_u32(vals.val[2], vals.val[1]);
    vals.val[3] = vaddq_u32(vals.val[3], vals.val[2]);

    // Now vals.val[3] contains the four local prefix sums:
    //   vals.val[3] = [s0=A+B+C+D, s1=E+F+G+H, 
    //                  s2=I+J+K+L, s3=M+N+O+P]

    // Compute prefix sum across the four local sums 
    uint32x4_t off = vextq_u32(zero, vals.val[3], 3);
    uint32x4_t ps = vaddq_u32(vals.val[3], off);       
    off = vextq_u32(zero, ps, 2);                      
    ps = vaddq_u32(ps, off);

    // Now ps contains cumulative sums across the four groups
    // Add the incoming carry from the previous 16-element block
    ps = vaddq_u32(ps, prev);

    // Prepare carry for next block: broadcast the last lane of ps
    prev = vdupq_laneq_u32(ps, 3);

    // The add vector to apply to the original lanes is the 
    // prefix up to previous group
    uint32x4_t add = vextq_u32(prev, ps, 3]);  

    // Apply carry/offset to each of the four transposed lanes
    vals.val[0] = vaddq_u32(vals.val[0], add);
    vals.val[1] = vaddq_u32(vals.val[1], add);
    vals.val[2] = vaddq_u32(vals.val[2], add);
    vals.val[3] = vaddq_u32(vals.val[3], add);

    // Store back the four lanes (interleaved)
    vst4q_u32(data + 16 * i, vals);
  }

  scalar_prefixsum_leftover(data, length, 16);
}

Apple M4プロセッサ（4.5 GHz）で実際に試してみましょう。

このSIMD方式はスカラー方式の約2.3倍高速でした。悪くない結果です。

ソースコードはGitHubで入手できます。

Appendix. Instrinsics