Commit 2883a7fd authored by Dion Häfner's avatar Dion Häfner

free MPI comm after use

parent c3cce72e
...@@ -284,31 +284,37 @@ def exchange_cyclic_boundaries(vs, arr): ...@@ -284,31 +284,37 @@ def exchange_cyclic_boundaries(vs, arr):
def _reduce(vs, arr, op, axis=None): def _reduce(vs, arr, op, axis=None):
if axis is None: if axis is None:
comm = rs.mpi_comm comm = rs.mpi_comm
disconnect_comm = False
else: else:
assert axis in (0, 1) assert axis in (0, 1)
pi = proc_rank_to_index(rst.proc_rank) pi = proc_rank_to_index(rst.proc_rank)
other_axis = 1 - axis other_axis = 1 - axis
comm = rs.mpi_comm.Split(pi[other_axis], rst.proc_rank) comm = rs.mpi_comm.Split(pi[other_axis], rst.proc_rank)
disconnect_comm = True
if np.isscalar(arr): try:
squeeze = True if np.isscalar(arr):
arr = np.array([arr]) squeeze = True
else: arr = np.array([arr])
squeeze = False else:
squeeze = False
arr = ascontiguousarray(arr) arr = ascontiguousarray(arr)
res = np.empty_like(arr) res = np.empty_like(arr)
comm.Allreduce( comm.Allreduce(
get_array_buffer(vs, arr), get_array_buffer(vs, arr),
get_array_buffer(vs, res), get_array_buffer(vs, res),
op=op op=op
) )
if squeeze: if squeeze:
res = res[0] res = res[0]
return res return res
finally:
if disconnect_comm:
comm.Disconnect()
@dist_context_only @dist_context_only
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment