/*
 * Copyright (C) CNRS, INRIA, Université Bordeaux 1, Télécom SudParis
 * See COPYING in top-level directory.
 */

#include <eztrace-lib/eztrace.h>
#include <eztrace-lib/eztrace_module.h>
#include <starpu.h>
#include <starpu_opencl.h>

#ifdef USE_MPI
#include <starpu_mpi.h>
#endif

#define CURRENT_MODULE starpu
DECLARE_CURRENT_MODULE;

/* set to 1 when all the hooks are set.
 * This is usefull in order to avoid recursive calls
 */
static int _starpu_initialized = 0;

// Internals
void (*lib_starpu_driver_start_job) (void *args, void *j, void *codelet_start, int rank, int profiling);
void (*lib_starpu_driver_end_job) (void * args, void *j, void *perf_arch, void *codelet_end, int rank, int profiling);

// API
void (*libstarpu_cublas_init) ();
void (*libstarpu_cublas_shutdown) ();
int (*libstarpu_init) (struct starpu_conf* conf);
void (*libstarpu_pause) ();
void (*libstarpu_resume) ();
void (*libstarpu_shutdown) ();
void (*libstarpu_data_unregister) (starpu_data_handle_t handle);
void (*libstarpu_data_unregister_no_coherency) (starpu_data_handle_t handle);
void (*libstarpu_data_unregister_submit) (starpu_data_handle_t handle);
void (*libstarpu_data_invalidate) (starpu_data_handle_t handle);
void (*libstarpu_data_invalidate_submit) (starpu_data_handle_t handle);
int (*libstarpu_data_acquire) (starpu_data_handle_t handle, enum starpu_data_access_mode mode);
int (*libstarpu_data_acquire_on_node) (starpu_data_handle_t handle, int node, enum starpu_data_access_mode mode);
int (*libstarpu_data_acquire_cb) (starpu_data_handle_t handle, enum starpu_data_access_mode mode, void (* callback), void *arg);
int (*libstarpu_data_acquire_on_node_cb) (starpu_data_handle_t handle, int node, enum starpu_data_access_mode mode, void (* callback), void *arg);
void (*libstarpu_data_release) (starpu_data_handle_t handle);
void (*libstarpu_data_release_on_node) (starpu_data_handle_t handle, int node);
int (*libstarpu_data_request_allocation) (starpu_data_handle_t handle, unsigned node);
uintptr_t (*libstarpu_malloc_on_node) (unsigned dst_node, size_t size);
void (*libstarpu_free_on_node) (unsigned dst_node, uintptr_t addr, size_t size);
#ifdef USE_MPI
int (*libstarpu_mpi_isend) (starpu_data_handle_t data_handle, starpu_mpi_req* req, int dest, int mpi_tag, MPI_Comm comm);
int (*libstarpu_mpi_irecv) (starpu_data_handle_t data_handle, starpu_mpi_req* req, int source, int mpi_tag, MPI_Comm comm);
int (*libstarpu_mpi_send) (starpu_data_handle_t data_handle, int dest, int mpi_tag, MPI_Comm comm);
int (*libstarpu_mpi_recv) (starpu_data_handle_t data_handle, int source, int mpi_tag, MPI_Comm comm, MPI_Status* status);
int (*libstarpu_mpi_isend_detached) (starpu_data_handle_t data_handle, int dest, int mpi_tag, MPI_Comm comm, void (* callback)(void *arg), void *arg);
int (*libstarpu_mpi_irecv_detached) (starpu_data_handle_t data_handle, int source, int mpi_tag, MPI_Comm comm, void (* callback)(void *arg), void *arg);
int (*libstarpu_mpi_wait) (starpu_mpi_req* req, MPI_Status* status);
int (*libstarpu_mpi_test) (starpu_mpi_req* req, int* flag, MPI_Status* status);
int (*libstarpu_mpi_barrier) (MPI_Comm comm);
int (*libstarpu_mpi_init) (int* argc, char*** argv, int initialize_mpi);
int (*libstarpu_mpi_initialize) ();
int (*libstarpu_mpi_initialize_extended) (int* rank, int* world_size);
int (*libstarpu_mpi_shutdown) ();
void (*libstarpu_mpi_get_data_on_node) (MPI_Comm comm, starpu_data_handle_t data_handle, int node);
void (*libstarpu_mpi_get_data_on_node_detached) (MPI_Comm comm, starpu_data_handle_t data_handle, int node, void (* callback)(void *arg), void *arg);
void (*libstarpu_mpi_redux_data) (MPI_Comm comm, starpu_data_handle_t data_handle);
int (*libstarpu_mpi_scatter_detached) (starpu_data_handle_t* data_handles, int count, int root, MPI_Comm comm, void (* scallback)(void *arg), void *arg,void (*rcallback)(void *), void *rarg);
int (*libstarpu_mpi_gather_detached) (starpu_data_handle_t* data_handles, int count, int root, MPI_Comm comm, void (* scallback)(void *arg), void *arg, void (*rcallback)(void *), void *rarg);
int (*libstarpu_mpi_isend_detached_unlock_tag) (starpu_data_handle_t data_handle, int dest, int mpi_tag, MPI_Comm comm, starpu_tag_t tag);
int (*libstarpu_mpi_irecv_detached_unlock_tag) (starpu_data_handle_t data_handle, int source, int mpi_tag, MPI_Comm comm, starpu_tag_t tag);
int (*libstarpu_mpi_isend_array_detached_unlock_tag) (unsigned array_size, starpu_data_handle_t* data_handle, int* dest, int* mpi_tag, MPI_Comm* comm, starpu_tag_t tag);
int (*libstarpu_mpi_irecv_array_detached_unlock_tag) (unsigned array_size, starpu_data_handle_t* data_handle, int* source, int* mpi_tag, MPI_Comm* comm, starpu_tag_t tag);
void (*libstarpu_mpi_comm_amounts_retrieve) (size_t* comm_amounts);
void (*libstarpu_mpi_cache_flush) (MPI_Comm comm, starpu_data_handle_t data_handle);
void (*libstarpu_mpi_cache_flush_all_data) (MPI_Comm comm);
#ifndef HAVE_STARPU_MPI_DATA_REGISTER_COMM
void (*libstarpu_mpi_data_register) (starpu_data_handle_t data_handle, int tag, int rank);
#else
void (*libstarpu_mpi_data_register_comm) (starpu_data_handle_t data_handle, int tag, int rank, MPI_Comm comm);
#endif
#endif // USE_MPI

void (*libstarpu_sched_ctx_add_workers) (int* workerids_ctx, unsigned nworkers_ctx, unsigned sched_ctx_id);
void (*libstarpu_sched_ctx_remove_workers) (int* workerids_ctx, unsigned nworkers_ctx, unsigned sched_ctx_id);
void (*libstarpu_sched_ctx_delete) (unsigned sched_ctx_id);
void (*libstarpu_sched_ctx_set_inheritor) (unsigned sched_ctx_id, unsigned inheritor);
void (*libstarpu_sched_ctx_set_context) (unsigned* sched_ctx_id);
unsigned (*libstarpu_sched_ctx_get_context) ();
void (*libstarpu_sched_ctx_stop_task_submission) ();
void (*libstarpu_sched_ctx_finished_submit) (unsigned sched_ctx_id);
struct starpu_worker_collection * (*libstarpu_sched_ctx_create_worker_collection) (unsigned sched_ctx_id, enum starpu_worker_collection_type type);
void (*libstarpu_sched_ctx_delete_worker_collection) (unsigned sched_ctx_id);
struct starpu_worker_collection * (*libstarpu_sched_ctx_get_worker_collection) (unsigned sched_ctx_id);
int (*libstarpu_prefetch_task_input_on_node) (struct starpu_task* task, unsigned node);
int (*libstarpu_malloc) (void** A, size_t dim);
int (*libstarpu_free) (void* A);
int (*libstarpu_malloc_flags) (void** A, size_t dim, int flags);
int (*libstarpu_free_flags) (void* A, size_t dim, int flags);
int (*libstarpu_tag_wait) (starpu_tag_t id);
int (*libstarpu_tag_wait_array) (unsigned ntags, starpu_tag_t* id);
void (*libstarpu_tag_notify_from_apps) (starpu_tag_t id);
void (*libstarpu_task_init) (struct starpu_task* task);
void (*libstarpu_task_clean) (struct starpu_task* task);
struct starpu_task * (*libstarpu_task_create) ();
void (*libstarpu_task_destroy) (struct starpu_task* task);
int (*libstarpu_task_submit) (struct starpu_task* task);
int (*libstarpu_task_submit_to_ctx) (struct starpu_task* task, unsigned sched_ctx_id);
int (*libstarpu_task_wait) (struct starpu_task* task);
int (*libstarpu_task_wait_for_all) ();
int (*libstarpu_task_wait_for_all_in_ctx) (unsigned sched_ctx_id);
int (*libstarpu_task_wait_for_no_ready) ();
int (*libstarpu_task_nready) ();
int (*libstarpu_task_nsubmitted) ();
struct starpu_task * (*libstarpu_task_dup) (struct starpu_task* task);
void (*libstarpu_task_bundle_create) (starpu_task_bundle_t* bundle);
int (*libstarpu_task_bundle_insert) (starpu_task_bundle_t bundle, struct starpu_task* task);
int (*libstarpu_task_bundle_remove) (starpu_task_bundle_t bundle, struct starpu_task* task);
void (*libstarpu_task_bundle_close) (starpu_task_bundle_t bundle);
void (*libstarpu_create_sync_task) (starpu_tag_t sync_tag, unsigned ndeps, starpu_tag_t* deps, void (* callback)(void *), void *callback_arg);
void (*libstarpu_execute_on_each_worker) (void (* func)(void *), void *arg, uint32_t where);
void (*libstarpu_execute_on_each_worker_ex) (void (* func)(void *), void *arg, uint32_t where, const char * name);
void (*libstarpu_execute_on_specific_workers) (void (* func)(void*), void * arg, unsigned num_workers, unsigned * workers, const char * name);
int (*libstarpu_data_cpy) (starpu_data_handle_t dst_handle, starpu_data_handle_t src_handle, int asynchronous, void (* callback_func)(void*), void *callback_arg);

int starpu_run_job_id = -1;

static void starpu_register_functions() {
  if(starpu_run_job_id < 0 ) {
    starpu_run_job_id = ezt_otf2_register_function("StarPU job");
  }
}

#define STARPU_ENTER(event_id) do {					\
    if(event_id<0) {							\
      starpu_register_functions();					\
      eztrace_assert(event_id >= 0);					\
    }									\
    EZTRACE_SHOULD_TRACE(OTF2_EvtWriter_Enter(evt_writer,		\
					      NULL,			\
					      ezt_get_timestamp(),	\
					      event_id));		\
  }while(0)

#define STARPU_LEAVE(event_id) do {					\
    if(event_id<0) {							\
      starpu_register_functions();					\
      eztrace_assert(event_id >= 0);					\
    }									\
    EZTRACE_SHOULD_TRACE(OTF2_EvtWriter_Leave(evt_writer,		\
					      NULL,			\
					      ezt_get_timestamp(),	\
					      event_id));		\
  }while(0)

void _starpu_driver_start_job (void *args, void *j, void *codelet_start, int rank, int profiling) {
    FUNCTION_ENTRY_WITH_ARGS (args, j, codelet_start, rank, profiling);
    lib_starpu_driver_start_job (args, j, codelet_start, rank, profiling);
    FUNCTION_EXIT;

    STARPU_ENTER(starpu_run_job_id);
}

void _starpu_driver_end_job (void* args, void* j, void* perf_arch, void *codelet_end, int rank, int profiling) {
  STARPU_LEAVE(starpu_run_job_id);

  FUNCTION_ENTRY_WITH_ARGS (args, j, perf_arch, codelet_end, rank, profiling)
  lib_starpu_driver_end_job (args, j, perf_arch, codelet_end, rank, profiling);
  FUNCTION_EXIT;
}


void starpu_cublas_init () {
  FUNCTION_ENTRY;
  libstarpu_cublas_init ();
  FUNCTION_EXIT;
}

void starpu_cublas_shutdown () {
  FUNCTION_ENTRY;
  libstarpu_cublas_shutdown ();
  FUNCTION_EXIT;
}

int starpu_init (struct starpu_conf* conf) {
  FUNCTION_ENTRY;
  int ret = libstarpu_init (conf);
  FUNCTION_EXIT;
  return ret;
}

void starpu_pause () {
  FUNCTION_ENTRY;
  libstarpu_pause ();
  FUNCTION_EXIT;
}

void starpu_resume () {
  FUNCTION_ENTRY;
  libstarpu_resume ();
  FUNCTION_EXIT;
}

void starpu_shutdown () {
  FUNCTION_ENTRY;
  libstarpu_shutdown ();
  FUNCTION_EXIT;
}

void starpu_data_unregister (starpu_data_handle_t handle) {
  FUNCTION_ENTRY_WITH_ARGS(handle);
  libstarpu_data_unregister (handle);
  FUNCTION_EXIT;
}

void starpu_data_unregister_no_coherency (starpu_data_handle_t handle) {
  FUNCTION_ENTRY_WITH_ARGS(handle);
  libstarpu_data_unregister_no_coherency (handle);
  FUNCTION_EXIT;
}

void starpu_data_unregister_submit (starpu_data_handle_t handle) {
  FUNCTION_ENTRY_WITH_ARGS(handle);
  libstarpu_data_unregister_submit (handle);
  FUNCTION_EXIT; 
}

void starpu_data_invalidate (starpu_data_handle_t handle) {
  FUNCTION_ENTRY_WITH_ARGS(handle);
  libstarpu_data_invalidate (handle);
  FUNCTION_EXIT;
}

void starpu_data_invalidate_submit (starpu_data_handle_t handle) {
  FUNCTION_ENTRY_WITH_ARGS(handle);
  libstarpu_data_invalidate_submit (handle);
  FUNCTION_EXIT;
}

int starpu_data_acquire (starpu_data_handle_t handle, enum starpu_data_access_mode mode) {
  FUNCTION_ENTRY_WITH_ARGS(handle, mode);
  int ret = libstarpu_data_acquire (handle, mode);
  FUNCTION_EXIT;
  return ret;
}

int starpu_data_acquire_on_node (starpu_data_handle_t handle, int node, enum starpu_data_access_mode mode) {
  FUNCTION_ENTRY_WITH_ARGS(handle, node, mode);
  int ret = libstarpu_data_acquire_on_node (handle, node, mode);
  FUNCTION_EXIT;
  return ret;
}

int starpu_data_acquire_cb (starpu_data_handle_t handle, enum starpu_data_access_mode mode, void (* callback)(void *), void *arg) {
  FUNCTION_ENTRY_WITH_ARGS(handle, mode, callback, arg);
  int ret = libstarpu_data_acquire_cb (handle, mode, callback, arg);
  FUNCTION_EXIT;
  return ret;
}

int starpu_data_acquire_on_node_cb (starpu_data_handle_t handle, int node, enum starpu_data_access_mode mode, void (* callback)(void *), void *arg) {
  FUNCTION_ENTRY_WITH_ARGS(handle, node, mode, callback, arg);
  int ret = libstarpu_data_acquire_on_node_cb (handle, node, mode, callback, arg);
  FUNCTION_EXIT;
  return ret;
}

void starpu_data_release (starpu_data_handle_t handle) {
  FUNCTION_ENTRY_WITH_ARGS(handle);
  libstarpu_data_release (handle);
  FUNCTION_EXIT;
}

void starpu_data_release_on_node (starpu_data_handle_t handle, int node) {
  FUNCTION_ENTRY_WITH_ARGS(handle, node);
  libstarpu_data_release_on_node (handle, node);
  FUNCTION_EXIT;
}

int starpu_data_request_allocation (starpu_data_handle_t handle, unsigned node) {
  FUNCTION_ENTRY_WITH_ARGS(handle, node);
  int ret = libstarpu_data_request_allocation (handle, node);
  FUNCTION_EXIT;
  return ret;
}

uintptr_t starpu_malloc_on_node (unsigned dst_node, size_t size) {
  FUNCTION_ENTRY_WITH_ARGS(dst_node, size);
  uintptr_t ret = libstarpu_malloc_on_node (dst_node, size);
  FUNCTION_EXIT;
  return ret;
}

void starpu_free_on_node (unsigned dst_node, uintptr_t addr, size_t size) {
  FUNCTION_ENTRY_WITH_ARGS(dst_node, addr, size);
  libstarpu_free_on_node (dst_node, addr, size);
  FUNCTION_EXIT;
}

#ifdef USE_MPI

int starpu_mpi_isend (starpu_data_handle_t data_handle, starpu_mpi_req* req, int dest, int mpi_tag, MPI_Comm comm) {
  FUNCTION_ENTRY_WITH_ARGS(data_handle, req, dest, mpi_tag, comm);
  int ret = libstarpu_mpi_isend (data_handle, req, dest, mpi_tag, comm);
  FUNCTION_EXIT;
  return ret;
}

int starpu_mpi_irecv (starpu_data_handle_t data_handle, starpu_mpi_req* req, int source, int mpi_tag, MPI_Comm comm) {
  FUNCTION_ENTRY_WITH_ARGS(data_handle, req, source, mpi_tag);
  int ret = libstarpu_mpi_irecv (data_handle, req, source, mpi_tag, comm);
  FUNCTION_EXIT;
  return ret;
}

int starpu_mpi_send (starpu_data_handle_t data_handle, int dest, int mpi_tag, MPI_Comm comm) {
  FUNCTION_ENTRY_WITH_ARGS(data_handle, dest, mpi_tag, comm);
  int ret = libstarpu_mpi_send (data_handle, dest, mpi_tag, comm);
  FUNCTION_EXIT;
  return ret;
}

int starpu_mpi_recv (starpu_data_handle_t data_handle, int source, int mpi_tag, MPI_Comm comm, MPI_Status* status) {
  FUNCTION_ENTRY_WITH_ARGS(data_handle, source, mpi_tag, comm);
  int ret = libstarpu_mpi_recv (data_handle, source, mpi_tag, comm, status);
  FUNCTION_EXIT;
  return ret;
}

int starpu_mpi_isend_detached (starpu_data_handle_t data_handle, int dest, int mpi_tag, MPI_Comm comm, void (* callback)(void *arg), void *arg) {
  FUNCTION_ENTRY_WITH_ARGS(data_handle, dest, mpi_tag, comm);
  int ret = libstarpu_mpi_isend_detached (data_handle, dest, mpi_tag, comm, callback, arg);
  FUNCTION_EXIT;
  return ret;
}

int starpu_mpi_irecv_detached (starpu_data_handle_t data_handle, int source, int mpi_tag, MPI_Comm comm, void (* callback)(void *arg), void *arg) {
  FUNCTION_ENTRY_WITH_ARGS(data_handle, source, mpi_tag, comm);
  int ret = libstarpu_mpi_irecv_detached (data_handle, source, mpi_tag, comm, callback, arg);
  FUNCTION_EXIT;
  return ret;
}

int starpu_mpi_wait (starpu_mpi_req* req, MPI_Status* status) {
  FUNCTION_ENTRY_WITH_ARGS(req, status);
  int ret = libstarpu_mpi_wait (req, status);
  FUNCTION_EXIT;
  return ret;
}

int starpu_mpi_test (starpu_mpi_req* req, int* flag, MPI_Status* status) {
  FUNCTION_ENTRY_WITH_ARGS(req, flag, status);
  int ret = libstarpu_mpi_test (req, flag, status);
  FUNCTION_EXIT;
  return ret;
}

int starpu_mpi_barrier (MPI_Comm comm) {
  FUNCTION_ENTRY_WITH_ARGS(comm);
  int ret = libstarpu_mpi_barrier (comm);
  FUNCTION_EXIT;
  return ret;
}

int starpu_mpi_init (int* argc, char*** argv, int initialize_mpi) {
  FUNCTION_ENTRY_WITH_ARGS(argc, argv, initialize_mpi);
  int ret = libstarpu_mpi_init (argc, argv, initialize_mpi);
  FUNCTION_EXIT;
  return ret;
}

int starpu_mpi_initialize () {
  FUNCTION_ENTRY;
  int ret = libstarpu_mpi_initialize ();
  FUNCTION_EXIT;
  return ret;
}

int starpu_mpi_initialize_extended (int* rank, int* world_size) {
  FUNCTION_ENTRY_WITH_ARGS(rank, world_size);
  int ret = libstarpu_mpi_initialize_extended (rank, world_size);
  FUNCTION_EXIT;
  return ret;
}

int starpu_mpi_shutdown () {
  FUNCTION_ENTRY;
  int ret = libstarpu_mpi_shutdown ();
  FUNCTION_EXIT;
  return ret;
}

void starpu_mpi_get_data_on_node (MPI_Comm comm, starpu_data_handle_t data_handle, int node) {
  FUNCTION_ENTRY_WITH_ARGS(comm, data_handle, node);
  libstarpu_mpi_get_data_on_node (comm, data_handle, node);
  FUNCTION_EXIT;
}

void starpu_mpi_get_data_on_node_detached (MPI_Comm comm, starpu_data_handle_t data_handle, int node, void (* callback)(void *arg), void *arg) {
  FUNCTION_ENTRY_WITH_ARGS(comm, data_handle, node);
  libstarpu_mpi_get_data_on_node_detached (comm, data_handle, node, callback, arg);
  FUNCTION_EXIT;
}

void starpu_mpi_redux_data (MPI_Comm comm, starpu_data_handle_t data_handle) {
  FUNCTION_ENTRY_WITH_ARGS(comm, data_handle);
  libstarpu_mpi_redux_data (comm, data_handle);
  FUNCTION_EXIT;
}

int starpu_mpi_scatter_detached (starpu_data_handle_t* data_handles, int count, int root, MPI_Comm comm, void (* scallback)(void *arg), void *arg, void (*rcallback)(void *), void *rarg) {
  FUNCTION_ENTRY_WITH_ARGS(data_handles, count, root, comm);
  int ret = libstarpu_mpi_scatter_detached (data_handles, count, root, comm, scallback, arg, rcallback, rarg);
  FUNCTION_EXIT;
  return ret;
}

int starpu_mpi_gather_detached (starpu_data_handle_t* data_handles, int count, int root, MPI_Comm comm, void (* scallback)(void *arg), void *arg, void (*rcallback)(void *), void *rarg) {
  FUNCTION_ENTRY_WITH_ARGS(data_handles, count, root, comm);
  int ret = libstarpu_mpi_gather_detached (data_handles, count, root, comm, scallback, arg, rcallback, rarg);
  FUNCTION_EXIT;
  return ret;
}

int starpu_mpi_isend_detached_unlock_tag (starpu_data_handle_t data_handle, int dest, int mpi_tag, MPI_Comm comm, starpu_tag_t tag) {
  FUNCTION_ENTRY_WITH_ARGS(data_handle, dest, mpi_tag, comm);
  int ret = libstarpu_mpi_isend_detached_unlock_tag (data_handle, dest, mpi_tag, comm, tag);
  FUNCTION_EXIT;
  return ret;
}

int starpu_mpi_irecv_detached_unlock_tag (starpu_data_handle_t data_handle, int source, int mpi_tag, MPI_Comm comm, starpu_tag_t tag) {
  FUNCTION_ENTRY_WITH_ARGS(data_handle, source, mpi_tag, comm);
  int ret = libstarpu_mpi_irecv_detached_unlock_tag (data_handle, source, mpi_tag, comm, tag);
  FUNCTION_EXIT;
  return ret;
}

int starpu_mpi_isend_array_detached_unlock_tag (unsigned array_size, starpu_data_handle_t* data_handle, int* dest, int* mpi_tag, MPI_Comm* comm, starpu_tag_t tag) {
  FUNCTION_ENTRY_WITH_ARGS(array_size, data_handle, dest, mpi_tag);
  int ret = libstarpu_mpi_isend_array_detached_unlock_tag (array_size, data_handle, dest, mpi_tag, comm, tag);
  FUNCTION_EXIT;
  return ret;
}

int starpu_mpi_irecv_array_detached_unlock_tag (unsigned array_size, starpu_data_handle_t* data_handle, int* source, int* mpi_tag, MPI_Comm* comm, starpu_tag_t tag) {
  FUNCTION_ENTRY_WITH_ARGS(array_size, data_handle, source, mpi_tag);
  int ret = libstarpu_mpi_irecv_array_detached_unlock_tag (array_size, data_handle, source, mpi_tag, comm, tag);
  FUNCTION_EXIT;
  return ret;
}

void starpu_mpi_comm_amounts_retrieve (size_t* comm_amounts) {
  FUNCTION_ENTRY_WITH_ARGS(comm_amounts);
  libstarpu_mpi_comm_amounts_retrieve (comm_amounts);
  FUNCTION_EXIT;
}

void starpu_mpi_cache_flush (MPI_Comm comm, starpu_data_handle_t data_handle) {
  FUNCTION_ENTRY_WITH_ARGS(comm, data_handle);
  libstarpu_mpi_cache_flush (comm, data_handle);
  FUNCTION_EXIT;
}

void starpu_mpi_cache_flush_all_data (MPI_Comm comm) {
  FUNCTION_ENTRY_WITH_ARGS(comm);
  libstarpu_mpi_cache_flush_all_data (comm);
  FUNCTION_EXIT;
}

#ifndef HAVE_STARPU_MPI_DATA_REGISTER_COMM
void starpu_mpi_data_register (starpu_data_handle_t data_handle, int tag, int rank) {
  FUNCTION_ENTRY;
  libstarpu_mpi_data_register (data_handle, tag, rank);
  FUNCTION_EXIT;
}
#else
void starpu_mpi_data_register_comm(starpu_data_handle_t data_handle, int tag, int rank, MPI_Comm comm) {
  FUNCTION_ENTRY;
  libstarpu_mpi_data_register_comm (data_handle, tag, rank, comm);
  FUNCTION_EXIT;
}
#endif
#endif



void starpu_sched_ctx_add_workers (int* workerids_ctx, unsigned nworkers_ctx, unsigned sched_ctx_id) {
  FUNCTION_ENTRY_WITH_ARGS(workerids_ctx, nworkers_ctx, sched_ctx_id);
  libstarpu_sched_ctx_add_workers (workerids_ctx, nworkers_ctx, sched_ctx_id);
  FUNCTION_EXIT;
}

void starpu_sched_ctx_remove_workers (int* workerids_ctx, unsigned nworkers_ctx, unsigned sched_ctx_id) {
  FUNCTION_ENTRY_WITH_ARGS(workerids_ctx, nworkers_ctx, sched_ctx_id);
  libstarpu_sched_ctx_remove_workers (workerids_ctx, nworkers_ctx, sched_ctx_id);
  FUNCTION_EXIT;
}

void starpu_sched_ctx_delete (unsigned sched_ctx_id) {
  FUNCTION_ENTRY_WITH_ARGS(sched_ctx_id);
  libstarpu_sched_ctx_delete (sched_ctx_id);
  FUNCTION_EXIT;
}

void starpu_sched_ctx_set_inheritor (unsigned sched_ctx_id, unsigned inheritor) {
  FUNCTION_ENTRY_WITH_ARGS(sched_ctx_id, inheritor);
  libstarpu_sched_ctx_set_inheritor (sched_ctx_id, inheritor);
  FUNCTION_EXIT;
}

void starpu_sched_ctx_set_context (unsigned* sched_ctx_id) {
  FUNCTION_ENTRY_WITH_ARGS(sched_ctx_id);
  libstarpu_sched_ctx_set_context (sched_ctx_id);
  FUNCTION_EXIT;
}

unsigned starpu_sched_ctx_get_context () {
  FUNCTION_ENTRY;
  unsigned ret = libstarpu_sched_ctx_get_context ();
  FUNCTION_EXIT;
  return ret;
}

void starpu_sched_ctx_stop_task_submission () {
  FUNCTION_ENTRY;
  libstarpu_sched_ctx_stop_task_submission ();
  FUNCTION_EXIT;
}

void starpu_sched_ctx_finished_submit (unsigned sched_ctx_id) {
  FUNCTION_ENTRY_WITH_ARGS(sched_ctx_id);
  libstarpu_sched_ctx_finished_submit (sched_ctx_id);
  FUNCTION_EXIT;
}

struct starpu_worker_collection * starpu_sched_ctx_create_worker_collection (unsigned sched_ctx_id, enum starpu_worker_collection_type type) {
  FUNCTION_ENTRY_WITH_ARGS(sched_ctx_id, type);
  struct starpu_worker_collection * ret = libstarpu_sched_ctx_create_worker_collection (sched_ctx_id, type);
  FUNCTION_EXIT(ret);
  return ret;
}

void starpu_sched_ctx_delete_worker_collection (unsigned sched_ctx_id) {
  FUNCTION_ENTRY_WITH_ARGS(sched_ctx_id);
  libstarpu_sched_ctx_delete_worker_collection (sched_ctx_id);
  FUNCTION_EXIT;
}

struct starpu_worker_collection * starpu_sched_ctx_get_worker_collection (unsigned sched_ctx_id) {
  FUNCTION_ENTRY_WITH_ARGS(sched_ctx_id);
  struct starpu_worker_collection * ret = libstarpu_sched_ctx_get_worker_collection (sched_ctx_id);
  FUNCTION_EXIT;
  return ret;
}

int starpu_prefetch_task_input_on_node (struct starpu_task* task, unsigned node) {
  FUNCTION_ENTRY_WITH_ARGS(task, node);
  int ret = libstarpu_prefetch_task_input_on_node (task, node);
  FUNCTION_EXIT;
  return ret;
}

int starpu_malloc (void** A, size_t dim) {
  FUNCTION_ENTRY_WITH_ARGS(A, dim);
  int ret = libstarpu_malloc (A, dim);
  FUNCTION_EXIT;
  return ret;
}

int starpu_free (void* A) {
  FUNCTION_ENTRY_WITH_ARGS(A);
  int ret = libstarpu_free (A);
  FUNCTION_EXIT;
  return ret;
}

int starpu_malloc_flags (void** A, size_t dim, int flags) {
  FUNCTION_ENTRY_WITH_ARGS(A, dim, flags);
  int ret = libstarpu_malloc_flags (A, dim, flags);
  FUNCTION_EXIT;
  return ret;
}

int starpu_free_flags (void* A, size_t dim, int flags) {
  FUNCTION_ENTRY_WITH_ARGS(A, dim, flags);
  int ret = libstarpu_free_flags (A, dim, flags);
  FUNCTION_EXIT;
  return ret;
}

int starpu_tag_wait (starpu_tag_t id) {
  FUNCTION_ENTRY_WITH_ARGS(id);
  int ret = libstarpu_tag_wait (id);
  FUNCTION_EXIT;
  return ret;
}

int starpu_tag_wait_array (unsigned ntags, starpu_tag_t* id) {
  FUNCTION_ENTRY_WITH_ARGS(ntags, id);
  int ret = libstarpu_tag_wait_array (ntags, id);
  FUNCTION_EXIT;
  return ret;
}

void starpu_tag_notify_from_apps (starpu_tag_t id) {
  FUNCTION_ENTRY_WITH_ARGS(id);
  libstarpu_tag_notify_from_apps (id);
  FUNCTION_EXIT;
}

void starpu_task_init (struct starpu_task* task) {
  FUNCTION_ENTRY_WITH_ARGS(task);
  libstarpu_task_init (task);
  FUNCTION_EXIT;
}

void starpu_task_clean (struct starpu_task* task) {
  FUNCTION_ENTRY_WITH_ARGS(task);
  libstarpu_task_clean (task);
  FUNCTION_EXIT;
}

struct starpu_task * starpu_task_create () {
  FUNCTION_ENTRY;
  struct starpu_task * ret = libstarpu_task_create ();
  FUNCTION_EXIT;
  return ret;
}

void starpu_task_destroy (struct starpu_task* task) {
  FUNCTION_ENTRY_WITH_ARGS(task);
  libstarpu_task_destroy (task);
  FUNCTION_EXIT;
}

#undef starpu_task_submit
int starpu_task_submit (struct starpu_task* task) {
  FUNCTION_ENTRY_WITH_ARGS(task);
  int ret = libstarpu_task_submit (task);
  FUNCTION_EXIT;
  return ret;
}

int starpu_task_submit_to_ctx (struct starpu_task* task, unsigned sched_ctx_id) {
  FUNCTION_ENTRY_WITH_ARGS(task, sched_ctx_id);
  int ret = libstarpu_task_submit_to_ctx (task, sched_ctx_id);
  FUNCTION_EXIT;
  return ret;
}

int starpu_task_wait (struct starpu_task* task) {
  FUNCTION_ENTRY_WITH_ARGS(task);
  int ret = libstarpu_task_wait (task);
  FUNCTION_EXIT;
  return ret;
}

int starpu_task_wait_for_all () {
  FUNCTION_ENTRY;
  int ret = libstarpu_task_wait_for_all ();
  FUNCTION_EXIT;
  return ret;
}

int starpu_task_wait_for_all_in_ctx (unsigned sched_ctx_id) {
  FUNCTION_ENTRY_WITH_ARGS(sched_ctx_id);
  int ret = libstarpu_task_wait_for_all_in_ctx (sched_ctx_id);
  FUNCTION_EXIT;
  return ret;
}

int starpu_task_wait_for_no_ready () {
  FUNCTION_ENTRY;
  int ret = libstarpu_task_wait_for_no_ready ();
  FUNCTION_EXIT;
  return ret;
}

int starpu_task_nready () {
  FUNCTION_ENTRY;
  int ret = libstarpu_task_nready ();
  FUNCTION_EXIT;
  return ret;
}

int starpu_task_nsubmitted () {
  FUNCTION_ENTRY;
  int ret = libstarpu_task_nsubmitted ();
  FUNCTION_EXIT;
  return ret;
}

struct starpu_task * starpu_task_dup (struct starpu_task* task) {
  FUNCTION_ENTRY_WITH_ARGS(task);
  struct starpu_task * ret = libstarpu_task_dup (task);
  FUNCTION_EXIT_WITH_ARGS(ret);
  return ret;
}

void starpu_task_bundle_create (starpu_task_bundle_t* bundle) {
  FUNCTION_ENTRY_WITH_ARGS(bundle);
  libstarpu_task_bundle_create (bundle);
  FUNCTION_EXIT;
}

int starpu_task_bundle_insert (starpu_task_bundle_t bundle, struct starpu_task* task) {
  FUNCTION_ENTRY_WITH_ARGS(bundle, task);
  int ret = libstarpu_task_bundle_insert (bundle, task);
  FUNCTION_EXIT;
  return ret;
}

int starpu_task_bundle_remove (starpu_task_bundle_t bundle, struct starpu_task* task) {
  FUNCTION_ENTRY_WITH_ARGS(bundle, task);
  int ret = libstarpu_task_bundle_remove (bundle, task);
  FUNCTION_EXIT;
  return ret;
}

void starpu_task_bundle_close (starpu_task_bundle_t bundle) {
  FUNCTION_ENTRY_WITH_ARGS(bundle);
  libstarpu_task_bundle_close (bundle);
  FUNCTION_EXIT;
}

void starpu_create_sync_task (starpu_tag_t sync_tag, unsigned ndeps, starpu_tag_t* deps, void (* callback)(void *), void *callback_arg) {
  FUNCTION_ENTRY_WITH_ARGS(sync_tag, ndeps, deps);
  libstarpu_create_sync_task (sync_tag, ndeps, deps, callback, callback_arg);
  FUNCTION_EXIT;
}

void starpu_execute_on_each_worker (void (* func)(void *), void *arg, uint32_t where) {
  FUNCTION_ENTRY_WITH_ARGS(func, arg, where);
  libstarpu_execute_on_each_worker (func, arg, where);
  FUNCTION_EXIT;
}

void starpu_execute_on_each_worker_ex (void (* func)(void *), void *arg, uint32_t where, const char * name) {
  FUNCTION_ENTRY_WITH_ARGS(func, arg, where, name);
  libstarpu_execute_on_each_worker_ex (func, arg, where, name);
  FUNCTION_EXIT;
}

void starpu_execute_on_specific_workers (void (* func)(void*), void * arg, unsigned num_workers, unsigned * workers, const char * name) {
  FUNCTION_ENTRY_WITH_ARGS(func, arg, num_workers, workers, name);
  libstarpu_execute_on_specific_workers (func, arg, num_workers, workers, name);
  FUNCTION_EXIT;
}

int starpu_data_cpy (starpu_data_handle_t dst_handle, starpu_data_handle_t src_handle, int asynchronous, void (* callback_func)(void*), void *callback_arg) {
  FUNCTION_ENTRY_WITH_ARGS(dst_handle, src_handle, asynchronous, callback_func, callback_arg);
  int ret = libstarpu_data_cpy (dst_handle, src_handle, asynchronous, callback_func, callback_arg);
  FUNCTION_EXIT;
  return ret;
}


PPTRACE_START_INTERCEPT_FUNCTIONS(starpu)
INTERCEPT3("_starpu_driver_start_job", lib_starpu_driver_start_job)
INTERCEPT3("_starpu_driver_end_job", lib_starpu_driver_end_job)
INTERCEPT3("starpu_cublas_init", libstarpu_cublas_init)
INTERCEPT3("starpu_cublas_shutdown", libstarpu_cublas_shutdown)
INTERCEPT3("starpu_init", libstarpu_init)
INTERCEPT3("starpu_pause", libstarpu_pause)
INTERCEPT3("starpu_resume", libstarpu_resume)
INTERCEPT3("starpu_shutdown", libstarpu_shutdown)
INTERCEPT3("starpu_data_unregister", libstarpu_data_unregister)
INTERCEPT3("starpu_data_unregister_no_coherency", libstarpu_data_unregister_no_coherency)
INTERCEPT3("starpu_data_unregister_submit", libstarpu_data_unregister_submit)
INTERCEPT3("starpu_data_invalidate", libstarpu_data_invalidate)
INTERCEPT3("starpu_data_invalidate_submit", libstarpu_data_invalidate_submit)
INTERCEPT3("starpu_data_acquire", libstarpu_data_acquire)
INTERCEPT3("starpu_data_acquire_on_node", libstarpu_data_acquire_on_node)
INTERCEPT3("starpu_data_acquire_cb", libstarpu_data_acquire_cb)
INTERCEPT3("starpu_data_acquire_on_node_cb", libstarpu_data_acquire_on_node_cb)
INTERCEPT3("starpu_data_release", libstarpu_data_release)
INTERCEPT3("starpu_data_release_on_node", libstarpu_data_release_on_node)
INTERCEPT3("starpu_data_request_allocation", libstarpu_data_request_allocation)
INTERCEPT3("starpu_malloc_on_node", libstarpu_malloc_on_node)
INTERCEPT3("starpu_free_on_node", libstarpu_free_on_node)
#ifdef USE_MPI
INTERCEPT3("starpu_mpi_isend", libstarpu_mpi_isend)
INTERCEPT3("starpu_mpi_irecv", libstarpu_mpi_irecv)
INTERCEPT3("starpu_mpi_send", libstarpu_mpi_send)
INTERCEPT3("starpu_mpi_recv", libstarpu_mpi_recv)
INTERCEPT3("starpu_mpi_isend_detached", libstarpu_mpi_isend_detached)
INTERCEPT3("starpu_mpi_irecv_detached", libstarpu_mpi_irecv_detached)
INTERCEPT3("starpu_mpi_wait", libstarpu_mpi_wait)
INTERCEPT3("starpu_mpi_test", libstarpu_mpi_test)
INTERCEPT3("starpu_mpi_barrier", libstarpu_mpi_barrier)
INTERCEPT3("starpu_mpi_init", libstarpu_mpi_init)
INTERCEPT3("starpu_mpi_initialize", libstarpu_mpi_initialize)
INTERCEPT3("starpu_mpi_initialize_extended", libstarpu_mpi_initialize_extended)
INTERCEPT3("starpu_mpi_shutdown", libstarpu_mpi_shutdown)
INTERCEPT3("starpu_mpi_get_data_on_node", libstarpu_mpi_get_data_on_node)
INTERCEPT3("starpu_mpi_get_data_on_node_detached", libstarpu_mpi_get_data_on_node_detached)
INTERCEPT3("starpu_mpi_redux_data", libstarpu_mpi_redux_data)
INTERCEPT3("starpu_mpi_scatter_detached", libstarpu_mpi_scatter_detached)
INTERCEPT3("starpu_mpi_gather_detached", libstarpu_mpi_gather_detached)
INTERCEPT3("starpu_mpi_isend_detached_unlock_tag", libstarpu_mpi_isend_detached_unlock_tag)
INTERCEPT3("starpu_mpi_irecv_detached_unlock_tag", libstarpu_mpi_irecv_detached_unlock_tag)
INTERCEPT3("starpu_mpi_isend_array_detached_unlock_tag", libstarpu_mpi_isend_array_detached_unlock_tag)
INTERCEPT3("starpu_mpi_irecv_array_detached_unlock_tag", libstarpu_mpi_irecv_array_detached_unlock_tag)
INTERCEPT3("starpu_mpi_comm_amounts_retrieve", libstarpu_mpi_comm_amounts_retrieve)
INTERCEPT3("starpu_mpi_cache_flush", libstarpu_mpi_cache_flush)
INTERCEPT3("starpu_mpi_cache_flush_all_data", libstarpu_mpi_cache_flush_all_data)
#ifndef H3VE_STARPU_MPI_DATA_REGISTER_COMM
INTERCEPT3("starpu_mpi_data_register", libstarpu_mpi_data_register)
#else	 3
INTERCEPT3("starpu_mpi_data_register_comm", libstarpu_mpi_data_register_comm)
#endif	 3
#endif
INTERCEPT3("starpu_sched_ctx_add_workers", libstarpu_sched_ctx_add_workers)
INTERCEPT3("starpu_sched_ctx_remove_workers", libstarpu_sched_ctx_remove_workers)
INTERCEPT3("starpu_sched_ctx_delete", libstarpu_sched_ctx_delete)
INTERCEPT3("starpu_sched_ctx_set_inheritor", libstarpu_sched_ctx_set_inheritor)
INTERCEPT3("starpu_sched_ctx_set_context", libstarpu_sched_ctx_set_context)
INTERCEPT3("starpu_sched_ctx_get_context", libstarpu_sched_ctx_get_context)
INTERCEPT3("starpu_sched_ctx_stop_task_submission", libstarpu_sched_ctx_stop_task_submission)
INTERCEPT3("starpu_sched_ctx_finished_submit", libstarpu_sched_ctx_finished_submit)
INTERCEPT3("starpu_sched_ctx_create_worker_collection", libstarpu_sched_ctx_create_worker_collection)
INTERCEPT3("starpu_sched_ctx_delete_worker_collection", libstarpu_sched_ctx_delete_worker_collection)
INTERCEPT3("starpu_sched_ctx_get_worker_collection", libstarpu_sched_ctx_get_worker_collection)
INTERCEPT3("starpu_prefetch_task_input_on_node", libstarpu_prefetch_task_input_on_node)
INTERCEPT3("starpu_malloc", libstarpu_malloc)
INTERCEPT3("starpu_free", libstarpu_free)
INTERCEPT3("starpu_malloc_flags", libstarpu_malloc_flags)
INTERCEPT3("starpu_free_flags", libstarpu_free_flags)
INTERCEPT3("starpu_tag_wait", libstarpu_tag_wait)
INTERCEPT3("starpu_tag_wait_array", libstarpu_tag_wait_array)
INTERCEPT3("starpu_task_init", libstarpu_task_init)
INTERCEPT3("starpu_task_clean", libstarpu_task_clean)
INTERCEPT3("starpu_task_create", libstarpu_task_create)
INTERCEPT3("starpu_task_destroy", libstarpu_task_destroy)
INTERCEPT3("starpu_task_submit", libstarpu_task_submit)
INTERCEPT3("starpu_task_submit_to_ctx", libstarpu_task_submit_to_ctx)
INTERCEPT3("starpu_task_wait", libstarpu_task_wait)
INTERCEPT3("starpu_task_wait_for_all", libstarpu_task_wait_for_all)
INTERCEPT3("starpu_task_wait_for_all_in_ctx", libstarpu_task_wait_for_all_in_ctx)
INTERCEPT3("starpu_task_wait_for_no_ready", libstarpu_task_wait_for_no_ready)
INTERCEPT3("starpu_task_nready", libstarpu_task_nready)
INTERCEPT3("starpu_task_nsubmitted", libstarpu_task_nsubmitted)
INTERCEPT3("starpu_task_dup", libstarpu_task_dup)
INTERCEPT3("starpu_task_bundle_create", libstarpu_task_bundle_create)
INTERCEPT3("starpu_task_bundle_insert", libstarpu_task_bundle_insert)
INTERCEPT3("starpu_task_bundle_remove", libstarpu_task_bundle_remove)
INTERCEPT3("starpu_task_bundle_close", libstarpu_task_bundle_close)
INTERCEPT3("starpu_create_sync_task", libstarpu_create_sync_task)
INTERCEPT3("starpu_execute_on_each_worker", libstarpu_execute_on_each_worker)
INTERCEPT3("starpu_execute_on_each_worker_ex", libstarpu_execute_on_each_worker_ex)
INTERCEPT3("starpu_execute_on_specific_workers", libstarpu_execute_on_specific_workers)
INTERCEPT3("starpu_data_cpy", libstarpu_data_cpy)
INTERCEPT3("starpu_tag_notify_from_apps", libstarpu_tag_notify_from_apps)
PPTRACE_END_INTERCEPT_FUNCTIONS(starpu)

static void init_starpu() {
  INSTRUMENT_FUNCTIONS(starpu);

  if (eztrace_autostart_enabled())
    eztrace_start();
}

static void finalize_starpu() {
  if(_starpu_initialized) {
    _starpu_initialized = 0;
    eztrace_stop();
  }
}

static void _starpu_finalize (void) __attribute__ ((destructor));
static void _starpu_finalize (void) {
  finalize_starpu();
}

static void _starpu_init (void) __attribute__ ((constructor));
/* Initialize the current library */
static void _starpu_init (void) {
  EZT_REGISTER_MODULE(starpu, "Module for StarPU functions",
		      init_starpu, finalize_starpu);
}
