Skip to content

Commit 13b509e

Browse files
authored
Merge pull request #7349 from IntelLabs/issue7344
Allow arbitrary walk-back in reduction nodes to find inplace_binop.
2 parents e6c4bd1 + eeac2fa commit 13b509e

2 files changed

Lines changed: 27 additions & 8 deletions

File tree

‎numba/parfors/parfor.py‎

Lines changed: 3 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -3566,15 +3566,10 @@ def get_reduction_init(nodes):
35663566
Get initial value for known reductions.
35673567
Currently, only += and *= are supported.
35683568
"""
3569-
require(len(nodes) >=1)
3570-
# there could be an extra assignment after the reduce node
3569+
require(len(nodes) >= 1)
3570+
# there could be multiple extra assignments after the reduce node
35713571
# See: test_reduction_var_reuse
3572-
if isinstance(nodes[-1].value, ir.Var):
3573-
require(len(nodes) >=2)
3574-
require(nodes[-2].target.name == nodes[-1].value.name)
3575-
acc_expr = nodes[-2].value
3576-
else:
3577-
acc_expr = nodes[-1].value
3572+
acc_expr = list(filter(lambda x: isinstance(x.value, ir.Expr), nodes))[-1].value
35783573
require(isinstance(acc_expr, ir.Expr) and acc_expr.op=='inplace_binop')
35793574
if acc_expr.fn == operator.iadd or acc_expr.fn == operator.isub:
35803575
return 0, acc_expr.fn

‎numba/tests/test_parfors.py‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1564,6 +1564,30 @@ def test_impl(n):
15641564
return acc
15651565
self.check(test_impl, 16)
15661566

1567+
def test_non_identity_initial(self):
1568+
# issue #7344
1569+
def test_impl(A, cond):
1570+
s = 1
1571+
for i in prange(A.shape[0]):
1572+
if cond[i]:
1573+
s += 1
1574+
return s
1575+
self.check(test_impl, np.ones(10), np.ones(10).astype('bool'))
1576+
1577+
def test_if_not_else_reduction(self):
1578+
# issue #7344
1579+
def test_impl(A, cond):
1580+
s = 1
1581+
t = 10
1582+
for i in prange(A.shape[0]):
1583+
if cond[i]:
1584+
s += 1
1585+
t += 1
1586+
else:
1587+
s += 2
1588+
return s + t
1589+
self.check(test_impl, np.ones(10), np.ones(10).astype('bool'))
1590+
15671591
def test_two_d_array_reduction_reuse(self):
15681592
def test_impl(n):
15691593
shp = (13, 17)

0 commit comments

Comments
 (0)