TERRYのブログ

最近は競技プログラミングにお熱。

AHC061で学ぶ、明日使えないビット演算 ~pdep/pextを添えて~

AtCoder Heuristic Contest 061お疲れ様でした。相対スコア2,230,541,945,022点で5位でした。*1

今回のAHCはモンテカルロ法を使っている人が多かったかなと思いますが、これは高速化が非常に重要でした。

今回の問題設定だと色々な処理をビット演算で高速化できるので、いくつか紹介してみようかなと思います。

ビジュアライザ (seed=0, score=293,457点)

問題概要

陣取りゲームでAI相手に圧倒的点数差で勝ってね。

atcoder.jp

明日使えるかもしれないビット演算

まずは使いどころがちょっとはありそうなやつから行きます。

BitBoard

今回のようなグリッドグラフでBFSをするとき、普通は2次元配列を持ってキューに各マスを突っ込んでいくと思います。ただそれでは遅いので、もっと高速な方法を考えます。

いきなり2次元だと難しいので、1次元で考えます。64bit整数を64要素のbool配列だと思うと、BFSの1ステップは元のビット列とそれを左右にビットシフトしたもののbitwise ORで表現できます。

# 以下、Pythonっぽい擬似コードで書きます
# 2進数か10進数かは特に断らないので雰囲気で読んでください

# i=3マスからスタートとする(0-indexed)
visited = 00001000

# visited = 00001000 | 00010000 | 00000100
# visited = 00011100
visited = visited | (visited << 1) | (visited >> 1)

壁などがあって進入できないマスがあっても大丈夫です。進入可能なマスとのbitwise ANDを取ってマスクしてあげれば良いです。*2

# 0マス目と4マス目に壁があるとする
# 壁のあるマスを1、壁のないマスを0とする
wall_mask= 00010001
visited = 00001000

# visited = (00001000 | 00010000 | 00000100)  & 11101110
# visited = 00011100 & 11101110
# visited = 00001100
visited = (visited | (visited << 1) | (visited >> 1)) & ~wall_mask

BFSの終了判定としては、 visited が変化しなくなったら終了という形にすると良いでしょう。

mask = 11101110
visited = 00001000

while True:
    new_visited = (visited | (visited << 1) | (visited >> 1)) & mask

    if new_visited == visited:
        break

     visited = new_visited

H\ \timesW マスの2次元グリッドになっても考え方は同じです。注意点として、何も工夫せずにやるとマス (i, W-1) とマス (i+1,0) との間で双方向にワープが可能になってしまいます。番兵*3を用いるか、以下のようにビットシフト前にマスクしましょう。

# 以下のような縦2マス×横3マスの2次元グリッドを考える
# ...
# ..#
# 今回はマスの番号を以下のように付ける
# 210
# 543
H = 2
W = 3
wall_mask = 11001000
visited = 00000001

# 上下左右方向への移動が可能なマスを1としたものを***_maskとする
left_mask = 00011011
right_mask = 00110110

while True:
    left = (visited & left_mask) << 1
    right = (visited & right_mask) >> 1
    down = visited << W
    up = visited >> W
    
    new_visited = (visited | left | right | down | up) & ~wall_mask 

    if new_visited == visited:
        break

    visited = new_visited

popcnt命令

Population Count 命令です。人口を数える?という感じですが、これはビット列の 1 の数を数える命令です。

src = 01101101
pcnt = popcnt(src)

# 5と出力される
print(pcnt)

使い途はいろいろありますが、先程のBFSと合わせて今回の問題に使うなら「次のターンの合法手の数をカウントする」とかですかね。例えばこんな感じで書きます。

# x: 訪問済み, .: 未訪問, #: 壁
#
# .xx
# ..#
H = 2
W = 3
wall_mask = 11001000
visited = 00000011
left_mask = 00011011
right_mask = 00110110

left = (visited & left_mask) << 1
right = (visited & right_mask) >> 1
down = visited << W
up = visited >> W

# BFSの1ステップ分と、未訪問のマス集合との bitwise AND が合法手となる
# 未訪問のマス集合は visited の bitwise NOT で求められる
# legal_hands = 00010100
bfs_step = (left | right | down | up) & ~wall_mask 
legal_hands = bfs_step & ~visited

# 合法手は2つと表示される
legal_cnt = popcnt(legal_hands)
print(legal_cnt)

popcnt 命令は、C++20以降なら std::popcount()*4 、Rustでは x.count_ones() で実行できます。*5

#include <bit>
#include <cstdint>

uint64_t x = 0b01001001;
int pcnt = std::popcount(x);
let x: u64 = 0b01001001;
let pcnt = x.count_ones();

コンパイル結果について

AtCoderのC++ユーザーはこの節を飛ばして良いです。Rustユーザーは読んでおいた方が良いかも?

注意点として、 std::popcount()x.count_ones() が popcnt 命令に変換されるかどうかはコンパイラが判断します。 popcnt 命令に対応していないCPUも存在し、そのようなCPUでもクラッシュしないようなアセンブリを吐き出す必要があるためです。2026/3/30時点におけるAtCoderのC++23 (GCC 15.2.0)には -march=native フラグが立っており、コンパイルするCPUの持つ命令セット全ての利用が許可されるのでほぼ間違いなく popcnt 命令が使用されるはずですが、Rustではそのようなフラグが指定されていないので popcnt 命令を使ってくれるとは限りません。*6

実際、私のローカルPC(Ryzen 9 7950X CPU)で cargo-show-asm を使用してアセンブリを表示させたところ、以下のようになりました。 popcnt 命令の姿は見えないですし、定数 6148914691236517205 は popcnt の有名アルゴリズムでおなじみの定数 0x5555555555555555 のことなので、ソフトウェア実装に展開されていますね。

mov rax, rdi
shr rax
movabs rcx, 6148914691236517205
and rcx, rax
sub rdi, rcx
movabs rax, 3689348814741910323
mov rcx, rdi
and rcx, rax
shr rdi, 2
and rax, rdi
add rax, rcx
mov rcx, rax
shr rcx, 4
add rcx, rax
movabs rdx, 1085102592571150095
and rdx, rcx
movabs rax, 72340172838076673
imul rax, rdx
shr rax, 56
ret

この問題を素直に解消するには、Rustコンパイラに -C target-cpu=native を渡してあげるとよいです。プロジェクト直下に .cargo/config.toml を置き、以下のように記述します。

[build]
rustflags = ["-C", "target-cpu=native"]

もしくは環境変数で渡してあげても良いです。

RUSTFLAGS='-C target-cpu=native' cargo build --release

このようにすると無事popcnt命令を使ってくれました。嬉しいですね。

popcnt rax, rdi
ret

しかし、AtCoder環境上のコンパイル方法を変えることはできません。C++のようにプリプロセッサでコンパイルオプションを上書きすることもできません。仕方がないのでこうします。*7

// target_featureを指定することで、実行環境依存の命令セットが使用可能になる
// CPUが命令セットに対応するかは呼び出し元が保証する必要があるため、関数にunsafeが付く
#[target_feature(enable = "popcnt")]
unsafe fn popcnt(x: u64) -> u32 { 
    x.count_ones()
}

バッドノウハウっぽいですが、 target_feature を指定することでコンパイラに各命令セットの使用を許可することができます。popcnt以外にも色々な値が指定できるので、必要に応じてリファレンスを参照してください。*8*9*10

(この情報、明日使えるのか……?)

明日使えないビット演算

ここからは明日使えません。

tzcnt命令

Trailing Zero Count命令です。整数を2進数で見たとき、最下位の 1 までに下位側から何個 0 が続いているかを数える命令です。兄弟に先頭側の 0 の個数を数える lzcnt 命令(Leading Zero Count)もいます。

C++20以降では std::countr_zero() 、Rustでは trailing_zeros() を使います。

#include <bit>
#include <cstdint>

uint64_t x = 0b01001000;
int tz = std::countr_zero(x);
let x: u64 = 0b01001000;
let tz = x.trailing_zeros();

使用にはCPUが BMI1 命令セットに対応している必要があります。ちなみにtzcntとよく似た命令として bsf 命令(Bit Scan Forward)もあり、これはかなり古くからあります。tzcntとの違いは入力が0のときに結果が未定義である点です。*11

これだけだと使い方が謎ですが、次のpdep命令とあわせて説明します。

pdep命令

Parallel Bits Deposit 命令です。なんのこっちゃ。

元のビット列 src とマスクビット列 mask を使って、srci bit目の 0/1maski 番目の 1 の場所に持っていきます。といっても文字で説明するのは限界があるため、図を見た方が早いでしょう。

pdep命令

こちらの記事も参考にしてみてください。

qiita.com

これが何に使えるかというと、整数集合から一様ランダムに1つ選ぶ時に使えます。普通にやると候補を一旦配列に詰める必要がありますが、pdep命令を使えば全てビット列上で完結できます。

# 候補となる要素に対応した場所が1となっているビット列(1が1個以上あると仮定)
candidates = 0110001110101001

# 1の数を数えて、0以上candidates_count未満の乱数を発生させる
candidates_count = popcnt(candidates)
k = random(0, candidates_count - 1)

# kだけ左ビットシフトしてからpdepすることで、candidatesのk個目の1のところのみのビットを立てる
shifted_k = 1 << k
shifted_result = pdep(shifted_k, candidates)

# 一番下の1までの0の個数をtzcntで数えることで、選ばれた要素がわかる
result = tzcnt(shifted_result)
print(result)

整数集合から一様ランダムに1つ選ぶ

pdepに相当する高レベルな関数は用意されておらず、intrinsics(組み込み関数)として呼び出す必要があります。C++20以降/Rustともに _pdep_u32 または _pdep_u64 で呼び出します。

#include <immintrin.h>
#include <cstdint>

uint64_t x = 0b0000000011010010;
uint64_t mask = 0b0110001110101001;
uint64_t pdep = _pdep_u64(x, mask);
use core::arch::x86_64::*;

#[target_feature(enable = "bmi2")]
unsafe fn sample() {
    let x: u64 = 0b0000000011010010;
    let mask: u64 = 0b0110001110101001;
    let pdep = _pdep_u64(x, mask);
    // do something
}

なお、pdep命令や次で説明するpext命令を使用する場合は、CPUがBMI2命令セットに対応している必要があります。AtCoder上のCPUなら基本大丈夫ですが、Zen~Zen2以前のCPUでは「一応対応しているがびっくりするほど遅い」みたいなことになっていたりするので注意が必要です。

pext命令

Parallel Bits Extract 命令です。pdep命令の逆のことができます。

pext命令

これも使い所がよく分からない命令だと思いますが、今回の命令だとグリッド上の陣地の簡易連結判定の高速化に使えました。以前のAHCであったように、陣地のうちあるマスが関節点であるかどうかは周囲 3\times3 マスを見ることで簡易的に判定することができ、これは 2^9 パターンについて事前に辞書に詰めておくことで高速に求めることができます。このとき、周囲 3\times3 マスが陣地となっているかどうかは普通にやるとループなどで求める必要がありますが、pext命令を使うと1命令で求めることができます。各マスについて、周囲 3\times 3 マスがどれか予め求めて、ビット列として保存しておくと良いです。

ちなみに、bitboard表現として番兵を使っていない場合、外周マスの扱いが厄介です。pextで集めようとしても、対応するbitがbitboard上に存在しないためです。一方で今回の問題は 10\times 10 グリッドで、これに番兵を入れて 12\times 12 にすると64bit整数2つに収まらなくなるので、それもやりたくありません。仕方がないので、bitboard上のマスだけをpextで集めた後、pdepでばら撒くことにします。何を言っているか分からないと思うので、コードを見てください。

# 以下のような2x4グリッドを考える
# 3210
# 7654
# マス0の周囲3x3マスで盤面内のマスは0, 1, 4, 5の4つ
dict_a[0] = 00110011

# マス0の周囲3x3マスについて以下のように名前を付ける(マス0がeにあるとする)
# cba
# fed
# ihg
# 周囲3x3マスのうち、盤面内のマスはe,f,h,iの4つ
dict_b[0] = 110110000

# マス0周りでpextをすることで、e,f,h,iがそれぞれ0,1,2,3bit目に集まる
# territory[i]はプレイヤーiの陣地をbitboardで持ったもの
x = pext(territory, dict_a[0])

# pdepを使うことで、e,f,h,iをそれぞれ4,5,7,8bit目に持っていく
y = pdep(x, dict_b[0])

# このyで辞書を引くことで、連結性を簡易的に判定する
is_articulation = dict_3x3[y]

……と書いておいて何なのですが、

  • bitwise OR/bitshiftは高速な上、ループアンローリングすればある程度複数ポートで並列実行が可能
  • pdep/pextは単純なビット演算と比較するとやや重い*12

あたりを考慮すると、ループでやるやつに比べてちゃんと速くなっているかは謎です。よい子の皆さんはちゃんと性能計測をしてから使いましょう。*13

おわりに

というわけで、明日使えないビット演算の紹介をしてみました。

pdep/pextのように、一見どこに使えるか見当も付かない命令の活用方法を見出すのは楽しいですね。皆さんも是非頭をひねって活用方法を見つけてあげてください!

*1:記事出すの大遅刻ですみません。色々忙しくて……。

*2:ちなみに現代のCPUでは、命令間の依存関係がない限り、単純なビット演算は複数のポートを使って複数個並列で処理できることが多いです。uops.infoというサイトに各命令ごとのレイテンシ・スループットがまとめられているので、興味のある方は見てみてください。

*3:縦H+2マス、横W+2マスのグリッドと見なして、外周を壁だと思うと良いです。実装は楽ですが、ビットが余分に必要になるのが玉に瑕です。

*4:C++17以前では builtin_popcount() または builtin_popcountll() を使います。

*5:Pythonについては詳しくないのですが、Pythonでのpopcntのやり方を探すより先にLLMに頼んで他の言語に書き換えてもらうのが先な気がします。

*6:AtCoder環境上ではコンパイルするマシンと実行するマシンのマイクロアーキテクチャが同一であるという保証がされていないことからこのような判断になったと記憶していますが、C++がそれを無視しているのでRustもこのフラグを付けて良い気がしてきました。次回言語アップデート時に提案したい……。

*7:is_x86_feature_detected!() を使って動的に分岐する方法もありますが、どうせAtCoder上でしか動かさないし、ほぼ間違いなく対応命令セットを持っているはずなので静的に解決します。

*8:target_feature を付け忘れてもフォールバックコードが動くのが罠で、「なんか遅いなあ……」と首をかしげることになります(一敗)

*9:このコードから作られたバイナリを非対応CPU上で動かすと未定義動作になるので気を付けてください(一敗)

*10:target_feature を指定した関数は、呼び出し元がそのfeatureをサポートしない場合はインライン化が行われないため、呼び出し元にも同じfeatureのサポートを要求しないと遅くなる可能性があることに注意してください(一敗)

*11:CPUがBMI1に対応していなくてもコンパイラくんが分岐なしで良い感じのアセンブリにしてくれるっぽいので、そんなにオーバーヘッドは大きくないです。

*12:pdep/pextはレイテンシ3/スループット1となっているCPUが多いようです。

*13:だってpext命令の活用方法が見付かって嬉しかったんだもん……。