We write numpy-module to accelerate math functions using SIMD instructions

The numpy and scipy packages provide excellent options for quickly solving various computational problems. The concept of universal functions (ufunc), working with both scalar values ​​and arrays of various dimensions, allows us to achieve high performance while maintaining the intrinsic Python language simplicity and elegance. A universal function is usually used to perform a single operation on a large array of data, which is ideal for optimization using SIMD instructions , but I could not find a complete solution based on free software that allows you to use SIMD for calculating mathematical functions like numpy sine, cosine and exponent. I didn’t want to implement the algorithms for calculating these functions from scratch, but fortunately there were several free libraries in the C language on the Internet. Overcoming laziness , I decided to write my own numpy-module, offering universal functions for sine, cosine and exponential functions. For details and test results, welcome under cat.

A bit about SIMD instructions

SIMD instructions allow you to simultaneously perform the same set of operations on several numbers written in one register. Thus, you can process several numbers at once in one clock cycle and potentially increase productivity by several times.
For example, the Advanced Vector Extensions (AVX) set of SIMD instructions allows you to perform operations with 256-bit registers, each of which can include eight 32-bit floating-point numbers (single precision numbers) or four 64-bit (double precision numbers) . The set of operations is rather modest, mainly it is addition, subtraction, multiplication and division. Accurate and fast implementation of trigonometric functions using these operations is not a trivial task and, in order not to reinvent the wheel, you should use some kind of ready-made library. In addition to the proprietary Intel MKL (which already knows how to work with numpy), there were only three options ( one , two , three ).
The first option is a header file with a very modest set of functions, with almost no documentation and tests. The second option is the C ++ library vecmathlib, which for some reason I stubbornly refused to compile, despite using the recommended GCC-4.7 compiler. Option promising, but still seems raw. The third option, the SLEEF library, was found thanks to vecmathlib, which uses its code base. I immediately liked the variant with the simplicity and clarity of the code, as well as with an abundance of tests.

Motivation test

To get enough motivation to write a module, and at the same time to deal with using SLEEF, I decided to compare the speed of calculating the sine in the “C” language when using SLEEF with the standard math.h. Naturally, we are talking about the element-by-element calculation of the sine for a large data array.
Unfortunately, there is practically no documentation and examples in SLEEF, but there are quite clearly written tests, so it was easy to understand the use of the library. The source code SLEEF consists of four directories: java , purec , simd and tester . In addition, there is a README file with a brief description of the library and a general Makefile pulling the Makefile from the listed directories. Naturally, I was most interested in the simd directory, which, as the name suggests, contained functions optimized using SIMD instructions.
From the Makefile in the simd directory, simd clear that 4 variants of SIMD instructions are supported: SSE2 , AVX , AVX2 and FMA4 . The function prototypes are defined in the sleefsimd.h file, and the required set of instructions is selected when compiling with the -DENABLE_SSE2 , -DENABLE_AVX , -DENABLE_AVX2 or -DENABLE_FMA4 . The makefile collects executables for testing functions using each of the instruction sets: iutsse2 , iutavx , iutavx2 or iutfma4 . These files are called from the universal tester program (from the tester directory) and execute the commands received from the tester. The implementation of the commands is in the iut.c file, from where the use of the library becomes obvious.

For double precision numbers ( double ), you need to define an array of length VECTLENDP , fill in the arguments of the function of interest and pass it to the vloadu function, which will copy them to the required place for using SIMD and return the value of type vdouble . We pass the xsin value to the vdouble function, which calculates the sine value for all VECLENDP arguments at once and returns vdouble again. The result is unpacked into an array from double using the vstoreu function.
For those who want to check the SLEEF on my machine, I cite the full source code of the program, which I wrote to assess the potential acceleration from using SIMD with the help of SLEEF.

The most advanced set of commands supported on my computer is AVX, so I compiled a program (recorded in the file simd/speedtest.c in the SLEEF sources) with the following command:

 gcc -O3 -Wall -Wno-unused -Wno-attributes -DENABLE_AVX -mavx speedtest.c sleefsimddp.c sleefsimdsp.c -o speedtest -lm 

I expected the acceleration about 4 times, but the result exceeded all my expectations:

 Testing sin, 400000000 values Finish sin, spent time = 14.95 sec Testing xsin Finish xsin, spent time = 1.31 sec Speed ratio: 11.412214 Max absolute delta: 5.55112e-17, relative delta 1.58441e-16 

Acceleration is more than 10 times , with a relative error of calculation of less than 2 · 10 ^-16 (the order of accuracy of the double itself), on one processor core! Of course, in a real application, the acceleration will be less, but the motivation for writing your numpy module is already enough.
')

A few words about universal functions

In numpy, data is represented as multidimensional arrays. Universal functions work elementwise with arrays of any dimension, and in the case of several parameters, their dimension may not coincide. The parameters of the universal function are first reduced to the same dimension in accordance with special rules (this is called Broadcasting ), and then the necessary calculations are performed element by element. The output is the largest array of dimensions.
For example, the same add function (which is automatically used when using the "+" operator for numpy arrays) allows you to add both two numbers or one-dimensional arrays, and add a number to an array or a one-dimensional array to a two-dimensional one.

More information on numpy in English can be found in the official documentation , in Russian - for example, here .

We write our numpy-module with a universal function and SIMD instructions

Numpy and skipy have a fairly convenient API and good documentation, in which for those who want to write their own numpy-module with a universal function there is a corresponding tutorial . First, we write a C-function, in a linear cycle, which calculates the value of the mathematical function of interest from a scalar argument:

Pointers to input and output numpy passes to us in the args array. In our case, the function has one input and one output, so the address of the input data is args[0] , the output is args[1] . The number of elements is passed to dimensions[0] . To navigate through the input data, you need to increment the pointer by steps[0] , on weekends - by steps[1] (it is important that the pointer is of type char , since steps specify values ​​in bytes). Unfortunately, I was not able to find in the numpy documentation that the values ​​in steps should be equal to the sizes of the corresponding data types, although the experiment showed that on my system for arrays of nonzero dimension this rule is fulfilled. In case of its violation, the calculations will be slower since additional copying of elements into and out of the tmp array is required.
The same universal function in numpy can work with different data types, but a separate C function is written for each data type. When registering a universal function, we indicate the supported types and for each combination of input and output types we pass a pointer to a C function that will work with this combination:

 static PyUFuncGenericFunction funcs[] = {&double_xsin}; static char types[] = {NPY_DOUBLE, NPY_DOUBLE}; static void *data[] = {NULL}; 

In the types array, both input and output types are funcs , so it is longer than the funcs and data arrays. The array of data pointers allows for each type combination to specify its own additional parameter, which will be passed to C-Functions as the last argument to void* data . In particular, this can be used to implement different universal functions with one C-function.
To register our universal function, you need to call PyUFunc_FromFuncAndData and pass the arrays described above ( funcs , data and types ), as well as the number of input and output arguments of the universal function, the number of supported type combinations, the name of the function in the module and the documentation string.

To build the module, I used the standard setup.py from the documentation , replacing the module name and adding the SLEEF library C files, compiler and linker flags. I saved the above source code for the module next to setup.py in a file called avxmath.c , renamed the sleef directory simd from the source code sleef and also put it next to setup.py .

To compile without installing to the system, you need to run the python setup.py build_ext --inplace command python setup.py build_ext --inplace , the result of the successful execution of which should be a ready-made module in the avxmath.so file. Now you can check the performance of our function. Run python in the same directory as avxmath.so , and check:

 >>> from numpy import array, pi >>> import avxmath >>> avxmath.sin(0) 0.0 >>> avxmath.sin(pi) 1.2246467991473532e-16 >>> avxmath.sin([0, pi/2, pi, 3*pi/2, 2*pi]) array([ 0.00000000e+00, 1.00000000e+00, 1.22464680e-16, -1.00000000e+00, -2.44929360e-16]) >>> 

Making sure that the avxmath module avxmath imported and works without errors, you can do a little test of the performance and accuracy of the new function.

So, we received an acceleration of more than 6 times with a calculation accuracy no worse than 3 · 10 ^-16 ! Replacing the calls of the xsin function xsin a simple copying of memory, it is easy to verify that the acceleration failed 10 times due to the fact that about 1 second of the 2.26 seconds of execution we received was spent on overhead. Similarly, replacing the xsin function xsin the usual sine from math.h , we find that the computation times using avxmath.sin and numpy.sin in our test will coincide with high accuracy.
Thus, using SIMD instructions, one can achieve a significant acceleration of calculations performed using numpy and scipy, simply by replacing the import of normal functions with optimized ones. And of course the source code is somewhat expanded compared with this article avxmath module available on Github by reference .
Upd : Do not use the sine of SLEEF for argument values ​​of the order of 1e10 or more (see comment )