Parallel numba rime

HISTORY.rst

@@ -4,6 +4,7 @@ History
 
 X.Y.Z (YYYY-MM-DD)
 ------------------
+* Parallel Numba RIME Implementation (:pr:`186`)
 * Update classifiers and correct license in setup.py to BSD3
 
 0.2.4 (2020-05-29)
@@ -18,6 +19,7 @@ X.Y.Z (YYYY-MM-DD)
 
 0.2.3 (2020-05-14)
 ------------------
+
 * Fix incorrect SPI calculation and make predict defaults MeqTree equivalent (:pr:`189`)
 * Depend on pytest-flake8 >= 1.0.6 (:pr:`187`, :pr:`188`)
 * MeqTrees Comparison Script Updates (:pr:`160`)

africanus/config.py

@@ -0,0 +1,20 @@
+from donfig import Config
+
+
+class AfricanusConfig(Config):
+    def numba_parallel(self, key):
+        value = self.get(key, False)
+
+        if value is False:
+            return {'parallel': False}
+        elif value is True:
+            return {'parallel': True}
+        elif isinstance(value, dict):
+            value['parallel'] = True
+            return value
+        else:
+            raise TypeError("key %s (%s) is not a bool or a dict",
+                            key, value)
+
+
+config = AfricanusConfig("africanus")

africanus/conftest.py

@@ -1,9 +1,38 @@
 # -*- coding: utf-8 -*-
 
+import importlib
+import pytest
 
 from africanus.util.testing import mark_in_pytest
 
 
+@pytest.fixture
+def cfg_parallel(request):
+    """ Performs parallel configuration setting and module reloading """
+    from africanus.config import config
+
+    module, cfg = request.param
+
+    assert isinstance(cfg, dict) and len(cfg) == 1
+
+    # Get module object, because importlib.reload doesn't take strings
+    mod = importlib.import_module(module)
+
+    with config.set(cfg):
+        importlib.reload(mod)
+
+        cfg = cfg.copy().popitem()[1]
+
+        if isinstance(cfg, dict):
+            yield cfg['parallel']
+        elif isinstance(cfg, bool):
+            yield cfg
+        else:
+            raise TypeError("Unhandled cfg type %s" % type(cfg))
+
+    importlib.reload(mod)
+
+
 # content of conftest.py
 def pytest_configure(config):
     mark_in_pytest(True)

africanus/install/requirements.py

@@ -0,0 +1,47 @@
+# -*- coding: utf-8 -*-
+
+
+# NOTE(sjperkins)
+# Non standard library imports should be avoided,
+# or should fail gracefully as functionality
+# in these modules is called by setup.py
+import os
+
+# requirements
+on_rtd = os.environ.get('READTHEDOCS') == 'True'
+
+# Basic requirements containing no C extensions.
+# This is necessary for building on RTD
+requirements = ['appdirs >= 1.4.3',
+                'decorator',
+                'donfig']
+
+if not on_rtd:
+    requirements += [
+        # astropy breaks with numpy 1.15.3
+        # https://github.com/astropy/astropy/issues/7943
+        'numpy >= 1.14.0, != 1.15.3',
+        'numba >= 0.46.0']
+
+extras_require = {
+    'cuda': ['cupy >= 5.0.0', 'jinja2 >= 2.10'],
+    'dask': ['dask[array] >= 1.1.0'],
+    'jax': ['jax == 0.1.27', 'jaxlib == 0.1.14'],
+    'scipy': ['scipy >= 1.0.0'],
+    'astropy': ['astropy >= 3.0'],
+    'python-casacore': ['python-casacore >= 3.2.0'],
+    'testing': ['pytest', 'flaky', 'pytest-flake8']
+}
+
+_non_cuda_extras = [er for n, er in extras_require.items() if n != "cuda"]
+_all_extras = extras_require.values()
+
+extras_require['complete'] = sorted(set(sum(_non_cuda_extras, [])))
+extras_require['complete-cuda'] = sorted(set(sum(_all_extras, [])))
+
+setup_requirements = []
+test_requirements = (extras_require['testing'] +
+                     extras_require['astropy'] +
+                     extras_require['python-casacore'] +
+                     extras_require['dask'] +
+                     extras_require['scipy'])

africanus/model/shape/gaussian_shape.py

@@ -3,12 +3,18 @@
 
 import numpy as np
 
+from africanus.config import config
 from africanus.util.docs import DocstringTemplate
 from africanus.util.numba import generated_jit
 from africanus.constants import c as lightspeed
 
+cfg = config.numba_parallel('model.shape.gaussian.parallel')
+parallel = cfg.get('parallel', False)
+axes = cfg.get('axes', set(('source', 'row')) if parallel else ())
 
-@generated_jit(nopython=True, nogil=True, cache=True)
+
+@generated_jit(nopython=True, nogil=True,
+               cache=not parallel, parallel=parallel)
 def gaussian(uvw, frequency, shape_params):
     # https://en.wikipedia.org/wiki/Full_width_at_half_maximum
     fwhm = 2.0 * np.sqrt(2.0 * np.log(2.0))
@@ -18,7 +24,16 @@ def gaussian(uvw, frequency, shape_params):
     dtype = np.result_type(*(np.dtype(a.dtype.name) for
                              a in (uvw, frequency, shape_params)))
 
+    from numba import prange, set_num_threads, get_num_threads
+    srange = prange if parallel and 'source' in axes else range
+    rrange = prange if parallel and 'row' in axes else range
+    threads = cfg.get("threads", None) if parallel else None
+
     def impl(uvw, frequency, shape_params):
+        if parallel and threads is not None:
+            prev_threads = get_num_threads()
+            set_num_threads(threads)
+
         nsrc = shape_params.shape[0]
         nrow = uvw.shape[0]
         nchan = frequency.shape[0]
@@ -30,15 +45,15 @@ def impl(uvw, frequency, shape_params):
         for f in range(frequency.shape[0]):
             scaled_freq[f] = frequency[f] * gauss_scale
 
-        for s in range(shape_params.shape[0]):
+        for s in srange(shape_params.shape[0]):
             emaj, emin, angle = shape_params[s]
 
             # Convert to l-projection, m-projection, ratio
             el = emaj * np.sin(angle)
             em = emaj * np.cos(angle)
             er = emin / (1.0 if emaj == 0.0 else emaj)
 
-            for r in range(uvw.shape[0]):
+            for r in rrange(uvw.shape[0]):
                 u, v, w = uvw[r]
 
                 u1 = (u*em - v*el)*er
@@ -50,6 +65,9 @@ def impl(uvw, frequency, shape_params):
 
                     shape[s, r, f] = np.exp(-(fu1*fu1 + fv1*fv1))
 
+        if parallel and threads is not None:
+            set_num_threads(prev_threads)
+
         return shape
 
     return impl

africanus/model/shape/tests/test_gaussian_shape.py

@@ -5,13 +5,22 @@
 from numpy.testing import assert_array_almost_equal
 import pytest
 
-from africanus.model.shape import gaussian as np_gaussian
 
-
-def test_gauss_shape():
+@pytest.mark.parametrize("cfg_parallel", [
+    ("africanus.model.shape.gaussian_shape",
+        {"model.shape.gaussian.parallel": True}),
+    ("africanus.model.shape.gaussian_shape", {
+        "model.shape.gaussian.parallel": {'threads': 2}}),
+    ("africanus.model.shape.gaussian_shape",
+        {"model.shape.gaussian.parallel": False}),
+    ], ids=["parallel", "parallel-2", "serial"], indirect=True)
+def test_gauss_shape(cfg_parallel):
+    from africanus.model.shape.gaussian_shape import gaussian as np_gaussian
     row = 10
     chan = 16
 
+    assert np_gaussian.targetoptions['parallel'] == cfg_parallel
+
     shape_params = np.array([[.4, .3, .2],
                              [.4, .3, .2]])
     uvw = np.random.random((row, 3))
@@ -24,6 +33,7 @@ def test_gauss_shape():
 
 def test_dask_gauss_shape():
     da = pytest.importorskip('dask.array')
+    from africanus.model.shape import gaussian as np_gaussian
     from africanus.model.shape.dask import gaussian as da_gaussian
 
     row_chunks = (5, 5)