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

       @generated from sparse_hip/blas/magma_zmatrixtools_gpu.hip.cpp, normal z -> s, Tue Apr  7 19:00:07 2026

*/
#include "magmasparse_internal.h"

#define PRECISION_s

#define SWAP(a, b)  { tmp = a; a = b; b = tmp; }



__global__ void 
magma_svalinit_kernel(  
    const magma_int_t num_el, 
    magmaFloat_ptr dval) 
{
    int k = blockDim.x * blockIdx.x + threadIdx.x;
    float zero = MAGMA_S_MAKE(0.0, 0.0);
    if (k < num_el) {
        dval[k] = zero;
    }
}


/**
    Purpose
    -------
    
    Initializes a device array with zero. 

    Arguments
    ---------

    @param[in]
    num_el      magma_int_t
                size of array

    @param[in,out]
    dval        magmaFloat_ptr
                array to initialize
                
    @param[in]
    queue       magma_queue_t
                Queue to execute in.
                
    @ingroup magmasparse_sgegpuk
    ********************************************************************/

extern "C" magma_int_t
magma_svalinit_gpu(
    magma_int_t num_el,
    magmaFloat_ptr dval,
    magma_queue_t queue)
{
    int blocksize1 = 128;
    int blocksize2 = 1;

    int dimgrid1 = magma_ceildiv(num_el, blocksize1);
    int dimgrid2 = 1;
    int dimgrid3 = 1;
    dim3 grid(dimgrid1, dimgrid2, dimgrid3);
    dim3 block(blocksize1, blocksize2, 1);
    hipLaunchKernelGGL(magma_svalinit_kernel, dim3(grid), dim3(block), 0, queue->hip_stream() , num_el, dval);

    return MAGMA_SUCCESS;
}




__global__ void 
magma_sindexinit_kernel(  
    const magma_int_t num_el, 
    magmaIndex_ptr dind) 
{
    int k = blockDim.x * blockIdx.x + threadIdx.x;
    if (k < num_el) {
        dind[k] = 0;
    }
}


/**
    Purpose
    -------
    
    Initializes a device array with zero. 

    Arguments
    ---------

    @param[in]
    num_el      magma_int_t
                size of array

    @param[in,out]
    dind        magmaIndex_ptr
                array to initialize
                
    @param[in]
    queue       magma_queue_t
                Queue to execute in.
                
    @ingroup magmasparse_sgegpuk
    ********************************************************************/

extern "C" magma_int_t
magma_sindexinit_gpu(
    magma_int_t num_el,
    magmaIndex_ptr dind,
    magma_queue_t queue)
{
    int blocksize1 = 128;
    int blocksize2 = 1;

    int dimgrid1 = magma_ceildiv(num_el, blocksize1);
    int dimgrid2 = 1;
    int dimgrid3 = 1;
    dim3 grid(dimgrid1, dimgrid2, dimgrid3);
    dim3 block(blocksize1, blocksize2, 1);
    hipLaunchKernelGGL(magma_sindexinit_kernel, dim3(grid), dim3(block), 0, queue->hip_stream() , num_el, dind);

    return MAGMA_SUCCESS;
}


/***************************************************************************//**
    Purpose
    -------
    Generates a matrix $U = A \cup B$. If both matrices have a nonzero value
    in the same location, the value of A is used.
    
    This is the GPU version of the operation.

    Arguments
    ---------

    @param[in]
    A           magma_s_matrix
                Input matrix 1.

    @param[in]
    B           magma_s_matrix
                Input matrix 2.

    @param[out]
    U           magma_s_matrix*
                $U = A \cup B$. If both matrices have a nonzero value
                in the same location, the value of A is used.

    @param[in]
    queue       magma_queue_t
                Queue to execute in.

    @ingroup magmasparse_saux
*******************************************************************************/

extern "C" magma_int_t
magma_scsr_sort_gpu(
    magma_s_matrix *A,
    magma_queue_t queue)
{   
    magma_int_t info = 0;
    hipsparseHandle_t handle=NULL;
    hipsparseMatDescr_t descrA=NULL;
    
    magmaFloat_ptr tmp=NULL, csrVal_sorted=NULL;
    char* pBuffer;
    int *P;
    size_t pBufferSizeInBytes;
    
    CHECK_CUSPARSE( hipsparseCreate( &handle ));
    CHECK_CUSPARSE( hipsparseSetStream( handle, queue->hip_stream() ));
    CHECK_CUSPARSE( hipsparseCreateMatDescr( &descrA ));
    CHECK_CUSPARSE( hipsparseSetMatType( descrA, 
        HIPSPARSE_MATRIX_TYPE_GENERAL ));
    CHECK_CUSPARSE( hipsparseSetMatDiagType( descrA, 
        HIPSPARSE_DIAG_TYPE_NON_UNIT ));
    CHECK_CUSPARSE( hipsparseSetMatIndexBase( descrA, 
        HIPSPARSE_INDEX_BASE_ZERO ));
    
    CHECK(magma_smalloc(&csrVal_sorted, A->nnz));
   
    // step 1: allocate buffer
    hipsparseXcsrsort_bufferSizeExt(handle, A->num_rows, A->num_cols, 
        A->nnz, A->drow, A->dcol, &pBufferSizeInBytes);
    hipMalloc( &pBuffer, sizeof(char)* pBufferSizeInBytes);
    
    // step 2: setup permutation vector P to identity
    hipMalloc( (void**)&P, sizeof(int)*A->nnz);
    hipsparseCreateIdentityPermutation(handle, A->nnz, P);
    
    // step 3: sort CSR format
    hipsparseXcsrsort(handle, A->num_rows, A->num_cols, A->nnz, 
        descrA, A->drow, A->dcol, P, pBuffer);
    
    // step 4: gather sorted csrVal
#if CUDA_VERSION >= 12000
    cusparseSpVecDescr_t vec_permutation;
    cusparseDnVecDescr_t vec_values;
    CHECK_CUSPARSE( cusparseCreateSpVec(&vec_permutation, A->nnz, A->nnz,
                                        P, csrVal_sorted,
                                        CUSPARSE_INDEX_32I,
                                        HIPSPARSE_INDEX_BASE_ZERO, HIPBLAS_C_64F) );
    CHECK_CUSPARSE( cusparseCreateDnVec(&vec_values, A->nnz, A->dval, HIPBLAS_C_64F) );
    CHECK_CUSPARSE( cusparseGather(handle, vec_values, vec_permutation) );
    
    // destroy matrix/vector descriptors
    CHECK_CUSPARSE( cusparseDestroySpVec(vec_permutation) );
    CHECK_CUSPARSE( cusparseDestroyDnVec(vec_values) );
#else
    hipsparseSgthr(handle, A->nnz, (float*)A->dval, (float*)csrVal_sorted, P, 
        HIPSPARSE_INDEX_BASE_ZERO);
#endif

    SWAP(A->dval, csrVal_sorted);
    
cleanup:
    hipsparseDestroyMatDescr( descrA );
    hipsparseDestroy( handle );
    magma_free(csrVal_sorted);

    return info;
}
