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 3.12.0 ドキュメント

しかしご存知の通り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演算の組み合わせだけで実現する効率的なアルゴリズムが知られています。詳細なアルゴリズムの説明は

Hamming weight - Wikipedia

などを参照ください。以下のコードは実際に合計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

*1:numexprというライブラリを使用するとこれらの問題をいい感じに解決してさらに高速化できるのですが、残念ながらnumexprはbit演算がサポートされていないため今回の目的では使用できません。

*2:CuPyでの計測時はcp.cuda.Stream.null.synchronize()まで含めること

*3:ローカル環境からビルドすることもできますが、環境構築が大変です

*4:NumPyにこだわらなければCythonなどを使う方法もありそうです

STAGの備忘録

みんなブログを書いている、書いていないのは俺だけ

NumPyを用いた配列に対するelementwiseなbitcountの実装