numpy_popcount

numpy/popcount/.gitignore

@@ -0,0 +1,7 @@
+/build
+/numpy_popcount/
+*.egg-info
+*.so
+*.pyc
+*.pyo
+__pycache__

numpy/popcount/__init__.py

@@ -0,0 +1 @@
+from .numpy_popcount import popcount

numpy/popcount/numpy_popcount.c

@@ -0,0 +1,87 @@
+#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
+#include "Python.h"
+#include "math.h"
+#include "numpy/ndarraytypes.h"
+#include "numpy/ufuncobject.h"
+#include "numpy/npy_3kcompat.h"
+
+static PyMethodDef PopcountMethods[] = {
+        {NULL, NULL, 0, NULL}
+};
+
+#define CONCAT2(A, B) A##B
+#define CONCAT(A, B) CONCAT2(A, B)
+
+#define MAKE_POPCOUNT(prefix, ty) \
+    static void CONCAT(prefix,_popcount)(char** args, const npy_intp* dimensions, \
+            const npy_intp* steps, void* data) { \
+        npy_intp i; \
+        npy_intp n = dimensions[0]; \
+        char* in = args[0]; \
+        char* out = args[1]; \
+        npy_intp in_step = steps[0], out_step = steps[1]; \
+        ty tmp; \
+        for (i = 0; i < n; i++) { \
+            tmp = *(ty*)in; \
+            *((ty*)out) = __builtin_popcountl((unsigned long) tmp); \
+            in += in_step; \
+            out += out_step; \
+        } \
+    }
+
+MAKE_POPCOUNT(u8, uint8_t)
+MAKE_POPCOUNT(u16, uint16_t)
+MAKE_POPCOUNT(u32, uint32_t)
+MAKE_POPCOUNT(u64, uint64_t)
+
+PyUFuncGenericFunction funcs[4] = {
+    &u8_popcount,
+    &u16_popcount,
+    &u32_popcount,
+    &u64_popcount
+};
+
+static char types[8] = {
+    NPY_UINT8, NPY_UINT8,
+    NPY_UINT16, NPY_UINT16,
+    NPY_UINT32, NPY_UINT32,
+    NPY_UINT64, NPY_UINT64,
+};
+
+static void* data[4] = {NULL, NULL, NULL, NULL};
+
+static struct PyModuleDef moduledef = {
+    PyModuleDef_HEAD_INIT,
+    "numpy_popcount",
+    NULL,
+    -1,
+    PopcountMethods,
+    NULL,
+    NULL,
+    NULL,
+    NULL
+};
+
+PyMODINIT_FUNC PyInit_numpy_popcount(void) {
+    PyObject* m;
+    PyObject* numpy_popcount;
+    PyObject* d;
+    m = PyModule_Create(&moduledef);
+    if (!m) {
+        return NULL;
+    }
+
+    import_array();
+    import_umath();
+
+    numpy_popcount = PyUFunc_FromFuncAndData(funcs, data, types, 4, 1, 1,
+            PyUFunc_None, "popcount",
+            "numpy_popcount doc", 0);
+
+    d = PyModule_GetDict(m);
+
+    PyDict_SetItemString(d, "popcount", numpy_popcount);
+    Py_DECREF(numpy_popcount);
+
+    return m;
+}

numpy/popcount/setup.py

@@ -0,0 +1,14 @@
+def configuration(parent_package='', top_path=None):
+    import numpy
+    from numpy.distutils.misc_util import Configuration
+
+    config = Configuration('numpy_popcount',
+                           parent_package,
+                           top_path)
+    config.add_extension('numpy_popcount', ['numpy_popcount.c'])
+
+    return config
+
+if __name__ == "__main__":
+    from numpy.distutils.core import setup
+    setup(configuration=configuration)