add multithreading tests

Author: vtavana

CHANGES.rst

@@ -79,7 +79,8 @@ transform improves multi-core utilization which may offset the performance loss
 
 Added :code:`scipy.fft` backend, see #42. Fixed #46.
 
-```
+```rst
+.. code-block:: python
 Python 3.7.5 (default, Nov 23 2019, 04:02:01)
 Type 'copyright', 'credits' or 'license' for more information
 IPython 7.11.1 -- An enhanced Interactive Python. Type '?' for help.

mkl_fft/_scipy_fft.py

@@ -71,7 +71,8 @@ def get_max_threads_count(self):
 
 class _workers_data:
     def __init__(self, workers=None):
-        if workers:
+        if workers is not None:
+            workers = _workers(workers)
             self.workers_ = workers
         else:
             self.workers_ = _cpu_max_threads_count().get_cpu_count()
@@ -86,6 +87,22 @@ def workers(self, workers_val):
         self.workerks_ = operator.index(workers_val)
 
 
+def _workers(workers):
+    _cpu_count = os.cpu_count()
+    if workers < 0:
+        if workers >= -_cpu_count:
+            workers += 1 + _cpu_count
+        else:
+            raise ValueError(
+                f"workers value out of range; got {workers}, must not be"
+                f" less than {-_cpu_count}"
+            )
+    elif workers == 0:
+        raise ValueError("workers must not be zero")
+
+    return workers
+
+
 _workers_global_settings = contextvars.ContextVar(
     "scipy_backend_workers", default=_workers_data()
 )

mkl_fft/tests/from_scipy/__pycache__/test_from_scipy.cpython-310-pytest-8.3.5.pyc

@@ -1,7 +1,7 @@
 # This file includes tests from scipy.fft module:
 # https://github.com/scipy/scipy/blob/main/scipy/fft/tests/test_basic.py
 
-# TODO: remove when hfft functions are added
+# TODO: remove when hfft* functions are added
 # pylint: disable=no-member
 
 import multiprocessing

mkl_fft/tests/from_scipy/__pycache__/test_multithreading.cpython-310-pytest-8.3.5.pyc

@@ -0,0 +1,104 @@
+# This file includes tests from scipy.fft module:
+# https://github.com/scipy/scipy/blob/main/scipy/fft/tests/test_multithreading.py.py
+
+import multiprocessing
+import os
+
+import numpy as np
+import pytest
+from numpy.testing import assert_allclose
+
+import mkl_fft.interfaces.scipy_fft as fft
+
+
+@pytest.fixture(scope="module")
+def x():
+    return np.random.randn(512, 128)  # Must be large enough to qualify for mt
+
+
+@pytest.mark.parametrize(
+    "func",
+    [
+        fft.fft,
+        fft.ifft,
+        fft.fft2,
+        fft.ifft2,
+        fft.fftn,
+        fft.ifftn,
+        fft.rfft,
+        fft.irfft,
+        fft.rfft2,
+        fft.irfft2,
+        fft.rfftn,
+        fft.irfftn,
+        # TODO: fft.hfft, fft.ihfft, fft.hfft2, fft.ihfft2, fft.hfftn, fft.ihfftn,
+        # TODO: fft.dct, fft.idct, fft.dctn, fft.idctn,
+        # TODO: fft.dst, fft.idst, fft.dstn, fft.idstn,
+    ],
+)
+@pytest.mark.parametrize("workers", [2, -1])
+def test_threaded_same(x, func, workers):
+    expected = func(x, workers=1)
+    actual = func(x, workers=workers)
+    assert_allclose(actual, expected)
+
+
+def _mt_fft(x):
+    return fft.fft(x, workers=2)
+
+
+@pytest.mark.slow
+def test_mixed_threads_processes(x):
+    # Test that the fft threadpool is safe to use before & after fork
+
+    expect = fft.fft(x, workers=2)
+
+    with multiprocessing.Pool(2) as p:
+        res = p.map(_mt_fft, [x for _ in range(4)])
+
+    for r in res:
+        assert_allclose(r, expect)
+
+    fft.fft(x, workers=2)
+
+
+def test_invalid_workers(x):
+    cpus = os.cpu_count()
+
+    fft.ifft([1], workers=-cpus)
+
+    with pytest.raises(ValueError, match="workers must not be zero"):
+        fft.fft(x, workers=0)
+
+    with pytest.raises(ValueError, match="workers value out of range"):
+        fft.ifft(x, workers=-cpus - 1)
+
+
+def test_set_get_workers():
+    cpus = os.cpu_count()
+    # scipy default is 1 but mkl_fft default is max number of threads
+    assert fft.get_workers() == cpus
+    with fft.set_workers(4):
+        assert fft.get_workers() == 4
+
+        with fft.set_workers(-1):
+            assert fft.get_workers() == cpus
+
+        assert fft.get_workers() == 4
+
+    # scipy default is 1 but mkl_fft default is max number of threads
+    assert fft.get_workers() == cpus
+
+    with fft.set_workers(-cpus):
+        assert fft.get_workers() == 1
+
+
+def test_set_workers_invalid():
+
+    with pytest.raises(ValueError):  # , match='workers must not be zero'):
+        with fft.set_workers(0):
+            pass
+
+    with pytest.raises(ValueError):  # , match='workers value out of range'):
+        with fft.set_workers(-os.cpu_count() - 1):
+            pass