diff options
| -rw-r--r-- | doc/release/upcoming_changes/19355.new_feature.rst | 13 | ||||
| -rw-r--r-- | numpy/__init__.pyi | 1 | ||||
| -rw-r--r-- | numpy/core/_add_newdocs_scalars.py | 19 | ||||
| -rw-r--r-- | numpy/core/include/numpy/npy_math.h | 11 | ||||
| -rw-r--r-- | numpy/core/src/multiarray/scalartypes.c.src | 52 | ||||
| -rw-r--r-- | numpy/core/src/npymath/npy_math_internal.h.src | 86 | ||||
| -rw-r--r-- | numpy/core/tests/test_scalar_methods.py | 18 |
7 files changed, 198 insertions, 2 deletions
diff --git a/doc/release/upcoming_changes/19355.new_feature.rst b/doc/release/upcoming_changes/19355.new_feature.rst new file mode 100644 index 000000000..cfa50b7a1 --- /dev/null +++ b/doc/release/upcoming_changes/19355.new_feature.rst @@ -0,0 +1,13 @@ +``bit_count`` to compute the number of 1-bits in an integer +----------------------------------------------------------- + +Computes the number of 1-bits in the absolute value of the input. +This works on all the numpy integer types. Analogous to the builtin +``int.bit_count`` or ``popcount`` in C++. + +.. code-block:: python + + >>> np.uint32(1023).bit_count() + 10 + >>> np.int32(-127).bit_count() + 7 diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi index 463d4a713..229b6b916 100644 --- a/numpy/__init__.pyi +++ b/numpy/__init__.pyi @@ -2750,6 +2750,7 @@ class integer(number[_NBit1]): # type: ignore ) -> int: ... def tolist(self) -> int: ... def is_integer(self) -> L[True]: ... + def bit_count(self: _ScalarType) -> int: ... def __index__(self) -> int: ... __truediv__: _IntTrueDiv[_NBit1] __rtruediv__: _IntTrueDiv[_NBit1] diff --git a/numpy/core/_add_newdocs_scalars.py b/numpy/core/_add_newdocs_scalars.py index 8773d6c96..94859a9d5 100644 --- a/numpy/core/_add_newdocs_scalars.py +++ b/numpy/core/_add_newdocs_scalars.py @@ -290,3 +290,22 @@ for float_name in ('half', 'single', 'double', 'longdouble'): >>> np.{float_name}(3.2).is_integer() False """)) + +for int_name in ('int8', 'uint8', 'int16', 'uint16', 'int32', 'uint32', + 'int64', 'uint64', 'int64', 'uint64', 'int64', 'uint64'): + # Add negative examples for signed cases by checking typecode + add_newdoc('numpy.core.numerictypes', int_name, ('bit_count', + f""" + {int_name}.bit_count() -> int + + Computes the number of 1-bits in the absolute value of the input. + Analogous to the builtin `int.bit_count` or ``popcount`` in C++. + + Examples + -------- + >>> np.{int_name}(127).bit_count() + 7""" + + (f""" + >>> np.{int_name}(-127).bit_count() + 7 + """ if dtype(int_name).char.islower() else ""))) diff --git a/numpy/core/include/numpy/npy_math.h b/numpy/core/include/numpy/npy_math.h index b1e6363e3..bead0dc14 100644 --- a/numpy/core/include/numpy/npy_math.h +++ b/numpy/core/include/numpy/npy_math.h @@ -150,6 +150,17 @@ NPY_INPLACE npy_long npy_lshiftl(npy_long a, npy_long b); NPY_INPLACE npy_longlong npy_rshiftll(npy_longlong a, npy_longlong b); NPY_INPLACE npy_longlong npy_lshiftll(npy_longlong a, npy_longlong b); +NPY_INPLACE uint8_t npy_popcountuhh(npy_ubyte a); +NPY_INPLACE uint8_t npy_popcountuh(npy_ushort a); +NPY_INPLACE uint8_t npy_popcountu(npy_uint a); +NPY_INPLACE uint8_t npy_popcountul(npy_ulong a); +NPY_INPLACE uint8_t npy_popcountull(npy_ulonglong a); +NPY_INPLACE uint8_t npy_popcounthh(npy_byte a); +NPY_INPLACE uint8_t npy_popcounth(npy_short a); +NPY_INPLACE uint8_t npy_popcount(npy_int a); +NPY_INPLACE uint8_t npy_popcountl(npy_long a); +NPY_INPLACE uint8_t npy_popcountll(npy_longlong a); + /* * C99 double math funcs */ diff --git a/numpy/core/src/multiarray/scalartypes.c.src b/numpy/core/src/multiarray/scalartypes.c.src index 0d52211a8..bbbc5bfa2 100644 --- a/numpy/core/src/multiarray/scalartypes.c.src +++ b/numpy/core/src/multiarray/scalartypes.c.src @@ -219,6 +219,27 @@ gentype_multiply(PyObject *m1, PyObject *m2) } /**begin repeat + * #TYPE = BYTE, UBYTE, SHORT, USHORT, INT, UINT, + * LONG, ULONG, LONGLONG, ULONGLONG# + * #type = npy_byte, npy_ubyte, npy_short, npy_ushort, npy_int, npy_uint, + * npy_long, npy_ulong, npy_longlong, npy_ulonglong# + * #c = hh, uhh, h, uh,, u, l, ul, ll, ull# + * #Name = Byte, UByte, Short, UShort, Int, UInt, + * Long, ULong, LongLong, ULongLong# + * #convert = Long*8, LongLong*2# + */ +static PyObject * +@type@_bit_count(PyObject *self) +{ + @type@ scalar = PyArrayScalar_VAL(self, @Name@); + uint8_t count = npy_popcount@c@(scalar); + PyObject *result = PyLong_From@convert@(count); + + return result; +} +/**end repeat**/ + +/**begin repeat * * #name = positive, negative, absolute, invert, int, float# */ @@ -2316,8 +2337,7 @@ static PyMethodDef @name@type_methods[] = { /**end repeat**/ /**begin repeat - * #name = byte, short, int, long, longlong, ubyte, ushort, - * uint, ulong, ulonglong, timedelta, cdouble# + * #name = timedelta, cdouble# */ static PyMethodDef @name@type_methods[] = { /* for typing; requires python >= 3.9 */ @@ -2328,6 +2348,23 @@ static PyMethodDef @name@type_methods[] = { }; /**end repeat**/ +/**begin repeat + * #name = byte, ubyte, short, ushort, int, uint, + * long, ulong, longlong, ulonglong# + */ +static PyMethodDef @name@type_methods[] = { + /* for typing; requires python >= 3.9 */ + {"__class_getitem__", + (PyCFunction)numbertype_class_getitem, + METH_CLASS | METH_O, NULL}, + {"bit_count", + (PyCFunction)npy_@name@_bit_count, + METH_NOARGS, NULL}, + {NULL, NULL, 0, NULL} /* sentinel */ +}; +/**end repeat**/ + + /************* As_mapping functions for void array scalar ************/ static Py_ssize_t @@ -4105,6 +4142,17 @@ initialize_numeric_types(void) /**end repeat**/ /**begin repeat + * #name = byte, short, int, long, longlong, + * ubyte, ushort, uint, ulong, ulonglong# + * #Name = Byte, Short, Int, Long, LongLong, + * UByte, UShort, UInt, ULong, ULongLong# + */ + + Py@Name@ArrType_Type.tp_methods = @name@type_methods; + + /**end repeat**/ + + /**begin repeat * #name = half, float, double, longdouble# * #Name = Half, Float, Double, LongDouble# */ diff --git a/numpy/core/src/npymath/npy_math_internal.h.src b/numpy/core/src/npymath/npy_math_internal.h.src index cae84befe..dd2424db8 100644 --- a/numpy/core/src/npymath/npy_math_internal.h.src +++ b/numpy/core/src/npymath/npy_math_internal.h.src @@ -55,6 +55,29 @@ */ #include "npy_math_private.h" +/* Magic binary numbers used by bit_count + * For type T, the magic numbers are computed as follows: + * Magic[0]: 01 01 01 01 01 01... = (T)~(T)0/3 + * Magic[1]: 0011 0011 0011... = (T)~(T)0/15 * 3 + * Magic[2]: 00001111 00001111... = (T)~(T)0/255 * 15 + * Magic[3]: 00000001 00000001... = (T)~(T)0/255 + * + * Counting bits set, in parallel + * Based on: http://graphics.stanford.edu/~seander/bithacks.html#CountBitsSetParallel + * + * Generic Algorithm for type T: + * a = a - ((a >> 1) & (T)~(T)0/3); + * a = (a & (T)~(T)0/15*3) + ((a >> 2) & (T)~(T)0/15*3); + * a = (a + (a >> 4)) & (T)~(T)0/255*15; + * c = (T)(a * ((T)~(T)0/255)) >> (sizeof(T) - 1) * CHAR_BIT; +*/ + +static const npy_uint8 MAGIC8[] = {0x55u, 0x33u, 0x0Fu, 0x01u}; +static const npy_uint16 MAGIC16[] = {0x5555u, 0x3333u, 0x0F0Fu, 0x0101u}; +static const npy_uint32 MAGIC32[] = {0x55555555ul, 0x33333333ul, 0x0F0F0F0Ful, 0x01010101ul}; +static const npy_uint64 MAGIC64[] = {0x5555555555555555ull, 0x3333333333333333ull, 0x0F0F0F0F0F0F0F0Full, 0x0101010101010101ull}; + + /* ***************************************************************************** ** BASIC MATH FUNCTIONS ** @@ -814,3 +837,66 @@ npy_rshift@u@@c@(npy_@u@@type@ a, npy_@u@@type@ b) } /**end repeat1**/ /**end repeat**/ + + +#define __popcnt32 __popcnt +/**begin repeat + * + * #type = ubyte, ushort, uint, ulong, ulonglong# + * #STYPE = BYTE, SHORT, INT, LONG, LONGLONG# + * #c = hh, h, , l, ll# + */ +#undef TO_BITS_LEN +#if 0 +/**begin repeat1 + * #len = 8, 16, 32, 64# + */ +#elif NPY_BITSOF_@STYPE@ == @len@ + #define TO_BITS_LEN(X) X##@len@ +/**end repeat1**/ +#endif + + +NPY_INPLACE uint8_t +npy_popcount_parallel@c@(npy_@type@ a) +{ + a = a - ((a >> 1) & (npy_@type@) TO_BITS_LEN(MAGIC)[0]); + a = ((a & (npy_@type@) TO_BITS_LEN(MAGIC)[1])) + ((a >> 2) & (npy_@type@) TO_BITS_LEN(MAGIC)[1]); + a = (a + (a >> 4)) & (npy_@type@) TO_BITS_LEN(MAGIC)[2]; + return (npy_@type@) (a * (npy_@type@) TO_BITS_LEN(MAGIC)[3]) >> ((NPY_SIZEOF_@STYPE@ - 1) * CHAR_BIT); +} + +NPY_INPLACE uint8_t +npy_popcountu@c@(npy_@type@ a) +{ +/* use built-in popcount if present, else use our implementation */ +#if (defined(__clang__) || defined(__GNUC__)) && NPY_BITSOF_@STYPE@ >= 32 + return __builtin_popcount@c@(a); +#elif defined(_MSC_VER) && NPY_BITSOF_@STYPE@ >= 16 + /* no builtin __popcnt64 for 32 bits */ + #if defined(_WIN64) || (defined(_WIN32) && NPY_BITSOF_@STYPE@ != 64) + return TO_BITS_LEN(__popcnt)(a); + /* split 64 bit number into two 32 bit ints and return sum of counts */ + #elif (defined(_WIN32) && NPY_BITSOF_@STYPE@ == 64) + npy_uint32 left = (npy_uint32) (a>>32); + npy_uint32 right = (npy_uint32) a; + return __popcnt32(left) + __popcnt32(right); + #endif +#else + return npy_popcount_parallel@c@(a); +#endif +} +/**end repeat**/ + +/**begin repeat + * + * #type = byte, short, int, long, longlong# + * #c = hh, h, , l, ll# + */ +NPY_INPLACE uint8_t +npy_popcount@c@(npy_@type@ a) +{ + /* Return popcount of abs(a) */ + return npy_popcountu@c@(a < 0 ? -a : a); +} +/**end repeat**/ diff --git a/numpy/core/tests/test_scalar_methods.py b/numpy/core/tests/test_scalar_methods.py index 6077c8f75..eef4c1433 100644 --- a/numpy/core/tests/test_scalar_methods.py +++ b/numpy/core/tests/test_scalar_methods.py @@ -183,3 +183,21 @@ def test_class_getitem_38(cls: Type[np.number]) -> None: match = "Type subscription requires python >= 3.9" with pytest.raises(TypeError, match=match): cls[Any] + + +class TestBitCount: + # derived in part from the cpython test "test_bit_count" + + @pytest.mark.parametrize("itype", np.sctypes['int']+np.sctypes['uint']) + def test_small(self, itype): + for a in range(max(np.iinfo(itype).min, 0), 128): + msg = f"Smoke test for {itype}({a}).bit_count()" + assert itype(a).bit_count() == bin(a).count("1"), msg + + def test_bit_count(self): + for exp in [10, 17, 63]: + a = 2**exp + assert np.uint64(a).bit_count() == 1 + assert np.uint64(a - 1).bit_count() == exp + assert np.uint64(a ^ 63).bit_count() == 7 + assert np.uint64((a - 1) ^ 510).bit_count() == exp - 8 |