Skip to content

Commit

Permalink
Add RMM to amps files
Browse files Browse the repository at this point in the history
  • Loading branch information
@gartavanis
gartavanis committed Dec 12, 2024
1 parent a02c041 commit f1fd9a4
Show file tree
Hide file tree
Showing 5 changed files with 24 additions and 95 deletions.
24 changes: 4 additions & 20 deletions pfsimulator/amps/cuda/amps.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1060,23 +1060,7 @@ static inline void amps_cuda_error(cudaError_t err, const char *file, int line)
}

#ifdef PARFLOW_HAVE_RMM
#include <rmm/rmm_api.h>
/**
* @brief RMM error handling.
*
* If error detected, print error message and exit.
*
* @param expr RMM error (of type rmmError_t) [IN]
*/
#define RMM_ERRCHK(err) (amps_rmm_error(err, __FILE__, __LINE__))
static inline void amps_rmm_error(rmmError_t err, const char *file, int line)
{
if (err != RMM_SUCCESS)
{
printf("\n\n%s in %s at line %d\n", rmmGetErrorString(err), file, line);
exit(1);
}
}
#include "rmm_wrapper.h"
#endif

/*--------------------------------------------------------------------------
Expand Down Expand Up @@ -1171,7 +1155,7 @@ static inline void *_amps_talloc_cuda(size_t size)
void *ptr = NULL;

#ifdef PARFLOW_HAVE_RMM
RMM_ERRCHK(rmmAlloc(&ptr, size, 0, __FILE__, __LINE__));
ptr = rmmAlloc(size);
#else
CUDA_ERRCHK(cudaMallocManaged((void**)&ptr, size, cudaMemAttachGlobal));
// CUDA_ERRCHK(cudaHostAlloc((void**)&ptr, size, cudaHostAllocMapped));
Expand All @@ -1195,7 +1179,7 @@ static inline void *_amps_ctalloc_cuda(size_t size)
void *ptr = NULL;

#ifdef PARFLOW_HAVE_RMM
RMM_ERRCHK(rmmAlloc(&ptr, size, 0, __FILE__, __LINE__));
ptr = rmmAlloc(ptr);
#else
CUDA_ERRCHK(cudaMallocManaged((void**)&ptr, size, cudaMemAttachGlobal));
// CUDA_ERRCHK(cudaHostAlloc((void**)&ptr, size, cudaHostAllocMapped));
Expand All @@ -1216,7 +1200,7 @@ static inline void *_amps_ctalloc_cuda(size_t size)
static inline void _amps_tfree_cuda(void *ptr)
{
#ifdef PARFLOW_HAVE_RMM
RMM_ERRCHK(rmmFree(ptr, 0, __FILE__, __LINE__));
rmmFree(ptr);
#else
CUDA_ERRCHK(cudaFree(ptr));
// CUDA_ERRCHK(cudaFreeHost(ptr));
Expand Down
26 changes: 4 additions & 22 deletions pfsimulator/amps/mpi1/amps.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@
#include <cuda_runtime.h>
#endif
#ifdef PARFLOW_HAVE_RMM
#include <rmm/rmm_api.h>
#include "rmm_wrapper.h"
#endif

/*
Expand Down Expand Up @@ -1081,24 +1081,6 @@ static inline void amps_cuda_error(cudaError_t err, const char *file, int line)
}
#endif // PARFLOW_HAVE_CUDA

#ifdef PARFLOW_HAVE_RMM
/**
* @brief RMM error handling
*
* If error detected, print error message and exit.
*
* @param expr RMM error (of type rmmError_t) [IN]
*/
#define RMM_ERRCHK(err) (amps_rmm_error(err, __FILE__, __LINE__))
static inline void amps_rmm_error(rmmError_t err, const char *file, int line)
{
if (err != RMM_SUCCESS)
{
printf("\n\n%s in %s at line %d\n", rmmGetErrorString(err), file, line);
exit(1);
}
}
#endif // PARFLOW_HAVE_RMM

/*--------------------------------------------------------------------------
* Define static unified memory allocation routines for devices
Expand Down Expand Up @@ -1139,7 +1121,7 @@ static inline void *_amps_talloc_device(size_t size)
void *ptr = NULL;

#ifdef PARFLOW_HAVE_RMM
RMM_ERRCHK(rmmAlloc(&ptr, size, 0, __FILE__, __LINE__));
ptr = rmmAlloc(size);
#elif defined(PARFLOW_HAVE_KOKKOS)
ptr = kokkosUVMAlloc(size);
#elif defined(PARFLOW_HAVE_CUDA)
Expand All @@ -1165,7 +1147,7 @@ static inline void *_amps_ctalloc_device(size_t size)
void *ptr = NULL;

#ifdef PARFLOW_HAVE_RMM
RMM_ERRCHK(rmmAlloc(&ptr, size, 0, __FILE__, __LINE__));
ptr = rmmAlloc(size);
#elif defined(PARFLOW_HAVE_KOKKOS)
ptr = kokkosUVMAlloc(size);
#elif defined(PARFLOW_HAVE_CUDA)
Expand Down Expand Up @@ -1193,7 +1175,7 @@ static inline void *_amps_ctalloc_device(size_t size)
static inline void _amps_tfree_device(void *ptr)
{
#ifdef PARFLOW_HAVE_RMM
RMM_ERRCHK(rmmFree(ptr, 0, __FILE__, __LINE__));
rmmFree(ptr);
#elif defined(PARFLOW_HAVE_KOKKOS)
kokkosUVMFree(ptr);
#elif defined(PARFLOW_HAVE_CUDA)
Expand Down
27 changes: 4 additions & 23 deletions pfsimulator/amps/oas3/amps.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@
#include <cuda_runtime.h>
#endif
#ifdef PARFLOW_HAVE_RMM
#include <rmm/rmm_api.h>
#include "rmm_wrapper.h"
#endif

/*
Expand Down Expand Up @@ -1112,25 +1112,6 @@ static inline void amps_cuda_error(cudaError_t err, const char *file, int line)
}
#endif // PARFLOW_HAVE_CUDA

#ifdef PARFLOW_HAVE_RMM
/**
* @brief RMM error handling
*
* If error detected, print error message and exit.
*
* @param expr RMM error (of type rmmError_t) [IN]
*/
#define RMM_ERRCHK(err) (amps_rmm_error(err, __FILE__, __LINE__))
static inline void amps_rmm_error(rmmError_t err, const char *file, int line)
{
if (err != RMM_SUCCESS)
{
printf("\n\n%s in %s at line %d\n", rmmGetErrorString(err), file, line);
exit(1);
}
}
#endif // PARFLOW_HAVE_RMM

/*--------------------------------------------------------------------------
* Define static unified memory allocation routines for devices
*--------------------------------------------------------------------------*/
Expand Down Expand Up @@ -1170,7 +1151,7 @@ static inline void *_amps_talloc_device(size_t size)
void *ptr = NULL;

#ifdef PARFLOW_HAVE_RMM
RMM_ERRCHK(rmmAlloc(&ptr, size, 0, __FILE__, __LINE__));
ptr = rmmAlloc(size);
#elif defined(PARFLOW_HAVE_KOKKOS)
ptr = kokkosUVMAlloc(size);
#elif defined(PARFLOW_HAVE_CUDA)
Expand All @@ -1196,7 +1177,7 @@ static inline void *_amps_ctalloc_device(size_t size)
void *ptr = NULL;

#ifdef PARFLOW_HAVE_RMM
RMM_ERRCHK(rmmAlloc(&ptr, size, 0, __FILE__, __LINE__));
ptr = rmmAlloc(size);
#elif defined(PARFLOW_HAVE_KOKKOS)
ptr = kokkosUVMAlloc(size);
#elif defined(PARFLOW_HAVE_CUDA)
Expand Down Expand Up @@ -1224,7 +1205,7 @@ static inline void *_amps_ctalloc_device(size_t size)
static inline void _amps_tfree_device(void *ptr)
{
#ifdef PARFLOW_HAVE_RMM
RMM_ERRCHK(rmmFree(ptr, 0, __FILE__, __LINE__));
rmmFree(ptr);
#elif defined(PARFLOW_HAVE_KOKKOS)
kokkosUVMFree(ptr);
#elif defined(PARFLOW_HAVE_CUDA)
Expand Down
27 changes: 4 additions & 23 deletions pfsimulator/amps/smpi/amps.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@
#include <cuda.h>
#include <cuda_runtime.h>
#ifdef PARFLOW_HAVE_RMM
#include <rmm/rmm_api.h>
#include "rmm_wrapper.h"
#endif
#endif

Expand Down Expand Up @@ -1088,25 +1088,6 @@ static inline void amps_cuda_error(cudaError_t err, const char *file, int line)
}
}

#ifdef PARFLOW_HAVE_RMM
/**
* @brief RMM error handling.
*
* If error detected, print error message and exit.
*
* @param expr RMM error (of type rmmError_t) [IN]
*/
#define RMM_ERRCHK(err) (amps_rmm_error(err, __FILE__, __LINE__))
static inline void amps_rmm_error(rmmError_t err, const char *file, int line)
{
if (err != RMM_SUCCESS)
{
printf("\n\n%s in %s at line %d\n", rmmGetErrorString(err), file, line);
exit(1);
}
}
#endif

/*--------------------------------------------------------------------------
* Define static unified memory allocation routines for CUDA
*--------------------------------------------------------------------------*/
Expand All @@ -1126,7 +1107,7 @@ static inline void *_amps_talloc_cuda(size_t size)
void *ptr = NULL;

#ifdef PARFLOW_HAVE_RMM
RMM_ERRCHK(rmmAlloc(&ptr, size, 0, __FILE__, __LINE__));
ptr = rmmAlloc(size);
#else
CUDA_ERRCHK(cudaMallocManaged((void**)&ptr, size, cudaMemAttachGlobal));
// CUDA_ERRCHK(cudaHostAlloc((void**)&ptr, size, cudaHostAllocMapped));
Expand All @@ -1150,7 +1131,7 @@ static inline void *_amps_ctalloc_cuda(size_t size)
void *ptr = NULL;

#ifdef PARFLOW_HAVE_RMM
RMM_ERRCHK(rmmAlloc(&ptr, size, 0, __FILE__, __LINE__));
ptr = rmmAlloc(size);
#else
CUDA_ERRCHK(cudaMallocManaged((void**)&ptr, size, cudaMemAttachGlobal));
// CUDA_ERRCHK(cudaHostAlloc((void**)&ptr, size, cudaHostAllocMapped));
Expand All @@ -1171,7 +1152,7 @@ static inline void *_amps_ctalloc_cuda(size_t size)
static inline void _amps_tfree_cuda(void *ptr)
{
#ifdef PARFLOW_HAVE_RMM
RMM_ERRCHK(rmmFree(ptr, 0, __FILE__, __LINE__));
rmmFree(ptr);
#else
CUDA_ERRCHK(cudaFree(ptr));
// CUDA_ERRCHK(cudaFreeHost(ptr));
Expand Down
15 changes: 8 additions & 7 deletions pfsimulator/parflow_exe/main.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -176,14 +176,15 @@ int main(int argc, char *argv [])
/*-----------------------------------------------------------------------
* Initialize RMM pool allocator
*-----------------------------------------------------------------------*/
#ifdef PARFLOW_HAVE_RMM
//#ifdef PARFLOW_HAVE_RMM
// RMM
rmmOptions_t rmmOptions;
rmmOptions.allocation_mode = (rmmAllocationMode_t)(PoolAllocation | CudaManagedMemory);
rmmOptions.initial_pool_size = 1; // size = 0 initializes half the device memory
rmmOptions.enable_logging = false;
RMM_ERR(rmmInitialize(&rmmOptions));
#endif // PARFLOW_HAVE_RMM
// rmmOptions_t rmmOptions;
// rmmOptions.allocation_mode = (rmmAllocationMode_t)(PoolAllocation | CudaManagedMemory);
// rmmOptions.initial_pool_size = 1; // size = 0 initializes half the device memory
// rmmOptions.enable_logging = false;
// RMM_ERR(rmmInitialize(&rmmOptions));
rmmInit();
//#endif // PARFLOW_HAVE_RMM

wall_clock_time = amps_Clock();

Expand Down

0 comments on commit f1fd9a4

Please sign in to comment.