/******************************************************************************
 * Copyright (c) 1998 Lawrence Livermore National Security, LLC and other
 * HYPRE Project Developers. See the top-level COPYRIGHT file for details.
 *
 * SPDX-License-Identifier: (Apache-2.0 OR MIT)
 ******************************************************************************/

#include "_hypre_utilities.h"
#include "_hypre_utilities.hpp"

#if defined(__linux__) && defined(HYPRE_USING_FPE_TRAP) && !defined(HYPRE_LONG_DOUBLE)
#ifndef __USE_GNU
#define __USE_GNU
#endif
#include <fenv.h>
#endif

#if defined(HYPRE_USING_CUDA) && defined(HYPRE_DEBUG)
static __global__ void
hypre_GPUKernelCompileFlagSafetyCheck(hypre_DeviceItem &item,
                                      hypre_int        *cuda_arch_compile)
{
   HYPRE_UNUSED_VAR(item);

#if defined(__CUDA_ARCH__)
   cuda_arch_compile[0] = __CUDA_ARCH__;
#else
   HYPRE_UNUSED_VAR(cuda_arch_compile);
#endif
}

static HYPRE_Int
cuda_compile_flag_check_local(void)
{
   HYPRE_Int device;
   hypre_GetDevice(&device);

   struct cudaDeviceProp props;
   HYPRE_CUDA_CALL(cudaGetDeviceProperties(&props, device));
   hypre_int cuda_arch_actual = props.major * 100 + props.minor * 10;
   hypre_int cuda_arch_compile = -1;
   dim3 gDim(1, 1, 1), bDim(1, 1, 1);

   hypre_int *cuda_arch_compile_d = NULL;
   HYPRE_CUDA_CALL(cudaMalloc(&cuda_arch_compile_d, sizeof(hypre_int)));
   HYPRE_CUDA_CALL(cudaMemcpy(cuda_arch_compile_d, &cuda_arch_compile, sizeof(hypre_int),
                              cudaMemcpyHostToDevice));
   HYPRE_GPU_LAUNCH(hypre_GPUKernelCompileFlagSafetyCheck, gDim, bDim, cuda_arch_compile_d);
   HYPRE_CUDA_CALL(cudaMemcpy(&cuda_arch_compile, cuda_arch_compile_d, sizeof(hypre_int),
                              cudaMemcpyDeviceToHost));
   HYPRE_CUDA_CALL(cudaFree(cuda_arch_compile_d));

   if (cuda_arch_actual < cuda_arch_compile)
   {
      char msg[256];

      if (-1 == cuda_arch_compile)
      {
         hypre_sprintf(msg, "hypre error: no proper cuda_arch found");
      }
      else
      {
         hypre_sprintf(msg,
                       "hypre error: Compile arch %d ('--generate-code arch=compute_%d') is greater than device arch %d",
                       cuda_arch_compile, cuda_arch_compile / 10, cuda_arch_actual);
      }

      hypre_error_w_msg(1, msg);
      hypre_ParPrintf(hypre_MPI_COMM_WORLD, "%s\n", msg);
      hypre_assert(0);
   }

   return hypre_error_flag;
}
#endif

/* global variable _hypre_handle:
 * Outside this file, do NOT access it directly,
 * but use hypre_handle() instead (see handle.h) */
#if defined (hypre_DEFINE_GLOBAL)
hypre_Handle *_hypre_handle = NULL;
#endif

/* accessor to the global ``_hypre_handle'' */
hypre_Handle*
hypre_handle(void)
{
   if (!_hypre_handle)
   {
      hypre_error_w_msg(HYPRE_ERROR_GENERIC,
                        "ERROR - _hypre_handle is not initialized. Calling HYPRE_Initialize(). All HYPRE_* or hypre_* function calls should occur between HYPRE_Initialize() and HYPRE_Finalize().\n");
      HYPRE_Initialize();
   }

   return _hypre_handle;
}

hypre_Handle*
hypre_HandleCreate(void)
{
   /* Note: this allocation is done directly with calloc in order to
      avoid a segmentation fault when building with HYPRE_USING_UMPIRE_HOST */
   hypre_Handle *hypre_handle_ = (hypre_Handle*) calloc(1, sizeof(hypre_Handle));

   hypre_HandleLogLevel(hypre_handle_)          = 0;
   hypre_HandleLogLevelSaved(hypre_handle_)     = 0;
   hypre_HandleMemoryLocation(hypre_handle_)    = HYPRE_MEMORY_DEVICE;

#if defined(HYPRE_USING_GPU) || defined(HYPRE_USING_DEVICE_OPENMP)
   hypre_HandleDeviceData(hypre_handle_)        = hypre_DeviceDataCreate();
   hypre_HandleDefaultExecPolicy(hypre_handle_) = HYPRE_EXEC_DEVICE;
#else
   hypre_HandleDefaultExecPolicy(hypre_handle_) = HYPRE_EXEC_HOST;
#endif

   return hypre_handle_;
}

HYPRE_Int
hypre_HandleDestroy(hypre_Handle *hypre_handle_)
{
   if (!hypre_handle_)
   {
      return hypre_error_flag;
   }

   hypre_TFree(hypre_HandleStructCommRecvBuffer(hypre_handle_), HYPRE_MEMORY_DEVICE);
   hypre_TFree(hypre_HandleStructCommSendBuffer(hypre_handle_), HYPRE_MEMORY_DEVICE);
#if defined(HYPRE_USING_GPU)
   hypre_DeviceDataDestroy(hypre_HandleDeviceData(hypre_handle_));
   hypre_HandleDeviceData(hypre_handle_) = NULL;
#endif

   /* Deallocate error messages in error handler */
   hypre_error_handler_clear_messages();

   /* Note: Directly using free since this variable was allocated with calloc */
   free((void*) hypre_handle_);

   return hypre_error_flag;
}

HYPRE_Int
hypre_SetDevice(hypre_int device_id, hypre_Handle *hypre_handle_)
{
#if defined(HYPRE_USING_DEVICE_OPENMP)
   omp_set_default_device(device_id);

#elif defined(HYPRE_USING_CUDA)
   HYPRE_CUDA_CALL( cudaSetDevice(device_id) );
   hypre_HandleDevice(hypre_handle_) = device_id;

#elif defined(HYPRE_USING_HIP)
   HYPRE_HIP_CALL( hipSetDevice(device_id) );
   hypre_HandleDevice(hypre_handle_) = device_id;

#elif defined(HYPRE_USING_SYCL)
   if (hypre_handle_)
   {
      if (!hypre_HandleDevice(hypre_handle_))
      {
         /* Note: this enforces "explicit scaling," i.e. we treat each tile of a multi-tile GPU as a separate device */
         sycl::platform platform(sycl::gpu_selector_v);
         auto gpu_devices = platform.get_devices(sycl::info::device_type::gpu);
         hypre_int n_devices = 0;
         hypre_GetDeviceCount(&n_devices);
         if (device_id >= n_devices)
         {
            hypre_error_w_msg(HYPRE_ERROR_GENERIC,
                              "ERROR: SYCL device-ID exceed the number of devices on-node\n");
         }

         hypre_int local_n_devices = 0;
         hypre_int i;
         for (i = 0; i < gpu_devices.size(); i++)
         {
            if (local_n_devices == device_id)
            {
               hypre_HandleDevice(hypre_handle_) = new sycl::device(gpu_devices[i]);
            }
            local_n_devices++;
         }
      }
      hypre_DeviceDataDeviceMaxWorkGroupSize(hypre_HandleDeviceData(hypre_handle_)) =
         hypre_DeviceDataDevice(hypre_HandleDeviceData(
                                   hypre_handle_))->get_info<sycl::info::device::max_work_group_size>();
   }
#else
   HYPRE_UNUSED_VAR(device_id);
   HYPRE_UNUSED_VAR(hypre_handle_);
#endif

   return hypre_error_flag;
}

HYPRE_Int
hypre_GetDeviceMaxShmemSize(hypre_int  device_id,
                            hypre_int *max_size_ptr,
                            hypre_int *max_size_optin_ptr)
{
   hypre_int max_size = 0, max_size_optin = 0;

#if defined(HYPRE_USING_GPU)
   hypre_Handle *handle = hypre_handle();

   if (!hypre_HandleDeviceMaxShmemPerBlockInited(handle))
   {
      if (device_id == -1)
      {
         hypre_GetDevice(&device_id);
      }

#if defined(HYPRE_USING_CUDA)
      HYPRE_CUDA_CALL(cudaDeviceGetAttribute(
                         &max_size,
                         cudaDevAttrMaxSharedMemoryPerBlock,
                         device_id));

      HYPRE_CUDA_CALL(cudaDeviceGetAttribute(
                         &max_size_optin,
                         cudaDevAttrMaxSharedMemoryPerBlockOptin,
                         device_id));

#elif defined(HYPRE_USING_HIP)
      HYPRE_HIP_CALL(hipDeviceGetAttribute(
                        &max_size,
                        hipDeviceAttributeMaxSharedMemoryPerBlock,
                        device_id));

#elif defined(HYPRE_USING_SYCL)
      auto device = *hypre_HandleDevice(hypre_handle());
      max_size = device.get_info<sycl::info::device::local_mem_size>();
#endif
      hypre_HandleDeviceMaxShmemPerBlock(handle)[0] = max_size;
      hypre_HandleDeviceMaxShmemPerBlock(handle)[1] = max_size_optin;

      hypre_HandleDeviceMaxShmemPerBlockInited(handle) = 1;
   }

   if (max_size_ptr)
   {
      *max_size_ptr = hypre_HandleDeviceMaxShmemPerBlock(handle)[0];
   }

   if (max_size_optin_ptr)
   {
      *max_size_optin_ptr = hypre_HandleDeviceMaxShmemPerBlock(handle)[1];
   }
#else /* not HYPRE_USING_GPU */
   HYPRE_UNUSED_VAR(device_id);

   if (max_size_ptr)
   {
      *max_size_ptr = max_size;
   }

   if (max_size_optin_ptr)
   {
      *max_size_optin_ptr = max_size_optin;
   }
#endif

   return hypre_error_flag;
}

/* Note: it doesn't return device_id in hypre_Handle->hypre_DeviceData,
 *       calls API instead. But these two should match at all times
 */
HYPRE_Int
hypre_GetDevice(hypre_int *device_id)
{
#if defined(HYPRE_USING_DEVICE_OPENMP)
   *device_id = omp_get_default_device();

#elif defined(HYPRE_USING_CUDA)
   HYPRE_CUDA_CALL( cudaGetDevice(device_id) );

#elif defined(HYPRE_USING_HIP)
   HYPRE_HIP_CALL( hipGetDevice(device_id) );

#elif defined(HYPRE_USING_SYCL)
   /* Note - no sycl call to get which device is setup for use (if the user has already setup a device at all)
    * Assume the rank/device binding below */
   HYPRE_Int my_id;
   hypre_int n_devices;
   hypre_MPI_Comm_rank(hypre_MPI_COMM_WORLD, &my_id);
   hypre_GetDeviceCount(&n_devices);
   (*device_id) = my_id % n_devices;

#else
   *device_id = 0;
#endif

   return hypre_error_flag;
}

HYPRE_Int
hypre_GetDeviceCount(hypre_int *device_count)
{
#if defined(HYPRE_USING_DEVICE_OPENMP)
   *device_count = omp_get_num_devices();

#elif defined(HYPRE_USING_CUDA)
   HYPRE_CUDA_CALL( cudaGetDeviceCount(device_count) );

#elif defined(HYPRE_USING_HIP)
   HYPRE_HIP_CALL( hipGetDeviceCount(device_count) );

#elif defined(HYPRE_USING_SYCL)
   (*device_count) = 0;
   sycl::platform platform(sycl::gpu_selector_v);
   auto const& gpu_devices = platform.get_devices(sycl::info::device_type::gpu);
   HYPRE_Int i;
   for (i = 0; i < gpu_devices.size(); i++)
   {
      (*device_count)++;
   }

#else
   *device_count = 0;
#endif

   return hypre_error_flag;
}

HYPRE_Int
hypre_GetDeviceLastError(void)
{
#if defined(HYPRE_USING_CUDA)
   HYPRE_CUDA_CALL( cudaGetLastError() );

#elif defined(HYPRE_USING_HIP)
   HYPRE_HIP_CALL( hipGetLastError() );

#elif defined(HYPRE_USING_SYCL)
   try
   {
      hypre_HandleComputeStream(hypre_handle())->wait_and_throw();
   }
   catch (sycl::exception const& e)
   {
      std::cout << "Caught synchronous SYCL exception:\n"
                << e.what() << std::endl;
   }
#endif

   return hypre_error_flag;
}

/******************************************************************************
 *
 * hypre initialization
 *
 *****************************************************************************/

HYPRE_Int
HYPRE_DeviceInitialize(void)
{
#if defined(HYPRE_USING_GPU) || defined(HYPRE_USING_DEVICE_OPENMP)
   hypre_Handle *handle = hypre_handle();

#if !defined(HYPRE_USING_SYCL)
   /* With sycl, cannot call hypre_GetDeviceLastError() until after device and queue setup */
   hypre_GetDeviceLastError();
#endif

   /* Notice: the cudaStream created is specific to the device
    * that was in effect when you created the stream.
    * So, we should first set the device and create the streams
    */
   hypre_int device_id;
   hypre_GetDevice(&device_id);
   hypre_SetDevice(device_id, handle);

   hypre_GetDeviceMaxShmemSize(device_id, NULL, NULL);

#if defined(HYPRE_USING_DEVICE_MALLOC_ASYNC)
   cudaMemPool_t mempool;
   cudaDeviceGetDefaultMemPool(&mempool, device_id);
   uint64_t threshold = UINT64_MAX;
   cudaMemPoolSetAttribute(mempool, cudaMemPoolAttrReleaseThreshold, &threshold);
#endif

   /* To include the cost of creating streams/cudahandles in HYPRE_Init */
   /* If not here, will be done at the first use */
#if defined(HYPRE_USING_CUDA_STREAMS)
   hypre_HandleComputeStream(handle);
#endif

   /* A separate stream for prefetching */
   //hypre_HandleCudaPrefetchStream(handle);

#if defined(HYPRE_USING_CUBLAS)
   hypre_HandleCublasHandle(handle);
#endif

#if defined(HYPRE_USING_CUSPARSE) || defined(HYPRE_USING_ROCSPARSE)
   hypre_HandleCusparseHandle(handle);
#endif

#if defined(HYPRE_USING_CURAND) || defined(HYPRE_USING_ROCRAND)
   hypre_HandleCurandGenerator(handle);
#endif

#if defined(HYPRE_USING_CUSOLVER) || defined(HYPRE_USING_ROCSOLVER)
   hypre_HandleVendorSolverHandle(handle);
#endif

   /* Check if cuda arch flags in compiling match the device */
#if defined(HYPRE_USING_CUDA) && defined(HYPRE_DEBUG)
   cuda_compile_flag_check_local();
#endif

   /* Check if OS supports UVM */
#if defined(HYPRE_USING_CUDA) && defined(HYPRE_USING_UNIFIED_MEMORY)
   {
      struct cudaDeviceProp deviceProp;

      HYPRE_CUDA_CALL( cudaGetDevice(&device_id) );
      HYPRE_CUDA_CALL( cudaGetDeviceProperties(&deviceProp, device_id) );

      if (deviceProp.managedMemory == 1 && deviceProp.concurrentManagedAccess == 1)
      {
         hypre_HandleDeviceUVM(handle) = 1;
      }
   }

#endif

#endif /* if defined(HYPRE_USING_GPU) || defined(HYPRE_USING_DEVICE_OPENMP) */

   return hypre_error_flag;
}

HYPRE_Int
HYPRE_Initialize(void)
{
   /* Return if the hypre library is in initialized state */
   if (hypre_Initialized())
   {
      return hypre_error_flag;
   }

#if defined(HYPRE_USING_MEMORY_TRACKER)
   if (!_hypre_memory_tracker)
   {
      _hypre_memory_tracker = hypre_MemoryTrackerCreate();
   }
#endif

   if (!_hypre_handle)
   {
      _hypre_handle = hypre_HandleCreate();
   }

#if defined(HYPRE_USING_DEVICE_OPENMP)
   HYPRE_OMPOffloadOn();
#endif

#if defined(HYPRE_USING_UMPIRE)
   hypre_UmpireInit(_hypre_handle);
#endif

#if defined(HYPRE_USING_MAGMA)
   hypre_MagmaInitialize();
#endif

#if defined(__linux__) && defined(HYPRE_USING_FPE_TRAP) && !defined(HYPRE_LONG_DOUBLE)
   feenableexcept(FE_INVALID | FE_DIVBYZERO | FE_OVERFLOW);
#endif

   /* Update library state */
   hypre_SetInitialized();

   return hypre_error_flag;
}

/******************************************************************************
 *
 * hypre finalization
 *
 *****************************************************************************/

HYPRE_Int
HYPRE_Finalize(void)
{
   /* Return if the hypre library has already been finalized */
   if (hypre_Finalized())
   {
      return hypre_error_flag;
   }

   HYPRE_Int exec_device = hypre_HandleDefaultExecPolicy(_hypre_handle) == HYPRE_EXEC_DEVICE;

#if defined(HYPRE_USING_UMPIRE)
   hypre_UmpireFinalize(_hypre_handle);
#endif

#if defined(HYPRE_USING_MAGMA)
   hypre_MagmaFinalize();
#endif

#if defined(HYPRE_USING_SYCL)
   /* With sycl, cannot call hypre_GetDeviceLastError() after destroying the handle, so do it here */
   if (exec_device)
   {
      hypre_GetDeviceLastError();
   }
#endif

   hypre_HandleDestroy(_hypre_handle);
   _hypre_handle = NULL;

#if !defined(HYPRE_USING_SYCL)
   if (exec_device)
   {
      hypre_GetDeviceLastError();
   }
#endif

#if defined(HYPRE_USING_MEMORY_TRACKER)
   hypre_PrintMemoryTracker(hypre_total_bytes, hypre_peak_bytes, hypre_current_bytes,
                            hypre_memory_tracker_print, hypre_memory_tracker_filename);

   hypre_MemoryTrackerDestroy(_hypre_memory_tracker);
   _hypre_memory_tracker = NULL;
#endif

   /* Update library state */
   hypre_SetFinalized();

   return hypre_error_flag;
}

HYPRE_Int
HYPRE_PrintDeviceInfo(void)
{
#if defined(HYPRE_USING_GPU)
   hypre_int dev = 0;
#endif

#if defined(HYPRE_USING_CUDA)
   struct cudaDeviceProp deviceProp;

   HYPRE_CUDA_CALL( cudaGetDevice(&dev) );
   HYPRE_CUDA_CALL( cudaGetDeviceProperties(&deviceProp, dev) );
   {
      hypre_printf("Running on \"%s\", Comp. Capability: %d.%d, Total VRAM: %.2f GiB, ",
                   deviceProp.name, deviceProp.major, deviceProp.minor,
                   ((HYPRE_Real) deviceProp.totalGlobalMem) / (HYPRE_Real) (1 << 30));
   }

#elif defined(HYPRE_USING_HIP)
   hipDeviceProp_t deviceProp;

   HYPRE_HIP_CALL( hipGetDevice(&dev) );
   HYPRE_HIP_CALL( hipGetDeviceProperties(&deviceProp, dev) );
   {
      const char *gfx_name = deviceProp.gcnArchName[0] ? deviceProp.gcnArchName : "unknown";
      hypre_printf("Running on \"%s\", %s, Total VRAM: %.2f GiB, ",
                   deviceProp.name, gfx_name,
                   ((HYPRE_Real) deviceProp.totalGlobalMem) / (HYPRE_Real) (1 << 30));
   }

#elif defined(HYPRE_USING_SYCL)
   auto device = *hypre_HandleDevice(hypre_handle());
   auto p_name = device.get_platform().get_info<sycl::info::platform::name>();
   hypre_printf("Platform Name: %s\n", p_name.c_str());
   auto p_version = device.get_platform().get_info<sycl::info::platform::version>();
   hypre_printf("Platform Version: %s\n", p_version.c_str());
   auto d_name = device.get_info<sycl::info::device::name>();
   hypre_printf("Device Name: %s\n", d_name.c_str());
   auto total_vram = device.get_info<sycl::info::device::global_mem_size>();
   hypre_printf("Total VRAM: %.2f GiB\n",
                ((HYPRE_Real) total_vram) / (HYPRE_Real) (1 << 30));
   auto max_work_group = device.get_info<sycl::info::device::max_work_group_size>();
   hypre_printf("Max Work Groups: %d\n", max_work_group);
   auto max_compute_units = device.get_info<sycl::info::device::max_compute_units>();
   hypre_printf("Max Compute Units: %d\n", max_compute_units);
#endif

#if defined(HYPRE_USING_GPU)
   hypre_int max_size = 0, max_size_optin = 0;
   hypre_GetDeviceMaxShmemSize(dev, &max_size, &max_size_optin);
   hypre_printf("MaxShmem: %d KiB, MaxShmemOptin: %d KiB\n",
                max_size / 1024, max_size_optin / 1024);
#endif

   return hypre_error_flag;
}

/*--------------------------------------------------------------------------
 * HYPRE_MemoryPrintUsage
 *--------------------------------------------------------------------------*/

HYPRE_Int
HYPRE_MemoryPrintUsage(MPI_Comm    comm,
                       HYPRE_Int   level,
                       const char *function,
                       HYPRE_Int   line)
{
   return hypre_MemoryPrintUsage(comm, level, function, line);
}

/******************************************************************************
 *
 * HYPRE memory location
 *
 *****************************************************************************/

HYPRE_Int
HYPRE_SetMemoryLocation(HYPRE_MemoryLocation memory_location)
{
   hypre_HandleMemoryLocation(hypre_handle()) = memory_location;

   return hypre_error_flag;
}

HYPRE_Int
HYPRE_GetMemoryLocation(HYPRE_MemoryLocation *memory_location)
{
   *memory_location = hypre_HandleMemoryLocation(hypre_handle());

   return hypre_error_flag;
}

/******************************************************************************
 *
 * HYPRE execution policy
 *
 *****************************************************************************/

HYPRE_Int
HYPRE_SetExecutionPolicy(HYPRE_ExecutionPolicy exec_policy)
{
   hypre_HandleDefaultExecPolicy(hypre_handle()) = exec_policy;

   return hypre_error_flag;
}

HYPRE_Int
HYPRE_GetExecutionPolicy(HYPRE_ExecutionPolicy *exec_policy)
{
   *exec_policy = hypre_HandleDefaultExecPolicy(hypre_handle());

   return hypre_error_flag;
}

const char*
HYPRE_GetExecutionPolicyName(HYPRE_ExecutionPolicy exec_policy)
{
   switch (exec_policy)
   {
      case HYPRE_EXEC_HOST:
         return "Host";

      case HYPRE_EXEC_DEVICE:
#if defined(HYPRE_USING_GPU) || defined(HYPRE_USING_DEVICE_OPENMP)
#if defined(HYPRE_USING_CUDA)
         return "Device (CUDA)";
#elif defined(HYPRE_USING_HIP)
         return "Device (HIP)";
#elif defined(HYPRE_USING_SYCL)
         return "Device (SYCL)";
#else
         return "Device (OpenMP)";
#endif
#else
         return "Host";
#endif
      case HYPRE_EXEC_UNDEFINED:
      default:
         return "Undefined";
   }
}
