MODULE cuda_fortran
CONTAINS
    ATTRIBUTES(global) SUBROUTINE dia_hsb_kernel(surf, e3t_n, surf_ini, e3t_ini, &
        &  tsn, hc_loc_ini, sc_loc_ini, tmask, zwrkv, zwrkh, zwrks, zwrk, jpi, jpj, jpk)
        IMPLICIT NONE
    	REAL(kind=8), DIMENSION(:) :: surf, e3t_n, surf_ini, e3t_ini, tsn, &
                                    & hc_loc_ini, sc_loc_ini, tmask, zwrkv, &
                                    & zwrkh, zwrks, zwrk
        REAL(kind=8), SHARED       :: sdata(*)
        INTEGER, VALUE             :: jpi, jpj, jpk
        INTEGER                    :: i, si, ti, globsize
        
        globsize = jpi*jpj*jpk

        i  = blockDim%x * (blockIdx%x-1) + threadIdx%x
        ti = threadIdx%x
        si = MOD(i, jpk)+1 !Shared memory indexing

        !Compute volume (zwrk), volume variation (zwrkv), heat deviation (zwrkh)
        !and salinity deviation (zwrks)
        IF ( i .le. globsize ) THEN
            sdata(ti) = surf(si) * e3t_n(i)
            zwrkv(i) = ( sdata(ti) - surf_ini(si) * e3t_ini(i) ) * tmask(i) * surf(si)
            zwrkh(i) = ( sdata(ti) * tsn(i) - surf_ini(si) * hc_loc_ini(i) ) * tmask(i) * surf(si)
            zwrks(i) = ( sdata(ti) * tsn(globsize + i) - surf_ini(si) * sc_loc_ini(i) ) &
                    * tmask(i) * surf(si)
            zwrk(i) = sdata(ti) * tmask(i) * surf(si)
        END IF
    END SUBROUTINE dia_hsb_kernel

    !CUDA reduction (IMPORTANT: in the original code every value is multiplied by a mask not implemented on this version!!)
    ATTRIBUTES(global) SUBROUTINE cuda_reduction(invec, iin, sz)
        IMPLICIT NONE
        REAL(kind=8), DIMENSION(:) :: invec, iin                !Values and accumulated error
        REAL(kind=8), SHARED       :: sdata(512), idata(512)    !shared arrays
        INTEGER, VALUE             :: sz                        !array size
        INTEGER                    :: tid, i, stride, idx
        COMPLEX(kind=8)            :: stmp, atmp                !used for the Knuth's trick

        tid = threadIdx%x
        i = blockDim%x*2 * (blockIdx%x - 1) + threadIdx%x

        !Init shared memory and first addition
        IF (i .le. (sz-blockDim%x)) THEN    
            stmp = CMPLX(invec(i), iin(i), 8)
            atmp = CMPLX(invec(i+blockDim%x), iin(i+blockDim%x), 8)
            CALL DDPDD_d(stmp, atmp)
            sdata(tid) = REAL(atmp)
            idata(tid) = AIMAG(atmp)
        ELSE IF (i .le. sz) THEN
            sdata(tid) = invec(i)
            idata(tid) = iin(i)
        END IF
        CALL syncthreads()
        
        !Reduction per block
        stride = min(blockDim%x/2, sz)
        DO WHILE (stride > 0)
            IF( tid <= stride) THEN
                atmp = CMPLX(sdata(tid + stride), idata(tid+stride))
                CALL DDPDD_d(CMPLX(sdata(tid), idata(tid), 8), atmp)
                sdata(tid) = REAL(atmp)
                idata(tid) = AIMAG(atmp)
            END IF
            stride = rshift(stride, 1)
            CALL syncthreads()
        END DO
    
        !Write result and Knuth's trick error back to Global memory
        IF ( tid .eq. 1 ) THEN
            invec(blockIdx%x) = sdata(tid)
            iin(blockIdx%x) = idata(tid)
        END IF

    END SUBROUTINE cuda_reduction
    
    !Knuth's trick
    ATTRIBUTES(device) SUBROUTINE DDPDD_d(ydda, yddb)
        IMPLICIT NONE
        COMPLEX(kind=8), INTENT(in   ) :: ydda              !Scalar to add
        COMPLEX(kind=8), INTENT(inout) :: yddb              !Total sum
        REAL(kind=8)                   :: zerr, zt1, zt2

        zt1  = REAL(ydda) + REAL(yddb)
        zerr = zt1 - REAL(ydda)
        zt2  = ( (REAL(yddb) - zerr) + (REAL(ydda) - (zt1 - zerr)) ) &
              & + AIMAG(ydda) + AIMAG(yddb)
        !The result is t1 + t2 after normalization
        yddb = CMPLX( zt1 + zt2, zt2 - ((zt1 + zt2) - zt1), 8 )
    END SUBROUTINE DDPDD_d
    
END MODULE

program diahsb
    USE cuda_fortran
    USE cudafor
    IMPLICIT NONE 
    !Equals to a single MPI process size. Modificable.
    INTEGER, SAVE :: jpi=87, jpj=105, jpk=75, globsize=87*105*75
    INTEGER       :: jp_tem, jp_sal, localsize
    REAL(kind=8)  :: v_sumr, h_sumr, s_sumr, tv_sumr            !Store the results
    INTEGER       ::  ji, jj, jk, i, iterator, piterator        !Dummy variables

    !Data for the diagnostic
    REAL(kind=8), DIMENSION(:), ALLOCATABLE, PINNED :: surf, surf_ini, e3t_n, & 
                                                    & e3t_ini, tmask, hc_loc_ini, &
                                                    & sc_loc_ini, zwrkv, zwrkh, zwrks, &
                                                    & zwrk, tsn

    ! Device data
    REAL(kind=8), DEVICE, DIMENSION(:), ALLOCATABLE :: d_surf, d_surf_ini
    REAL(kind=8), DEVICE, DIMENSION(:), ALLOCATABLE :: d_e3t_n, d_e3t_ini, d_tmask, &
                                                     & d_hc_loc_ini, d_sc_loc_ini
    REAL(kind=8), DEVICE, DIMENSION(:), ALLOCATABLE :: d_zwrkv, d_zwrkh, d_zwrks, d_zwrk
    REAL(kind=8), DEVICE, DIMENSION(:), ALLOCATABLE :: d_tsn
    REAL(kind=8), DEVICE, DIMENSION(:), ALLOCATABLE :: d_iin, d_iinv, d_iinh, d_iins

    ! dim3 variables
    type(dim3) :: dimGrid, dimBlock

    !Allocate Host data
    allocate(surf(jpi*jpj), surf_ini(jpi*jpj))
    allocate(e3t_n(globsize), e3t_ini(globsize), tmask(globsize), hc_loc_ini(globsize), &
           & sc_loc_ini(globsize), zwrkv(globsize), zwrkh(globsize), zwrks(globsize), zwrk(globsize))
    allocate(tsn(globsize*2))

    ! Initialization with random values
    print *, "Initialize matrixs"
    call RANDOM_NUMBER(surf)
    call RANDOM_NUMBER(surf_ini)
    call RANDOM_NUMBER(e3t_n) 
    call RANDOM_NUMBER(e3t_ini)
    call RANDOM_NUMBER(tmask)
    call RANDOM_NUMBER(tsn)
    call RANDOM_NUMBER(hc_loc_ini)
    call RANDOM_NUMBER(sc_loc_ini)
    
    !Allocate GPU data
    allocate(d_surf(jpi*jpj), d_surf_ini(jpi*jpj))
    allocate(d_e3t_n(globsize), d_e3t_ini(globsize), d_tmask(globsize), &
           & d_hc_loc_ini(globsize), d_sc_loc_ini(globsize))
    allocate(d_tsn(globsize*2))
    allocate(d_iin(globsize), d_iinv(globsize), d_iinh(globsize), d_iins(globsize))

    !To store partial results
    allocate (d_zwrkv(globsize), d_zwrkh(globsize), d_zwrks(globsize), d_zwrk(globsize))

    !Send data to GPU
    d_surf       = surf
    d_surf_ini   = surf_ini
    d_e3t_n      = e3t_n
    d_e3t_ini    = e3t_ini
    d_tmask      = tmask
    d_hc_loc_ini = hc_loc_ini
    d_sc_loc_ini = sc_loc_ini
    d_tsn        = tsn
            
    !Create the grid and block dimensions
    dimBlock = dim3(512,1,1)
    dimGrid  = dim3(ceiling(real(globsize)/dimBlock%x), 1, 1) 
    !Call dia_hsb GPU kernel
    CALL dia_hsb_kernel<<<dimGrid, dimBlock, 8*dimBlock%x>>> &
        & (d_surf, d_e3t_n, d_surf_ini, d_e3t_ini, d_tsn, d_hc_loc_ini, &
        & d_sc_loc_ini, d_tmask, d_zwrkv, d_zwrkh, d_zwrks, d_zwrk, jpi, jpj, jpk)

    !GPU Rreduction
    print *, "Reduction:"
    iterator  = globsize
    piterator = iterator
    d_iin  = 0 !Used to store the imaginary part for the Knuth's trick
    d_iinv = 0
    d_iinh = 0
    d_iins = 0
    DO WHILE ( iterator .gt. 1 )
        print *, iterator, "Remaining elements"
        dimBlock = dim3(512, 1, 1)
        dimGrid  = dim3(ceiling(real(ceiling(real(iterator)/dimBlock%x))/2), 1, 1)
        CALL cuda_reduction<<< dimGrid, dimBlock >>>(d_zwrk, d_iin, globsize)
        CALL cuda_reduction<<< dimGrid, dimBlock >>>(d_zwrkv, d_iinv, globsize)
        CALL cuda_reduction<<< dimGrid, dimBlock >>>(d_zwrkh, d_iinh, globsize)
        CALL cuda_reduction<<< dimGrid, dimBlock >>>(d_zwrks, d_iins, globsize)
        piterator = iterator
        iterator = dimGrid%x
        d_zwrk ((iterator + 1):piterator) = 0
        d_zwrkv((iterator + 1):piterator) = 0
        d_zwrkh((iterator + 1):piterator) = 0
        d_zwrks((iterator + 1):piterator) = 0
    END DO
    i = cudaMemcpy(tv_sumr, d_zwrk, 1)
    i = cudaMemcpy(v_sumr, d_zwrkv, 1)
    i = cudaMemcpy(h_sumr, d_zwrkh, 1)
    i = cudaMemcpy(s_sumr, d_zwrks, 1)
 
    print *, "Total volume variation:",     v_sumr
    print *, "Total heat deviation:",       h_sumr   
    print *, "Total salinity deviation:",   s_sumr
    print *, "Total volume:",               tv_sumr

    !Free Host memory
    deallocate(surf, surf_ini)
    deallocate(e3t_n, e3t_ini, tmask, hc_loc_ini, sc_loc_ini)
    deallocate(tsn)

    !Free Device memory
    deallocate(d_surf, d_surf_ini, d_e3t_n, d_e3t_ini, &
             & d_tmask, d_hc_loc_ini, d_sc_loc_ini, d_tsn)
    deallocate(d_zwrkv, d_zwrkh, d_zwrks, d_zwrk)

    print *, "Exiting..."
END
