<gh_stars>0
from mpi4py import MPI
import mpi4py
import numpy as np
from utils import GetProcessTasks, GetTasksOffset

#
# Fonctions de communication parallèle
#

def SendSubLayer(loaded_weights, loaded_bias, loaded_func, entry):
	comm = MPI.COMM_WORLD
	world_size = comm.Get_size()
	rank = comm.Get_rank()

	n_neuron = len(loaded_weights)
	comm.bcast(n_neuron, root=0)

	comm.bcast(len(entry), root=0)

	process_tasks = GetProcessTasks(n_neuron, len(entry))
	tasks_offset = GetTasksOffset(process_tasks)

	# print(loaded_weights)

	weights = np.zeros(process_tasks[rank])
	comm.Scatterv([loaded_weights, process_tasks, tasks_offset, MPI.DOUBLE], weights, root=0)

	process_tasks = GetProcessTasks(n_neuron, 1)
	tasks_offset = GetTasksOffset(process_tasks)

	bias = np.zeros(process_tasks[rank])
	comm.Scatterv([loaded_bias, process_tasks, tasks_offset, MPI.DOUBLE], bias, root=0)

	func = None
	func = comm.bcast(loaded_func, root=0)

	return [weights, bias, func]

def RecvSubLayer():
	comm = MPI.COMM_WORLD
	world_size = comm.Get_size()
	rank = comm.Get_rank()

	n_neuron = comm.bcast(None, root=0)
	segment_size = int(n_neuron / world_size)
	remaining_tasks = n_neuron - (world_size) * segment_size

	entry_size = comm.bcast(None, root=0)

	process_tasks = segment_size
	if rank <= remaining_tasks :
		process_tasks += 1

	weights = np.zeros(process_tasks * entry_size)
	comm.Scatterv([None, None, None, MPI.DOUBLE], weights, root=0)

	bias = np.zeros(process_tasks)
	comm.Scatterv([None, None, None, MPI.DOUBLE], bias, root=0)

	func = None
	func = comm.bcast(None, root=0)

	return [weights, bias, func]

# def GatherResult(result, total_tasks):
#     comm = MPI.COMM_WORLD
#     world_size = comm.Get_size()
#     rank = comm.Get_rank()

#     total_tasks = comm.bcast(total_tasks, root=0)

#     res = np.zeros(total_tasks)

#     print(rank, res, total_tasks)
#     print(rank, result, total_tasks)

#     process_tasks = GetProcessTasks(total_tasks, 1)
#     tasks_offset = GetTasksOffset(process_tasks)

#     comm.Gatherv(result,[res,process_tasks,tasks_offset,MPI.DOUBLE], root=0)

#     return res

def GatherResult(res, total_tasks):
	comm = MPI.COMM_WORLD
	rank = comm.Get_rank()
	world_size = comm.size

	total_tasks = comm.bcast(total_tasks, root=0)
	segment_size = int(total_tasks / world_size)

	remaining_tasks = total_tasks - (world_size) * segment_size

	if rank == 0:
		process_task = [segment_size] * world_size
		for process_id in range(world_size):
			if remaining_tasks >= process_id and process_id > 0:
				process_task[process_id] += 1
	
	if remaining_tasks >= rank and rank > 0:
		segment_size += 1

	if rank == 0:
		for process_id in range(1,world_size):
			for neuron in range(process_task[process_id]):
				res.append(comm.recv(None, process_id))
	else:
		for neuron in range(segment_size):
			comm.send(res[neuron], 0)

	final_res = []

	for neuron in range(total_tasks):
		if rank == 0:
			final_res.append(comm.bcast(res[neuron], root=0))
		else:
			final_res.append(comm.bcast(None, root=0))

	return final_res
from mpi4py import MPI
import mpi4py
import numpy as np
from utils import GetProcessTasks, GetTasksOffset

#
# Fonctions de communication parallèle
#

def SendSubLayer(loaded_weights, loaded_bias, loaded_func, entry):
	comm = MPI.COMM_WORLD
	world_size = comm.Get_size()
	rank = comm.Get_rank()

	n_neuron = len(loaded_weights)
	comm.bcast(n_neuron, root=0)

	comm.bcast(len(entry), root=0)

	process_tasks = GetProcessTasks(n_neuron, len(entry))
	tasks_offset = GetTasksOffset(process_tasks)

	# print(loaded_weights)

	weights = np.zeros(process_tasks[rank])
	comm.Scatterv([loaded_weights, process_tasks, tasks_offset, MPI.DOUBLE], weights, root=0)

	process_tasks = GetProcessTasks(n_neuron, 1)
	tasks_offset = GetTasksOffset(process_tasks)

	bias = np.zeros(process_tasks[rank])
	comm.Scatterv([loaded_bias, process_tasks, tasks_offset, MPI.DOUBLE], bias, root=0)

	func = None
	func = comm.bcast(loaded_func, root=0)

	return [weights, bias, func]

def RecvSubLayer():
	comm = MPI.COMM_WORLD
	world_size = comm.Get_size()
	rank = comm.Get_rank()

	n_neuron = comm.bcast(None, root=0)
	segment_size = int(n_neuron / world_size)
	remaining_tasks = n_neuron - (world_size) * segment_size

	entry_size = comm.bcast(None, root=0)

	process_tasks = segment_size
	if rank <= remaining_tasks :
		process_tasks += 1

	weights = np.zeros(process_tasks * entry_size)
	comm.Scatterv([None, None, None, MPI.DOUBLE], weights, root=0)

	bias = np.zeros(process_tasks)
	comm.Scatterv([None, None, None, MPI.DOUBLE], bias, root=0)

	func = None
	func = comm.bcast(None, root=0)

	return [weights, bias, func]

# def GatherResult(result, total_tasks):
#     comm = MPI.COMM_WORLD
#     world_size = comm.Get_size()
#     rank = comm.Get_rank()

#     total_tasks = comm.bcast(total_tasks, root=0)

#     res = np.zeros(total_tasks)

#     print(rank, res, total_tasks)
#     print(rank, result, total_tasks)

#     process_tasks = GetProcessTasks(total_tasks, 1)
#     tasks_offset = GetTasksOffset(process_tasks)

#     comm.Gatherv(result,[res,process_tasks,tasks_offset,MPI.DOUBLE], root=0)

#     return res

def GatherResult(res, total_tasks):
	comm = MPI.COMM_WORLD
	rank = comm.Get_rank()
	world_size = comm.size

	total_tasks = comm.bcast(total_tasks, root=0)
	segment_size = int(total_tasks / world_size)

	remaining_tasks = total_tasks - (world_size) * segment_size

	if rank == 0:
		process_task = [segment_size] * world_size
		for process_id in range(world_size):
			if remaining_tasks >= process_id and process_id > 0:
				process_task[process_id] += 1
	
	if remaining_tasks >= rank and rank > 0:
		segment_size += 1

	if rank == 0:
		for process_id in range(1,world_size):
			for neuron in range(process_task[process_id]):
				res.append(comm.recv(None, process_id))
	else:
		for neuron in range(segment_size):
			comm.send(res[neuron], 0)

	final_res = []

	for neuron in range(total_tasks):
		if rank == 0:
			final_res.append(comm.bcast(res[neuron], root=0))
		else:
			final_res.append(comm.bcast(None, root=0))

	return final_res
# # Fonctions de communication parallèle # # print(loaded_weights) # def GatherResult(result, total_tasks): #     comm = MPI.COMM_WORLD #     world_size = comm.Get_size() #     rank = comm.Get_rank() #     total_tasks = comm.bcast(total_tasks, root=0) #     res = np.zeros(total_tasks) #     print(rank, res, total_tasks) #     print(rank, result, total_tasks) #     process_tasks = GetProcessTasks(total_tasks, 1) #     tasks_offset = GetTasksOffset(process_tasks) #     comm.Gatherv(result,[res,process_tasks,tasks_offset,MPI.DOUBLE], root=0) #     return res