diff --git a/sdc/tests/test_prange_utils.py b/sdc/tests/test_prange_utils.py index 5ae0f629b..beac0982a 100644 --- a/sdc/tests/test_prange_utils.py +++ b/sdc/tests/test_prange_utils.py @@ -68,6 +68,8 @@ def _get_chunks_data(self): Chunk(start=3, stop=6), Chunk(start=6, stop=9), ] + yield (0, 5), [] + yield (5, 0), [] def _check_get_chunks(self, args, expected_chunks): pyfunc = get_chunks diff --git a/sdc/utilities/prange_utils.py b/sdc/utilities/prange_utils.py index 0f24d34ea..4a44fe249 100644 --- a/sdc/utilities/prange_utils.py +++ b/sdc/utilities/prange_utils.py @@ -47,11 +47,15 @@ def get_pool_size(): @sdc_register_jitable def get_chunks(size, pool_size): + chunks = [] + + if size < 1 or pool_size < 1: + return chunks + pool_size = min(pool_size, size) chunk_size = size // pool_size overload_size = size % pool_size - chunks = [] for i in range(pool_size): start = i * chunk_size + min(i, overload_size) stop = (i + 1) * chunk_size + min(i + 1, overload_size)