Numpy bitcount benchmark

benchmark.py

import timeit
import numpy as np

# Naive bitcount implementation without popcount
m1 = np.int64(0x5555555555555555)
m2 = np.int64(0x3333333333333333)
m3 = np.int64(0x0F0F0F0F0F0F0F0F)
m4 = np.int64(0x0101010101010101)

mask = np.int64(-1)
# TODO - precompute type specific hashes
s55 = np.int64(m1 & mask)  # Add more digits for 128bit support
s33 = np.int64(m2 & mask)
s0F = np.int64(m3 & mask)
s01 = np.int64(m4 & mask)
num_bytes_64 = 8

def bit_count64(arr):
    # Make the values type-agnostic (as long as it's integers)
    # baseline
    #
    # -arr = arr - ((arr >> 1) & s55)
    # +arr -= ((arr >> 1) & s55)
    #
    # before
    #     5106   32.987    0.006   32.987    0.006 hamming.py:27(bit_count64)
    #     5106   32.436    0.006   32.436    0.006 hamming.py:26(bit_count64)
    #     5106   35.277    0.007   35.277    0.007 hamming.py:26(bit_count64)

    # after
    #     5106   26.593    0.005   26.593    0.005 hamming.py:26(bit_count64)
    #     5106   26.308    0.005   26.308    0.005 hamming.py:28(bit_count64)
    #
    # reduce copies by subtract in place

    # assert arr.dtype == np.int64

    arr -= ((arr >> 1) & s55)
    arr = (arr & s33) + ((arr >> 2) & s33)

    arr += (arr >> 4)
    arr &= s0F
    arr *= s01
    arr >>= (8 * (num_bytes_64 - 1))

    return arr


# Random array of 64bit integers
arr_size = (10000, 100)
max_val = np.iinfo(np.int64).max
arr = np.random.randint(0, max_val, arr_size, dtype=np.int64)

# Benchmark np.bitwise_count
print(timeit.timeit(lambda: np.bitwise_count(arr), number=1000))

# Benchmark
print(timeit.timeit(lambda: bit_count64(arr), number=1000))

assert np.all(bit_count64(arr) == np.bitwise_count(arr))