Python Accelerations

Here's a whiteboard diagram we came up with which seemed to communicate our current thinking on how we are proposing doing Python performance acceleration options on large vector operations.

It's kind of "choose your own adventure". For any given vector computation, one can go as far to the right as one is capable. One can write a numpy implementation and call it good for now. One can go further and do a numba implementation. Or go further and do a ctypes compiled c extension. Or go a bit further and make this faster with OpenMP. Or go even further and make it faster with a SIMD implementation.

Some notes:

• At the end of the day, from numpy to simd, these are all just ways of operating on contiguous arrays of numbers. All these solutions operate on the same datatypes really.
• numba works on for loops, not numpy operations.
• In theory, if you have GPU hardware, you can go from numba to numba+cuda with just the cuda extension to numba, a decorator and "a little setup code". We have yet to truly demonstrate this for ourselves though.
• numba doesn't really start to show its worth unless you use prange.
• Going from numpy to numba or ctypes requires starting to thinking about datatypes.
• Developing a ctypes extension requires a compiler. (Numba has a built in compiler.)
• Speeding up a C implementation with OpenMP is truly simple.
• OpenMP naturally uses all the cores it finds, without recompiling, which is very useful for making sure one is leveraging all available hardware, but also means that performance gains tested & recorded will depend on what machine it is being run on.
• Note that the benefits of parallelization can be mostly eliminated if there is a critical section (e.g., updating a common counter) somewhere in the for loop being parallelized. Though this can usually be fixed by using thread-local storage and then adding an accumulator section at the end. GPT can help with this.
• SIMD acceleration requires a certain low-level programming comfort. Working code can all be provided by GPT/Copilot, but it can take some trial and error.
• We expect to provide SIMD implementations for AVX2, AVX512 and SSE. One can compile all of these together but one can only test on the hardware one has. Testing all 3 implementations requires getting a bit fancy, though presumably GHA can support this.
• We are actively working on a reference, stress test which demonstrates performance numbers on each of these implementations.

Examples

Let's look at some sample code for a function that ages everyone in our model. We have some expansion slots for the "not yet born" with negative age so we co a condition check in addition to an addition, which makes our code at least not totally trivial.

Naive Python

def naive_one(data: np.ndarray) -> None:
    one_day = np.float32(1.0 / 365)
    for i in range(data.shape[0]):
        if data[i] >= 0:
            data[i] += one_day
    return

Numpy

def vectorized(data: np.ndarray) -> None:
    one_day = np.float32(1.0 / 365)
    data[data >= 0] += one_day
    return

Numba

@nb.njit((nb.float32[:],), parallel=True, nogil=True)
def numbaized(data: np.ndarray) -> None:
    one_day = np.float32(1.0 / 365)
    for i in nb.****prange****(data.shape[0]):
        if data[i] >= 0:
            data[i] += one_day
    return

Compiled C (ctypes extension)

Python part:

# 'import' the module
update_ages_lib = ctypes.CDLL('./update_ages.so')
# define the interface to the function (datatypes)
update_ages_lib.update_ages_vanilla.argtypes = [
    ctypes.c_size_t,  # start_idx
    ctypes.c_size_t,  # stop_idx
    np.ctypeslib.ndpointer(dtype=np.float32, flags='C_CONTIGUOUS')
]
# Invoke
def compiled_c(data: np.ndarray) -> None:
    update_ages_lib.update_ages_vanilla(
        0,
        len( data ),
        data
    )

C Part:

const float one_day = 1.0f/365.0f;
void update_ages_vanilla(unsigned long int start_idx, unsigned long int stop_idx, float *ages) {
    for (size_t i = start_idx; i <= stop_idx; i++) {
        if( ages[i] >= 0 )
        {
            ages[i] += one_day;
        }
    }
}

C with OpenMP

#include <omp.h>

const float one_day = 1.0f/365.0f;
void update_ages_vanilla(unsigned long int start_idx, unsigned long int stop_idx, float *ages) {
    #pragma omp parallel for
    for (size_t i = start_idx; i <= stop_idx; i++) {
        if( ages[i] >= 0 )
        {
            ages[i] += one_day;
        }
    }
}

SIMD (avx2)

// avxv2
void update_ages(unsigned long int start_idx, unsigned long int stop_idx, float *ages) {
    #pragma omp parallel for
    for (size_t i = start_idx; i <= stop_idx; i += 8) {
        // Load age values into SIMD registers
        __m256 ages_vec = _mm256_loadu_ps(&ages[i]);

        // Create a vector with the one_day value replicated 8 times
        __m256 one_day_vec = _mm256_set1_ps(one_day);

        // Mask for elements greater than or equal to zero
        __m256 mask = _mm256_cmp_ps(ages_vec, _mm256_setzero_ps(), _CMP_GE_OQ);

        // Increment age values by one_day for elements greater than or equal to zero
        ages_vec = _mm256_blendv_ps(ages_vec, _mm256_add_ps(ages_vec, one_day_vec), mask);

        // Store the result back to memory
        _mm256_storeu_ps(&ages[i], ages_vec);
    }
}

SIMD (sse & avx2 & avx512 with run-time instruction set detection)

https://gist.github.com/clorton/16d8ca411c83058b31c6275a4e54f352