Created
August 15, 2023 14:46
-
-
Save softwaredoug/8210fa941f8b806b9eea83fd1e762d03 to your computer and use it in GitHub Desktop.
Numpy bitcount benchmark
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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)) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment