Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add del's to reductions. #5047

Closed
wants to merge 6 commits into from
Closed
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
15 changes: 14 additions & 1 deletion numba/npyufunc/parfor.py
Original file line number Diff line number Diff line change
Expand Up @@ -101,6 +101,7 @@ def _lower_parfor_parallel(lowerer, parfor):
# init reduction array allocation here.
nredvars = len(parfor_redvars)
redarrs = {}
redtoset_dels = []
if nredvars > 0:
# reduction arrays outer dimension equal to thread count
thread_count = get_thread_count()
Expand All @@ -114,7 +115,7 @@ def _lower_parfor_parallel(lowerer, parfor):
redarrvar_typ = redtyp_to_redarraytype(redvar_typ)
reddtype = redarrvar_typ.dtype
if config.DEBUG_ARRAY_OPT:
print("redvar_typ", redvar_typ, redarrvar_typ, reddtype, types.DType(reddtype))
print("redvar_typ", redvar, redvar_typ, redarrvar_typ, reddtype, types.DType(reddtype))

# If this is reduction over an array,
# the reduction array has just one added per-worker dimension.
Expand Down Expand Up @@ -204,6 +205,7 @@ def _lower_parfor_parallel(lowerer, parfor):
full_call = ir.Expr.call(full_func, [redshape_var, init_val_var], {}, loc=loc)
lowerer.fndesc.calltypes[full_call] = full_sig
redtoset = ir.Var(scope, mk_unique_var("redtoset"), loc)
redtoset_dels.append(redtoset)
redtoset_assign = ir.Assign(full_call, redtoset, loc)
typemap[redtoset.name] = redvar_typ
lowerer.lower_inst(redtoset_assign)
Expand Down Expand Up @@ -395,6 +397,7 @@ def _lower_parfor_parallel(lowerer, parfor):
# Add calltype back in for the expr with updated signature.
lowerer.fndesc.calltypes[rhs] = ct
lowerer.lower_inst(inst)

if isinstance(inst, ir.Assign) and name == inst.target.name:
break

Expand All @@ -418,10 +421,20 @@ def _lower_parfor_parallel(lowerer, parfor):
print("res_print2", res_print)
lowerer.lower_inst(res_print)

lowerer.lower_inst(ir.Del(oneelem.name, loc=loc))
lowerer.lower_inst(ir.Del(init_var.name, loc=loc))

for i in range(nredvars):
name = parfor_redvars[i]
for inst in parfor_reddict[name][1][:-1]:
if isinstance(typemap[inst.target.name], types.npytypes.Array):
lowerer.lower_inst(ir.Del(inst.target.name, loc=loc))

# Cleanup reduction variable
for v in redarrs.values():
lowerer.lower_inst(ir.Del(v.name, loc=loc))
for v in redtoset_dels:
lowerer.lower_inst(ir.Del(v.name, loc=loc))
# Restore the original typemap of the function that was replaced temporarily at the
# Beginning of this function.
lowerer.fndesc.typemap = orig_typemap
Expand Down