Repository: numba/numba

Issue:
[feature request] user defined priorities of automatic threading layer choice


<!--

Thanks for opening an issue! To help the Numba team handle your information
efficiently, please first ensure that there is no other issue present that
already describes the issue you have
(search at https://github.com/numba/numba/issues?&q=is%3Aissue).

For more general "how do I do X?" type questions, please speak to us in real
time on https://gitter.im/numba/numba or post to the Numba mailing list
https://groups.google.com/a/continuum.io/forum/#!forum/numba-users.

-->

## Feature request

<!--

Please include details of the feature you would like to see, why you would
like to see it/the use case

-->
In <https://numba.pydata.org/numba-doc/dev/user/threading-layer.html>, it is shown how to select the 3 available threading layers.

This feature request concerns setting relative priorities between this 3 threading layers. i.e. I still want Numba to automatically choose one for me, but I want to specify the priorities between them. For example if the first is not available, fall back to lower priority ones.

The priority I want to set is `'omp', 'tbb', 'workqueue'`. By the way, is the current behavior having the priority of `'tbb', 'omp', 'workqueue'`, i.e. `tbb` is chosen first, and if not available fall back to `omp`, etc.?


Discussion:
Yes, the default is `tbb, omp, workqueue`.
@stuartarchibald, you have thought about this more than I have.  
I think that this request is feasible and can be accommodated without breaking anyone relying on the current behaviour. @ickc would you perhaps be interested in providing a patch and tests? Thanks.

Resolving patch:
diff --git a/numba/core/config.py b/numba/core/config.py
--- a/numba/core/config.py
+++ b/numba/core/config.py
@@ -306,6 +306,11 @@ def avx_default():
         DISABLE_JIT = _readenv("NUMBA_DISABLE_JIT", int, 0)
 
         # choose parallel backend to use
+        THREADING_LAYER_PRIORITY = _readenv(
+            "NUMBA_THREADING_LAYER_PRIORITY",
+            lambda string: string.split(),
+            ['tbb', 'omp', 'workqueue'],
+        )
         THREADING_LAYER = _readenv("NUMBA_THREADING_LAYER", str, 'default')
 
         # CUDA Configs
diff --git a/numba/np/ufunc/parallel.py b/numba/np/ufunc/parallel.py
--- a/numba/np/ufunc/parallel.py
+++ b/numba/np/ufunc/parallel.py
@@ -422,7 +422,15 @@ def select_from_backends(backends):
                 return lib, backend
 
             t = str(config.THREADING_LAYER).lower()
-            namedbackends = ['tbb', 'omp', 'workqueue']
+            namedbackends = config.THREADING_LAYER_PRIORITY
+            if not (len(namedbackends) == 3 and
+                    set(namedbackends) == {'tbb', 'omp', 'workqueue'}):
+                raise ValueError(
+                    "THREADING_LAYER_PRIORITY invalid: %s. "
+                    "It must be a permutation of "
+                    "{'tbb', 'omp', 'workqueue'}"
+                    % namedbackends
+                )
 
             lib = None
             err_helpers = dict()


Test patch:
diff --git a/numba/tests/test_parallel_backend.py b/numba/tests/test_parallel_backend.py
--- a/numba/tests/test_parallel_backend.py
+++ b/numba/tests/test_parallel_backend.py
@@ -4,6 +4,7 @@
 Tests the parallel backend
 """
 import faulthandler
+import itertools
 import multiprocessing
 import os
 import random
@@ -521,6 +522,80 @@ def generate(cls):
 TestThreadingLayerSelection.generate()
 
 
+@skip_parfors_unsupported
+class TestThreadingLayerPriority(ThreadLayerTestHelper):
+
+    def each_env_var(self, env_var: str):
+        """Test setting priority via env var NUMBA_THREADING_LAYER_PRIORITY.
+
+        :return: threading_layer_priority, stderr
+            (containing ``@threading_layer@``)
+        """
+        env = os.environ.copy()
+        env['NUMBA_THREADING_LAYER'] = 'default'
+        env['NUMBA_THREADING_LAYER_PRIORITY'] = env_var
+
+        code = """import sys
+import numba
+
+# trigger threading layer decision
+# hence catching invalid THREADING_LAYER_PRIORITY
+@numba.jit(
+    'float64[::1](float64[::1], float64[::1])',
+    nopython=True,
+    parallel=True,
+)
+def plus(x, y):
+    return x + y
+
+print(' '.join(numba.config.THREADING_LAYER_PRIORITY))
+print("@%s@" % numba.threading_layer(), file=sys.stderr)
+"""
+        cmd = [
+            sys.executable,
+            '-c',
+            code,
+        ]
+        return self.run_cmd(cmd, env=env)
+
+    def test_valid_env_var(self):
+        default = ['tbb', 'omp', 'workqueue']
+        for p in itertools.permutations(default):
+            env_var = ' '.join(p)
+            threading_layer_priority, _ = self.each_env_var(env_var)
+            self.assertEqual(threading_layer_priority.strip(), env_var)
+
+    def test_invalid_env_var(self):
+        env_var = 'tbb omp workqueue notvalidhere'
+        with self.assertRaises(AssertionError) as raises:
+            self.each_env_var(env_var)
+        for msg in (
+            "THREADING_LAYER_PRIORITY invalid:",
+            "It must be a permutation of"
+        ):
+            self.assertIn(f"{msg}", str(raises.exception))
+
+    @skip_no_omp
+    def test_omp(self):
+        for env_var in ("omp tbb workqueue", "omp workqueue tbb"):
+            threading_layer_priority, out = self.each_env_var(env_var)
+            self.assertEqual(threading_layer_priority.strip(), env_var)
+            self.assertIn("@omp@", out)
+
+    @skip_no_tbb
+    def test_tbb(self):
+        for env_var in ("tbb omp workqueue", "tbb workqueue omp"):
+            threading_layer_priority, out = self.each_env_var(env_var)
+            self.assertEqual(threading_layer_priority.strip(), env_var)
+            self.assertIn("@tbb@", out)
+
+    def test_workqueue(self):
+        for env_var in ("workqueue tbb omp", "workqueue omp tbb"):
+            threading_layer_priority, out = self.each_env_var(env_var)
+            self.assertEqual(threading_layer_priority.strip(), env_var)
+            self.assertIn("@workqueue@", out)
+
+
 @skip_parfors_unsupported
 class TestMiscBackendIssues(ThreadLayerTestHelper):
     """