diff --git a/clifford/_bit_helpers.py b/clifford/_bit_helpers.py new file mode 100644 index 00000000..57138135 --- /dev/null +++ b/clifford/_bit_helpers.py @@ -0,0 +1,48 @@ +""" Helpers for efficient bit operations + +Bitmaps are used in clifford as representations of basis blades, hence the need +for these operations. These are only for internal use. +""" + +from typing import Iterator + +import numba +import numba.extending +import numba.types +import numba.config + + +def set_bit_indices(x: int) -> Iterator[int]: + """ Iterate over the indices of bits set to 1 in `x`, in ascending order """ + n = 0 + while x > 0: + if x & 1: + yield n + x = x >> 1 + n = n + 1 + + +@numba.extending.intrinsic +def __builtin_popcnt(tyctx, x): + """ Emulate clang and GCC's `__builtin_popcnt` """ + if isinstance(x, numba.types.Integer): + def impl(cgctx, builder, sig, args): + x, = args + return builder.ctpop(x) + sig = x(x) + return sig, impl + + +if numba.config.DISABLE_JIT: + def count_set_bits(bitmap: int) -> int: + """ Counts the number of bits set to 1 in bitmap """ + count = 0 + for i in set_bit_indices(bitmap): + count += 1 + return count + +else: + @numba.njit + def count_set_bits(x: int) -> int: + """ Counts the number of bits set to 1 in bitmap """ + return __builtin_popcnt(x) diff --git a/clifford/_layout_helpers.py b/clifford/_layout_helpers.py index 310c3d93..b8f8cd01 100644 --- a/clifford/_layout_helpers.py +++ b/clifford/_layout_helpers.py @@ -14,41 +14,7 @@ import operator from . import _numba_utils -import numba -import numba.extending -import numba.types -import numba.config - - -@numba.extending.intrinsic -def __builtin_popcnt(tyctx, x): - """ Emulate clang and GCC's `__builtin_popcnt` """ - if isinstance(x, numba.types.Integer): - def impl(cgctx, builder, sig, args): - x, = args - return builder.ctpop(x) - sig = x(x) - return sig, impl - - -if numba.config.DISABLE_JIT: - def count_set_bits(bitmap): - """ Counts the number of bits set to 1 in bitmap """ - bmp = bitmap - count = 0 - n = 1 - while bmp > 0: - if bmp & 1: - count += 1 - bmp = bmp >> 1 - n = n + 1 - return count - -else: - @numba.njit - def count_set_bits(x): - """ Counts the number of bits set to 1 in bitmap """ - return __builtin_popcnt(x) +from ._bit_helpers import count_set_bits, set_bit_indices @_numba_utils.njit @@ -192,15 +158,7 @@ def __init__(self, blade_ids: Sequence[IdT]): def bitmap_as_tuple(self, bitmap: int) -> Tuple[IdT]: """ Convert a bitmap representation into a tuple of ids. """ - bmp = bitmap - blade = [] - n = 0 - while bmp > 0: - if bmp & 1: - blade.append(self.values[n]) - bmp = bmp >> 1 - n = n + 1 - return tuple(blade) + return tuple(self.values[n] for n in set_bit_indices(bitmap)) def id_as_bitmap(self, id: IdT) -> int: """ Convert the id of a single vector into a bitmap representation. """ diff --git a/clifford/test/test_bit_helpers.py b/clifford/test/test_bit_helpers.py new file mode 100644 index 00000000..6b5e057a --- /dev/null +++ b/clifford/test/test_bit_helpers.py @@ -0,0 +1,16 @@ +""" Tests of clifford._bit_helpers """ +from clifford._bit_helpers import count_set_bits, set_bit_indices + + +def test_count_bits(): + assert count_set_bits(0b0) == 0 + assert count_set_bits(0b1) == 1 + assert count_set_bits(0b1001) == 2 + assert count_set_bits(0b1111) == 4 + + +def test_bit_indices(): + assert list(set_bit_indices(0b0)) == [] + assert list(set_bit_indices(0b1)) == [0] + assert list(set_bit_indices(0b101)) == [0, 2] + assert list(set_bit_indices(0b101010)) == [1, 3, 5]