SIMDとBit演算を用いたテトリスの高速化に関する考察
テトリス Advent Calendar 2024 12日目です。
テトリスAI向けのシミュレーション高速化に関する記事をゆるーく書きます。
限界まで高速化したい一定の知識のある開発者向けの内容を想定しています。
ぷよぷよAIで使われてる連鎖や落下などのシミュレーションはbitboardで高速化されているのを見て[1]、テトリスもbitboardで高速化できそうだと思ったので色々試してみました。
今回bitboardを使うにあたって、CPU拡張命令を使います。
使用する命令セットはAVX2・BMI2を用います。
盤面データの持ち方
通常のテトリスの盤面だと、10*40で400bit必要になりますが、今回はAVX2を使う上では、1列あたりのサイズにプレイに支障がなく大きすぎない4byte(=32bit)を使います。
すると幅は8列となりますが、足りないため、256bitを左右で2つ用意することで盤面5列+壁3列で表します。
画像で表すとこのようになります。
操作ミノ情報と操作
操作ミノの情報には、回転状態・現在位置・ミノの状態があります。
まず回転状態ですが、これは4種類なので2bitで管理すればオーバーフローとアンダーフローで1周させられif文を節約できます。
ミノの状態ですが、普通にテトリスを作る際には、ブロックがある位置を配列で管理する方法と4*4などで用意する方法の2種類が主にあると思います。
今回はbitをフル活用したいので演算がしやすい後者の方法を選択します。
例えば、Iミノの場合は
0000
1111
0000
0000
なので0000111100000000となります。
回転する際にはDeltaSwapという方法を用いて高速に処理します。
詳細はビットボードの回転がとても分かりやすかったです。
4*4の場合、MaskとDeltaはこのようになりました。
fn flip_horizontal_16(mut data: u16) -> u16 {
data = Self::delta_swap(data, 0x5555, 1);
data = Self::delta_swap(data, 0x3333, 2);
data
}
fn flip_vertical_16(mut data: u16) -> u16 {
data = Self::delta_swap(data, 0xF0F, 4);
data = Self::delta_swap(data, 0xFF, 8);
data
}
fn flip_diagonal_a1d4(mut data: u16) -> u16 {
data = Self::delta_swap(data, 0xA0A, 3);
data = Self::delta_swap(data, 0xCC, 6);
data
}
fn flip_diagonal_a4d1(mut data: u16) -> u16 {
data = Self::delta_swap(data, 0x505, 5);
data = Self::delta_swap(data, 0xCC, 10);
data
}
pub fn rotate_cw(&mut self) {
self.data = Self::flip_horizontal_16(Self::flip_diagonal_a1d4(self.data));
}
pub fn rotate_ccw(&mut self) {
self.data = Self::flip_vertical_16(Self::flip_diagonal_a1d4(self.data));
}
全ての回転状態をあらかじめ用意することもできますが、メモリやキャッシュのアクセス速度を考えると毎回計算した方が早そうだったのでこの手法を用いました。
また、I・Oミノの場合回転の軸が真ん中にあるため4*4でも大丈夫ですが、その他のミノは3*3のサイズで中心に軸が来るためにこのままだと上手く回転できません。
カレントミノのデータ生成は頻繁に行われるわけではないので、生成時に関数ポインタで4*4と3*3それぞれのDeltaSwapを渡しても良いですし、その後SRSの(0,0)回転を変えることでも行けそうです。
カレントミノの移動判定
ここは愚直に実装すると、操作ミノを左右の盤面用の2つの_m256iにして論理積を求めた結果が0かどうかで判定することができますが、激遅です。
この判定をするにあたって、論理積で判定するため2つ分の_m256iを作る方向は変わらないのですが、__m256iのデータの持ち方の都合上、分割したデータを跨いだ処理が難しいため、x座標を合わせるのにコストがかかります。（y座標はそれぞれの分割データをシフトするだけ）
そこで、4*4のミノを盤面4*16まで拡張し、シフトでx座標を合わせた後__m256iに変換することにより解決しました。
また、y座標合わせについてですが、_mm256_slli_epi32を用いることでそれぞれの列をシフトすることができるのですが、シフトする量が即値なので使えません。
gccのcppで実装した知り合いは通っていたのでRust側の仕様かと思ったのですが、Intelの公式ドキュメントを見ても即値なんですよね...?
まあそれぞれの列で独立した処理なので愚直に実装してもパイプラインでいい感じにやってくれそう。（要検証）
ラインクリアのマスク作成
各列の論理積を求めることで取得できます。
元の盤面と1列左にシフト盤面の論理積を幅分繰り返せば求められますが、シフトの代わりにshuffleを使うことで1度に(残りの個数/2)個分の比較ができるのでお得です。
これでできたフラグをいい感じにごにょごにょすると消えるラインのマスクができます。
クリアしたラインを詰める
詰めるには、拡張命令のPEXTとPDEPを使い処理します。
各列に対して処理を行うわけですが、各列のデータを取得するには一度メモリに移す方法と列データをそれぞれ直接レジスタに取得する方法がありますが、メモリにアクセスするのは遅いので後者の方法でやります。
操作ミノを設置する
操作ミノの移動判定のやり方で盤面に合わせてから、その論理和を使うことで簡単に取得できます。
おまけ
穴の数を取得する
穴の定義は上が1、下が0であるマスとします。
下にずらしたbitboardと元でNANDを取ることで上下で違う種類同士の数がわかります。…①
後は上が1であるという条件が必要なので、①の場所で元のbitboardが1であるものを取得することで、盤面上のすべての穴の数を計算できます。
まとめ
ということで拡張命令を用いて高速にシミュレーションするプログラムの実装・考察をしてみました。
自分で全て実装して確認したわけではないので、理論通りいかない部分もあるかもしれないです。
また、アルゴリズムにあたっては議論の余地が多く残っていると思うので、もっと高速なアルゴリズムがあれば是非教えていただけると幸いです。
今回実装したコードは諸般の事情により公開はありませんが、いつかテトリスAIを作る機会があれば一緒に公開するかもしれません。
最後に、この記事を作るにあたって、アルゴリズムや実装を知り合いのAlcearさんに手伝ってもらいました。
Discussion