File tree Expand file tree Collapse file tree
Expand file tree Collapse file tree Original file line number Diff line number Diff 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
Original file line number Diff line number Diff 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 )
You can’t perform that action at this time.
0 commit comments