競プロ備忘録

競プロerの備忘録

ABC91D - Two Sequences

どうしても解けなかった。ただ、SIMDで無理やり高速化できれば通る予感がしていて、実際解説を見るとSIMD高速化で通せる可能性に言及されている。
そんな訳なので、大して経験もないSIMDで突っ張って通してみた。

解法

凝ったことは何もせず単に二重ループを回し、問題の定義通りに足し算とXOR演算を行うコードを書く。
このままでは当然TLEなので、AVX2を有効化し、256bitのYMMレジスタに8個分の値を詰め、8並列で高速化を図る。8要素単位で計算を行うのでハンパが出ることがあるが、そればっかりはどうしようもないので別途計算する。

YMMレジスタへのLoadだが、これには_mm256_load_si256か、_mm256_loadu_si256が使えそう。前者はSourceのアドレスが32Byteでアラインメントされている必要があり、アラインメント制約を守らないとSIGSEGVで死ぬ。後者はそのような制約はない。
当然、アラインメント制約がついている分前者のほうがパフォーマンスが良いが、proconioで読み込んだVecが32Byteでアラインメントされている保証などあるはずがないので、32Byte Alignedが保証された配列を用意してコピーを行う前用意の必要があり、そのオーバーヘッドと比較してどうなのかはわからない。とりあえず今回は32Byte Alignedな配列を用意した。

逆に計算済みの値をYMMレジスタからStoreするには_mm256_store_si256か、_mm256_storeu_si256が使える。Load同様、前者はアラインメント制約付き、後者は制約なしである。
StoreはLoadと違い、気を遣えば今回は1回しかやらなくてよいので、どちらを使っても良い。ただ、気分的に前者を使った。

足し算には_mm256_add_epi32が使え、XOR演算には_mm256_xor_si256が使える。これはただやるだけでよい。

これだけだと全要素の組を計算することはできないので、何とか工夫して、 a _ {i} b _ {j}の組を全部計算しないといけない。

まず考えたのは、 bの要素を8回repeatしてソートし、8要素ずつ単に計算する方法。
これではダメなようで、最大ケースっぽいケースでいくつか落ちた。前用意が重すぎたのでは?と考えているが、原因はわからない。あるいは、レジスタのLoad回数がかなり多いため、その分のオーバーヘッドという可能性も大きい(というかこっちのほうが有力かも)。

次は bはそのまま回し、ループを展開する方法。YMMレジスタの一定間隔ごとの値を並べ替える命令(shufflepermuteというやつ)があるので、これを使って自力で値を並べ替えれば、ループが展開できる。

shuffle命令は256bitの幅のうち、上位下位それぞれ128bitずつの範囲内で値の並べ替えを実現できる。今回は32bitの値を並べ替えたいので_mm256_shuffle_epi32が使える。
一方permute命令は128bitの壁を越えて並べ替えが可能で、今回は32bitの値を並べ替えるために_mm256_permutevar8x32_epi32が使える。
permuteのほうが便利ではあるが、shuffleは128bitの壁がある分permuteよりは性能が良い(良かったはず…)ので、なるべくshuffleを使って並べ替える。すると、

 (7, 6, 5, 4, 3, 2, 1, 0) > (4, 7, 6, 5, 0, 3, 2, 1) > (5, 4, 7, 6, 1, 0, 3, 2) > (6, 5, 4, 7, 2, 1, 0, 3)

となるようにshuffleを利用して上下128bit内の並べ替えを行い、その後1回だけpermuteを行って

 (6, 5, 4, 7, 2, 1, 0, 3) > (3, 2, 1, 0, 7, 6, 5, 4)

と並べ替え、その後は先述のshuffleと同じ並べ替えをやれば、全通りの足し算とXOR演算が実現できる。

結果的には、ここまでの工夫でTLに間に合わせることができた。(実装例:Submission #37118454 - AtCoder Beginner Contest 091)
自身が以前解説通りに書いたC++コードよりほんのちょっと遅い程度の速度なので、まあまあ満足するレベルでは動作していると思う。

Load命令をアラインメント制約なしのものに書き換えたり、並べ替えを全部permuteで実装したりするとどうなるのかちょっと気になるところではあるが、試してない。
また、32Byte Alignedな配列を固定長で用意しているからなのか、ローカル環境でコードを回すとスタックオーバーフローに出くわしてしまった。なぜかAtCoderのジャッジ環境では正常に動くので、スタックサイズをいじれば動くのだろうと思うが、コンパイラオプションをいじってもスタックオーバーフローし続けているので、何が何やら。あんまり追及するのも面倒なので放置して、結局提出デバッグを繰り返すことになった。最終的には数回の提出でACをもえらたので、結果オーライではある。

追記

結局Loadをアラインメント制約なしのものに書き換えて投げてみたが、若干低速にはなったものの余裕をもってACできた。
ローカル環境でもスタックオーバーフローしなくなり、コードも短くなったのでとても良い。
(_mm256_loadu_si256を使った実装例:Submission #37119461 - AtCoder Beginner Contest 091)