/*
 * -----------------------------------------------------------------
 * Programmer(s): Cody J. Balos @ LLNL
 * -----------------------------------------------------------------
 * SUNDIALS Copyright Start
 * Copyright (c) 2002-2024, Lawrence Livermore National Security
 * and Southern Methodist University.
 * All rights reserved.
 *
 * See the top-level LICENSE and NOTICE files for details.
 *
 * SPDX-License-Identifier: BSD-3-Clause
 * SUNDIALS Copyright End
 * -----------------------------------------------------------------
 * This is the header file is for the cuSPARSE implementation of the
 * SUNMATRIX module.
 * -----------------------------------------------------------------
 */

#include <stdio.h>
#include <stdlib.h>
#include <sunmatrix/sunmatrix_cusparse.h>
#include <sunmemory/sunmemory_cuda.h>

#include "cusparse_kernels.cuh"
#include "sundials/sundials_errors.h"
#include "sundials/sundials_math.h"
#include "sundials_cuda.h"
#include "sundials_cusparse.h"
#include "sundials_debug.h"

/* Use the namespace for the kernels */
using namespace sundials::cuda;
using namespace sundials::sunmatrix_cusparse;

/* Constants */
#define ZERO SUN_RCONST(0.0)
#define ONE  SUN_RCONST(1.0)

/* Private function prototypes */
static sunbooleantype SMCompatible_cuSparse(SUNMatrix, SUNMatrix);
static SUNMatrix SUNMatrix_cuSparse_NewEmpty(SUNContext sunctx);
#if CUDART_VERSION >= 11000
static cusparseStatus_t CreateSpMatDescr(SUNMatrix, cusparseSpMatDescr_t*);
#endif

#if CUDART_VERSION >= 12000
#define SPMV_ALG CUSPARSE_SPMV_CSR_ALG1
#else
#define SPMV_ALG CUSPARSE_MV_ALG_DEFAULT
#endif

/* Macros for handling the different function names based on precision */
#if defined(SUNDIALS_DOUBLE_PRECISION)
#define cusparseXcsrmv cusparseDcsrmv
#define CUDA_R_XF      CUDA_R_64F
#elif defined(SUNDIALS_SINGLE_PRECISION)
#define cusparseXcsrmv cusparseScsrmv
#define CUDA_R_XF      CUDA_R_32F
#endif

/* Content accessor macros */
#define SMCU_CONTENT(A)      ((SUNMatrix_Content_cuSparse)(A->content))
#define SMCU_ROWS(A)         (SMCU_CONTENT(A)->M)
#define SMCU_COLUMNS(A)      (SMCU_CONTENT(A)->N)
#define SMCU_NNZ(A)          (SMCU_CONTENT(A)->NNZ)
#define SMCU_NBLOCKS(A)      (SMCU_CONTENT(A)->nblocks)
#define SMCU_BLOCKROWS(A)    (SMCU_CONTENT(A)->blockrows)
#define SMCU_BLOCKCOLS(A)    (SMCU_CONTENT(A)->blockcols)
#define SMCU_BLOCKNNZ(A)     (SMCU_CONTENT(A)->blocknnz)
#define SMCU_NP(A)           (SMCU_CONTENT(A)->NP)
#define SMCU_SPARSETYPE(A)   (SMCU_CONTENT(A)->sparse_type)
#define SMCU_OWNMATD(A)      (SMCU_CONTENT(A)->own_matd)
#define SMCU_DATA(A)         (SMCU_CONTENT(A)->data)
#define SMCU_DATAp(A)        ((sunrealtype*)SMCU_CONTENT(A)->data->ptr)
#define SMCU_INDEXVALS(A)    (SMCU_CONTENT(A)->colind)
#define SMCU_INDEXPTRS(A)    (SMCU_CONTENT(A)->rowptrs)
#define SMCU_INDEXVALSp(A)   ((int*)SMCU_CONTENT(A)->colind->ptr)
#define SMCU_INDEXPTRSp(A)   ((int*)SMCU_CONTENT(A)->rowptrs->ptr)
#define SMCU_MEMHELP(A)      (SMCU_CONTENT(A)->mem_helper)
#define SMCU_MATDESCR(A)     (SMCU_CONTENT(A)->mat_descr)
#define SMCU_CUSPHANDLE(A)   (SMCU_CONTENT(A)->cusp_handle)
#define SMCU_FIXEDPATTERN(A) (SMCU_CONTENT(A)->fixed_pattern)
#define SMCU_EXECPOLICY(A)   (SMCU_CONTENT(A)->exec_policy)

/* ------------------------------------------------------------------
 * Default execution policy definition.
 *
 * This policy tries to help us leverage the structure of the matrix.
 * It will choose block sizes which are a multiple of the warp size,
 * and it will choose a grid size to such that all work elements are
 * covered.
 * ------------------------------------------------------------------ */

class SUNCuSparseMatrixExecPolicy : public ExecPolicy
{
public:
  SUNCuSparseMatrixExecPolicy(const cudaStream_t stream = 0)
    : ExecPolicy(stream)
  {}

  SUNCuSparseMatrixExecPolicy(const SUNCuSparseMatrixExecPolicy& ex)
    : ExecPolicy(ex.stream_)
  {}

  virtual size_t gridSize(size_t numWorkElements, size_t blockDim = 0) const
  {
    return (numWorkElements + blockDim - 1) / blockDim;
  }

  virtual size_t blockSize(size_t numWorkElements = 0, size_t gridDim = 0) const
  {
    return (
      max_block_size(WARP_SIZE * (numWorkElements + WARP_SIZE - 1) / WARP_SIZE));
  }

  virtual const cudaStream_t* stream() const { return (&stream_); }

  virtual ExecPolicy* clone() const
  {
    return (static_cast<ExecPolicy*>(new SUNCuSparseMatrixExecPolicy(*this)));
  }

  static size_t max_block_size(int val)
  {
    return ((val > MAX_BLOCK_SIZE) ? MAX_BLOCK_SIZE : val);
  }
};

SUNCuSparseMatrixExecPolicy DEFAULT_EXEC_POLICY;

/* ------------------------------------------------------------------
 * Constructors.
 * ------------------------------------------------------------------ */

SUNMatrix SUNMatrix_cuSparse_NewCSR(int M, int N, int NNZ,
                                    cusparseHandle_t cusp, SUNContext sunctx)
{
  SUNMemory d_colind, d_rowptr, d_values;
  int alloc_fail = 0;

  /* return with NULL matrix on illegal input */
  if ((M <= 0) || (N <= 0) || (NNZ < 0))
  {
    SUNDIALS_DEBUG_PRINT("ERROR in SUNMatrix_NewCSR_cuSparse: illegal value(s) "
                         "for M, N, or NNZ\n");
    return (NULL);
  }

  SUNMatrix A = SUNMatrix_cuSparse_NewEmpty(sunctx);
  if (A == NULL)
  {
    SUNDIALS_DEBUG_PRINT("ERROR in SUNMatrix_NewCSR_cuSparse: "
                         "SUNMatrix_cuSparse_NewEmpty returned NULL\n");
    return (NULL);
  }

  SMCU_MEMHELP(A) = SUNMemoryHelper_Cuda(sunctx);
  if (SMCU_MEMHELP(A) == NULL)
  {
    SUNDIALS_DEBUG_PRINT("ERROR in SUNMatrix_NewCSR_cuSparse: "
                         "SUNMemoryHelper_Cuda returned NULL\n");
    SUNMatDestroy(A);
    return (NULL);
  }

  /* Allocate device memory for the matrix */
  alloc_fail += SUNMemoryHelper_Alloc(SMCU_MEMHELP(A), &d_colind,
                                      sizeof(int) * NNZ, SUNMEMTYPE_DEVICE,
                                      nullptr);
  alloc_fail += SUNMemoryHelper_Alloc(SMCU_MEMHELP(A), &d_rowptr,
                                      sizeof(int) * (M + 1), SUNMEMTYPE_DEVICE,
                                      nullptr);
  alloc_fail += SUNMemoryHelper_Alloc(SMCU_MEMHELP(A), &d_values,
                                      sizeof(sunrealtype) * NNZ,
                                      SUNMEMTYPE_DEVICE, nullptr);
  if (alloc_fail)
  {
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(A), d_colind, nullptr);
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(A), d_rowptr, nullptr);
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(A), d_values, nullptr);
    SUNMatDestroy(A);
    return (NULL);
  }

  /* Choose sensible defaults */
  cusparseStatus_t cusparse_status = CUSPARSE_STATUS_SUCCESS;
  cusparseMatDescr_t mat_descr;
  cusparse_status = cusparseCreateMatDescr(&mat_descr);
  if (!SUNDIALS_CUSPARSE_VERIFY(cusparse_status))
  {
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(A), d_colind, nullptr);
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(A), d_rowptr, nullptr);
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(A), d_values, nullptr);
    SUNMatDestroy(A);
    return (NULL);
  }

  cusparse_status = cusparseSetMatType(mat_descr, CUSPARSE_MATRIX_TYPE_GENERAL);
  if (!SUNDIALS_CUSPARSE_VERIFY(cusparse_status))
  {
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(A), d_colind, nullptr);
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(A), d_rowptr, nullptr);
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(A), d_values, nullptr);
    cusparseDestroyMatDescr(mat_descr);
    SUNMatDestroy(A);
    return (NULL);
  }

  cusparse_status = cusparseSetMatIndexBase(mat_descr, CUSPARSE_INDEX_BASE_ZERO);
  if (!SUNDIALS_CUSPARSE_VERIFY(cusparse_status))
  {
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(A), d_colind, nullptr);
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(A), d_rowptr, nullptr);
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(A), d_values, nullptr);
    cusparseDestroyMatDescr(mat_descr);
    SUNMatDestroy(A);
    return (NULL);
  }

  cudaStream_t stream;
  if (!SUNDIALS_CUSPARSE_VERIFY(cusparseGetStream(cusp, &stream)))
  {
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(A), d_colind, nullptr);
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(A), d_rowptr, nullptr);
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(A), d_values, nullptr);
    cusparseDestroyMatDescr(mat_descr);
    SUNMatDestroy(A);
    return (NULL);
  }

  /* Fill the content */
  SMCU_CONTENT(A)->M              = M;
  SMCU_CONTENT(A)->N              = N;
  SMCU_CONTENT(A)->NNZ            = NNZ;
  SMCU_CONTENT(A)->nblocks        = 1;
  SMCU_CONTENT(A)->blockrows      = M;
  SMCU_CONTENT(A)->blockcols      = N;
  SMCU_CONTENT(A)->blocknnz       = NNZ;
  SMCU_CONTENT(A)->own_matd       = SUNTRUE;
  SMCU_CONTENT(A)->matvec_issetup = SUNFALSE;
  SMCU_CONTENT(A)->fixed_pattern  = SUNFALSE;
  SMCU_CONTENT(A)->sparse_type    = SUNMAT_CUSPARSE_CSR;
  SMCU_CONTENT(A)->colind         = d_colind;
  SMCU_CONTENT(A)->rowptrs        = d_rowptr;
  SMCU_CONTENT(A)->data           = d_values;
  SMCU_CONTENT(A)->mat_descr      = mat_descr;
  SMCU_CONTENT(A)->cusp_handle    = cusp;
  SMCU_CONTENT(A)->exec_policy = DEFAULT_EXEC_POLICY.clone_new_stream(stream);

#if CUDART_VERSION >= 11000
  cusparseSpMatDescr_t spmat_descr;
  if (!SUNDIALS_CUSPARSE_VERIFY(CreateSpMatDescr(A, &spmat_descr)))
  {
    SUNMatDestroy(A);
    return (NULL);
  }
  SMCU_CONTENT(A)->spmat_descr = spmat_descr;
  SMCU_CONTENT(A)->dBufferMem  = NULL;
  SMCU_CONTENT(A)->bufferSize  = 0;
  SMCU_CONTENT(A)->vecX        = NULL;
  SMCU_CONTENT(A)->vecY        = NULL;
#endif

  return A;
}

SUNMatrix SUNMatrix_cuSparse_MakeCSR(cusparseMatDescr_t mat_descr, int M, int N,
                                     int NNZ, int* rowptrs, int* colind,
                                     sunrealtype* data, cusparseHandle_t cusp,
                                     SUNContext sunctx)
{
  /* return with NULL matrix on illegal input */
  if ((M <= 0) || (N <= 0) || (NNZ < 0))
  {
    SUNDIALS_DEBUG_PRINT("ERROR in SUNMatrix_MakeCSR_cuSparse: illegal "
                         "value(s) for M, N, or NNZ\n");
    return (NULL);
  }

  if ((rowptrs == NULL) || (colind == NULL) || (data == NULL))
  {
    SUNDIALS_DEBUG_PRINT("ERROR in SUNMatrix_MakeCSR_cuSparse: rowptrs, "
                         "colind, or data is NULL\n");
    return (NULL);
  }

  if (cusparseGetMatIndexBase(mat_descr) != CUSPARSE_INDEX_BASE_ZERO)
  {
    SUNDIALS_DEBUG_PRINT(
      "ERROR in SUNMatrix_MakeCSR_cuSparse: the cusparseMatDescr_t must have "
      "index base CUSPARSE_INDEX_BASE_ZERO\n");
    return (NULL);
  }

  SUNMatrix A = SUNMatrix_cuSparse_NewEmpty(sunctx);
  if (A == NULL)
  {
    SUNDIALS_DEBUG_PRINT("ERROR in SUNMatrix_MakeCSR_cuSparse: "
                         "SUNMatrix_cuSparse_NewEmpty returned NULL\n");
    return (NULL);
  }

  SMCU_MEMHELP(A) = SUNMemoryHelper_Cuda(sunctx);
  if (SMCU_MEMHELP(A) == NULL)
  {
    SUNDIALS_DEBUG_PRINT("ERROR in SUNMatrix_NewCSR_cuSparse: "
                         "SUNMemoryHelper_Cuda returned NULL\n");
    SUNMatDestroy(A);
    return (NULL);
  }

  cudaStream_t stream;
  if (!SUNDIALS_CUSPARSE_VERIFY(cusparseGetStream(cusp, &stream)))
  {
    SUNMatDestroy(A);
    return (NULL);
  }

  /* Fill content */
  SMCU_CONTENT(A)->M              = M;
  SMCU_CONTENT(A)->N              = N;
  SMCU_CONTENT(A)->NNZ            = NNZ;
  SMCU_CONTENT(A)->nblocks        = 1;
  SMCU_CONTENT(A)->blockrows      = M;
  SMCU_CONTENT(A)->blockcols      = N;
  SMCU_CONTENT(A)->blocknnz       = NNZ;
  SMCU_CONTENT(A)->own_matd       = SUNFALSE;
  SMCU_CONTENT(A)->matvec_issetup = SUNFALSE;
  SMCU_CONTENT(A)->fixed_pattern  = SUNFALSE;
  SMCU_CONTENT(A)->sparse_type    = SUNMAT_CUSPARSE_CSR;
  SMCU_CONTENT(A)->colind      = SUNMemoryHelper_Wrap(SMCU_MEMHELP(A), colind,
                                                      SUNMEMTYPE_DEVICE);
  SMCU_CONTENT(A)->rowptrs     = SUNMemoryHelper_Wrap(SMCU_MEMHELP(A), rowptrs,
                                                      SUNMEMTYPE_DEVICE);
  SMCU_CONTENT(A)->data        = SUNMemoryHelper_Wrap(SMCU_MEMHELP(A), data,
                                                      SUNMEMTYPE_DEVICE);
  SMCU_CONTENT(A)->mat_descr   = mat_descr;
  SMCU_CONTENT(A)->cusp_handle = cusp;

  SMCU_CONTENT(A)->exec_policy = DEFAULT_EXEC_POLICY.clone_new_stream(stream);

  if (SMCU_CONTENT(A)->colind == NULL || SMCU_CONTENT(A)->rowptrs == NULL ||
      SMCU_CONTENT(A)->data == NULL)
  {
    SUNMatDestroy(A);
    return (NULL);
  }

#if CUDART_VERSION >= 11000
  cusparseSpMatDescr_t spmat_descr;
  if (!SUNDIALS_CUSPARSE_VERIFY(CreateSpMatDescr(A, &spmat_descr)))
  {
    SUNMatDestroy(A);
    return (NULL);
  }
  SMCU_CONTENT(A)->spmat_descr = spmat_descr;
  SMCU_CONTENT(A)->dBufferMem  = NULL;
  SMCU_CONTENT(A)->bufferSize  = 0;
  SMCU_CONTENT(A)->vecX        = NULL;
  SMCU_CONTENT(A)->vecY        = NULL;
#endif

  return (A);
}

SUNMatrix SUNMatrix_cuSparse_NewBlockCSR(int nblocks, int blockrows,
                                         int blockcols, int blocknnz,
                                         cusparseHandle_t cusp, SUNContext sunctx)
{
  SUNMemory d_colind, d_rowptr, d_values;
  int M, N, NNZ;
  int alloc_fail = 0;

  /* Return with NULL matrix on illegal input */
  if (blockrows != blockcols)
  {
    SUNDIALS_DEBUG_PRINT("ERROR in SUNMatrix_cuSparse_NewBlockCSR: matrix must "
                         "be square for the BCSR format\n");
    return (NULL);
  }

  M   = nblocks * blockrows;
  N   = M;
  NNZ = nblocks * blocknnz;

  /* Return with NULL matrix on illegal input */
  if ((M <= 0) || (N <= 0) || (NNZ < 0))
  {
    SUNDIALS_DEBUG_PRINT("ERROR in SUNMatrix_cuSparse_NewBlockCSR: illegal "
                         "value(s) for M, N, or NNZ\n");
    return (NULL);
  }

  /* Allocate the SUNMatrix object */
  SUNMatrix A = SUNMatrix_cuSparse_NewEmpty(sunctx);
  if (A == NULL)
  {
    SUNDIALS_DEBUG_PRINT("ERROR in SUNMatrix_cuSparse_NewBlockCSR: "
                         "SUNMatrix_cuSparse_NewEmpty returned NULL\n");
    return (NULL);
  }

  SMCU_MEMHELP(A) = SUNMemoryHelper_Cuda(sunctx);
  if (SMCU_MEMHELP(A) == NULL)
  {
    SUNDIALS_DEBUG_PRINT("ERROR in SUNMatrix_NewCSR_cuSparse: "
                         "SUNMemoryHelper_Cuda returned NULL\n");
    SUNMatDestroy(A);
    return (NULL);
  }

  /* Allocate device memory for the matrix */
  alloc_fail += SUNMemoryHelper_Alloc(SMCU_MEMHELP(A), &d_colind,
                                      sizeof(int) * blocknnz, SUNMEMTYPE_DEVICE,
                                      nullptr);
  alloc_fail += SUNMemoryHelper_Alloc(SMCU_MEMHELP(A), &d_rowptr,
                                      sizeof(int) * (blockrows + 1),
                                      SUNMEMTYPE_DEVICE, nullptr);
  alloc_fail += SUNMemoryHelper_Alloc(SMCU_MEMHELP(A), &d_values,
                                      sizeof(sunrealtype) * blocknnz * nblocks,
                                      SUNMEMTYPE_DEVICE, nullptr);
  if (alloc_fail)
  {
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(A), d_colind, nullptr);
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(A), d_rowptr, nullptr);
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(A), d_values, nullptr);
    SUNMatDestroy(A);
    return (NULL);
  }

  /* Choose sensible defaults */
  cusparseStatus_t cusparse_status = CUSPARSE_STATUS_SUCCESS;
  cusparseMatDescr_t mat_descr;
  cusparse_status = cusparseCreateMatDescr(&mat_descr);
  if (!SUNDIALS_CUSPARSE_VERIFY(cusparse_status))
  {
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(A), d_colind, nullptr);
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(A), d_rowptr, nullptr);
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(A), d_values, nullptr);
    SUNMatDestroy(A);
    return (NULL);
  }

  cusparse_status = cusparseSetMatType(mat_descr, CUSPARSE_MATRIX_TYPE_GENERAL);
  if (!SUNDIALS_CUSPARSE_VERIFY(cusparse_status))
  {
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(A), d_colind, nullptr);
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(A), d_rowptr, nullptr);
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(A), d_values, nullptr);
    cusparseDestroyMatDescr(mat_descr);
    SUNMatDestroy(A);
    return (NULL);
  }

  cusparse_status = cusparseSetMatIndexBase(mat_descr, CUSPARSE_INDEX_BASE_ZERO);
  if (!SUNDIALS_CUSPARSE_VERIFY(cusparse_status))
  {
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(A), d_colind, nullptr);
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(A), d_rowptr, nullptr);
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(A), d_values, nullptr);
    cusparseDestroyMatDescr(mat_descr);
    SUNMatDestroy(A);
    return (NULL);
  }

  cudaStream_t stream;
  if (!SUNDIALS_CUSPARSE_VERIFY(cusparseGetStream(cusp, &stream)))
  {
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(A), d_colind, nullptr);
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(A), d_rowptr, nullptr);
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(A), d_values, nullptr);
    cusparseDestroyMatDescr(mat_descr);
    SUNMatDestroy(A);
    return (NULL);
  }

  /* Fill the content */
  SMCU_CONTENT(A)->M              = M;
  SMCU_CONTENT(A)->N              = N;
  SMCU_CONTENT(A)->NNZ            = NNZ;
  SMCU_CONTENT(A)->nblocks        = nblocks;
  SMCU_CONTENT(A)->blockrows      = blockrows;
  SMCU_CONTENT(A)->blockcols      = blockrows;
  SMCU_CONTENT(A)->blocknnz       = blocknnz;
  SMCU_CONTENT(A)->own_matd       = SUNTRUE;
  SMCU_CONTENT(A)->matvec_issetup = SUNFALSE;
  SMCU_CONTENT(A)->cusp_handle    = cusp;
  SMCU_CONTENT(A)->fixed_pattern  = SUNFALSE;
  SMCU_CONTENT(A)->sparse_type    = SUNMAT_CUSPARSE_BCSR;
  SMCU_CONTENT(A)->colind         = d_colind;
  SMCU_CONTENT(A)->rowptrs        = d_rowptr;
  SMCU_CONTENT(A)->data           = d_values;
  SMCU_CONTENT(A)->mat_descr      = mat_descr;
  SMCU_CONTENT(A)->exec_policy = DEFAULT_EXEC_POLICY.clone_new_stream(stream);

#if CUDART_VERSION >= 11000
  cusparseSpMatDescr_t spmat_descr;
  if (!SUNDIALS_CUSPARSE_VERIFY(CreateSpMatDescr(A, &spmat_descr)))
  {
    SUNMatDestroy(A);
    return (NULL);
  }
  SMCU_CONTENT(A)->spmat_descr = spmat_descr;
  SMCU_CONTENT(A)->dBufferMem  = NULL;
  SMCU_CONTENT(A)->bufferSize  = 0;
  SMCU_CONTENT(A)->vecX        = NULL;
  SMCU_CONTENT(A)->vecY        = NULL;
#endif

  return (A);
}

/* ------------------------------------------------------------------
 * Implementation specific routines.
 * ------------------------------------------------------------------ */

int SUNMatrix_cuSparse_SparseType(SUNMatrix A) { return SMCU_SPARSETYPE(A); }

int SUNMatrix_cuSparse_Rows(SUNMatrix A) { return SMCU_ROWS(A); }

int SUNMatrix_cuSparse_Columns(SUNMatrix A) { return SMCU_COLUMNS(A); }

int SUNMatrix_cuSparse_NNZ(SUNMatrix A) { return SMCU_NNZ(A); }

int* SUNMatrix_cuSparse_IndexPointers(SUNMatrix A)
{
  if (SUNMatGetID(A) == SUNMATRIX_CUSPARSE) { return (SMCU_INDEXPTRSp(A)); }
  else { return (NULL); }
}

int* SUNMatrix_cuSparse_IndexValues(SUNMatrix A)
{
  if (SUNMatGetID(A) == SUNMATRIX_CUSPARSE) { return (SMCU_INDEXVALSp(A)); }
  else { return (NULL); }
}

sunrealtype* SUNMatrix_cuSparse_Data(SUNMatrix A)
{
  if (SUNMatGetID(A) == SUNMATRIX_CUSPARSE) { return (SMCU_DATAp(A)); }
  else { return (NULL); }
}

int SUNMatrix_cuSparse_NumBlocks(SUNMatrix A) { return SMCU_NBLOCKS(A); }

int SUNMatrix_cuSparse_BlockRows(SUNMatrix A) { return SMCU_BLOCKROWS(A); }

int SUNMatrix_cuSparse_BlockColumns(SUNMatrix A) { return SMCU_BLOCKCOLS(A); }

int SUNMatrix_cuSparse_BlockNNZ(SUNMatrix A) { return SMCU_BLOCKNNZ(A); }

sunrealtype* SUNMatrix_cuSparse_BlockData(SUNMatrix A, int blockidx)
{
  sunrealtype* matdata;
  int offset;

  if (SUNMatGetID(A) != SUNMATRIX_CUSPARSE) { return (NULL); }

  if (blockidx >= SMCU_NBLOCKS(A)) { return (NULL); }

  matdata = SMCU_DATAp(A);
  offset  = SMCU_BLOCKNNZ(A) * blockidx;

  return (&matdata[offset]);
}

cusparseMatDescr_t SUNMatrix_cuSparse_MatDescr(SUNMatrix A)
{
  if (SUNMatGetID(A) == SUNMATRIX_CUSPARSE) { return (SMCU_MATDESCR(A)); }
  else { return (NULL); }
}

SUNErrCode SUNMatrix_cuSparse_SetFixedPattern(SUNMatrix A, sunbooleantype yesno)
{
  if (SUNMatGetID(A) != SUNMATRIX_CUSPARSE) { return (SUN_ERR_ARG_WRONGTYPE); }

  SMCU_FIXEDPATTERN(A) = yesno;

  return SUN_SUCCESS;
}

SUNErrCode SUNMatrix_cuSparse_SetKernelExecPolicy(SUNMatrix A,
                                                  SUNCudaExecPolicy* exec_policy)
{
  if (SUNMatGetID(A) != SUNMATRIX_CUSPARSE) { return (SUN_ERR_ARG_WRONGTYPE); }

  /* Reset to the default policy if the new one is NULL */
  delete SMCU_EXECPOLICY(A);
  if (exec_policy) { SMCU_EXECPOLICY(A) = exec_policy->clone(); }
  else
  {
    SMCU_EXECPOLICY(A) =
      DEFAULT_EXEC_POLICY.clone_new_stream(*SMCU_EXECPOLICY(A)->stream());
  }

  return SUN_SUCCESS;
}

SUNErrCode SUNMatrix_cuSparse_CopyToDevice(SUNMatrix dA, sunrealtype* h_data,
                                           int* h_idxptrs, int* h_idxvals)
{
  int retval;
  SUNMemory _h_data, _h_idxptrs, _h_idxvals;
  const cudaStream_t* stream;
  int nidxvals, nidxptrs;

  if (SUNMatGetID(dA) != SUNMATRIX_CUSPARSE) { return (SUN_ERR_ARG_WRONGTYPE); }

  stream = SMCU_EXECPOLICY(dA)->stream();

  if (h_data != NULL)
  {
    _h_data = SUNMemoryHelper_Wrap(SMCU_MEMHELP(dA), h_data, SUNMEMTYPE_HOST);
    retval = SUNMemoryHelper_CopyAsync(SMCU_MEMHELP(dA), SMCU_DATA(dA), _h_data,
                                       SMCU_NNZ(dA) * sizeof(sunrealtype),
                                       (void*)stream);
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(dA), _h_data, nullptr);
    if (retval != 0) { return (SUN_ERR_OP_FAIL); }
  }

  switch (SMCU_SPARSETYPE(dA))
  {
  case SUNMAT_CUSPARSE_CSR:
    nidxptrs = SMCU_ROWS(dA) + 1;
    nidxvals = SMCU_NNZ(dA);
    break;
  case SUNMAT_CUSPARSE_BCSR:
    nidxptrs = SMCU_BLOCKROWS(dA) + 1;
    nidxvals = SMCU_BLOCKNNZ(dA);
    break;
  default:
    SUNDIALS_DEBUG_PRINT(
      "ERROR in SUNMatrix_cuSparse_CopyToDevice: unrecognized sparse type\n");
    return (SUN_ERR_ARG_INCOMPATIBLE);
  }

  if (h_idxptrs != NULL)
  {
    _h_idxptrs = SUNMemoryHelper_Wrap(SMCU_MEMHELP(dA), h_idxptrs,
                                      SUNMEMTYPE_HOST);
    retval     = SUNMemoryHelper_CopyAsync(SMCU_MEMHELP(dA), SMCU_INDEXPTRS(dA),
                                           _h_idxptrs, nidxptrs * sizeof(int),
                                           (void*)stream);
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(dA), _h_idxptrs, nullptr);
    if (retval != 0) { return (SUN_ERR_OP_FAIL); }
  }

  if (h_idxvals != NULL)
  {
    _h_idxvals = SUNMemoryHelper_Wrap(SMCU_MEMHELP(dA), h_idxvals,
                                      SUNMEMTYPE_HOST);
    retval     = SUNMemoryHelper_CopyAsync(SMCU_MEMHELP(dA), SMCU_INDEXVALS(dA),
                                           _h_idxvals, nidxvals * sizeof(int),
                                           (void*)stream);
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(dA), _h_idxvals, nullptr);
    if (retval != 0) { return (SUN_ERR_OP_FAIL); }
  }

  return SUN_SUCCESS;
}

SUNErrCode SUNMatrix_cuSparse_CopyFromDevice(SUNMatrix dA, sunrealtype* h_data,
                                             int* h_idxptrs, int* h_idxvals)
{
  int retval;
  SUNMemory _h_data, _h_idxptrs, _h_idxvals;
  const cudaStream_t* stream;
  int nidxvals, nidxptrs;

  if (SUNMatGetID(dA) != SUNMATRIX_CUSPARSE) { return (SUN_ERR_ARG_WRONGTYPE); }

  stream = SMCU_EXECPOLICY(dA)->stream();

  if (h_data != NULL)
  {
    _h_data = SUNMemoryHelper_Wrap(SMCU_MEMHELP(dA), h_data, SUNMEMTYPE_HOST);
    retval = SUNMemoryHelper_CopyAsync(SMCU_MEMHELP(dA), _h_data, SMCU_DATA(dA),
                                       SMCU_NNZ(dA) * sizeof(sunrealtype),
                                       (void*)stream);
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(dA), _h_data, nullptr);
    if (retval != 0) { return (SUN_ERR_OP_FAIL); }
  }

  switch (SMCU_SPARSETYPE(dA))
  {
  case SUNMAT_CUSPARSE_CSR:
    nidxptrs = SMCU_ROWS(dA) + 1;
    nidxvals = SMCU_NNZ(dA);
  case SUNMAT_CUSPARSE_BCSR:
    nidxptrs = SMCU_BLOCKROWS(dA) + 1;
    nidxvals = SMCU_BLOCKNNZ(dA);
  }

  if (h_idxptrs != NULL)
  {
    _h_idxptrs = SUNMemoryHelper_Wrap(SMCU_MEMHELP(dA), h_idxptrs,
                                      SUNMEMTYPE_HOST);
    retval     = SUNMemoryHelper_CopyAsync(SMCU_MEMHELP(dA), _h_idxptrs,
                                           SMCU_INDEXPTRS(dA),
                                           nidxptrs * sizeof(int), (void*)stream);
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(dA), _h_idxptrs, nullptr);
    if (retval != 0) { return (SUN_ERR_OP_FAIL); }
  }

  if (h_idxvals != NULL)
  {
    _h_idxvals = SUNMemoryHelper_Wrap(SMCU_MEMHELP(dA), h_idxvals,
                                      SUNMEMTYPE_HOST);
    retval     = SUNMemoryHelper_CopyAsync(SMCU_MEMHELP(dA), _h_idxvals,
                                           SMCU_INDEXVALS(dA),
                                           nidxvals * sizeof(int), (void*)stream);
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(dA), _h_idxvals, nullptr);
    if (retval != 0) { return (SUN_ERR_OP_FAIL); }
  }

  return SUN_SUCCESS;
}

/*
 * -----------------------------------------------------------------
 * implementation of matrix operations
 * -----------------------------------------------------------------
 */

SUNMatrix_ID SUNMatGetID_cuSparse(SUNMatrix A) { return (SUNMATRIX_CUSPARSE); }

/* Returns a new matrix allocated to have the same structure as A,
   but it does not copy any nonzeros, column vals, or row pointers. */
SUNMatrix SUNMatClone_cuSparse(SUNMatrix A)
{
  SUNMatrix B;

  switch (SMCU_SPARSETYPE(A))
  {
  case SUNMAT_CUSPARSE_CSR:
    B = SUNMatrix_cuSparse_NewCSR(SMCU_ROWS(A), SMCU_COLUMNS(A), SMCU_NNZ(A),
                                  SMCU_CUSPHANDLE(A), A->sunctx);
    break;
  case SUNMAT_CUSPARSE_BCSR:
    B = SUNMatrix_cuSparse_NewBlockCSR(SMCU_NBLOCKS(A), SMCU_BLOCKROWS(A),
                                       SMCU_BLOCKCOLS(A), SMCU_BLOCKNNZ(A),
                                       SMCU_CUSPHANDLE(A), A->sunctx);
    break;
  default:
    SUNDIALS_DEBUG_PRINT(
      "ERROR in SUNMatClone_cuSparse: sparse type not recognized\n");
    B = NULL;
  }

  SMCU_FIXEDPATTERN(B) = SMCU_FIXEDPATTERN(A);
  delete SMCU_EXECPOLICY(B);
  SMCU_EXECPOLICY(B) = SMCU_EXECPOLICY(A)->clone();

  return (B);
}

/* Deallocates the SUNMatrix object and all data it owns */
void SUNMatDestroy_cuSparse(SUNMatrix A)
{
  if (A == NULL) { return; }

  /* free content */
  if (A->content != NULL)
  {
    if (SMCU_MEMHELP(A))
    {
      SUNMemoryHelper_Dealloc(SMCU_MEMHELP(A), SMCU_DATA(A), nullptr);
      SUNMemoryHelper_Dealloc(SMCU_MEMHELP(A), SMCU_INDEXPTRS(A), nullptr);
      SUNMemoryHelper_Dealloc(SMCU_MEMHELP(A), SMCU_INDEXVALS(A), nullptr);
    }
    else
    {
      SUNDIALS_DEBUG_PRINT(
        "WARNING in SUNMatDestroy_cuSparse: mem_helper was NULL when trying "
        "to "
        "dealloc data, this could result in a memory leak\n");
    }

    if (SMCU_OWNMATD(A))
    {
      /* free cusparseMatDescr_t */
      SUNDIALS_CUSPARSE_VERIFY(cusparseDestroyMatDescr(SMCU_MATDESCR(A)));
    }

#if CUDART_VERSION >= 11000
    SUNDIALS_CUSPARSE_VERIFY(cusparseDestroyDnVec(SMCU_CONTENT(A)->vecX));
    SUNDIALS_CUSPARSE_VERIFY(cusparseDestroyDnVec(SMCU_CONTENT(A)->vecY));
    SUNDIALS_CUSPARSE_VERIFY(cusparseDestroySpMat(SMCU_CONTENT(A)->spmat_descr));
    SUNMemoryHelper_Dealloc(SMCU_MEMHELP(A), SMCU_CONTENT(A)->dBufferMem,
                            nullptr);
#endif

    if (SMCU_EXECPOLICY(A))
    {
      delete SMCU_EXECPOLICY(A);
      SMCU_EXECPOLICY(A) = NULL;
    }

    SUNMemoryHelper_Destroy(SMCU_MEMHELP(A));

    /* free content struct */
    free(A->content);
    A->content = NULL;
  }

  /* free ops and matrix */
  if (A->ops)
  {
    free(A->ops);
    A->ops = NULL;
  }
  free(A);
  A = NULL;

  return;
}

/* Performs A_ij = 0 */
SUNErrCode SUNMatZero_cuSparse(SUNMatrix A)
{
  cudaError_t cuerr;
  cudaStream_t stream;

  stream = *SMCU_EXECPOLICY(A)->stream();

  /* set all data to zero */
  cuerr = cudaMemsetAsync(SMCU_DATAp(A), 0, SMCU_NNZ(A) * sizeof(sunrealtype),
                          stream);
  if (!SUNDIALS_CUDA_VERIFY(cuerr)) { return (SUN_ERR_OP_FAIL); }

  /* set all rowptrs to zero unless the sparsity pattern is fixed */
  if (!SMCU_FIXEDPATTERN(A))
  {
    cuerr = cudaMemsetAsync(SMCU_INDEXPTRSp(A), 0,
                            (SMCU_BLOCKROWS(A) + 1) * sizeof(int), stream);
    if (!SUNDIALS_CUDA_VERIFY(cuerr)) { return (SUN_ERR_OP_FAIL); }

    /* set all colind to zero */
    cuerr = cudaMemsetAsync(SMCU_INDEXVALSp(A), 0,
                            SMCU_BLOCKNNZ(A) * sizeof(int), stream);
    if (!SUNDIALS_CUDA_VERIFY(cuerr)) { return (SUN_ERR_OP_FAIL); }
  }

  return SUN_SUCCESS;
}

/* Copies the nonzeros, column vals, and row pointers into dst */
SUNErrCode SUNMatCopy_cuSparse(SUNMatrix src, SUNMatrix dst)
{
  int retval;
  const cudaStream_t* stream;

  /* Verify that src and dst are compatible */
  if (!SMCompatible_cuSparse(src, dst)) { return (SUN_ERR_ARG_INCOMPATIBLE); }

  stream = SMCU_EXECPOLICY(src)->stream();

  /* Ensure that dst is allocated with at least as
     much memory as we have nonzeros in src */
  if (SMCU_NNZ(dst) < SMCU_NNZ(src))
  {
    SUNDIALS_DEBUG_PRINT("ERROR in SUNMatCopy_cuSparse: the destination matrix "
                         "has less nonzeros than the source\n");
    return (SUN_ERR_ARG_INCOMPATIBLE);
  }

  /* Zero out dst so that copy works correctly */
  if (SUNMatZero_cuSparse(dst) != SUN_SUCCESS)
  {
    SUNDIALS_DEBUG_PRINT(
      "ERROR in SUNMatCopy_cuSparse: SUNMatZero_cuSparse failed\n");
    return (SUN_ERR_OP_FAIL);
  }

  /* Copy the data over */
  retval = SUNMemoryHelper_CopyAsync(SMCU_MEMHELP(src), SMCU_DATA(dst),
                                     SMCU_DATA(src),
                                     SMCU_NNZ(src) * sizeof(sunrealtype),
                                     (void*)stream);
  if (retval) { return (SUN_ERR_OP_FAIL); }

  /* Copy the row pointers over */
  retval = SUNMemoryHelper_CopyAsync(SMCU_MEMHELP(src), SMCU_INDEXPTRS(dst),
                                     SMCU_INDEXPTRS(src),
                                     (SMCU_BLOCKROWS(src) + 1) * sizeof(int),
                                     (void*)stream);
  if (retval) { return (SUN_ERR_OP_FAIL); }

  /* Copy the column indices over */
  retval = SUNMemoryHelper_CopyAsync(SMCU_MEMHELP(src), SMCU_INDEXVALS(dst),
                                     SMCU_INDEXVALS(src),
                                     SMCU_BLOCKNNZ(src) * sizeof(int),
                                     (void*)stream);
  if (retval) { return (SUN_ERR_OP_FAIL); }

  return SUN_SUCCESS;
}

/* Performs A = cA + I. Requires the diagonal to be allocated already. */
SUNErrCode SUNMatScaleAddI_cuSparse(sunrealtype c, SUNMatrix A)
{
  unsigned threadsPerBlock, gridSize;
  cudaStream_t stream = *SMCU_EXECPOLICY(A)->stream();

  switch (SMCU_SPARSETYPE(A))
  {
  case SUNMAT_CUSPARSE_CSR:
    /* Choose the grid size to be the number of rows in the matrix,
        and then choose threadsPerBlock to be a multiple of the warp size
        that results in enough threads to have one per 2 columns. */
    threadsPerBlock =
      SUNMAX(1, SMCU_EXECPOLICY(A)->blockSize(SMCU_COLUMNS(A) / 2));
    gridSize =
      SUNMAX(1, SMCU_EXECPOLICY(A)->gridSize(SMCU_ROWS(A) * SMCU_COLUMNS(A) / 2,
                                             threadsPerBlock));
    scaleAddIKernelCSR<sunrealtype, int>
      <<<gridSize, threadsPerBlock, 0, stream>>>(SMCU_ROWS(A), c, SMCU_DATAp(A),
                                                 SMCU_INDEXPTRSp(A),
                                                 SMCU_INDEXVALSp(A));
    break;
  case SUNMAT_CUSPARSE_BCSR:
    /* Choose the grid size to be the number of blocks in the matrix,
         and then choose threadsPerBlock to be a multiple of the warp size
         that results in enough threads to have one per row of the block. */
    threadsPerBlock = SMCU_EXECPOLICY(A)->blockSize(SMCU_BLOCKROWS(A));
    gridSize = SMCU_EXECPOLICY(A)->gridSize(SMCU_NBLOCKS(A) * SMCU_BLOCKROWS(A),
                                            threadsPerBlock);
    scaleAddIKernelBCSR<sunrealtype, int>
      <<<gridSize, threadsPerBlock, 0, stream>>>(SMCU_BLOCKROWS(A),
                                                 SMCU_NBLOCKS(A),
                                                 SMCU_BLOCKNNZ(A), c,
                                                 SMCU_DATAp(A),
                                                 SMCU_INDEXPTRSp(A),
                                                 SMCU_INDEXVALSp(A));
    break;
  default:
    SUNDIALS_DEBUG_PRINT(
      "ERROR in SUNMatScaleAddI_cuSparse: sparse type not recognized\n");
    return (SUN_ERR_ARG_INCOMPATIBLE);
  }

#ifdef SUNDIALS_DEBUG_CUDA_LASTERROR
  cudaDeviceSynchronize();
  if (!SUNDIALS_CUDA_VERIFY(cudaGetLastError())) return (SUN_ERR_OP_FAIL);
#endif

  return SUN_SUCCESS;
}

/* Performs A = cA + B */
SUNErrCode SUNMatScaleAdd_cuSparse(sunrealtype c, SUNMatrix A, SUNMatrix B)
{
  cudaStream_t stream;
  unsigned threadsPerBlock, gridSize;

  if (!SMCompatible_cuSparse(A, B))
  {
    SUNDIALS_DEBUG_PRINT(
      "ERROR in SUNMatScaleAdd_cuSparse: SUNMatScaleAdd_cuSparse failed\n");
    return (SUN_ERR_ARG_INCOMPATIBLE);
  }

  stream = *SMCU_EXECPOLICY(A)->stream();

  switch (SMCU_SPARSETYPE(A))
  {
  case SUNMAT_CUSPARSE_CSR:
    /* Choose the grid size to be the number of rows in the matrix,
        and then choose threadsPerBlock to be a multiple of the warp size
        that results in enough threads to have one per 2 columns. */
    threadsPerBlock =
      SUNMAX(1, SMCU_EXECPOLICY(A)->blockSize(SMCU_COLUMNS(A) / 2));
    gridSize =
      SUNMAX(1, SMCU_EXECPOLICY(A)->gridSize(SMCU_ROWS(A) * SMCU_COLUMNS(A) / 2,
                                             threadsPerBlock));
    scaleAddKernelCSR<sunrealtype, int>
      <<<gridSize, threadsPerBlock, 0, stream>>>(SMCU_NNZ(A), c, SMCU_DATAp(A),
                                                 SMCU_DATAp(B));
    break;
  case SUNMAT_CUSPARSE_BCSR:
    /* Choose the grid size to be the number of blocks in the matrix,
         and then choose threadsPerBlock to be a multiple of the warp size
         that results in enough threads to have one per row of the block. */
    threadsPerBlock = SMCU_EXECPOLICY(A)->blockSize(SMCU_BLOCKROWS(A));
    gridSize = SMCU_EXECPOLICY(A)->gridSize(SMCU_NBLOCKS(A) * SMCU_BLOCKROWS(A),
                                            threadsPerBlock);
    scaleAddKernelCSR<sunrealtype, int>
      <<<gridSize, threadsPerBlock, 0, stream>>>(SMCU_NNZ(A), c, SMCU_DATAp(A),
                                                 SMCU_DATAp(B));
    break;
  default:
    SUNDIALS_DEBUG_PRINT(
      "ERROR in SUNMatScaleAdd_cuSparse: sparse type not recognized\n");
    return (SUN_ERR_ARG_INCOMPATIBLE);
  }

#ifdef SUNDIALS_DEBUG_CUDA_LASTERROR
  cudaDeviceSynchronize();
  if (!SUNDIALS_CUDA_VERIFY(cudaGetLastError())) return (SUN_ERR_OP_FAIL);
#endif

  return SUN_SUCCESS;
}

/* Setup buffers needed for Matvec */
SUNErrCode SUNMatMatvecSetup_cuSparse(SUNMatrix A)
{
#if CUDART_VERSION >= 11000
  sunrealtype placeholder[1];
  const sunrealtype one = ONE;

  /* Check if setup has already been done */
  if (!(SMCU_CONTENT(A)->matvec_issetup))
  {
    SUNDIALS_CUSPARSE_VERIFY(cusparseCreateDnVec(&SMCU_CONTENT(A)->vecX,
                                                 SMCU_COLUMNS(A), placeholder,
                                                 CUDA_R_XF));
    SUNDIALS_CUSPARSE_VERIFY(cusparseCreateDnVec(&SMCU_CONTENT(A)->vecY,
                                                 SMCU_ROWS(A), placeholder,
                                                 CUDA_R_XF));

    SUNDIALS_CUSPARSE_VERIFY(
      cusparseSpMV_bufferSize(SMCU_CUSPHANDLE(A), CUSPARSE_OPERATION_NON_TRANSPOSE,
                              &one, SMCU_CONTENT(A)->spmat_descr,
                              SMCU_CONTENT(A)->vecX, &one, SMCU_CONTENT(A)->vecY,
                              CUDA_R_XF, SPMV_ALG, &SMCU_CONTENT(A)->bufferSize));

    if (SUNMemoryHelper_Alloc(SMCU_MEMHELP(A), &SMCU_CONTENT(A)->dBufferMem,
                              SMCU_CONTENT(A)->bufferSize, SUNMEMTYPE_DEVICE,
                              nullptr))
      return (SUN_ERR_OP_FAIL);
  }
#endif
  SMCU_CONTENT(A)->matvec_issetup = SUNTRUE;
  return SUN_SUCCESS;
}

/* Perform y = Ax */
SUNErrCode SUNMatMatvec_cuSparse(SUNMatrix A, N_Vector x, N_Vector y)
{
  /* Verify that the dimensions of A, x, and y agree */
  if ((SMCU_COLUMNS(A) != N_VGetLength(x)) || (SMCU_ROWS(A) != N_VGetLength(y)))
  {
    SUNDIALS_DEBUG_PRINT(
      "ERROR in SUNMatMatvec_cuSparse: dimensions do not agree\n");
    return (SUN_ERR_ARG_INCOMPATIBLE);
  }

  sunrealtype* d_xdata = N_VGetDeviceArrayPointer(x);
  sunrealtype* d_ydata = N_VGetDeviceArrayPointer(y);

  if (SMCU_SPARSETYPE(A) == SUNMAT_CUSPARSE_CSR)
  {
    const sunrealtype one = ONE;

    /* Zero result vector */
    N_VConst(ZERO, y);

#if CUDART_VERSION >= 11000
    {
      /* Setup matvec if it has not been done yet */
      if (!SMCU_CONTENT(A)->matvec_issetup && SUNMatMatvecSetup_cuSparse(A))
      {
        return (SUN_ERR_OP_FAIL);
      }

      SUNDIALS_CUSPARSE_VERIFY(
        cusparseDnVecSetValues(SMCU_CONTENT(A)->vecX, d_xdata));
      SUNDIALS_CUSPARSE_VERIFY(
        cusparseDnVecSetValues(SMCU_CONTENT(A)->vecY, d_ydata));

      SUNDIALS_CUSPARSE_VERIFY(
        cusparseSpMV(SMCU_CUSPHANDLE(A), CUSPARSE_OPERATION_NON_TRANSPOSE, &one,
                     SMCU_CONTENT(A)->spmat_descr, SMCU_CONTENT(A)->vecX, &one,
                     SMCU_CONTENT(A)->vecY, CUDA_R_XF, SPMV_ALG,
                     SMCU_CONTENT(A)->dBufferMem->ptr));
    }
#else
    SUNDIALS_CUSPARSE_VERIFY(
      cusparseXcsrmv(SMCU_CUSPHANDLE(A), CUSPARSE_OPERATION_NON_TRANSPOSE,
                     SMCU_ROWS(A), SMCU_COLUMNS(A), SMCU_NNZ(A), &one,
                     SMCU_MATDESCR(A), SMCU_DATAp(A), SMCU_INDEXPTRSp(A),
                     SMCU_INDEXVALSp(A), d_xdata, &one, d_ydata));
#endif
  }
  else if (SMCU_SPARSETYPE(A) == SUNMAT_CUSPARSE_BCSR)
  {
    cudaStream_t stream;
    unsigned gridSize, threadsPerBlock;

    stream = *SMCU_EXECPOLICY(A)->stream();

    /* Choose the grid size to be the number of blocks in the matrix,
       and then choose threadsPerBlock to be a multiple of the warp size
       that results in enough threads to have one per row of the block. */
    threadsPerBlock = SMCU_EXECPOLICY(A)->blockSize(SMCU_COLUMNS(A) / 2);
    gridSize = SMCU_EXECPOLICY(A)->gridSize(SMCU_ROWS(A) * SMCU_COLUMNS(A) / 2,
                                            threadsPerBlock);
    matvecBCSR<sunrealtype, int>
      <<<gridSize, threadsPerBlock, 0, stream>>>(SMCU_BLOCKROWS(A),
                                                 SMCU_NBLOCKS(A),
                                                 SMCU_BLOCKNNZ(A), SMCU_DATAp(A),
                                                 SMCU_INDEXPTRSp(A),
                                                 SMCU_INDEXVALSp(A), d_xdata,
                                                 d_ydata);
  }
  else
  {
    SUNDIALS_DEBUG_PRINT(
      "ERROR in SUNMatMatvec_cuSparse: sparse type not recognized\n");
    return (SUN_ERR_ARG_INCOMPATIBLE);
  }

#ifdef SUNDIALS_DEBUG_CUDA_LASTERROR
  cudaDeviceSynchronize();
  if (!SUNDIALS_CUDA_VERIFY(cudaGetLastError())) return (SUN_ERR_OP_FAIL);
#endif

  return SUN_SUCCESS;
}

/*
 * =================================================================
 * private functions
 * =================================================================
 */

/* -----------------------------------------------------------------
 * Function to check compatibility of two sparse SUNMatrix objects
 */
static sunbooleantype SMCompatible_cuSparse(SUNMatrix A, SUNMatrix B)
{
  /* both matrices must be sparse */
  if ((SUNMatGetID(A) != SUNMATRIX_CUSPARSE) ||
      (SUNMatGetID(B) != SUNMATRIX_CUSPARSE))
  {
    return (SUNFALSE);
  }

  /* both matrices must have the same shape and sparsity type */
  if (SMCU_ROWS(A) != SMCU_ROWS(B)) { return (SUNFALSE); }
  if (SMCU_COLUMNS(A) != SMCU_COLUMNS(B)) { return (SUNFALSE); }
  if (SMCU_SPARSETYPE(A) != SMCU_SPARSETYPE(B)) { return (SUNFALSE); }

  return (SUNTRUE);
}

/* -----------------------------------------------------------------
 * Function to create empty SUNMatrix with ops attached and
 * the content structure allocated.
 */
SUNMatrix SUNMatrix_cuSparse_NewEmpty(SUNContext sunctx)
{
  /* Create an empty matrix object */
  SUNMatrix A = NULL;
  A           = SUNMatNewEmpty(sunctx);
  if (A == NULL)
  {
    SUNDIALS_DEBUG_PRINT(
      "ERROR in SUNMatrix_cuSparse_NewEmpty: SUNMatNewEmpty failed\n");
    return (NULL);
  }

  /* Attach operations */
  A->ops->getid       = SUNMatGetID_cuSparse;
  A->ops->clone       = SUNMatClone_cuSparse;
  A->ops->destroy     = SUNMatDestroy_cuSparse;
  A->ops->zero        = SUNMatZero_cuSparse;
  A->ops->copy        = SUNMatCopy_cuSparse;
  A->ops->scaleadd    = SUNMatScaleAdd_cuSparse;
  A->ops->scaleaddi   = SUNMatScaleAddI_cuSparse;
  A->ops->matvecsetup = SUNMatMatvecSetup_cuSparse;
  A->ops->matvec      = SUNMatMatvec_cuSparse;

  /* Create content */
  SUNMatrix_Content_cuSparse content = NULL;
  content = (SUNMatrix_Content_cuSparse)malloc(sizeof *content);
  if (content == NULL)
  {
    SUNDIALS_DEBUG_PRINT(
      "ERROR in SUNMatrix_cuSparse_NewEmpty: failed to malloc content\n");
    SUNMatDestroy(A);
    return (NULL);
  }

  /* Attach content */
  A->content          = content;
  content->mem_helper = NULL;

  return (A);
}

#if CUDART_VERSION >= 11000
cusparseStatus_t CreateSpMatDescr(SUNMatrix A, cusparseSpMatDescr_t* spmat_descr)
{
  /* CUDA 11 introduced the "Generic API" and removed the cusparseXcsrmv that
    works on the old cusparseMatDescr_t and raw data arrays. However,
    cuSolverSp stuff requires the cusparseMatDescr_t still. So, we have to
    create this cusparseSpMatDescr_t *and* the cusparseMatDescr_t. */
  return (cusparseCreateCsr(spmat_descr, SMCU_ROWS(A), SMCU_COLUMNS(A),
                            SMCU_NNZ(A), SMCU_INDEXPTRSp(A), SMCU_INDEXVALSp(A),
                            SMCU_DATAp(A), CUSPARSE_INDEX_32I, CUSPARSE_INDEX_32I,
                            CUSPARSE_INDEX_BASE_ZERO, CUDA_R_XF));
}
#endif
