#include "hip/hip_runtime.h"
/*
    -- MAGMA (version 2.6.2) --
       Univ. of Tennessee, Knoxville
       Univ. of California, Berkeley
       Univ. of Colorado, Denver
       @date March 2022

       @author Ahmad Abdelfattah
*/

#include <hip/hip_runtime.h>    // for CUDA_VERSION
#include "magma_internal.h"

#if defined(MAGMA_HAVE_HIP)
#include <hip/hip_fp16.h>
#endif

#define BLK_X 32
#define BLK_Y 4

/******************************************************************************/
__device__
void hlag2s_device(
    int m, int n,
    magmaHalf_const_ptr A, int lda,
    float             *SA, int ldsa )
{
#if CUDA_VERSION >= 7500 || defined(MAGMA_HAVE_HIP)
    const int gtx = blockIdx.x * BLK_X + threadIdx.x;
    const int gty = blockIdx.y * BLK_Y + threadIdx.y;

    for(int j = 0; j < n; j+= gridDim.y) {
        const int gty_ = gty + j;
        for(int i = 0; i < m; i+= gridDim.x) {
            const int gtx_ = gtx + i;
            if(gtx_ < m && gty_ < n) {
                SA[gty_ * ldsa + gtx_] = __half2float( A[gty_ * lda + gtx_] );
            }
        }
    }
#endif
}

/******************************************************************************/
__global__
void hlag2s_kernel(
        int m, int n,
        magmaHalf_const_ptr dA, int lda,
        float             *dSA, int ldsa )
{
    hlag2s_device(m, n, dA, lda, dSA, ldsa);
}

/******************************************************************************/
__global__
void hlag2s_kernel_batched(
        int m, int n,
        magmaHalf const * const * dAarray, int lda,
        float**                  dSAarray, int ldsa )
{
    const int batchid = blockIdx.z;
    hlag2s_device(m, n, dAarray[batchid], lda, dSAarray[batchid], ldsa);
}

/******************************************************************************/
extern "C" void
magmablas_hlag2s(
    magma_int_t m, magma_int_t n,
    magmaHalf_const_ptr dA, magma_int_t lda,
    float             *dSA, magma_int_t ldsa,
    magma_queue_t queue )
{
    magma_int_t info = 0;
    if ( m < 0 )
        info = -1;
    else if ( n < 0 )
        info = -2;
    else if ( lda < max(1,m) )
        info = -4;
    else if ( ldsa < max(1,m) )
        info = -6;

    if (info != 0) {
        magma_xerbla( __func__, -(info) );
        return;
    }

    /* quick return */
    if ( m == 0 || n == 0 ) {
        return;
    }

    const int max_gridy = 65000; // the kernel can work with any gridx/gridy dimension
    dim3 threads( BLK_X, BLK_Y );
    dim3 grid(magma_ceildiv( m, BLK_X ), min(max_gridy, magma_ceildiv(n, BLK_Y)), 1);
    hipLaunchKernelGGL(hlag2s_kernel, dim3(grid), dim3(threads), 0, queue->hip_stream() ,  m, n, dA, lda, dSA, ldsa );
}


/******************************************************************************/
extern "C" void
magmablas_hlag2s_batched(
    magma_int_t m, magma_int_t n,
    magmaHalf const * const * dAarray, magma_int_t lda,
    float               **dSAarray, magma_int_t ldsa,
    magma_int_t batchCount, magma_queue_t queue )
{
    magma_int_t info = 0;
    if ( m < 0 )
        info = -1;
    else if ( n < 0 )
        info = -2;
    else if ( lda < max(1,m) )
        info = -4;
    else if ( ldsa < max(1,m) )
        info = -6;

    if (info != 0) {
        magma_xerbla( __func__, -(info) );
        return;
    }

    /* quick return */
    if ( m == 0 || n == 0 ) {
        return;
    }

    dim3 threads( BLK_X, BLK_Y );
    const int maxBatch = queue->get_maxBatch();
    for(int i = 0; i < batchCount; i+=maxBatch){
        magma_int_t batch = min(maxBatch, batchCount-i);
        dim3 grid(magma_ceildiv( m, BLK_X ), magma_ceildiv( n, BLK_Y ), batch);
        hipLaunchKernelGGL(hlag2s_kernel_batched, dim3(grid), dim3(threads), 0, queue->hip_stream() ,  m, n, dAarray + i, lda, dSAarray + i, ldsa);
    }
}
