Skip to content

Commit 069ad43

Browse files
authored
Merge pull request numba#9532 from IntelLabs/issue9529
Fix hoisting bug to exclude variables used in containers or calls.
2 parents 0a9a2cc + 6947ae1 commit 069ad43

3 files changed

Lines changed: 59 additions & 0 deletions

File tree

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
Fix parfor variable hoisting analysis.
2+
--------------------------------------
3+
4+
If a variable is used to build a container (e.g., tuple, list, map, set) or is passed as an
5+
argument to a call then conservatively assume it could escape the current iteration of
6+
the parfor and so should not be hoisted.

‎numba/parfors/parfor_lowering.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -922,6 +922,20 @@ def find_setitems_block(setitems, itemsset, block, typemap):
922922
elif isinstance(inst, parfor.Parfor):
923923
find_setitems_block(setitems, itemsset, inst.init_block, typemap)
924924
find_setitems_body(setitems, itemsset, inst.loop_body, typemap)
925+
elif isinstance(inst, ir.Assign):
926+
# If something of mutable type is given to a build_tuple or
927+
# used in a call then consider it unanalyzable and so
928+
# unavailable for hoisting.
929+
rhs = inst.value
930+
if isinstance(rhs, ir.Expr):
931+
if rhs.op in ["build_tuple", "build_list", "build_set", "build_map"]:
932+
for item in rhs.items:
933+
if getattr(typemap[item.name], "mutable", False):
934+
itemsset.add(item.name)
935+
elif rhs.op == "call":
936+
for item in list(rhs.args) + [x[1] for x in rhs.kws]:
937+
if getattr(typemap[item.name], "mutable", False):
938+
itemsset.add(item.name)
925939

926940
def find_setitems_body(setitems, itemsset, loop_body, typemap):
927941
"""

‎numba/tests/test_parfors.py‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,8 @@
2929
set_num_threads, get_num_threads, typeof)
3030
from numba.core import (types, errors, ir, rewrites,
3131
typed_passes, inline_closurecall, config, compiler, cpu)
32+
from numba.typed import Dict, List
33+
3234
from numba.extending import (overload_method, register_model,
3335
typeof_impl, unbox, NativeValue, models)
3436
from numba.core.registry import cpu_target
@@ -230,6 +232,10 @@ def copy_args(*args):
230232
new_args.append(copy.deepcopy(x))
231233
elif isinstance(x, list):
232234
new_args.append(x[:])
235+
elif isinstance(x, Dict):
236+
new_args.append(copy.copy(x))
237+
elif isinstance(x, List):
238+
new_args.append(copy.copy(x))
233239
else:
234240
raise ValueError('Unsupported argument type encountered')
235241
return tuple(new_args)
@@ -4082,6 +4088,11 @@ def test_impl(x, par, numthreads):
40824088
self.prange_tester(test_impl, x, par, 2)
40834089

40844090

4091+
@register_jitable
4092+
def test_call_hoisting_outcall(a,b):
4093+
return (a, b)
4094+
4095+
40854096
@skip_parfors_unsupported
40864097
class TestPrangeSpecific(TestPrangeBase):
40874098
""" Tests specific features/problems found under prange"""
@@ -4421,6 +4432,34 @@ def test_impl():
44214432

44224433
self.prange_tester(test_impl)
44234434

4435+
def test_tuple_hoisting(self):
4436+
# issue9529
4437+
def test_impl(inputs):
4438+
outputs = [(Dict.empty(key_type=types.int64, value_type=types.float64), np.zeros(1)) for _ in range(len(inputs))]
4439+
for i in range(len(inputs)):
4440+
y = inputs[i]
4441+
out = np.zeros(1)
4442+
out[0] = i
4443+
outputs[i] = (inputs[i], out)
4444+
return outputs[0][1][0]
4445+
4446+
N = config.NUMBA_NUM_THREADS + 1
4447+
self.prange_tester(test_impl, [Dict.empty(key_type=types.int64, value_type=types.float64) for i in range(N)])
4448+
4449+
def test_call_hoisting(self):
4450+
# issue9529
4451+
def test_impl(inputs):
4452+
outputs = [(Dict.empty(key_type=types.int64, value_type=types.float64), np.zeros(1)) for _ in range(len(inputs))]
4453+
for i in range(len(inputs)):
4454+
y = inputs[i]
4455+
out = np.zeros(1)
4456+
out[0] = i
4457+
outputs[i] = test_call_hoisting_outcall(inputs[i], out)
4458+
return outputs[0][1][0]
4459+
4460+
N = config.NUMBA_NUM_THREADS + 1
4461+
self.prange_tester(test_impl, [Dict.empty(key_type=types.int64, value_type=types.float64) for i in range(N)])
4462+
44244463
def test_record_array_setitem(self):
44254464
# issue6704
44264465
state_dtype = np.dtype([('var', np.int32)])

0 commit comments

Comments
 (0)