I have a written a c-extension for the numpy library which is used for computing a specific type of bincount. From the lack of a better name, let's call it fast_compiled and place the method signature in numpy/core/src/multiarray/multiarraymodule.c inside array_module_methods:
{"fast_compiled", (PyCFunction)arr_fast_compiled,
METH_VARARGS | METH_KEYWORDS, NULL},
And the actual implementation inside numpy/core/src/multiarray/compiled_base.c (and compiled_base.h):
NPY_NO_EXPORT PyObject *
arr_fast_compiled(PyObject *NPY_UNUSED(self), PyObject *args, PyObject *kwds)
{
PyObject *list_obj = NULL, *strides_obj = Py_None;
PyArrayObject *list_arr = NULL, *ans = NULL, *strides_arr = NULL;
npy_intp len, ans_size, total_size;
npy_intp i, j, k;
double *dans, *weights;
npy_intp* strides;
static char *kwlist[] = {"weights", "strides", NULL};
if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O",
kwlist, &list_obj, &strides_obj)) {
goto fail;
}
list_arr = (PyArrayObject *)PyArray_ContiguousFromAny(list_obj, NPY_DOUBLE, 2, 2);
if (list_arr == NULL) {
goto fail;
}
len = PyArray_DIM(list_arr, 0);
weights = (double *)PyArray_DATA(list_arr);
ans_size = 2*len-1;
ans = (PyArrayObject *)PyArray_ZEROS(1, &ans_size, NPY_DOUBLE, 0);
if (ans == NULL) {
goto fail;
}
dans = (double *)PyArray_DATA(ans);
NPY_BEGIN_ALLOW_THREADS;
if (strides_obj == Py_None) {
for (i = 0; i < len; ++i) {
k = i * len;
for (j = i; j < i + len; ++j, ++k) {
dans[j] += weights[k];
}
}
Py_DECREF(list_arr);
}
else {
total_size = len*len;
strides_arr = (PyArrayObject *)PyArray_ContiguousFromAny(
strides_obj, NPY_INTP, 1, 1);
strides = (npy_intp *)PyArray_DATA(strides_arr);
for (i = 0; i < total_size; ++i) {
dans[strides[i]] += weights[i];
}
Py_DECREF(list_arr);
Py_DECREF(strides_arr);
}
NPY_END_ALLOW_THREADS;
return (PyObject *)ans;
fail:
Py_XDECREF(list_arr);
Py_XDECREF(strides_arr);
Py_XDECREF(ans);
return NULL;
}
The method takes one required positional argument weights and one optional keyword argument strides. Depending on if strides is specified, it will use a different (equivalent) way of computing the answer.
I am curious to why precomputing the strides and specifying it as the keyword argument is slower than computing the stride in a nested for-loop. I.e. Why is this:
for (i = 0; i < total_size; ++i) {
dans[strides[i]] += weights[i];
}
Slower than this:
for (i = 0; i < len; ++i) {
k = i * len;
for (j = i; j < i + len; ++j, ++k) {
dans[j] += weights[k];
}
}
Here is how I computed my benchmark:
import numpy as np
import perfplot
def fast_compiled(args):
A, _ = args
return np.fast_compiled(A)
def fast_compiled_strides(args):
A, strides = args
return np.fast_compiled(A, strides=strides)
def setup(n):
A = np.random.normal(size=(n, n))
strides = np.arange(n*n)
strides = np.lib.stride_tricks.sliding_window_view(strides, (n,))
strides = strides[:n]
strides = strides.flatten() # make sure it is continous
return A, strides
perfplot.show(
setup=setup,
kernels=[fast_compiled, fast_compiled_strides],
n_range=[2 ** k for k in range(3, 15)],
xlabel='n',
relative_to=0,
)

