Multiprocess performance of linear algebra and FFTs in Python

python

data analysis

signal processing

Benchmarking NumPy backends (Intel MKL) and multiprocessing for linear-algebra and FFT workloads, and when spreading work across cores actually pays off.

Author

Anderson M. Amaral

Published

January 16, 2023

Numpy may use several backends for improving the calculation speed. Which one is currently active? Also, let’s see how faster can we improve the numpy calculations using more optimized routines. In particular, it should be noticed that after installing Intel’s MKL optimized numpy there was a significant improvement in the speed. While the number of cores used in the computation must be set before loading numpy (as in the examples below), the examples below do not use all cores at once.

Basic implementation

# See http://mitrocketscience.blogspot.com/2018/11/automatic-mulit-threading-with-python.html

# Also notice that I've installed Intel's mkl as suggested at
# https://www.intel.com/content/www/us/en/developer/articles/technical/using-intel-distribution-for-python-with-anaconda.html

import os

NTHREADS = '12' # Value set as string

# Export system variables before loading numpy
for i in ['OMP_NUM_THREADS', 'OPENBLAS_NUM_THREADS', 'MKL_NUM_THREADS']:#,
          #'VECLIB_MAXIMUM_THREADS', 'NUMEXPR_NUM_THREADS']:
    os.environ[i] = NTHREADS

import numpy as np

# Defining two big matrices for testing

matrix_shape = (2048, 2048)
a = np.random.random(matrix_shape) + 1j * np.random.random(matrix_shape)
b = np.random.random(matrix_shape) + 1j * np.random.random(matrix_shape)

np.show_config()

blas_mkl_info:
    libraries = ['mkl_rt']
    library_dirs = ['C:/Users/Anderson/anaconda3\\Library\\lib']
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = ['C:/Users/Anderson/anaconda3\\Library\\include']
blas_opt_info:
    libraries = ['mkl_rt']
    library_dirs = ['C:/Users/Anderson/anaconda3\\Library\\lib']
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = ['C:/Users/Anderson/anaconda3\\Library\\include']
lapack_mkl_info:
    libraries = ['mkl_rt']
    library_dirs = ['C:/Users/Anderson/anaconda3\\Library\\lib']
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = ['C:/Users/Anderson/anaconda3\\Library\\include']
lapack_opt_info:
    libraries = ['mkl_rt']
    library_dirs = ['C:/Users/Anderson/anaconda3\\Library\\lib']
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = ['C:/Users/Anderson/anaconda3\\Library\\include']

Basic numpy (after adding Intel’s MKL) using single core

While I don’t have the record for the benchmarks below before adding Intel’s MKL, some of the operations below were already significantly improved at single-core level with Intel’s optimizations.

# How long does it take to calculate the dot product?
%timeit np.dot(a, b)

1.27 s ± 4.66 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Looking at Windows’ resources monitor, it can be clearly seen that only a single core is working at a time.

%timeit np.fft.fft(a)

37.1 ms ± 449 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit np.fft.ifft(a)

38.9 ms ± 590 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit -n 100 a * b

17.3 ms ± 1.18 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)

pyfftw

pyfftw is an interface to FFTW, which claims to be a very fast FFT library. Below we can see how fast does it perform without optimization. It should be noticed that pyfftw states that significant improvements can be seen by tuning the library calls.

import pyfftw

pyfftw.config.NUM_THREADS = 1

%timeit -n 100 pyfftw.interfaces.numpy_fft.fft(a)

63.3 ms ± 1.4 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)

so it seems that Intel’s optimized numpy is faster than pyfftw.

Numpy and pyfft

Now let’s consider increasing the MKL and pyfftw number of threads to 12 (the max in my current cpu) and running again the tests

%timeit np.dot(a, b)

322 ms ± 8.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

f"An improvement of {(1-322/1270) * 100:.2f}%"

'An improvement of 74.65%'

%timeit np.fft.fft(a)

33.4 ms ± 620 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

f"An improvement of {(1-33.4/37.1) * 100:.2f}%"

'An improvement of 9.97%'

%timeit np.fft.ifft(a)

35.1 ms ± 337 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

f"An improvement of {(1-35.1/38.9) * 100:.2f}%"

'An improvement of 9.77%'

%timeit -n 100 a * b

13.2 ms ± 242 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

f"An improvement of {(1-13.2/17.3) * 100:.2f}%"

'An improvement of 23.70%'

pyfftw.config.NUM_THREADS = pyfftw.config.multiprocessing.cpu_count()

%timeit -n 100 pyfftw.interfaces.numpy_fft.fft(a)

46.4 ms ± 259 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

f"An improvement of {(1-46.4/63.3) * 100:.2f}%"

'An improvement of 26.70%'

Some operations as the dot product have become significantly faster, but the speed increase in other expressions was not so significant. Perhaps the array size is too small for significant improvements and the process creation/destruction overhead dominates in this case. pyfftw became faster, but the numpy’s fft remains faster.

Numba

Numba is a python library that often speeds up numerical calculations by using just-int-time compilation and avoiding the overhead of dynamic types in python scripts. Here we install it and perform some benchmarks. Unfortunately it does not support ffts, which would be a significant improvement for algorithms as split-step propagation method.

!pip install numba

Requirement already satisfied: numba in c:\users\anderson\anaconda3\lib\site-packages (0.54.1)
Requirement already satisfied: llvmlite<0.38,>=0.37.0rc1 in c:\users\anderson\anaconda3\lib\site-packages (from numba) (0.37.0)
Requirement already satisfied: numpy<1.21,>=1.17 in c:\users\anderson\anaconda3\lib\site-packages (from numba) (1.20.3)
Requirement already satisfied: setuptools in c:\users\anderson\anaconda3\lib\site-packages (from numba) (58.0.4)

from numba import jit, njit, objmode

%timeit np.dot(a, b) + np.exp(a * b **2)

586 ms ± 8.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

@njit
def example():
    return np.dot(a, b) + np.exp(a * b **2)

%timeit example()

396 ms ± 11.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

f"An improvement of {(1-396/586) * 100:.2f}%"

'An improvement of 32.42%'

def example():
    np.fft.fft2(a * np.conj(b))*np.conj(b)
    np.fft.ifft2(a * b)*b
    return 0

%timeit example()

336 ms ± 3.62 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

"""
Since numba does not have native support for fft,
there is a workaround suggested at
https://github.com/numba/numba/issues/5864#issuecomment-690838747
"""

@jit(nopython=True)
def example2():
    d = a * np.conj(b)
    with objmode(d='complex128[:]'):
        d = np.fft.fft2(d)        
    d = d * np.conj(b)
    
    d = a * b
    with objmode(d='complex128[:]'):
        d = np.fft.ifft2(d)        
    d = d * b    
    return 0

%timeit example2()

277 ms ± 8.44 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

f"An improvement of {(1-277/336) * 100:.2f}%"

'An improvement of 17.56%'

@jit()
def example3():
    np.fft.fft2(a * np.conj(b))*np.conj(b)
    np.fft.ifft2(a * b)*b
    return 0

%timeit example3()

C:\Users\Anderson\AppData\Local\Temp/ipykernel_22088/3408064900.py:1: NumbaWarning: 
Compilation is falling back to object mode WITH looplifting enabled because Function "example3" failed type inference due to: Unknown attribute 'fft2' of type Module(<module 'numpy.fft' from 'C:\\Users\\Anderson\\anaconda3\\lib\\site-packages\\numpy\\fft\\__init__.py'>)

File "..\..\..\..\AppData\Local\Temp\ipykernel_22088\3408064900.py", line 3:
<source missing, REPL/exec in use?>

During: typing of get attribute at C:\Users\Anderson\AppData\Local\Temp/ipykernel_22088/3408064900.py (3)

File "..\..\..\..\AppData\Local\Temp\ipykernel_22088\3408064900.py", line 3:
<source missing, REPL/exec in use?>

  @jit
C:\Users\Anderson\anaconda3\lib\site-packages\numba\core\object_mode_passes.py:151: NumbaWarning: Function "example3" was compiled in object mode without forceobj=True.

File "..\..\..\..\AppData\Local\Temp\ipykernel_22088\3408064900.py", line 1:
<source missing, REPL/exec in use?>

  warnings.warn(errors.NumbaWarning(warn_msg,
C:\Users\Anderson\anaconda3\lib\site-packages\numba\core\object_mode_passes.py:161: NumbaDeprecationWarning: 
Fall-back from the nopython compilation path to the object mode compilation path has been detected, this is deprecated behaviour.

For more information visit https://numba.pydata.org/numba-doc/latest/reference/deprecation.html#deprecation-of-object-mode-fall-back-behaviour-when-using-jit

File "..\..\..\..\AppData\Local\Temp\ipykernel_22088\3408064900.py", line 1:
<source missing, REPL/exec in use?>

  warnings.warn(errors.NumbaDeprecationWarning(msg,

343 ms ± 10.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit example3()

344 ms ± 13.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

f"An improvement of {(1-344/336) * 100:.2f}%"

'An improvement of -2.38%'

Thus, while numba is able to deliver improvements around 20-30%, it’s lack of fft support slows down the computation significantly. In particular the nopython=True option is critical to have significant performance gains, otherwise it can even make the code slower as in the example 3 above.

Tensorflow

Since tensorflow has native support for ffts, some tests were also performed using the code below. I’ve tested it using my GPUless PC and also google colab’s cloud computing service. Since it’s not feasible to show the timeit results inline with so many configurations, the data is summarized at the end.

import tensorflow as tf

import tensorflow.experimental.numpy as tnp
tnp.experimental_enable_numpy_behavior()

# Gerando matrizes equivalentes ao teste anterior, no tensorflow

matrix_shape = (2048, 2048)

a = tf.random.uniform(matrix_shape, dtype=tf.dtypes.float64)
b = tf.random.uniform(matrix_shape, dtype=tf.dtypes.float64)

c = tf.cast(a, tf.dtypes.complex128) + 1j * tf.cast(b, tf.dtypes.complex128)
d = tf.cast(b**2 - a**2, tf.dtypes.complex128) + 1j * tf.cast(a*b, tf.dtypes.complex128)

a = c
b = d

# Shows if a gpu is available
tf.config.list_physical_devices()

tf.debugging.set_log_device_placement(True)

# Example 0
%timeit tf.experimental.numpy.dot(a, b)

tf.debugging.set_log_device_placement(True)

# Example 1
%timeit tf.signal.fft2d(a)

tf.debugging.set_log_device_placement(True)

# Example 2
%timeit tf.signal.ifft2d(a)

tf.debugging.set_log_device_placement(True)

# Example 3
%timeit -n 100 a * b

tf.debugging.set_log_device_placement(True)

# Example 4
%timeit tf.signal.fft( a * b) * b * tf.signal.ifft( a * b) * b

import pandas as pd

df = pd.DataFrame()

#times in ms

df['Home (CPU)'] = [783, 140, 173, 14.9, 218]
df['Colab (CPU)'] = [3460, 477, 605, 18.1, 503]
df['Colab (GPU-fastest)'] = [1.23, 0.266, 3.51, 0.093, 2.94]
df['Colab (GPU-slowest)'] = [1.23*14.63, 0.266*1515, 3.51, 0.093*97.98, 2.94]

df['Colab GPU % improvement'] = df['Colab (CPU)']*2./(df['Colab (GPU-fastest)']+df['Colab (GPU-slowest)']) * 100

df

	Home (CPU)	Colab (CPU)	Colab (GPU-fastest)	Colab (GPU-slowest)	Colab GPU % improvement
0	783.0	3460.0	1.230	17.99490	35994.985670
1	140.0	477.0	0.266	402.99000	236.574285
2	173.0	605.0	3.510	3.51000	17236.467236
3	14.9	18.1	0.093	9.11214	393.258549
4	218.0	503.0	2.940	2.94000	17108.843537

It’s worth to mention that the CPU available on colab performs slower than the one I have at home, but tensorflow shows a significant performance improvement when executed on the GPU. Also, when executing locally it can be seen that tensorflow code is very optimized to use all cores. While this strategy may not be optimal for a processor with a few cores, it’s certainly interesting for GPUs with their massive number of parallel processing cores.

Something to be aware of is that timeit on colab works differently, such that it reports the best time, while on jupyter lab it reports the mean time. When the calculation can buffer some sections of the calculus for speedup, the reported value may be an optimistic view of the true computation time. Is it time for a GPU?

Also, the timeit results from colab are inconsistent. Depending on the run, the values may become significantly larger or smaller. However, it seems reasonable to expect at least one order of magnitude speedup factor when using tensorflow+gpu for these calculations. However it should be possible to speedup up to 4 orders of magnitude, depending on the operations used. This would be very consistent with the performance gains of up to 1000x reported in

https://arxiv.org/abs/2010.08895

https://www.technologyreview.com/2020/10/30/1011435/ai-fourier-neural-network-cracks-navier-stokes-and-partial-differential-equations/

but without the need for NNs to solve these PDEs, but it’s just related with the TF’s improved speed when using gpus.

Pytorch

Notice that pytorch also has a native fft library, and it probably should be interesting to benchmark it as well.