NumPyを用いた配列に対するelementwiseなbitcountの実装
TL;DR
- PythonとNumPyだけで配列に対してelementwiseにbitcount処理を高速に行いたいよ
- 前計算したり、bit演算だけを使うアルゴリズムを使うとNumPyと相性よくそれなりに高速で動作したよ
- CuPy使うともっと速くなるよ、カーネル融合すると更に速くなるよ
- そもそも、次のNumPyのリリースで組み込みでbitcount処理がサポートされるみたいだから、それを使うのが一番楽で速いよ
1. 概要
プログラムのパフォーマンスを改善したいとき、bool型の配列を64要素ごとにまとめて整数型とみなしてまとめて処理することで高速化をしたい場合があると思います(いわゆるbitset高速化)。 bitsetを扱う場合、各種のbit演算(AND,OR,NOT,XOR,bitshift,bitcount)を行える必要があります。そして実際の場面では配列に格納された各要素に対してこのような操作を行うことが必要になります。
Pythonでこれを実現するにはどうすればよいでしょうか? Pythonでのint型に対してのbitcount処理自体はPython3.10にて組み込みでサポートされています。
しかしご存知の通りPythonでforループを使用すると遅いですので、配列に対してelementwiseに操作を行いたい場合はNumPyを使用することになります。
幸いNumPyでもbitwise演算はほとんど用意されています。
Binary operations — NumPy v1.26 Manual
しかし1が立っているbit数をカウントするbitcount処理だけはNumpy=1.26の時点ではまだ実装されていません。bitcount処理は結構使いどころがあり、例えば(AND,OR,NOT,XOR)などの演算を行った最後に 1が立っているbitの個数を集計したいときなどに使用します。
bitcount処理をnumpyで行いたい場合は自分で関数を書く必要があります。 この記事では、bitcount処理をNumPy配列に対してelementwiseに行う場合、どのように実装するのが一番速いのか検証します。
2. 比較した方法
この記事では以下の方法でそれぞれ性能評価しました。
- bool型配列に対する単純なsumで計算する方法(ベンチマーク)
- python組み込みのbitcount関数を使用する方法
- 前計算しておいたテーブルを活用する方法
- bitwise演算の組み合わせによる効率的なアルゴリズムを使用する方法
- (おまけ)Cupyを使用した高速化
- (おまけ)近日Numpyに追加される予定のbitcountのuniversal functionを使用する方法
3. 実験設定
長さ$n=200,000$のランダムな符号なし64bit整数の配列に対して、すべての要素のbitcountをした結果を返す処理の速度を比較します。
def elementwise_bitcount(arr:np.ndarray) -> np.ndarray: # 何らかのbitcountアルゴリズム return res #以下の処理時間を計測 arr = np.random.randint(low=0, hight=2**63-1, size=n, dtype=np.uint64) res = elementwise_bitcount(arr)
実験環境は以下の通りです。
- Python=3.11
- Numpy=1.26
- CPU: 12th Gen Intel(R) Core(TM) i9-129000K
- GPU: NVIDIA GeForce RTX 3080
4. 実験結果
まずは1~4までの方法で実際に処理時間の計測結果と分析を行います。5と6の方法については後述します。
1.bool型配列に対する単純なsumで計算する方法(ベンチマーク)
まずはベンチマークとして、各要素を64bit整数ではなく長さ64のbool型配列で表現した、$n×64$のbool型配列に対して、単純に足し算を行うことでbitcountを実現する方法を試してみます。単純に考えるとこれは各要素に対して64回の演算を行っているので効率は悪そうです。
def naive_bitcount(arr : np.ndarray) -> np.ndarray: res = np.sum(arr, axis=1) return res
処理時間の計測結果は1回あたり平均7455 μsでした。
2.Python組み込みのbitcount関数を使用する方法
Python組込みのbitcount関数をnumpyのvectorize関数でベクトル化して適用してみます。NumPyのvectorize関数は本質的にはpythonのforループを回しているだけなのでそこまで速いわけではないですが、組込みのbitcount関数が十分速いので期待できそうです。
f = np.vectorize(lambda x: x.bit_count()) def pybuiltin_bitcount(arr: np.ndarray): return f(arr)
処理時間の計測結果は1回あたり平均14877 μsでした。
むしろベンチマークの方法に比べて2倍遅くなってしまいました。単純にvectorizeするだけでは速くはならないようです。
3.前計算しておいたテーブルを活用する方法
次は予め全ての整数に対してbitcountの結果を前計算しておいて、その結果を毎回参照する方法です。とはいえ64bitで表現できる整数すべてについて前計算すると膨大なメモリが必要になります。
そこで16bit整数で表現できる0~65535までbitcountの結果を前計算しておき、64bit整数を16bitずつに分けてそれぞれのbitcountを足すという方法を取ります。(16bitで区切るのは単純にキリが良いからで、メモリが潤沢に使用できる場合だともっとたくさんの桁まで前計算してもよいかもしれません。)
またNumPyのfancy indexing機能で上手く書くことで、前計算した結果を配列の各要素に適用することができます。
# 16bitで表現できるすべての整数のbitcountの結果を前計算しておく precalc_16bit = np.array( [n.bit_count() for n in range(2**16 - 1)], dtype=np.uint8 ) def precalc_bitcount(arr: np.ndarray): # fancy indexingでarrの各要素を前計算した結果に置き換えた配列にする arr = precalc_16bit[arr] return arr.sum(axis = 1)
処理時間の計測結果は1回あたり平均3426 μsでした。
ベンチマークの方法に比べて2倍くらい速くなっています。16bitごとに分割した分、最後にsumを取る分の計算回数は増えますが、前計算によりbitcount計算が省略できたのと、NumPyで完結するよう処理できた部分が高速化に貢献していますね。
4.bitwise演算の組み合わせによる効率的なアルゴリズムを使用する方法
実はbitcount処理をbitwise演算の組み合わせだけで実現する効率的なアルゴリズムが知られています。詳細なアルゴリズムの説明は
などを参照ください。以下のコードは実際に合計17回のbit演算でuint64型の整数をbitcountした結果を返すことができます。
def efficient_bitcount(n: np.uint64) -> np.uint64: n -= (n >> 1) & 0x5555555555555555 n = (n & 0x3333333333333333) + ((n >> 2) & 0x3333333333333333) n = (n + (n >> 4)) & 0x0F0F0F0F0F0F0F0F n += n >> 8 n += n >> 16 n += n >> 32 return n & 0x7F
上記の関数はif文や参照を使用していないbit演算だけで構成されているので、引数のnをそのままNumPyの配列に変えてあげれば自動的に各要素で同じ計算を行ってくれるため、NumPyで完結する処理だけでelementwiseにbitcount処理を行うことができます。
結果は1回あたり平均3292 μsでした。
ベンチマークの方法に比べて2倍くらい速くなっており、また前計算で求める方法よりも少し速くなっています。bool型配列の場合と比べて配列サイズが1/64になっているとはいえ、bitcointのために17回の演算を行っているため思ったより高速化できていない印象です。
これは単純に計算回数の問題だけではなくNumPyの性質も関係してきます。この方法では17回の計算の度にNumPyの関数の呼び出し、それぞれの計算の中間結果の保存のためのメモリ確保といった色々なオーバーヘッドが発生しています。*1
5. まとめ.Part1
方法 | 処理時間(μs) | 高速化倍率 |
---|---|---|
1. ベンチマーク | 7455 | 1.0x |
2. python組み込みのbitcount関数を使用する方法 | 14877 | 0.5x |
3. 前計算しておいたテーブルを活用する方法 | 3426 | 2.17x |
4. bitwise演算の組み合わせによる効率的なアルゴリズムを使用する方法 | 3292 | 2.26x |
というわけで現在のPython3.11,NumPy1.26の時点では4番の方法で計算するのが最も速いという結果になりました。メモリが潤沢に使える場合は3の方法で前計算する整数の範囲を更に広げることでより高速化できるかもしれません。ともあれ、ベンチマークの方法よりは大体2倍程度は速くなるので試してみる価値はあると思います。
6.おまけ
上記の2倍程度の高速化では満足できない人のために更に頑張ってみましょう。
CuPyを使用した高速化.Part1
先ほど、4番の方法はNumPyで完結する処理だけで書けることを述べました。そこでNumPyの処理をCuPyで書き換えることでGPUを活用することができます。特に今回のような単純なelementwiseな演算については並列化の恩恵をかなり受けることが可能なので、高速化が期待できます。CuPyはインストールが若干面倒ですが、インストールができてしまえば今回の場合ほとんど既存のコードに手を加えることなくCuPy化が可能です。
import cupy as cp def efficient_bitcount(arr: cp.uint64) -> cp.uint64: arr -= (arr >> 1) & 0x5555555555555555 arr = (arr & 0x3333333333333333) + ((arr >> 2) & 0x3333333333333333) arr = (arr + (arr >> 4)) & 0x0F0F0F0F0F0F0F0F arr += arr >> 8 arr += arr >> 16 arr += arr >> 32 return arr & 0x7F
処理時間の計測結果は1回あたり平均194 μsでした。*2
そもそも配列のサイズが大きめのでGPUのほうが有利なのは当然ですが、CPUの場合と比較して16倍の高速化を達成しています。
CuPyを使用した高速化.Part2
CuPyを使用することで劇的に高速化できることはわかりましたが、先ほど述べたようにCuPyもNumPyと同様に計算の度に関数の呼び出し、それぞれの計算の中間結果の保存のためのメモリ確保といった色々なオーバーヘッドが発生します。CuPyにはカーネル融合という機能が用意されており、これを使うことで更に高速化することができます。詳細は以下を参照
CuPy カーネル融合の拡張 - Preferred Networks Research & Development
User-Defined Kernels — CuPy 12.2.0 documentation
今回の場合は以下のように自作のbitcount用のカーネルを作成し呼び出します。
cp_kernel_fusion_bitcount = cp.ElementwiseKernel( "uint64 x", "uint64 z", """ z -= (x >> 1) & 0x5555555555555555; z = (z & 0x3333333333333333) + ((z >> 2) & 0x3333333333333333); z = (z + (z >> 4)) & 0x0F0F0F0F0F0F0F0F; z += z >> 8; z += z >> 16; z += z >> 32; """, "bitcount_kernel", ) cp_kernel_fusion_bitcount(x)
処理時間の計測結果は1回あたり平均22 μsでした。
カーネル融合をしないCuPyの方法より10倍以上速くなりました。今回のようにCuPyでの演算を何回も行う場合はカーネル融合の恩恵が大きいようです。
近日Numpyに追加される予定のbitcountのuniversal functionを使用する方法
CuPyとカーネル融合を使って高速化できることはわかりましたが、CPUとNumPyだけではこれ以上高速化は無理そうでしょうか。 繰り返しになりますが、現状の方法のボトルネックはbitcount時のNumPy配列に対する17回の演算です。これがCuPyのカーネル融合のようにNumPy内部で計算されるようになれば、高速化できそうです。
実は現在開発中のNumPyではbicount処理がNumPy組込みのuniversal functionとして実装される予定です! numpy.org
開発中のNumPyの機能をお手軽に試すには、
GitHub - numpy/numpy: The fundamental package for scientific computing with Python.
で用意されているGihub Codespacesの環境でNumPyをソースからビルドすることで使用することができます。*3
Codespacesに入ったらnumpyのリポジトリに入り
pip install .
でmainブランチにあるNumPyをビルドします。ビルドができたらnumpyというフォルダができていると思うのでその中でPythonを起動することでmainブランチにあるNumPyの関数が使えるようになります。
実際にCodespaces内で以下のコードを実行して、4の方法と比較してみます。
import numpy as np #開発中のNumpyをimport np.bitwise_count(arr)
処理時間の計測結果は同じくCodespaces内で実行したbitwise演算の組み合わせによる効率的なアルゴリズムを使用する方法と比べて、10倍ほど速くなりました。(実験環境がローカルと異なるので倍率で比較しています。)
よって新たなNumPyがリリースされれば、bitwise_countを使うのが一番楽に高速化できるということになります。
まとめ.Part2
今回試した方法をすべてまとめました。GPUが使えて配列のサイズが大きい場合にはCuPyのカーネル融合による方法が一番速いです。 NumPyだけの場合、次のリリースで追加されるbitwise_countを使うのが速いですが、リリースされるまでは3や4の方法を使うのが良さそうです。*4
方法 | 処理時間(μs) | 高速化倍率 |
---|---|---|
1. ベンチマーク | 7455 | 1.0x |
2. python組み込みのbitcount関数を使用する方法 | 14877 | 0.5x |
3. 前計算しておいたテーブルを活用する方法 | 3426 | 2.17x |
4. bitwise演算の組み合わせによる効率的なアルゴリズムを使用する方法 | 3292 | 2.26x |
5.1 CuPy | 194 | 38.4x |
5.2 CuPy(カーネル融合) | 22 | 338x |
6. numpy.bitwise_count | 329(推定) | 22.6x |