633 lines
23 KiB
Python
633 lines
23 KiB
Python
|
# -*- coding: utf-8 -*-
|
||
|
|
from __future__ import print_function, absolute_import, division
|
||
|
|
|
||
|
|
import sys
|
||
|
|
import os
|
||
|
|
import re
|
||
|
|
import multiprocessing
|
||
|
|
import unittest
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
from numba import (njit, set_num_threads, get_num_threads, prange, config,
|
||
|
|
threading_layer, guvectorize)
|
||
|
|
from numba.np.ufunc.parallel import get_thread_id
|
||
|
|
from numba.core.errors import TypingError
|
||
|
|
from numba.tests.support import TestCase, skip_parfors_unsupported, tag
|
||
|
|
from numba.tests.test_parallel_backend import TestInSubprocess
|
||
|
|
|
||
|
|
|
||
|
|
class TestNumThreads(TestCase):
|
||
|
|
_numba_parallel_test_ = False
|
||
|
|
|
||
|
|
def setUp(self):
|
||
|
|
# Make sure the num_threads is set to the max. This also makes sure
|
||
|
|
# the threads are launched.
|
||
|
|
set_num_threads(config.NUMBA_NUM_THREADS)
|
||
|
|
|
||
|
|
def check_mask(self, expected, result):
|
||
|
|
# There's no guarantee that TBB will use a full mask worth of
|
||
|
|
# threads if it deems it inefficient to do so
|
||
|
|
if threading_layer() == 'tbb':
|
||
|
|
self.assertTrue(np.all(result <= expected))
|
||
|
|
elif threading_layer() in ('omp', 'workqueue'):
|
||
|
|
np.testing.assert_equal(expected, result)
|
||
|
|
else:
|
||
|
|
assert 0, 'unreachable'
|
||
|
|
|
||
|
|
@skip_parfors_unsupported
|
||
|
|
def test_set_num_threads_type(self):
|
||
|
|
|
||
|
|
@njit
|
||
|
|
def foo():
|
||
|
|
set_num_threads('wrong_type')
|
||
|
|
|
||
|
|
expected = "The number of threads specified must be an integer"
|
||
|
|
for fn, errty in ((foo, TypingError), (foo.py_func, TypeError)):
|
||
|
|
with self.assertRaises(errty) as raises:
|
||
|
|
fn()
|
||
|
|
self.assertIn(expected, str(raises.exception))
|
||
|
|
|
||
|
|
@skip_parfors_unsupported
|
||
|
|
@unittest.skipIf(config.NUMBA_NUM_THREADS < 2, "Not enough CPU cores")
|
||
|
|
def _test_set_num_threads_basic(self):
|
||
|
|
max_threads = config.NUMBA_NUM_THREADS
|
||
|
|
|
||
|
|
self.assertEqual(get_num_threads(), max_threads)
|
||
|
|
set_num_threads(2)
|
||
|
|
self.assertEqual(get_num_threads(), 2)
|
||
|
|
set_num_threads(max_threads)
|
||
|
|
self.assertEqual(get_num_threads(), max_threads)
|
||
|
|
|
||
|
|
with self.assertRaises(ValueError):
|
||
|
|
set_num_threads(0)
|
||
|
|
|
||
|
|
with self.assertRaises(ValueError):
|
||
|
|
set_num_threads(max_threads + 1)
|
||
|
|
|
||
|
|
@skip_parfors_unsupported
|
||
|
|
@unittest.skipIf(config.NUMBA_NUM_THREADS < 2, "Not enough CPU cores")
|
||
|
|
def _test_set_num_threads_basic_jit(self):
|
||
|
|
max_threads = config.NUMBA_NUM_THREADS
|
||
|
|
|
||
|
|
@njit
|
||
|
|
def get_n():
|
||
|
|
return get_num_threads()
|
||
|
|
|
||
|
|
self.assertEqual(get_n(), max_threads)
|
||
|
|
set_num_threads(2)
|
||
|
|
self.assertEqual(get_n(), 2)
|
||
|
|
set_num_threads(max_threads)
|
||
|
|
self.assertEqual(get_n(), max_threads)
|
||
|
|
|
||
|
|
@njit
|
||
|
|
def set_get_n(n):
|
||
|
|
set_num_threads(n)
|
||
|
|
return get_num_threads()
|
||
|
|
|
||
|
|
self.assertEqual(set_get_n(2), 2)
|
||
|
|
self.assertEqual(set_get_n(max_threads), max_threads)
|
||
|
|
|
||
|
|
@skip_parfors_unsupported
|
||
|
|
@unittest.skipIf(config.NUMBA_NUM_THREADS < 2, "Not enough CPU cores")
|
||
|
|
def _test_set_num_threads_basic_guvectorize(self):
|
||
|
|
max_threads = config.NUMBA_NUM_THREADS
|
||
|
|
|
||
|
|
@guvectorize(['void(int64[:])'],
|
||
|
|
'(n)',
|
||
|
|
nopython=True,
|
||
|
|
target='parallel')
|
||
|
|
def get_n(x):
|
||
|
|
x[:] = get_num_threads()
|
||
|
|
|
||
|
|
x = np.zeros((5000000,), dtype=np.int64)
|
||
|
|
get_n(x)
|
||
|
|
np.testing.assert_equal(x, max_threads)
|
||
|
|
set_num_threads(2)
|
||
|
|
x = np.zeros((5000000,), dtype=np.int64)
|
||
|
|
get_n(x)
|
||
|
|
np.testing.assert_equal(x, 2)
|
||
|
|
set_num_threads(max_threads)
|
||
|
|
x = np.zeros((5000000,), dtype=np.int64)
|
||
|
|
get_n(x)
|
||
|
|
np.testing.assert_equal(x, max_threads)
|
||
|
|
|
||
|
|
@guvectorize(['void(int64[:])'],
|
||
|
|
'(n)',
|
||
|
|
nopython=True,
|
||
|
|
target='parallel')
|
||
|
|
def set_get_n(n):
|
||
|
|
set_num_threads(n[0])
|
||
|
|
n[:] = get_num_threads()
|
||
|
|
|
||
|
|
x = np.zeros((5000000,), dtype=np.int64)
|
||
|
|
x[0] = 2
|
||
|
|
set_get_n(x)
|
||
|
|
np.testing.assert_equal(x, 2)
|
||
|
|
x = np.zeros((5000000,), dtype=np.int64)
|
||
|
|
x[0] = max_threads
|
||
|
|
set_get_n(x)
|
||
|
|
np.testing.assert_equal(x, max_threads)
|
||
|
|
|
||
|
|
@skip_parfors_unsupported
|
||
|
|
@unittest.skipIf(config.NUMBA_NUM_THREADS < 2, "Not enough CPU cores")
|
||
|
|
def _test_set_num_threads_outside_jit(self):
|
||
|
|
|
||
|
|
# Test set_num_threads outside a jitted function
|
||
|
|
set_num_threads(2)
|
||
|
|
|
||
|
|
@njit(parallel=True)
|
||
|
|
def test_func():
|
||
|
|
x = 5
|
||
|
|
buf = np.empty((x,))
|
||
|
|
for i in prange(x):
|
||
|
|
buf[i] = get_num_threads()
|
||
|
|
return buf
|
||
|
|
|
||
|
|
@guvectorize(['void(int64[:])'],
|
||
|
|
'(n)',
|
||
|
|
nopython=True,
|
||
|
|
target='parallel')
|
||
|
|
def test_gufunc(x):
|
||
|
|
x[:] = get_num_threads()
|
||
|
|
|
||
|
|
out = test_func()
|
||
|
|
np.testing.assert_equal(out, 2)
|
||
|
|
|
||
|
|
x = np.zeros((5000000,), dtype=np.int64)
|
||
|
|
test_gufunc(x)
|
||
|
|
np.testing.assert_equal(x, 2)
|
||
|
|
|
||
|
|
@skip_parfors_unsupported
|
||
|
|
@unittest.skipIf(config.NUMBA_NUM_THREADS < 2, "Not enough CPU cores")
|
||
|
|
def _test_set_num_threads_inside_jit(self):
|
||
|
|
# Test set_num_threads inside a jitted function
|
||
|
|
@njit(parallel=True)
|
||
|
|
def test_func(nthreads):
|
||
|
|
x = 5
|
||
|
|
buf = np.empty((x,))
|
||
|
|
set_num_threads(nthreads)
|
||
|
|
for i in prange(x):
|
||
|
|
buf[i] = get_num_threads()
|
||
|
|
return buf
|
||
|
|
|
||
|
|
mask = 2
|
||
|
|
out = test_func(mask)
|
||
|
|
np.testing.assert_equal(out, mask)
|
||
|
|
|
||
|
|
@skip_parfors_unsupported
|
||
|
|
@unittest.skipIf(config.NUMBA_NUM_THREADS < 2, "Not enough CPU cores")
|
||
|
|
def _test_set_num_threads_inside_guvectorize(self):
|
||
|
|
# Test set_num_threads inside a jitted guvectorize function
|
||
|
|
@guvectorize(['void(int64[:])'],
|
||
|
|
'(n)',
|
||
|
|
nopython=True,
|
||
|
|
target='parallel')
|
||
|
|
def test_func(x):
|
||
|
|
set_num_threads(x[0])
|
||
|
|
x[:] = get_num_threads()
|
||
|
|
|
||
|
|
x = np.zeros((5000000,), dtype=np.int64)
|
||
|
|
mask = 2
|
||
|
|
x[0] = mask
|
||
|
|
test_func(x)
|
||
|
|
np.testing.assert_equal(x, mask)
|
||
|
|
|
||
|
|
@skip_parfors_unsupported
|
||
|
|
@unittest.skipIf(config.NUMBA_NUM_THREADS < 2, "Not enough CPU cores")
|
||
|
|
def _test_get_num_threads_truth_outside_jit(self):
|
||
|
|
|
||
|
|
for mask in range(2, min(6, config.NUMBA_NUM_THREADS + 1)):
|
||
|
|
set_num_threads(mask)
|
||
|
|
|
||
|
|
# a lot of work, hopefully will trigger "mask" count of threads to
|
||
|
|
# join the parallel region (for those backends with dynamic threads)
|
||
|
|
@njit(parallel=True)
|
||
|
|
def test_func():
|
||
|
|
x = 5000000
|
||
|
|
buf = np.empty((x,))
|
||
|
|
for i in prange(x):
|
||
|
|
buf[i] = get_thread_id()
|
||
|
|
return len(np.unique(buf)), get_num_threads()
|
||
|
|
|
||
|
|
out = test_func()
|
||
|
|
self.check_mask((mask, mask), out)
|
||
|
|
|
||
|
|
@guvectorize(['void(int64[:], int64[:])'],
|
||
|
|
'(n), (m)',
|
||
|
|
nopython=True,
|
||
|
|
target='parallel')
|
||
|
|
def test_gufunc(x, out):
|
||
|
|
x[:] = get_thread_id()
|
||
|
|
out[0] = get_num_threads()
|
||
|
|
|
||
|
|
# Reshape to force parallelism
|
||
|
|
x = np.full((5000000,), -1, dtype=np.int64).reshape((100, 50000))
|
||
|
|
out = np.zeros((1,), dtype=np.int64)
|
||
|
|
test_gufunc(x, out)
|
||
|
|
self.check_mask(mask, out)
|
||
|
|
self.check_mask(mask, len(np.unique(x)))
|
||
|
|
|
||
|
|
@skip_parfors_unsupported
|
||
|
|
@unittest.skipIf(config.NUMBA_NUM_THREADS < 2, "Not enough CPU cores")
|
||
|
|
def _test_get_num_threads_truth_inside_jit(self):
|
||
|
|
|
||
|
|
for mask in range(2, min(6, config.NUMBA_NUM_THREADS + 1)):
|
||
|
|
|
||
|
|
# a lot of work, hopefully will trigger "mask" count of threads to
|
||
|
|
# join the parallel region (for those backends with dynamic threads)
|
||
|
|
@njit(parallel=True)
|
||
|
|
def test_func():
|
||
|
|
set_num_threads(mask)
|
||
|
|
x = 5000000
|
||
|
|
buf = np.empty((x,))
|
||
|
|
for i in prange(x):
|
||
|
|
buf[i] = get_thread_id()
|
||
|
|
return len(np.unique(buf)), get_num_threads()
|
||
|
|
|
||
|
|
out = test_func()
|
||
|
|
self.check_mask((mask, mask), out)
|
||
|
|
|
||
|
|
@guvectorize(['void(int64[:], int64[:])'],
|
||
|
|
'(n), (m)',
|
||
|
|
nopython=True,
|
||
|
|
target='parallel')
|
||
|
|
def test_gufunc(x, out):
|
||
|
|
set_num_threads(mask)
|
||
|
|
x[:] = get_thread_id()
|
||
|
|
out[0] = get_num_threads()
|
||
|
|
|
||
|
|
# Reshape to force parallelism
|
||
|
|
x = np.full((5000000,), -1, dtype=np.int64).reshape((100, 50000))
|
||
|
|
out = np.zeros((1,), dtype=np.int64)
|
||
|
|
test_gufunc(x, out)
|
||
|
|
self.check_mask(mask, out)
|
||
|
|
self.check_mask(mask, len(np.unique(x)))
|
||
|
|
|
||
|
|
# this test can only run on OpenMP (providing OMP_MAX_ACTIVE_LEVELS is not
|
||
|
|
# set or >= 2) and TBB backends
|
||
|
|
@skip_parfors_unsupported
|
||
|
|
@unittest.skipIf(config.NUMBA_NUM_THREADS < 2, "Not enough CPU cores")
|
||
|
|
def _test_nested_parallelism_1(self):
|
||
|
|
if threading_layer() == 'workqueue':
|
||
|
|
self.skipTest("workqueue is not threadsafe")
|
||
|
|
|
||
|
|
# check that get_num_threads is ok in nesting
|
||
|
|
mask = config.NUMBA_NUM_THREADS - 1
|
||
|
|
|
||
|
|
N = config.NUMBA_NUM_THREADS
|
||
|
|
M = 2 * config.NUMBA_NUM_THREADS
|
||
|
|
|
||
|
|
@njit(parallel=True)
|
||
|
|
def child_func(buf, fid):
|
||
|
|
M, N = buf.shape
|
||
|
|
for i in prange(N):
|
||
|
|
buf[fid, i] = get_num_threads()
|
||
|
|
|
||
|
|
def get_test(test_type):
|
||
|
|
if test_type == 'njit':
|
||
|
|
def test_func(nthreads, py_func=False):
|
||
|
|
@njit(parallel=True)
|
||
|
|
def _test_func(nthreads):
|
||
|
|
acc = 0
|
||
|
|
buf = np.zeros((M, N))
|
||
|
|
set_num_threads(nthreads)
|
||
|
|
for i in prange(M):
|
||
|
|
local_mask = 1 + i % mask
|
||
|
|
# set threads in parent function
|
||
|
|
set_num_threads(local_mask)
|
||
|
|
if local_mask < N:
|
||
|
|
child_func(buf, local_mask)
|
||
|
|
acc += get_num_threads()
|
||
|
|
return acc, buf
|
||
|
|
if py_func:
|
||
|
|
return _test_func.py_func(nthreads)
|
||
|
|
else:
|
||
|
|
return _test_func(nthreads)
|
||
|
|
|
||
|
|
elif test_type == 'guvectorize':
|
||
|
|
def test_func(nthreads, py_func=False):
|
||
|
|
def _test_func(acc, buf, local_mask):
|
||
|
|
set_num_threads(nthreads)
|
||
|
|
# set threads in parent function
|
||
|
|
set_num_threads(local_mask[0])
|
||
|
|
if local_mask[0] < N:
|
||
|
|
child_func(buf, local_mask[0])
|
||
|
|
acc[0] += get_num_threads()
|
||
|
|
|
||
|
|
buf = np.zeros((M, N), dtype=np.int64)
|
||
|
|
acc = np.zeros((M, 1), dtype=np.int64)
|
||
|
|
local_mask = (1 + np.arange(M) % mask).reshape((M, 1))
|
||
|
|
sig = ['void(int64[:], int64[:, :], int64[:])']
|
||
|
|
layout = '(p), (n, m), (p)'
|
||
|
|
if not py_func:
|
||
|
|
_test_func = guvectorize(sig, layout, nopython=True,
|
||
|
|
target='parallel')(_test_func)
|
||
|
|
else:
|
||
|
|
_test_func = guvectorize(sig, layout,
|
||
|
|
forceobj=True)(_test_func)
|
||
|
|
_test_func(acc, buf, local_mask)
|
||
|
|
return acc, buf
|
||
|
|
|
||
|
|
return test_func
|
||
|
|
|
||
|
|
for test_type in ['njit', 'guvectorize']:
|
||
|
|
test_func = get_test(test_type)
|
||
|
|
got_acc, got_arr = test_func(mask)
|
||
|
|
exp_acc, exp_arr = test_func(mask, py_func=True)
|
||
|
|
np.testing.assert_equal(exp_acc, got_acc)
|
||
|
|
np.testing.assert_equal(exp_arr, got_arr)
|
||
|
|
|
||
|
|
# check the maths reconciles, guvectorize does not reduce, njit does
|
||
|
|
math_acc_exp = 1 + np.arange(M) % mask
|
||
|
|
if test_type == 'guvectorize':
|
||
|
|
math_acc = math_acc_exp.reshape((M, 1))
|
||
|
|
else:
|
||
|
|
math_acc = np.sum(math_acc_exp)
|
||
|
|
|
||
|
|
np.testing.assert_equal(math_acc, got_acc)
|
||
|
|
|
||
|
|
math_arr = np.zeros((M, N))
|
||
|
|
for i in range(1, N):
|
||
|
|
# there's branches on 1, ..., num_threads - 1
|
||
|
|
math_arr[i, :] = i
|
||
|
|
np.testing.assert_equal(math_arr, got_arr)
|
||
|
|
|
||
|
|
# this test can only run on OpenMP (providing OMP_MAX_ACTIVE_LEVELS is not
|
||
|
|
# set or >= 2) and TBB backends
|
||
|
|
@skip_parfors_unsupported
|
||
|
|
@unittest.skipIf(config.NUMBA_NUM_THREADS < 2, "Not enough CPU cores")
|
||
|
|
def _test_nested_parallelism_2(self):
|
||
|
|
if threading_layer() == 'workqueue':
|
||
|
|
self.skipTest("workqueue is not threadsafe")
|
||
|
|
|
||
|
|
# check that get_num_threads is ok in nesting
|
||
|
|
|
||
|
|
N = config.NUMBA_NUM_THREADS + 1
|
||
|
|
M = 4 * config.NUMBA_NUM_THREADS + 1
|
||
|
|
|
||
|
|
def get_impl(child_type, test_type):
|
||
|
|
|
||
|
|
if child_type == 'parallel':
|
||
|
|
child_dec = njit(parallel=True)
|
||
|
|
elif child_type == 'njit':
|
||
|
|
child_dec = njit(parallel=False)
|
||
|
|
elif child_type == 'none':
|
||
|
|
def child_dec(x):
|
||
|
|
return x
|
||
|
|
|
||
|
|
@child_dec
|
||
|
|
def child(buf, fid):
|
||
|
|
M, N = buf.shape
|
||
|
|
set_num_threads(fid) # set threads in child function
|
||
|
|
for i in prange(N):
|
||
|
|
buf[fid, i] = get_num_threads()
|
||
|
|
|
||
|
|
if test_type in ['parallel', 'njit', 'none']:
|
||
|
|
if test_type == 'parallel':
|
||
|
|
test_dec = njit(parallel=True)
|
||
|
|
elif test_type == 'njit':
|
||
|
|
test_dec = njit(parallel=False)
|
||
|
|
elif test_type == 'none':
|
||
|
|
def test_dec(x):
|
||
|
|
return x
|
||
|
|
|
||
|
|
@test_dec
|
||
|
|
def test_func(nthreads):
|
||
|
|
buf = np.zeros((M, N))
|
||
|
|
set_num_threads(nthreads)
|
||
|
|
for i in prange(M):
|
||
|
|
local_mask = 1 + i % mask
|
||
|
|
# when the threads exit the child functions they should
|
||
|
|
# have a TLS slot value of the local mask as it was set
|
||
|
|
# in child
|
||
|
|
if local_mask < config.NUMBA_NUM_THREADS:
|
||
|
|
child(buf, local_mask)
|
||
|
|
assert get_num_threads() == local_mask
|
||
|
|
return buf
|
||
|
|
else:
|
||
|
|
if test_type == 'guvectorize':
|
||
|
|
test_dec = guvectorize(['int64[:,:], int64[:]'],
|
||
|
|
'(n, m), (k)', nopython=True,
|
||
|
|
target='parallel')
|
||
|
|
elif test_type == 'guvectorize-obj':
|
||
|
|
test_dec = guvectorize(['int64[:,:], int64[:]'],
|
||
|
|
'(n, m), (k)', forceobj=True)
|
||
|
|
|
||
|
|
def test_func(nthreads):
|
||
|
|
@test_dec
|
||
|
|
def _test_func(buf, local_mask):
|
||
|
|
set_num_threads(nthreads)
|
||
|
|
# when the threads exit the child functions they should
|
||
|
|
# have a TLS slot value of the local mask as it was set
|
||
|
|
# in child
|
||
|
|
if local_mask[0] < config.NUMBA_NUM_THREADS:
|
||
|
|
child(buf, local_mask[0])
|
||
|
|
assert get_num_threads() == local_mask[0]
|
||
|
|
|
||
|
|
buf = np.zeros((M, N), dtype=np.int64)
|
||
|
|
local_mask = (1 + np.arange(M) % mask).reshape((M, 1))
|
||
|
|
_test_func(buf, local_mask)
|
||
|
|
return buf
|
||
|
|
|
||
|
|
return test_func
|
||
|
|
|
||
|
|
mask = config.NUMBA_NUM_THREADS - 1
|
||
|
|
|
||
|
|
res_arrays = {}
|
||
|
|
for test_type in ['parallel', 'njit', 'none',
|
||
|
|
'guvectorize', 'guvectorize-obj']:
|
||
|
|
for child_type in ['parallel', 'njit', 'none']:
|
||
|
|
if child_type == 'none' and test_type != 'none':
|
||
|
|
continue
|
||
|
|
set_num_threads(mask)
|
||
|
|
res_arrays[test_type, child_type] = get_impl(
|
||
|
|
child_type, test_type)(mask)
|
||
|
|
|
||
|
|
py_arr = res_arrays['none', 'none']
|
||
|
|
for arr in res_arrays.values():
|
||
|
|
np.testing.assert_equal(arr, py_arr)
|
||
|
|
|
||
|
|
# check the maths reconciles
|
||
|
|
math_arr = np.zeros((M, N))
|
||
|
|
# there's branches on modulo mask but only NUMBA_NUM_THREADS funcs
|
||
|
|
for i in range(1, config.NUMBA_NUM_THREADS):
|
||
|
|
math_arr[i, :] = i
|
||
|
|
|
||
|
|
np.testing.assert_equal(math_arr, py_arr)
|
||
|
|
|
||
|
|
# this test can only run on OpenMP (providing OMP_MAX_ACTIVE_LEVELS is not
|
||
|
|
# set or >= 2) and TBB backends
|
||
|
|
# This test needs at least 3 threads to run, N>=2 for the launch, M>=N+1 for
|
||
|
|
# the nested function
|
||
|
|
@skip_parfors_unsupported
|
||
|
|
@unittest.skipIf(config.NUMBA_NUM_THREADS < 3, "Not enough CPU cores")
|
||
|
|
def _test_nested_parallelism_3(self):
|
||
|
|
if threading_layer() == 'workqueue':
|
||
|
|
self.skipTest("workqueue is not threadsafe")
|
||
|
|
|
||
|
|
# check that the right number of threads are present in nesting
|
||
|
|
# this relies on there being a load of cores present
|
||
|
|
BIG = 1000000
|
||
|
|
|
||
|
|
@njit(parallel=True)
|
||
|
|
def work(local_nt): # arg is value 3
|
||
|
|
tid = np.zeros(BIG)
|
||
|
|
acc = 0
|
||
|
|
set_num_threads(local_nt) # set to 3 threads
|
||
|
|
for i in prange(BIG):
|
||
|
|
acc += 1
|
||
|
|
tid[i] = get_thread_id()
|
||
|
|
return acc, np.unique(tid)
|
||
|
|
|
||
|
|
@njit(parallel=True)
|
||
|
|
def test_func_jit(nthreads):
|
||
|
|
set_num_threads(nthreads) # set to 2 threads
|
||
|
|
lens = np.zeros(nthreads)
|
||
|
|
total = 0
|
||
|
|
for i in prange(nthreads):
|
||
|
|
my_acc, tids = work(nthreads + 1) # call with value 3
|
||
|
|
lens[i] = len(tids)
|
||
|
|
total += my_acc
|
||
|
|
return total, np.unique(lens)
|
||
|
|
|
||
|
|
NT = 2
|
||
|
|
expected_acc = BIG * NT
|
||
|
|
expected_thread_count = NT + 1
|
||
|
|
|
||
|
|
got_acc, got_tc = test_func_jit(NT)
|
||
|
|
self.assertEqual(expected_acc, got_acc)
|
||
|
|
self.check_mask(expected_thread_count, got_tc)
|
||
|
|
|
||
|
|
def test_guvectorize(nthreads):
|
||
|
|
@guvectorize(['int64[:], int64[:]'],
|
||
|
|
'(n), (n)',
|
||
|
|
nopython=True,
|
||
|
|
target='parallel')
|
||
|
|
def test_func_guvectorize(total, lens):
|
||
|
|
my_acc, tids = work(nthreads + 1)
|
||
|
|
lens[0] = len(tids)
|
||
|
|
total[0] += my_acc
|
||
|
|
|
||
|
|
total = np.zeros((nthreads, 1), dtype=np.int64)
|
||
|
|
lens = np.zeros(nthreads, dtype=np.int64).reshape((nthreads, 1))
|
||
|
|
|
||
|
|
test_func_guvectorize(total, lens)
|
||
|
|
# vectorize does not reduce, so total is summed
|
||
|
|
return total.sum(), np.unique(lens)
|
||
|
|
|
||
|
|
got_acc, got_tc = test_guvectorize(NT)
|
||
|
|
|
||
|
|
self.assertEqual(expected_acc, got_acc)
|
||
|
|
self.check_mask(expected_thread_count, got_tc)
|
||
|
|
|
||
|
|
@skip_parfors_unsupported
|
||
|
|
@unittest.skipIf(config.NUMBA_NUM_THREADS < 2, "Not enough CPU cores")
|
||
|
|
@unittest.skipIf(not sys.platform.startswith('linux'), "Linux only")
|
||
|
|
def _test_threadmask_across_fork(self):
|
||
|
|
forkctx = multiprocessing.get_context('fork')
|
||
|
|
|
||
|
|
@njit
|
||
|
|
def foo():
|
||
|
|
return get_num_threads()
|
||
|
|
|
||
|
|
def wrap(queue):
|
||
|
|
queue.put(foo())
|
||
|
|
|
||
|
|
mask = 1
|
||
|
|
self.assertEqual(foo(), config.NUMBA_NUM_THREADS)
|
||
|
|
set_num_threads(mask)
|
||
|
|
self.assertEqual(foo(), mask)
|
||
|
|
shared_queue = forkctx.Queue()
|
||
|
|
# check TLS slot inheritance in fork
|
||
|
|
p = forkctx.Process(target=wrap, args=(shared_queue,))
|
||
|
|
p.start()
|
||
|
|
p.join()
|
||
|
|
self.assertEqual(shared_queue.get(), mask)
|
||
|
|
|
||
|
|
def tearDown(self):
|
||
|
|
set_num_threads(config.NUMBA_NUM_THREADS)
|
||
|
|
|
||
|
|
@skip_parfors_unsupported
|
||
|
|
def _test_get_thread_id_not_parallel(self):
|
||
|
|
python_get_thread_id = get_thread_id()
|
||
|
|
check_array_size = 8
|
||
|
|
|
||
|
|
@njit(parallel=False)
|
||
|
|
def par_false(size):
|
||
|
|
njit_par_false_tid = get_thread_id()
|
||
|
|
res = np.ones(size)
|
||
|
|
for i in prange(size):
|
||
|
|
res[i] = get_thread_id()
|
||
|
|
return njit_par_false_tid, res
|
||
|
|
|
||
|
|
@njit(parallel=True)
|
||
|
|
def par_true(size):
|
||
|
|
njit_par_true_tid = get_thread_id()
|
||
|
|
res = np.ones(size)
|
||
|
|
for i in range(size):
|
||
|
|
res[i] = get_thread_id()
|
||
|
|
return njit_par_true_tid, res
|
||
|
|
|
||
|
|
self.assertEqual(python_get_thread_id, 0)
|
||
|
|
njit_par_false_tid, njit_par_false_arr = par_false(check_array_size)
|
||
|
|
self.assertEqual(njit_par_false_tid, 0)
|
||
|
|
np.testing.assert_equal(njit_par_false_arr, 0)
|
||
|
|
njit_par_true_tid, njit_par_true_arr = par_true(check_array_size)
|
||
|
|
self.assertEqual(njit_par_true_tid, 0)
|
||
|
|
np.testing.assert_equal(njit_par_true_arr, 0)
|
||
|
|
|
||
|
|
|
||
|
|
class TestNumThreadsBackends(TestInSubprocess, TestCase):
|
||
|
|
_class = TestNumThreads
|
||
|
|
_DEBUG = False
|
||
|
|
|
||
|
|
# 1 is mainly here to ensure tests skip correctly
|
||
|
|
num_threads = [i for i in [1, 2, 4, 8, 16] if i <= config.NUMBA_NUM_THREADS]
|
||
|
|
|
||
|
|
def run_test_in_separate_process(self, test, threading_layer, num_threads):
|
||
|
|
env_copy = os.environ.copy()
|
||
|
|
env_copy['NUMBA_THREADING_LAYER'] = str(threading_layer)
|
||
|
|
env_copy['NUMBA_NUM_THREADS'] = str(num_threads)
|
||
|
|
cmdline = [sys.executable, "-m", "numba.runtests", "-v", test]
|
||
|
|
return self.run_cmd(cmdline, env_copy)
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def _inject(cls, name, backend, backend_guard, num_threads):
|
||
|
|
themod = cls.__module__
|
||
|
|
thecls = cls._class.__name__
|
||
|
|
injected_method = '%s.%s.%s' % (themod, thecls, name)
|
||
|
|
|
||
|
|
def test_template(self):
|
||
|
|
o, e = self.run_test_in_separate_process(injected_method, backend,
|
||
|
|
num_threads)
|
||
|
|
if self._DEBUG:
|
||
|
|
print('stdout:\n "%s"\n stderr:\n "%s"' % (o, e))
|
||
|
|
# If the test was skipped in the subprocess, then mark this as a
|
||
|
|
# skipped test.
|
||
|
|
m = re.search(r"\.\.\. skipped '(.*?)'", e)
|
||
|
|
if m is not None:
|
||
|
|
self.skipTest(m.group(1))
|
||
|
|
self.assertIn('OK', e)
|
||
|
|
self.assertTrue('FAIL' not in e)
|
||
|
|
self.assertTrue('ERROR' not in e)
|
||
|
|
|
||
|
|
injected_test = "%s_%s_%s_threads" % (name[1:], backend, num_threads)
|
||
|
|
setattr(cls, injected_test,
|
||
|
|
tag('long_running')(backend_guard(test_template)))
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def generate(cls):
|
||
|
|
for name in cls._class.__dict__.copy():
|
||
|
|
for backend, backend_guard in cls.backends.items():
|
||
|
|
for num_threads in cls.num_threads:
|
||
|
|
if not name.startswith('_test_'):
|
||
|
|
continue
|
||
|
|
cls._inject(name, backend, backend_guard, num_threads)
|
||
|
|
|
||
|
|
|
||
|
|
TestNumThreadsBackends.generate()
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
unittest.main()
|