88 lines
2.0 KiB
C
88 lines
2.0 KiB
C
#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;
|
|
}
|