Precomputing strided access pattern to array gives worse performance?

Viewed 129

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,
)

Relative runtime Runtime

2 Answers

fast_compiled is faster than fast_compiled_strides because it works on contiguous data known at compile time enabling compilers to use SIMD instructions (eg. typically SSE on x86-like platforms or NEON on ARM ones). It should also be faster because of less data cache to retrieve from the L1 cache (more fetches are needed due to the indirection).

Indeed, dans[j] += weights[k] can be vectorized by loading m items of dans and m items of weights adding the m items using one instruction and storing the m items back in dans. This solution is efficient and cache friendly.

dans[strides[i]] += weights[i] cannot be efficiently vectorized on most mainstream hardware. The processor need to perform a costly gather from the memory hierarchy due to the indirection, then do the sum and then perform a scatter store which is also expensive. Even if strides would contain contiguous indices, the instructions are generally much more expensive than loading a contiguous block of data from memory. Moreover, compiler often fail to vectorize the code or just find that this is not worth using SIMD instruction in that case. As a result the generated code is likely a less efficient scalar code.

Actually, the performance difference between the two codes should be bigger on modern processors with good compilation flags. I suspect you only use SSE on a x86 processor here and so the speed up is close to 2 theoretically since 2 double-precision floating-point numbers can be computed in a row. However, using AVX/AVX-2 would lead to a speed up close to 4 theoretically (as 4 numbers can be computed in a row). Very recent Intel processors can even compute 8 double-precision floating-point numbers in a row. Note that computing simple-precision floating-point numbers can also results in a theoretical 2x speed up. The same apply for other architecture like ARM with NEON and SVE instruction sets or POWER. Since future processors will likely use wider SIMD registers (because of their efficiency), it is very important to write SIMD-friendly codes.

This is how the inner loop for fast_compiled looks like:

dans[j] += weights[k];

In assembly:

│   0x7ffff6c35260 <arr_fast_compiled+336>  movsd  (%rbx,%rdx,8),%xmm0                                                                                                                                            │
│   0x7ffff6c35265 <arr_fast_compiled+341>  addsd  0x0(%rbp,%rdx,8),%xmm0                                                                                                                                         │
│   0x7ffff6c3526b <arr_fast_compiled+347>  movsd  %xmm0,(%rbx,%rdx,8)                                                                                                                                            │
│   0x7ffff6c35270 <arr_fast_compiled+352>  add    $0x1,%rdx                                                                                                                                                      │
│   0x7ffff6c35274 <arr_fast_compiled+356>  cmp    %rcx,%rdx                                                                                                                                                      │
│   0x7ffff6c35277 <arr_fast_compiled+359>  jne    0x7ffff6c35260 <arr_fast_compiled+336>

The first 3 instructions are SIMD instructions from the SSE2 Instructions set.

This is how the loop body for fast_compiled_strides look like:

dans[strides[i]] += weights[i]; 

In assembly:

│   0x7ffff6c35330 <arr_fast_compiled+544>  mov    (%rsi,%rdx,8),%rcx                                                                                                                                     │
│   0x7ffff6c35334 <arr_fast_compiled+548>  lea    (%rbx,%rcx,8),%rcx                                                                                                                                     │
│   0x7ffff6c35338 <arr_fast_compiled+552>  movsd  (%rcx),%xmm0                                                                                                                                           │
│   0x7ffff6c3533c <arr_fast_compiled+556>  addsd  0x0(%rbp,%rdx,8),%xmm0                                                                                                                                 │
│   0x7ffff6c35342 <arr_fast_compiled+562>  add    $0x1,%rdx                                                                                                                                              │
│   0x7ffff6c35346 <arr_fast_compiled+566>  movsd  %xmm0,(%rcx)                                                                                                                                           │
│   0x7ffff6c3534a <arr_fast_compiled+570>  cmp    %rdx,%r14                                                                                                                                              │
│   0x7ffff6c3534d <arr_fast_compiled+573>  jne    0x7ffff6c35330 <arr_fast_compiled+544> 

It looks similar to fast_compiled, except that it needs to load the strides in the beginning. But more noteworthy, they seem to both use the same SIMD instructions movsd, addsd and movsd ?

I am running numpy with gdb in the same way as described in the numpy dev. setup guide, so I assume that all compilation flags are already set.

I am using intel cpu which should have access to SSE4.2 and AVX2:

sudo lshw -c cpu

Output:

  *-cpu                     
       description: CPU
       product: Intel(R) Core(TM) i5-4690K CPU @ 3.50GHz
       vendor: Intel Corp.
       physical id: e
       bus info: cpu@0
       version: Intel(R) Core(TM) i5-4690K CPU @ 3.50GHz
       slot: CPUSocket
       size: 3727MHz
       capacity: 3900MHz
       width: 64 bits
       clock: 100MHz
       capabilities: lm fpu fpu_exception wp vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp x86-64 constant_tsc arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf pni pclmulqdq dtes64 monitor ds_cpl vmx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm cpuid_fault invpcid_single pti ssbd ibrs ibpb stibp tpr_shadow vnmi flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid xsaveopt dtherm ida arat pln pts md_clear flush_l1d cpufreq
       configuration: cores=4 enabledcores=4 threads=4
Related