intermediate sum at nodes

Author: jishnub

src/ParallelUtilities.jl

@@ -169,25 +169,41 @@ get_nprocs_node(procs_used::Vector{<:Integer}=workers()) = get_nprocs_node(get_h
 function pmapsum(f::Function,iterable,args...;kwargs...)
 
 	procs_used = workers_active(iterable)
-	num_workers = length(procs_used)
-	hostnames = get_hostnames(procs_used)
-	nodes = get_nodes(hostnames)
-	pid_rank0_on_node = [procs_used[findfirst(isequal(node),hostnames)] for node in nodes]
+	num_workers = length(procs_used);
+	hostnames = get_hostnames(procs_used);
+	nodes = get_nodes(hostnames);
+	pid_rank0_on_node = [procs_used[findfirst(isequal(node),hostnames)] for node in nodes];
+
+	nprocs_node = get_nprocs_node(procs_used)
+	node_channels = Dict(node=>RemoteChannel(()->Channel{Any}(nprocs_node[node]),pid_node)
+		for (node,pid_node) in zip(nodes,pid_rank0_on_node))
 
 	# Worker at which final reduction takes place
 	p_final = first(pid_rank0_on_node)
 
-	sum_channel = RemoteChannel(()->Channel{Any}(100),p_final)
-	result = 0
+	sum_channel = RemoteChannel(()->Channel{Any}(length(pid_rank0_on_node)),p_final)
+	result = nothing
 
 	# Run the function on each processor and compute the sum at each node
-	@sync for (rank,p) in enumerate(procs_used)
+	@sync for (rank,(p,node)) in enumerate(zip(procs_used,hostnames))
 		@async begin
+			
 			iterable_on_proc = split_across_processors(iterable,num_workers,rank)
 			r = @spawnat p f(iterable_on_proc,args...;kwargs...)
-			@spawnat p put!(sum_channel,fetch(r))
+
+			node_remotechannel = node_channels[node]
+			np_node = nprocs_node[node]
+			
+			@spawnat p put!(node_remotechannel,fetch(r))
+
+			if p in pid_rank0_on_node
+				s = @spawnat p sum(take!(node_remotechannel) for i=1:np_node)
+				@spawnat p put!(sum_channel,fetch(s))
+			end
+
 			if p==p_final
-				result = @fetchfrom p_final sum(take!(sum_channel) for i=1:num_workers)			
+				result = @fetchfrom p_final sum(take!(sum_channel) 
+					for i=1:length(pid_rank0_on_node))			
 			end
 		end
 	end