/* Copyright 2008 BORDEAUX I UNIVERSITY & INRIA 
**
** This file is part of the PaStiX parallel sparse matrix package.
**
** This software is governed by the CeCILL-C license under French law
** and abiding by the rules of distribution of free software. You can
** use, modify and/or redistribute the software under the terms of the
** CeCILL-C license as circulated by CEA, CNRS and INRIA at the following
** URL: "http://www.cecill.info".
** 
** As a counterpart to the access to the source code and rights to copy,
** modify and redistribute granted by the license, users are provided
** only with a limited warranty and the software's author, the holder of
** the economic rights, and the successive licensors have only limited
** liability.
** 
** In this respect, the user's attention is drawn to the risks associated
** with loading, using, modifying and/or developing or reproducing the
** software by the user in light of its specific status of free software,
** that may mean that it is complicated to manipulate, and that also
** therefore means that it is reserved for developers and experienced
** professionals having in-depth computer knowledge. Users are therefore
** encouraged to load and test the software's suitability as regards
** their requirements in conditions enabling the security of their
** systems and/or data to be ensured and, more generally, to use and
** operate it in the same conditions as regards security.
** 
** The fact that you are presently reading this means that you have had
** knowledge of the CeCILL-C license and that you accept its terms.
*/
/**
 * @file starpu_register_data.c
 *
 * @author Xavier Lacoste
 * @precisions normal z -> s d c
 */

#include "starpu_defines.h"
#include "common_pastix.h"
#include "sopalin3d.h"
#include "starpu_register_data.h"
#include "sopalin_acces.h"

void starpu_fanin_init_cpu_func(void *descr[], void *cl_arg)
{
    PASTIX_FLOAT *L      = (PASTIX_FLOAT*)STARPU_MATRIX_GET_PTR(descr[0]);
    PASTIX_INT        stride = STARPU_MATRIX_GET_LD(descr[0]);
    PASTIX_INT        ncol   = STARPU_MATRIX_GET_NY(descr[0]);
    memset(L, 0, stride*ncol*sizeof(PASTIX_FLOAT));
}

#ifdef STARPU_USE_CUDA
void starpu_fanin_init_cuda_func(void *descr[], void *cl_arg)
{
    PASTIX_FLOAT *L      = (PASTIX_FLOAT*)STARPU_MATRIX_GET_PTR(descr[0]);
    PASTIX_INT        stride = STARPU_MATRIX_GET_LD(descr[0]);
    PASTIX_INT        ncol   = STARPU_MATRIX_GET_NY(descr[0]);
    cudaMemsetAsync(L, 0, stride*ncol*sizeof(PASTIX_FLOAT), starpu_cuda_get_local_stream());
}
#endif

static struct starpu_codelet starpu_fanin_init_codelet =
{
    .where = STARPU_CPU|STARPU_CUDA,
    .cpu_funcs = {starpu_fanin_init_cpu_func, NULL},
    //.cpu_funcs_name = {"starpu_fanin_init_cpu_func", NULL},
#ifdef STARPU_USE_CUDA
    .cuda_funcs = {starpu_fanin_init_cuda_func, NULL},
#  ifdef STARPU_WITH_ASYNC_CUDA
    .cuda_flags = {STARPU_CUDA_ASYNC},
#  endif
#endif
/* #ifdef STARPU_USE_OPENCL */
/*     .opencl_funcs = {init_opencl_func, NULL}, */
/* #endif */
    .modes = {STARPU_W},
    .nbuffers = 1,
    .name = "init",
};

int
starpu_register_fanin(Sopalin_Data_t * sopalin_data,
                      SolverMatrix * solvmtx,
                      starpu_data_handle_t  *** Lfanin_handle,
                      starpu_data_handle_t  *** Ufanin_handle) {
    PASTIX_INT clustnum;
    PASTIX_INT max_ftgtnbr;
    MALLOC_INTERN(*Lfanin_handle, solvmtx->clustnbr, starpu_data_handle_t*);
    if (Ufanin_handle != NULL)
        MALLOC_INTERN(*Ufanin_handle, solvmtx->clustnbr, starpu_data_handle_t*);
    for (clustnum = 0; clustnum < solvmtx->clustnbr; clustnum++) {
        PASTIX_INT coef = 1;
        SolverCblk *fanin;
        starpu_data_handle_t *Lhandle;
        starpu_data_handle_t *Uhandle;
        SolverCblk * ffanin = solvmtx->fcblktab[clustnum];
        SolverCblk * lfanin = ffanin + solvmtx->fcblknbr[clustnum];

        MALLOC_INTERN((*Lfanin_handle)[clustnum],
                      solvmtx->fcblknbr[clustnum], starpu_data_handle_t);
        if (Ufanin_handle != NULL) {
            MALLOC_INTERN((*Ufanin_handle)[clustnum],
                          solvmtx->fcblknbr[clustnum], starpu_data_handle_t);
            coef = 2;
        }
        for (fanin = ffanin; fanin < lfanin; fanin++) {
            PASTIX_INT faninnum = fcblk_getnum(solvmtx, fanin, clustnum);
            assert( clustnum == solvmtx->clustnum ||
                    solvmtx->clustnum == fanin->procdiag );
            assert(fanin->procdiag < solvmtx->clustnbr);

            Lhandle = &((*Lfanin_handle)[clustnum][faninnum]);
            if (Ufanin_handle != NULL) {
                Uhandle = &((*Ufanin_handle)[clustnum][faninnum]);
            }
            starpu_matrix_data_register(Lhandle, -1,
                                        (uintptr_t)NULL,
                                        (uint32_t)fanin->stride,
                                        (uint32_t)fanin->stride,
                                        cblk_colnbr(fanin),
                                        sizeof(PASTIX_FLOAT));

            if (Ufanin_handle != NULL) {
                starpu_matrix_data_register(Uhandle, -1,
                                            (uintptr_t)NULL,
                                            (uint32_t)fanin->stride,
                                            (uint32_t)fanin->stride,
                                            cblk_colnbr(fanin),
                                            sizeof(PASTIX_FLOAT));
            }
            starpu_mpi_data_register(*Lhandle,
                                     coef * ( solvmtx->gcblknbr +
                                              clustnum * solvmtx->gcblknbr ) +
                                     fanin->gcblknum,
                                     clustnum);
            if (clustnum == solvmtx->clustnum) {
                starpu_data_set_reduction_methods(*Lhandle, NULL,
                                                  &starpu_fanin_init_codelet);
            }
            if (Ufanin_handle != NULL) {
                starpu_mpi_data_register(*Uhandle,
                                         coef * ( solvmtx->gcblknbr +
                                                  clustnum * solvmtx->gcblknbr ) +
                                         solvmtx->gcblknbr +
                                         fanin->gcblknum,
                                         clustnum);
                if (clustnum == solvmtx->clustnum) {
                    starpu_data_set_reduction_methods(*Uhandle, NULL,
                                                      &starpu_fanin_init_codelet);
                }
            }
        }
    }
    return NO_ERR;
}

int
starpu_unregister_fanin( SolverMatrix * solvmtx,
                          starpu_data_handle_t  *** Lfanin_handle,
                          starpu_data_handle_t  *** Ufanin_handle) {
    PASTIX_INT clustnum;
    PASTIX_INT max_ftgtnbr;
    for (clustnum = 0; clustnum < solvmtx->clustnbr; clustnum++) {
        SolverCblk *fanin;
        SolverCblk * ffanin = solvmtx->fcblktab[clustnum];
        SolverCblk * lfanin = ffanin + solvmtx->fcblknbr[clustnum];
        /* Outgoing fanin are unregistred as soon as sent */
        if (clustnum != solvmtx->clustnum) {
            for (fanin = ffanin; fanin < lfanin; fanin++) {
                PASTIX_INT faninnum = fcblk_getnum(solvmtx, fanin, clustnum);
                starpu_data_unregister((*Lfanin_handle)[clustnum][faninnum]);
                if (Ufanin_handle != NULL) {
                    starpu_data_unregister((*Ufanin_handle)[clustnum][faninnum]);
                }
            }
        }
        memFree_null((*Lfanin_handle)[clustnum]);
        if (Ufanin_handle != NULL) {
            memFree_null((*Ufanin_handle)[clustnum]);
        }
    }
    memFree_null(*Lfanin_handle);
    if (Ufanin_handle != NULL)
        memFree_null(*Ufanin_handle);
    return NO_ERR;
}

int
starpu_register_halo( Sopalin_Data_t * sopalin_data,
                      SolverMatrix * datacode,
                       starpu_data_handle_t ** Lhalo_handle,
                       starpu_data_handle_t ** Uhalo_handle) {
    PASTIX_INT itercblk;
    MALLOC_INTERN(*Lhalo_handle, SOLV_HCBLKNBR, starpu_data_handle_t);
    if (Uhalo_handle) {
        MALLOC_INTERN(*Uhalo_handle, SOLV_HCBLKNBR, starpu_data_handle_t);
    }

    for (itercblk=0;itercblk<SOLV_HCBLKNBR;itercblk++) {
        SolverCblk * cblk = datacode->hcblktab + itercblk;
        ASSERT(SOLV_GCBLK2HALO(HCBLK_GCBLK(itercblk)) == itercblk,
               MOD_SOPALIN);
        assert(cblk->procdiag < datacode->clustnbr);
        starpu_matrix_data_register(&((*Lhalo_handle)[itercblk]), -1,
                                    (uintptr_t)NULL,
                                    (uint32_t)cblk->stride,
                                    (uint32_t)cblk->stride,
                                    cblk_colnbr(cblk),
                                    sizeof(PASTIX_FLOAT));
        starpu_mpi_data_register((*Lhalo_handle)[itercblk],
                                 cblk->gcblknum,
                                 cblk->procdiag);
        if (Uhalo_handle) {
            starpu_matrix_data_register(&((*Uhalo_handle)[itercblk]), -1,
                                        (uintptr_t)NULL,
                                        (uint32_t)cblk->stride,
                                        (uint32_t)cblk->stride,
                                        cblk_colnbr(cblk),
                                        sizeof(PASTIX_FLOAT));
            starpu_mpi_data_register((*Uhalo_handle)[itercblk],
                                     SOLV_GCBLKNBR + cblk->gcblknum,
                                     cblk->procdiag);
        }
    }
    return NO_ERR;
}

int
starpu_unregister_halo( SolverMatrix * datacode,
                         starpu_data_handle_t ** Lhalo_handle,
                         starpu_data_handle_t ** Uhalo_handle) {
    PASTIX_INT itercblk;
    for (itercblk=0;itercblk<SOLV_HCBLKNBR;itercblk++) {
        ASSERT(SOLV_GCBLK2HALO(HCBLK_GCBLK(itercblk)) == itercblk,
               MOD_SOPALIN);
        starpu_data_unregister((*Lhalo_handle)[itercblk]);
        if (Uhalo_handle) {
            starpu_data_unregister((*Uhalo_handle)[itercblk]);
        }
    }
    memFree_null(*Lhalo_handle);
    if (Uhalo_handle) {
        memFree_null(*Uhalo_handle);
    }
    return NO_ERR;
}

int
starpu_register_cblk( Sopalin_Data_t * sopalin_data,
                      SolverMatrix * datacode,
                       starpu_data_handle_t ** L_handle,
                       starpu_data_handle_t ** U_handle ) {
    PASTIX_INT itercblk;
    MALLOC_INTERN(*L_handle,     SYMB_CBLKNBR, starpu_data_handle_t);
    if (U_handle) {
        MALLOC_INTERN(*U_handle,    SYMB_CBLKNBR, starpu_data_handle_t);
    }

    for (itercblk=0;itercblk<SYMB_CBLKNBR;itercblk++) {
        SolverCblk * cblk = datacode->cblktab+itercblk;
        MALLOC_INTERN(cblk->coeftab, cblk->stride*cblk_colnbr(cblk), PASTIX_FLOAT);
        memset(cblk->coeftab, 0, cblk->stride*cblk_colnbr(cblk)*sizeof(PASTIX_FLOAT));
        starpu_matrix_data_register(&((*L_handle)[itercblk]), 0,
                                    (uintptr_t)cblk->coeftab,
                                    (uint32_t)cblk->stride,
                                    (uint32_t)cblk->stride,
                                    cblk_colnbr(cblk),
                                    sizeof(PASTIX_FLOAT));
        starpu_mpi_data_register((*L_handle)[itercblk],
                                 UPDOWN_LOC2GLOB(itercblk),
                                 SOLV_PROCNUM);
        if (U_handle) {
          MALLOC_INTERN(cblk->ucoeftab, cblk->stride*cblk_colnbr(cblk), PASTIX_FLOAT);
          memset(cblk->ucoeftab, 0, cblk->stride*cblk_colnbr(cblk)*sizeof(PASTIX_FLOAT));
            starpu_matrix_data_register(&((*U_handle)[itercblk]), 0,
                                        (uintptr_t)cblk->ucoeftab,
                                        (uint32_t)cblk->stride,
                                        (uint32_t)cblk->stride,
                                        cblk_colnbr(cblk),
                                        sizeof(PASTIX_FLOAT));
            starpu_mpi_data_register((*U_handle)[itercblk],
                                     SOLV_GCBLKNBR + UPDOWN_LOC2GLOB(itercblk),
                                     SOLV_PROCNUM);
        }
    }
    return NO_ERR;
}

int
starpu_unregister_cblk( SolverMatrix * datacode,
                         starpu_data_handle_t ** L_handle,
                         starpu_data_handle_t ** U_handle ) {
    PASTIX_INT itercblk;
    for (itercblk=0;itercblk<SYMB_CBLKNBR;itercblk++) {
        starpu_data_unregister((*L_handle)[itercblk]);
        if (U_handle) {
            starpu_data_unregister((*U_handle)[itercblk]);
        }
    }
    memFree_null(*L_handle);
    if (U_handle) {
        memFree_null(*U_handle);
    }

    return NO_ERR;
}

int
starpu_register_blocktab( Sopalin_Data_t        * sopalin_data,
                           starpu_data_handle_t ** blocktab_handles,
                           int                  ** blocktab) {
    SolverMatrix * datacode = sopalin_data->datacode;
    int iter;

    MALLOC_INTERN(*blocktab_handles, SOLV_PROCNBR, starpu_data_handle_t);
    for (iter = 0; iter < SOLV_PROCNBR; iter++) {
        int size;
        if (SOLV_PROCNUM == iter) {
            /* build blocktab */
            int iterblock;
            MALLOC_INTERN((*blocktab), (SYMB_BLOKNBR+SOLV_HBLOKNBR)*2, int);
            if (sopalin_data->sopar->iparm[IPARM_VERBOSE] > API_VERBOSE_NO)
                fprintf(stdout, "sizeof blocktab : %d integers\n",
                        (int)(2*(SYMB_BLOKNBR+SOLV_HBLOKNBR)));
            for (iterblock = 0; iterblock < SYMB_BLOKNBR; iterblock++) {
                (*blocktab)[2*iterblock]   = SYMB_FROWNUM(iterblock);
                (*blocktab)[2*iterblock+1] = SYMB_LROWNUM(iterblock);
            }
            for (iterblock = 0; iterblock < SOLV_HBLOKNBR; iterblock++) {
                (*blocktab)[2*(iterblock+SYMB_BLOKNBR)]   = HBLOCK_FROWNUM(iterblock);
                (*blocktab)[2*(iterblock+SYMB_BLOKNBR)+1] = HBLOCK_LROWNUM(iterblock);
            }
            size = 2*(SYMB_BLOKNBR + SOLV_HBLOKNBR);
            MPI_Bcast(&size, 1, MPI_INT, iter, sopalin_data->sopar->pastix_comm);
            starpu_vector_data_register((*blocktab_handles)+iter, 0,
                                        (uintptr_t)(*blocktab), size,
                                        sizeof(int));
        } else {
            MPI_Bcast(&size, 1, MPI_INT, iter, sopalin_data->sopar->pastix_comm);
            starpu_vector_data_register((*blocktab_handles)+iter, -1,
                                        (uintptr_t)NULL, size,
                                        sizeof(int));
        }

        starpu_mpi_data_register((*blocktab_handles)[iter], -20-iter, iter);
        /* starpu 1.1.x has quadratic algorithm to manage sequential consistency
         * this flag reduces the costs.
         */
        starpu_data_set_sequential_consistency_flag((*blocktab_handles)[iter], 0);
    }
    return NO_ERR;
}

int
starpu_unregister_blocktab( SolverMatrix          * datacode,
                             starpu_data_handle_t ** blocktab_handles,
                             int                  ** blocktab) {
    int iter;
    for (iter = 0; iter < SOLV_PROCNBR; iter++) {
        starpu_data_unregister((*blocktab_handles)[iter]);
    }
    memFree_null(*blocktab_handles);
    memFree_null(*blocktab);
    return NO_ERR;
}

int
starpu_unregister_work( SolverMatrix * datacode,
                         starpu_data_handle_t * WORK_handle ) {
    starpu_data_unregister(*WORK_handle);
    return NO_ERR;
}

int
starpu_register_work( Sopalin_Data_t * sopalin_data,
                      SolverMatrix * datacode,
                       starpu_data_handle_t * WORK_handle,
                       PASTIX_INT WORK_size ) {
    starpu_vector_data_register(WORK_handle, -1, (uintptr_t)NULL,
                                WORK_size,
                                sizeof(PASTIX_FLOAT));
    starpu_mpi_data_register(*WORK_handle, -3, SOLV_PROCNUM);
    return NO_ERR;
}

int
starpu_register_data( Sopalin_Data_t         * sopalin_data,
                       starpu_data_handle_t  ** L_handle,
                       starpu_data_handle_t  ** U_handle,
                       starpu_data_handle_t  ** Lhalo_handle,
                       starpu_data_handle_t  ** Uhalo_handle,
                       starpu_data_handle_t *** Lfanin_handle,
                       starpu_data_handle_t *** Ufanin_handle,
                       starpu_data_handle_t  ** blocktab_handles,
                       int                   ** blocktab,
                       starpu_data_handle_t   * WORK_handle,
                       PASTIX_INT             WORK_size )
{
    SolverMatrix       * datacode         = sopalin_data->datacode;
    starpu_register_cblk(sopalin_data, datacode, L_handle, U_handle);
    /* register halo */
    starpu_register_halo(sopalin_data, datacode,
                          Lhalo_handle,
                          Uhalo_handle);
    if (sopalin_data->sopar->iparm[IPARM_STARPU_FANIN] == API_YES) {
        /* register FANIN */
        starpu_register_fanin(sopalin_data, datacode,
                               Lfanin_handle,
                               Ufanin_handle);
    }
    starpu_register_blocktab(sopalin_data, blocktab_handles, blocktab);
    starpu_register_work(sopalin_data, datacode, WORK_handle, WORK_size);

    return NO_ERR;
}

int
starpu_unregister_data( Sopalin_Data_t         * sopalin_data,
                         starpu_data_handle_t  ** L_handle,
                         starpu_data_handle_t  ** U_handle,
                         starpu_data_handle_t  ** Lhalo_handle,
                         starpu_data_handle_t  ** Uhalo_handle,
                         starpu_data_handle_t *** Lfanin_handle,
                         starpu_data_handle_t *** Ufanin_handle,
                         starpu_data_handle_t  ** blocktab_handles,
                         int                   ** blocktab,
                         starpu_data_handle_t   * WORK_handle)
{
    SolverMatrix       * datacode         = sopalin_data->datacode;
    starpu_unregister_cblk(datacode, L_handle, U_handle);
    /* register halo */
    starpu_unregister_halo(datacode,
                            Lhalo_handle,
                            Uhalo_handle);
    if (sopalin_data->sopar->iparm[IPARM_STARPU_FANIN] == API_YES) {
        /* register FANIN */
        starpu_unregister_fanin(datacode,
                                 Lfanin_handle,
                                 Ufanin_handle);
    }
    starpu_unregister_blocktab(datacode, blocktab_handles, blocktab);
    starpu_unregister_work(datacode, WORK_handle);
    return NO_ERR;
}
