Skip to content
This repository was archived by the owner on Feb 2, 2024. It is now read-only.

Commit 1a2f68e

Browse files
authored
Move prange utils in separate PR (#612)
1 parent 43f658e commit 1a2f68e

File tree

1 file changed

+89
-0
lines changed

1 file changed

+89
-0
lines changed

sdc/utilities/prange_utils.py

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
# *****************************************************************************
2+
# Copyright (c) 2020, Intel Corporation All rights reserved.
3+
#
4+
# Redistribution and use in source and binary forms, with or without
5+
# modification, are permitted provided that the following conditions are met:
6+
#
7+
# Redistributions of source code must retain the above copyright notice,
8+
# this list of conditions and the following disclaimer.
9+
#
10+
# Redistributions in binary form must reproduce the above copyright notice,
11+
# this list of conditions and the following disclaimer in the documentation
12+
# and/or other materials provided with the distribution.
13+
#
14+
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
15+
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
16+
# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
17+
# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
18+
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
19+
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
20+
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
21+
# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
22+
# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
23+
# OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE,
24+
# EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
25+
# *****************************************************************************
26+
27+
28+
import numba
29+
import sdc
30+
31+
from typing import NamedTuple
32+
from sdc.utilities.utils import sdc_overload
33+
34+
35+
class Chunk(NamedTuple):
36+
start: int
37+
stop: int
38+
39+
40+
def get_pool_size():
41+
if sdc.config.config_use_parallel_overloads:
42+
return numba.config.NUMBA_NUM_THREADS
43+
44+
return 1
45+
46+
47+
@sdc_overload(get_pool_size)
48+
def get_pool_size_overload():
49+
pool_size = get_pool_size()
50+
51+
def get_pool_size_impl():
52+
return pool_size
53+
54+
return get_pool_size_impl
55+
56+
57+
def get_chunks(size, pool_size=0):
58+
if pool_size == 0:
59+
pool_size = get_pool_size()
60+
61+
chunk_size = (size - 1) // pool_size + 1
62+
63+
chunks = []
64+
for i in range(pool_size):
65+
start = min(i * chunk_size, size)
66+
stop = min((i + 1) * chunk_size, size)
67+
chunks.append(Chunk(start, stop))
68+
69+
return chunks
70+
71+
72+
@sdc_overload(get_chunks)
73+
def get_chunks_overload(size, pool_size=0):
74+
def get_chunks_impl(size, pool_size=0):
75+
if pool_size == 0:
76+
pool_size = get_pool_size()
77+
78+
chunk_size = (size - 1) // pool_size + 1
79+
80+
chunks = []
81+
for i in range(pool_size):
82+
start = min(i * chunk_size, size)
83+
stop = min((i + 1) * chunk_size, size)
84+
chunk = Chunk(start, stop)
85+
chunks.append(chunk)
86+
87+
return chunks
88+
89+
return get_chunks_impl

0 commit comments

Comments
 (0)