/**
 *
 * @file starpu/chameleon_starpu.h
 *
 * @copyright 2009-2014 The University of Tennessee and The University of
 *                      Tennessee Research Foundation. All rights reserved.
 * @copyright 2012-2021 Bordeaux INP, CNRS (LaBRI UMR 5800), Inria,
 *                      Univ. Bordeaux. All rights reserved.
 *
 ***
 *
 * @brief Chameleon StarPU runtime header
 *
 * @version 1.1.0
 * @author Mathieu Faverge
 * @author Cedric Castagnede
 * @author Florent Pruvost
 * @author Philippe Swartvagher
 * @author Samuel Thibault
 * @date 2021-03-16
 *
 */
#ifndef _chameleon_starpu_h_
#define _chameleon_starpu_h_

#include "chameleon/config.h"

/* StarPU options */
#cmakedefine HAVE_STARPU_IDLE_PREFETCH
#cmakedefine HAVE_STARPU_ITERATION_PUSH
#cmakedefine HAVE_STARPU_DATA_WONT_USE
#cmakedefine HAVE_STARPU_DATA_SET_COORDINATES
#cmakedefine HAVE_STARPU_MALLOC_ON_NODE_SET_DEFAULT_FLAGS
#cmakedefine HAVE_STARPU_DATA_SET_OOC_FLAG
#cmakedefine HAVE_STARPU_INTERFACE_COPY2D
#cmakedefine HAVE_STARPU_DATA_PEEK
#cmakedefine HAVE_STARPU_MPI_DATA_MIGRATE
#cmakedefine HAVE_STARPU_MPI_DATA_REGISTER
#cmakedefine HAVE_STARPU_MPI_COMM_RANK
#cmakedefine HAVE_STARPU_MPI_CACHED_RECEIVE
#cmakedefine HAVE_STARPU_MPI_COMM_GET_ATTR
#cmakedefine HAVE_STARPU_MPI_INIT_CONF
#cmakedefine HAVE_STARPU_MPI_WAIT_FOR_ALL
#cmakedefine HAVE_STARPU_MPI_INTERFACE_DATATYPE_NODE_REGISTER
#cmakedefine HAVE_STARPU_MPI_INTERFACE_DATATYPE_REGISTER

#if (!defined(HAVE_STARPU_MPI_INTERFACE_DATATYPE_NODE_REGISTER) && !defined(HAVE_STARPU_MPI_INTERFACE_DATATYPE_REGISTER)) && defined(CHAMELEON_USE_MPI_DATATYPES)
#error "This version of StarPU does not support MPI datatypes (Please compile with -DCHAMELEON_USE_MPI_DATATYPES=OFF)"
#endif

#if defined(CHAMELEON_USE_MPI)
#include <starpu_mpi.h>
#else
#include <starpu.h>
#endif

#include <starpu_profiling.h>

#if defined(CHAMELEON_USE_CUDA) && !defined(CHAMELEON_SIMULATION)
#include <starpu_scheduler.h>
#include <starpu_cuda.h>

#include <cublas.h>
#include <starpu_cublas.h>
#if defined(CHAMELEON_USE_CUBLAS_V2)
#include <cublas_v2.h>
#include <starpu_cublas_v2.h>
#endif
#endif

#if defined(CHAMELEON_SIMULATION)
# if !defined(STARPU_SIMGRID)
#  error "Starpu was not built with simgrid support (--enable-simgrid). Can not run Chameleon with simulation support."
# endif
#else
# if defined(STARPU_SIMGRID)
#  warning "Starpu was built with simgrid support. Better build Chameleon with simulation support (-DCHAMELEON_SIMULATION=YES)."
# endif
#endif

#include "control/common.h"
#include "runtime_codelets.h"
#include "runtime_profiling.h"
#include "runtime_codelet_profile.h"
#include "runtime_workspace.h"
#include "cham_tile_interface.h"

typedef struct starpu_conf starpu_conf_t;

/* Structure used to give some options during one request (procedure) */
typedef struct starpu_option_request_s {
    int workerid; // to force task execution on a specific workerid
} starpu_option_request_t;

/**/

/*
 * MPI Redefinitions
 */
#if defined(CHAMELEON_RUNTIME_SYNC)
#define TASK_SYNCHRONOUS , STARPU_TASK_SYNCHRONOUS, 1
#else
#define TASK_SYNCHRONOUS
#endif

#if defined(CHAMELEON_USE_MPI)
#undef STARPU_REDUX

#if defined(CHAMELEON_RUNTIME_SYNC)
#define rt_starpu_insert_task( _codelet_, ... )                         \
    starpu_mpi_insert_task( MPI_COMM_WORLD, (_codelet_), STARPU_TASK_SYNCHRONOUS, 1, ##__VA_ARGS__ )
#else
#define rt_starpu_insert_task( _codelet_, ... )                         \
    starpu_mpi_insert_task( MPI_COMM_WORLD, (_codelet_), ##__VA_ARGS__ )
#endif

#else

#if defined(CHAMELEON_RUNTIME_SYNC)
#define rt_starpu_insert_task( _codelet_, ... )                         \
    starpu_insert_task( (_codelet_), STARPU_TASK_SYNCHRONOUS, 1, ##__VA_ARGS__ )
#else
#define rt_starpu_insert_task( _codelet_, ... )                         \
    starpu_insert_task( (_codelet_), ##__VA_ARGS__ )
#endif

#endif

/*
 * cuBlasAPI v2 - StarPU enable the support for cublas handle
 */
#if defined(CHAMELEON_USE_CUDA) && defined(CHAMELEON_USE_CUBLAS_V2)
#define RUNTIME_getStream(_stream_)                             \
    cublasHandle_t _stream_ = starpu_cublas_get_local_handle();
#else
#define RUNTIME_getStream(_stream_)                             \
    cudaStream_t _stream_ = starpu_cuda_get_local_stream();     \
    cublasSetKernelStream( stream );

#endif

/*
 * Enable codelets names
 */
#if (STARPU_MAJOR_VERSION > 1) || ((STARPU_MAJOR_VERSION == 1) && (STARPU_MINOR_VERSION > 1))
#define CHAMELEON_CODELETS_HAVE_NAME
#endif

/**
 * Access to block pointer and leading dimension
 */
#define RTBLKADDR( desc, type, m, n ) ( (starpu_data_handle_t)RUNTIME_data_getaddr( desc, m, n ) )

void RUNTIME_set_reduction_methods(starpu_data_handle_t handle, cham_flttype_t dtyp);

#if defined(CHAMELEON_USE_MPI) && defined(HAVE_STARPU_MPI_CACHED_RECEIVE)
static inline int
chameleon_starpu_data_iscached(const CHAM_desc_t *A, int m, int n)
{
    int64_t mm = m + (A->i / A->mb);
    int64_t nn = n + (A->j / A->nb);

    starpu_data_handle_t *ptrtile = A->schedopt;
    ptrtile += ((int64_t)A->lmt) * nn + mm;

    if (!(*ptrtile)) {
        return 0;
    }

    return starpu_mpi_cached_receive(*ptrtile);
}

#define RUNTIME_ACCESS_WRITE_CACHED(A, Am, An) do {                 \
        if (chameleon_starpu_data_iscached(A, Am, An)) __chameleon_need_submit = 1; } while(0)

#else

#if defined(CHAMELEON_USE_MPI)
#warning "WAR dependencies need starpu_mpi_cached_receive support from StarPU 1.2.1 or greater"
#endif
#define RUNTIME_ACCESS_WRITE_CACHED(A, Am, An) do {} while (0)

#endif

#ifdef CHAMELEON_ENABLE_PRUNING_STATS

#define RUNTIME_PRUNING_STATS_BEGIN_ACCESS_DECLARATION \
    int __chameleon_exec = 0; \
    int __chameleon_changed = 0;

#define RUNTIME_PRUNING_STATS_ACCESS_W(A, Am, An) \
    if (chameleon_desc_islocal(A, Am, An)) \
        __chameleon_exec = 1;

#define RUNTIME_PRUNING_STATS_END_ACCESS_DECLARATION \
    RUNTIME_total_tasks++; \
    if (__chameleon_exec) \
        RUNTIME_exec_tasks++; \
    else if (__chameleon_need_submit) \
        RUNTIME_comm_tasks++; \
    else if (__chameleon_changed) \
        RUNTIME_changed_tasks++;

#define RUNTIME_PRUNING_STATS_RANK_CHANGED(rank) \
    int __chameleon_myrank; \
    RUNTIME_comm_rank(&__chameleon_myrank); \
    __chameleon_exec = (rank) == __chameleon_myrank; \
    __chameleon_changed = 1; \

#else
#define RUNTIME_PRUNING_STATS_BEGIN_ACCESS_DECLARATION
#define RUNTIME_PRUNING_STATS_ACCESS_W(A, Am, An)
#define RUNTIME_PRUNING_STATS_END_ACCESS_DECLARATION
#define RUNTIME_PRUNING_STATS_RANK_CHANGED(rank)
#endif

#define RUNTIME_BEGIN_ACCESS_DECLARATION        \
    RUNTIME_PRUNING_STATS_BEGIN_ACCESS_DECLARATION

#define RUNTIME_ACCESS_R(A, Am, An)

#define RUNTIME_ACCESS_W(A, Am, An)             \
    RUNTIME_PRUNING_STATS_ACCESS_W(A, Am, An);  \
    RUNTIME_ACCESS_WRITE_CACHED(A, Am, An)

#define RUNTIME_ACCESS_RW(A, Am, An)            \
    RUNTIME_PRUNING_STATS_ACCESS_W(A, Am, An);  \
    RUNTIME_ACCESS_WRITE_CACHED(A, Am, An)

#define RUNTIME_RANK_CHANGED(rank)              \
    RUNTIME_PRUNING_STATS_RANK_CHANGED(rank)

#define RUNTIME_END_ACCESS_DECLARATION          \
    RUNTIME_PRUNING_STATS_END_ACCESS_DECLARATION;

#endif /* _chameleon_starpu_h_ */
