/* -----------------------------------------------------------------------------
 * Programmer(s): David J. Gardner @ 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
 * -----------------------------------------------------------------------------
 * CUDA/HIP functions for exchange buffers
 * ---------------------------------------------------------------------------*/

#include "diffusion_2D.hpp"

#if defined(USE_HIP)
#include <hip/hip_runtime.h>
#define BLOCK_SIZE   256
#define BLOCK_SIZE_X 16
#define BLOCK_SIZE_Y 16
#elif defined(USE_CUDA)
#include <cuda_runtime.h>
#define BLOCK_SIZE   256
#define BLOCK_SIZE_X 16
#define BLOCK_SIZE_Y 16
#else
#error Define USE_CUDA or USE_HIP
#endif

// Pack exchange buffers kernel
__global__ void pack_buffers_kernel(const sunindextype nx_loc,
                                    const sunindextype ny_loc,
                                    const sunrealtype* u, sunrealtype* wbuf,
                                    sunrealtype* ebuf, sunrealtype* sbuf,
                                    sunrealtype* nbuf)
{
  // Thread ID
  int i = blockIdx.x * blockDim.x + threadIdx.x;

  // West and East faces
  if (i < ny_loc)
  {
    if (wbuf) { wbuf[i] = u[i * nx_loc]; }
    if (ebuf) { ebuf[i] = u[(i + 1) * nx_loc - 1]; }
  }

  // South and North faces
  if (i < nx_loc)
  {
    if (sbuf) { sbuf[i] = u[i]; }
    if (nbuf) { nbuf[i] = u[i + (ny_loc - 1) * nx_loc]; }
  }
}

// Pack exchange buffers
int UserData::pack_buffers(const N_Vector u)
{
  // Access data array
  const sunrealtype* uarray =
    N_VGetDeviceArrayPointer(N_VGetLocalVector_MPIPlusX(u));
  if (check_flag((void*)uarray, "N_VGetDeviceArrayPointer", 0)) return -1;

  sunindextype maxdim = max(nx_loc, ny_loc);
  dim3 block(BLOCK_SIZE);
  dim3 grid(ICEIL(maxdim, BLOCK_SIZE));

  pack_buffers_kernel<<<grid, block>>>(nx_loc, ny_loc, uarray, Wsend, Esend,
                                       Ssend, Nsend);

  return 0;
}

// Allocate exchange buffers
int UserData::allocate_buffers()
{
  if (HaveNbrW)
  {
#if defined(USE_HIP)
    hipMalloc(&Wrecv, ny_loc * sizeof(sunrealtype));
    hipMalloc(&Wsend, ny_loc * sizeof(sunrealtype));
#elif defined(USE_CUDA)
    cudaMalloc(&Wrecv, ny_loc * sizeof(sunrealtype));
    cudaMalloc(&Wsend, ny_loc * sizeof(sunrealtype));
#else
#error Define USE_CUDA or USE_HIP
#endif
  }

  if (HaveNbrE)
  {
#if defined(USE_HIP)
    hipMalloc(&Erecv, ny_loc * sizeof(sunrealtype));
    hipMalloc(&Esend, ny_loc * sizeof(sunrealtype));
#elif defined(USE_CUDA)
    cudaMalloc(&Erecv, ny_loc * sizeof(sunrealtype));
    cudaMalloc(&Esend, ny_loc * sizeof(sunrealtype));
#else
#error Define USE_CUDA or USE_HIP
#endif
  }

  if (HaveNbrS)
  {
#if defined(USE_HIP)
    hipMalloc(&Srecv, nx_loc * sizeof(sunrealtype));
    hipMalloc(&Ssend, nx_loc * sizeof(sunrealtype));
#elif defined(USE_CUDA)
    cudaMalloc(&Srecv, nx_loc * sizeof(sunrealtype));
    cudaMalloc(&Ssend, nx_loc * sizeof(sunrealtype));
#else
#error Define USE_CUDA or USE_HIP
#endif
  }

  if (HaveNbrN)
  {
#if defined(USE_HIP)
    hipMalloc(&Nrecv, nx_loc * sizeof(sunrealtype));
    hipMalloc(&Nsend, nx_loc * sizeof(sunrealtype));
#elif defined(USE_CUDA)
    cudaMalloc(&Nrecv, nx_loc * sizeof(sunrealtype));
    cudaMalloc(&Nsend, nx_loc * sizeof(sunrealtype));
#else
#error Define USE_CUDA or USE_HIP
#endif
  }

  return 0;
}

// Free exchange buffers
int UserData::free_buffers()
{
#if defined(USE_HIP)
  if (Wrecv != NULL) hipFree(Wrecv);
  if (Wsend != NULL) hipFree(Wsend);
  if (Erecv != NULL) hipFree(Erecv);
  if (Esend != NULL) hipFree(Esend);
  if (Srecv != NULL) hipFree(Srecv);
  if (Ssend != NULL) hipFree(Ssend);
  if (Nrecv != NULL) hipFree(Nrecv);
  if (Nsend != NULL) hipFree(Nsend);
#elif defined(USE_CUDA)
  if (Wrecv != NULL) cudaFree(Wrecv);
  if (Wsend != NULL) cudaFree(Wsend);
  if (Erecv != NULL) cudaFree(Erecv);
  if (Esend != NULL) cudaFree(Esend);
  if (Srecv != NULL) cudaFree(Srecv);
  if (Ssend != NULL) cudaFree(Ssend);
  if (Nrecv != NULL) cudaFree(Nrecv);
  if (Nsend != NULL) cudaFree(Nsend);
#else
#error Define USE_CUDA or USE_HIP
#endif

  return 0;
}
