starpu-max/mpi/src/mpi_failure_tolerance/starpu_mpi_checkpoint.c

/* StarPU --- Runtime system for heterogeneous multicore architectures.
 *
 * Copyright (C) 2014-2020  Université de Bordeaux, CNRS (LaBRI UMR 5800), Inria
 *
 * StarPU is free software; you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published by
 * the Free Software Foundation; either version 2.1 of the License, or (at
 * your option) any later version.
 *
 * StarPU is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 *
 * See the GNU Lesser General Public License in COPYING.LGPL for more details.
 */


#include <stdlib.h>

#include <mpi_failure_tolerance/starpu_mpi_checkpoint.h>
#include <mpi_failure_tolerance/starpu_mpi_checkpoint_template.h>
#include <mpi_failure_tolerance/starpu_mpi_checkpoint_package.h>
#include <mpi_failure_tolerance/starpu_mpi_ft_service_comms.h>
#include <starpu_mpi_private.h>
#include <mpi/starpu_mpi_mpi_backend.h> // Should be deduced at preprocessing (Nmad vs MPI)
#include "starpu_mpi_cache.h"

#define MAX_CP_TEMPLATE_NUMBER 32 // Arbitrary limit

starpu_pthread_mutex_t           cp_lib_mutex;
int                              my_rank;


extern struct _starpu_mpi_req* _starpu_mpi_isend_cache_aware(starpu_data_handle_t data_handle, int dest, starpu_mpi_tag_t data_tag, MPI_Comm comm, unsigned detached, unsigned sync, int prio, void (*callback)(void *), void *_arg, void (*alt_callback)(void *), void *_alt_arg, int sequential_consistency);
extern struct _starpu_mpi_req* _starpu_mpi_irecv_cache_aware(starpu_data_handle_t data_handle, int source, starpu_mpi_tag_t data_tag, MPI_Comm comm, unsigned detached, unsigned sync, void (*callback)(void *), void *_arg, void (*alt_callback)(void *), void *_alt_arg, int sequential_consistency, int is_internal_req, starpu_ssize_t count);



void _starpu_mpi_treat_ack_receipt_cb(void* _args)
{
	struct _starpu_mpi_cp_ack_arg_cb* arg = (struct _starpu_mpi_cp_ack_arg_cb*) _args;
	int ret;
	_STARPU_MPI_DEBUG(3, "ack msg recved id:%d inst:%d\n", arg->msg.checkpoint_id, arg->msg.checkpoint_instance);
	ret = _checkpoint_template_digest_ack_reception(arg->msg.checkpoint_id, arg->msg.checkpoint_instance);
	if (ret == 0) {
		free(arg);
	}
	else if (ret == -1)
	{
		STARPU_ABORT_MSG("Could not find CP template, cpid:%d - cpinst:%d\n", arg->msg.checkpoint_id, arg->msg.checkpoint_instance);
	}
}

void _arg_free(void* _args)
{
	struct _starpu_mpi_cp_ack_arg_cb* arg = (struct _starpu_mpi_cp_ack_arg_cb*) _args;
	_STARPU_MPI_DEBUG(3, "Ack send succeeded cpid:%d, cpinst:%d, dest:%d\n", arg->msg.checkpoint_id, arg->msg.checkpoint_instance, arg->rank);
	free(arg);
}

void _starpu_mpi_store_data_and_send_ack_cb(void* _args)
{
	struct _starpu_mpi_cp_ack_arg_cb* arg = (struct _starpu_mpi_cp_ack_arg_cb*) _args;
	if (STARPU_VALUE == arg->type) {
		// an handle has specifically been created, Let's get the value back, and unregister the handle
		arg->copy_handle = starpu_data_handle_to_pointer(arg->handle, STARPU_MAIN_RAM);
		starpu_data_unregister_submit(arg->handle);
	}
	checkpoint_package_data_add(arg->msg.checkpoint_id, arg->msg.checkpoint_instance, arg->rank, arg->tag, arg->type, arg->copy_handle, arg->count);
	_STARPU_MPI_DEBUG(3,"Send ack msg to %d: id=%d inst=%d\n", arg->rank, arg->msg.checkpoint_id, arg->msg.checkpoint_instance);
	_ft_service_msg_isend_cb((void*)&arg->msg, sizeof(struct _starpu_mpi_cp_ack_msg), arg->rank, _STARPU_MPI_TAG_CP_ACK, MPI_COMM_WORLD, _arg_free, _args);

}void _starpu_mpi_release_and_store_data_and_send_ack_cb(void* _args)
{
	struct _starpu_mpi_cp_ack_arg_cb* arg = (struct _starpu_mpi_cp_ack_arg_cb*) _args;
	starpu_data_release(arg->copy_handle);
	checkpoint_package_data_add(arg->msg.checkpoint_id, arg->msg.checkpoint_instance, arg->rank, arg->tag, arg->type, arg->copy_handle, arg->count);
	_STARPU_MPI_DEBUG(3,"Send ack msg to %d: id=%d inst=%d\n", arg->rank, arg->msg.checkpoint_id, arg->msg.checkpoint_instance);
	_ft_service_msg_isend_cb((void*)&arg->msg, sizeof(struct _starpu_mpi_cp_ack_msg), arg->rank, _STARPU_MPI_TAG_CP_ACK, MPI_COMM_WORLD, _arg_free, _args);
}

void _starpu_mpi_push_cp_ack_recv_cb(void* _args)
{
	struct _starpu_mpi_cp_ack_arg_cb* arg = (struct _starpu_mpi_cp_ack_arg_cb*) _args;
	if (STARPU_VALUE == arg->type)
	{
		free(starpu_data_handle_to_pointer(arg->handle, STARPU_MAIN_RAM));
		starpu_data_unregister_submit(arg->handle);
	}
	_STARPU_MPI_DEBUG(3, "Posting ack recv cb from %d\n", arg->rank);
	_ft_service_msg_irecv_cb((void*)&arg->msg, sizeof(struct _starpu_mpi_cp_ack_msg), arg->rank, _STARPU_MPI_TAG_CP_ACK, MPI_COMM_WORLD, _starpu_mpi_treat_ack_receipt_cb, _args);
}

void _starpu_mpi_cached_push_cp_ack_recv_cb(void* _args)
{
	struct _starpu_mpi_cp_ack_arg_cb* arg = (struct _starpu_mpi_cp_ack_arg_cb*) _args;
	if (STARPU_R == arg->type)
	{
		starpu_data_release(arg->handle);
	}
	_STARPU_MPI_DEBUG(3, "Posting ack recv cb from %d\n", arg->rank);
	_ft_service_msg_irecv_cb((void*)&arg->msg, sizeof(struct _starpu_mpi_cp_ack_msg), arg->rank, _STARPU_MPI_TAG_CP_ACK, MPI_COMM_WORLD, _starpu_mpi_treat_ack_receipt_cb, _args);
}

void _starpu_data_release_cb(void* _arg)
{
	starpu_data_release(_arg);
}


int starpu_mpi_submit_checkpoint_template(starpu_mpi_checkpoint_template_t cp_template)
{
	starpu_data_handle_t* handle;
	struct _starpu_mpi_checkpoint_tracker* tracker;
	struct _starpu_mpi_cp_ack_arg_cb* arg;
	void* cpy_ptr;
	struct _starpu_mpi_checkpoint_template_item* item;
	int current_instance;

	current_instance = increment_current_instance();
//	_starpu_mpi_checkpoint_template_create_instance_tracker(cp_template, cp_template->cp_id, cp_template->checkpoint_domain, current_instance);
	_starpu_mpi_checkpoint_post_cp_discard_recv(cp_template);

	item = _starpu_mpi_checkpoint_template_get_first_data(cp_template);
	while (item != _starpu_mpi_checkpoint_template_end(cp_template))
	{
		switch (item->type)
		{
			case STARPU_VALUE:
				// TODO: Maybe do not pass via starpu handles for external data, and need to reimplement mpi comm layer for
				arg = calloc(1, sizeof(struct _starpu_mpi_cp_ack_arg_cb));
				arg->tag = item->tag;
				arg->type = STARPU_VALUE;
				arg->count = item->count;
				arg->msg.checkpoint_id = cp_template->cp_id;
				arg->msg.checkpoint_instance = current_instance;
				if (item->backupped_by != -1)
				{
					cpy_ptr = malloc(item->count);
					memcpy(cpy_ptr, item->ptr, item->count);
					starpu_variable_data_register(&arg->handle, STARPU_MAIN_RAM, (uintptr_t)cpy_ptr, item->count);
					arg->rank = item->backupped_by;
					_STARPU_MPI_DEBUG(0, "Submit CP: sending external data:%d, tag:%ld, to :%d\n", (int)(*(int*)cpy_ptr), arg->tag, arg->rank);
					starpu_mpi_isend_detached_prio(arg->handle, arg->rank, arg->tag, 0, MPI_COMM_WORLD,
												   &_starpu_mpi_push_cp_ack_recv_cb, (void*)arg);
					// The callback needs to free the handle specially created for the send, and post ack recv
				}
				else if (item->backup_of != -1)
				{
					cpy_ptr = malloc(item->count);
					starpu_variable_data_register(&arg->handle, STARPU_MAIN_RAM, (uintptr_t)cpy_ptr, item->count);
					arg->rank = item->backup_of;
					_STARPU_MPI_DEBUG(0, "Submit CP: receiving external data tag:%ld, from :%d\n", arg->tag, arg->rank);
					starpu_mpi_irecv_detached(arg->handle, arg->rank, arg->tag, MPI_COMM_WORLD,
											  &_starpu_mpi_store_data_and_send_ack_cb, (void*)arg);
					// The callback needs to store the received data and post ack send
				}
				break;
			case STARPU_R:
				handle = (starpu_data_handle_t*)item->ptr;
				if (starpu_mpi_data_get_rank(*handle)==my_rank)
				{
					_STARPU_MPI_DEBUG(0, "Submit CP: sending starPU data to %d (tag %d)\n", item->backupped_by, (int)starpu_mpi_data_get_tag(*handle));
					arg = calloc(1, sizeof(struct _starpu_mpi_cp_ack_arg_cb));
					arg->rank = item->backupped_by;
					arg->handle = *handle;
					arg->tag = starpu_mpi_data_get_tag(*handle);
					arg->type = STARPU_R;
					arg->count = item->count;
					arg->msg.checkpoint_id = cp_template->cp_id;
					arg->msg.checkpoint_instance = current_instance;
					_starpu_mpi_isend_cache_aware(*handle, item->backupped_by, starpu_mpi_data_get_tag(*handle), MPI_COMM_WORLD, 1, 0, 0,
					                              &_starpu_mpi_push_cp_ack_recv_cb, (void*)arg, &_starpu_mpi_cached_push_cp_ack_recv_cb, (void*)arg, 1);
					// the callbacks need to post ack recv. The cache one needs to release the handle.

				}
				else if (item->backup_of == starpu_mpi_data_get_rank(*handle))
				{
					_STARPU_MPI_DEBUG(0, "Submit CP: receiving starPU data from %d (tag %d)\n", starpu_mpi_data_get_rank(*handle), (int)starpu_mpi_data_get_tag(*handle));
					arg = calloc(1, sizeof(struct _starpu_mpi_cp_ack_arg_cb));
					arg->rank = item->backup_of;
					arg->handle = *handle;
					arg->tag = starpu_mpi_data_get_tag(*handle);
					arg->type = STARPU_R;
					arg->count = item->count;
					arg->msg.checkpoint_id = cp_template->cp_id;
					arg->msg.checkpoint_instance = current_instance;
					_starpu_mpi_irecv_cache_aware(*handle, starpu_mpi_data_get_rank(*handle), starpu_mpi_data_get_tag(*handle), MPI_COMM_WORLD, 1, 0,
					                              NULL, NULL, &_starpu_data_release_cb, (void*)arg->handle, 1, 0, 1);
					// The callback needs to do nothing. The cached one must release the handle.
					starpu_data_dup_ro(&arg->copy_handle, arg->handle, 1);
					starpu_data_acquire_cb(arg->copy_handle, STARPU_R, _starpu_mpi_release_and_store_data_and_send_ack_cb, arg);
					// The callback need to store the data and post ack send.
				}
				break;
		}

		item = _starpu_mpi_checkpoint_template_get_next_data(cp_template, item);
	}

	return 0;
}

//
///**
// * receives param of type starpu_mpi_checkpoint_template_t
// * @param args
// * @return
// */
//void _starpu_mpi_checkpoint_ack_send_cb(void* args)
//{
//	starpu_mpi_checkpoint_template_t cp_template = (starpu_mpi_checkpoint_template_t) args;
//	starpu_pthread_mutex_lock(&cp_template->mutex);
//	cp_template->remaining_ack_awaited--;
//	starpu_pthread_mutex_unlock(&cp_template->mutex);
//}


//
//void _starpu_checkpoint_cached_data_send_copy_and_ack(void* _arg)
//{
//	struct _starpu_mpi_cp_ack_arg_cb* arg = (struct _starpu_mpi_cp_ack_arg_cb*) _arg;
//	starpu_data_register_same(&arg->copy_handle, arg->handle);
//	starpu_data_cpy(arg->copy_handle, arg->handle, 1, _starpu_mpi_push_cp_ack_recv_cb, _arg);
//	starpu_data_release(arg->handle);
//}
//
//void _starpu_checkpoint_data_send_copy_and_ack(void* _args)
//{
//	struct _starpu_mpi_cp_ack_arg_cb* arg = (struct _starpu_mpi_cp_ack_arg_cb*) _args;
//	starpu_data_register_same(&arg->copy_handle, arg->handle);
//	starpu_data_cpy(arg->copy_handle, arg->handle, 1, _starpu_mpi_push_cp_ack_recv_cb, _args);
//}
//
//void _starpu_mpi_treat_cache_ack_no_lock_cb(void* _args)
//{
//	starpu_mpi_checkpoint_template_t cp_template = (starpu_mpi_checkpoint_template_t)_args;
//	cp_template->remaining_ack_awaited--;
//}