SSEによる高速化ケーススタディ

  • 21
    いいね
  • 0
    コメント
この記事は最終更新日から1年以上が経過しています。

この記事はCompetitive Programming Advent Calendar 2014 22日目の記事として書かれました。

はじめに

 競技プログラミングの問題(主にTopCoder SRMやICPCなどの短期間マッチ)において、制約が甘いなどの理由からオーダ記法で一回り想定解より大きな計算量のアルゴリズムでも定数倍高速化することによって通ってしまうことがあります。一方で、最近のx86(互換)プロセッサにはほぼ確実に SSE (Streaming SIMD Extensions) と呼ばれる拡張命令セットが実装されており、アルゴリズムによってはベクトル計算をうまく使うことによってプログラムを数倍高速化することができます。
 最近のコンパイラでは適切なオプションを設定してやると、最適化の一部としてある程度自動でベクトル化を行うのですが、多くのジャッジサーバではそのようなオプションは有効になっていません。(自動ベクトル化についてはnodchipさんが3日目の記事で言及されています。)しかし、まったくSSEの命令が利用できないのかというとそういうことはなく、手動で使用する命令を指定することで利用することができてしまいます。
 この記事では、愚直解をベースに手動でSSEを利用した定数倍高速化を施すことでAcceptを取れてしまった問題をいくつかご紹介いたします。

ベクトル計算・SIMDって何? という方は、以下の記事のようなSIMDプログラミングの概要をあらかじめ眺めておいた方が良いかもしれません。
http://cell.fixstars.com/ps3linux/index.php/第2章_SIMDプログラミングの基礎
また、どのような命令があるか・この命令(組み込み関数)は何をするのかを調べるときには Intel Intrinsics Guide を参照するのがおすすめです。

平和協定

問題概要

 2次元平面上の点$(A_i, B_i)$が$N$個与えられる。それらのうち、$(A_i - A_j) \times (B_i - B_j)$が$S1$以上$S2$以下となる$i, j$のペアの数を求めよ。

  • $1 \le N \le 50000$
  • $1 \le A_i, B_i \le 50000$

解法

 愚直にすべての2点間の距離を求めます。割と素直にベクトル化しただけなので、命令リファレンスなどと突き合わせていけばなんとなく見えてくるのではないかと思います。(この問題の解答例は比較的しっかりコメントが入っています。)

解答例

Design Tutorial: Increase the Constraints

問題概要

 01のみからなる文字列$a$, $b$が与えられる。それに対するクエリとして、$p_1$, $p_2$, $\mathit{len}$の組が$q$個与えられる。各クエリに対して、$a$の$p_1$番目の文字から$\mathit{len}$文字を切り出した文字列と$b$の$p_2$番目の文字から$\mathit{len}$文字を切り出した文字列のハミング距離を求めよ。

  • $1 \le |a|, |b| \le 200000$
  • $1 \le q \le 400000$

解法

 愚直にクエリあたり$O(\mathit{len})$で解いていきます。まず、文字列が01のみしか含まないので1文字あたり1bitとなるようにパッキングします。パッキングしてやるとハミング距離は部分列同士のXORのうち、1となっているビットの数がクエリに対する答えになります。このとき、XORの部分は1/32や1/64のようなかなり小さい定数をかけることができるのでどうにかなるのですが、そのままだと1のビットを数える部分 (popcount) がボトルネックになってしまいます。そこで、ビットカウントにpopcnt命令を使用することでかなり高速に処理することができるようになります。
 その他、片方のビット列についてあらかじめ [0, 31] bit 分ずらしたパターンを生成しておき、ビットシフトをクエリ処理ループの外に逃がすなどのテクニックも使用しています。

解答例

Arithmetic Progressions

問題概要

 $N$要素の数からなる数列$A$から3つの要素$A_i$, $A_j$, $A_k$ $(1 \le i < j < k \le N)$ を選択する時に、$A_j - A_i = A_k - A_j$となっている$(i, j, k)$の組み合わせが何通りあるかを求めよ。

  • $3 \le N \le 100000$
  • $1 \le A_i \le 30000 (1 \le i \le N)$

解法

 まず、$j$を決め打ちした場合に$i$, $k$の選び方が何通りあるかを考えます。$[1, j)$番目の要素のヒストグラムを$H_l$, $(j, N]$番目の要素のヒストグラムを$H_r$としてやると、$j$に対応する$i$, $j$の組み合わせの数は$\sum_t {H_l(A_j+t) H_r(A_j-t)}$となります。想定解法だとここからさらに考察が必要なのですが、定数倍高速化でゴリ押しする場合はこのくらいで十分です。頑張ってこの総和の部分を早く計算しましょう。
 ヒストグラムの更新については$j$を決めるループ内で更新してやればOKです。このとき、どちらかのヒストグラムを逆順で保持しておくと総和を求める部分の処理がよりシンプルになります。ヒストグラムの更新ができたらあとはひたすらpmuldqやpaddqなどで総和の式を解いていきます。細かいテクニックとしては、ヒストグラムの終端に番兵としてベクトル長と同程度だけ0を追加しておくと端数処理が不要になって少しコードがシンプルになります。

解答例

XORAND

問題概要

 $N$個の要素からなる数列$A$と$Q$個のそれに対するクエリが与えられる。各クエリは$l$, $r$の2つの整数からなる。各クエリに対して、部分列$A_l, A_{l+1}, ... , A_r$を$k (l < k < r)$番目の要素を境に分割し、$(A_l \mathrm{\ AND\ } A_{l + 1} \mathrm{\ AND\ } ... \mathrm{\ AND\ } A_{k - 1}) - (A_k \mathrm{\ XOR\ } A_{k+1} \mathrm{\ XOR\ } ... \mathrm{\ XOR\ } A_r)$の値を最大化したときのその値を求めよ。

  • $1 \le N \le 100000$
  • $1 \le Q \le 100000$
  • $0 \le A_i < 2^{31}$
  • クエリの先読みは不可

解法

 とりあえずXORについては累積和をもっておけば区間を一発で結合できるので累積和テーブルを作ります。ANDの方は普通のテーブルではどうしようもないので毎回求めることにします。ただし、ANDで結合したものが変化するタイミングは高々31回しかないので、あるビットが次に0になるタイミングを控えておいて、ANDで結合した値が変化する位置を高速に探せるようにしておきます。そうすると、その値が変化するまでの区間内の$k$のうちXORで結合した値が最小になるものを探せばよくなります。あとは、これまでの問題と同様に最小値を探す部分のループをベクトル化すればOKです。
 ちなみに、ANDで結合した値とXORで結合した値を同じループで愚直に求める方法ではさすがに制限内に収めることができませんでした。

解答例

まとめ

 SSEを利用することでかなり定数倍高速化が達成でき、結果として時間制限内に愚直解(またはそれに近い解法)で答えが求まってしまった問題をいくつかご紹介しました。自動ベクトル化ほど手軽ではないうえにそもそも本当に間に合う保証がないため、他にやることがない場合か行けそうな手ごたえがある場合くらいにしか出番はないのですが、もしかしたら解ける問題がなくなった後に+1問分の点数を稼げるかもしれません。逆にジャッジ側の方にとっては頭の痛い問題だとは思いますが、ベクトル化が簡単そうだったら時間制限を長くして制約を大きくするなどの策を講じた方がいい場合もあるかもしれません。