xref: /petsc/src/ksp/ksp/tests/raja/memoryManager.hpp (revision 9dd11ecf0918283bb567d8b33a92f53ac4ea7840)
1 //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
2 // Copyright (c) 2016-21, Lawrence Livermore National Security, LLC
3 // and RAJA project contributors. See the RAJA/COPYRIGHT file for details.
4 //
5 // SPDX-License-Identifier: (BSD-3-Clause)
6 //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
7 
8 #pragma once
9 
10 #include "RAJA/RAJA.hpp"
11 
12 #if defined(RAJA_ENABLE_CUDA)
13   #include "RAJA/policy/cuda/raja_cudaerrchk.hpp"
14 #endif
15 
16 #if defined(RAJA_ENABLE_HIP)
17   #include "RAJA/policy/hip/raja_hiperrchk.hpp"
18 #endif
19 
20 /*
21   As RAJA does not manage memory we include a general purpose memory
22   manager which may be used to perform c++ style allocation/deallocation
23   or allocate/deallocate CUDA unified memory. The type of memory allocated
24   is dependent on how RAJA was configured.
25 */
26 namespace memoryManager
27 {
28 
29 #if defined(RAJA_ENABLE_SYCL)
30 static camp::resources::Resource *sycl_res;
31 #endif
32 
33 template <typename T>
allocate(RAJA::Index_type size)34 T *allocate(RAJA::Index_type size)
35 {
36   T *ptr;
37 #if defined(RAJA_ENABLE_CUDA)
38   cudaErrchk(cudaMallocManaged((void **)&ptr, sizeof(T) * size, cudaMemAttachGlobal));
39 #elif defined(RAJA_ENABLE_HIP)
40   hipErrchk(hipMalloc((void **)&ptr, sizeof(T) * size));
41 #elif defined(RAJA_ENABLE_SYCL)
42   ptr = sycl_res->allocate<T>(size);
43 #else
44   ptr = new T[size];
45 #endif
46   return ptr;
47 }
48 
49 template <typename T>
deallocate(T * & ptr)50 void deallocate(T *&ptr)
51 {
52   if (ptr) {
53 #if defined(RAJA_ENABLE_CUDA)
54     cudaErrchk(cudaFree(ptr));
55 #elif defined(RAJA_ENABLE_HIP)
56     hipErrchk(hipFree(ptr));
57 #elif defined(RAJA_ENABLE_SYCL)
58     sycl_res->deallocate(ptr);
59 #else
60     delete[] ptr;
61 #endif
62     ptr = nullptr;
63   }
64 }
65 
66 #if defined(RAJA_ENABLE_CUDA) || defined(RAJA_ENABLE_HIP) || defined(RAJA_ENABLE_SYCL)
67 template <typename T>
allocate_gpu(RAJA::Index_type size)68 T *allocate_gpu(RAJA::Index_type size)
69 {
70   T *ptr;
71   #if defined(RAJA_ENABLE_CUDA)
72   cudaErrchk(cudaMalloc((void **)&ptr, sizeof(T) * size));
73   #elif defined(RAJA_ENABLE_HIP)
74   hipErrchk(hipMalloc((void **)&ptr, sizeof(T) * size));
75   #elif defined(RAJA_ENABLE_SYCL)
76   auto qu = sycl_res->get<camp::resources::Sycl>().get_queue();
77   ptr     = cl::sycl::malloc_device<T>(size, *qu);
78   #endif
79   return ptr;
80 }
81 
82 template <typename T>
deallocate_gpu(T * & ptr)83 void deallocate_gpu(T *&ptr)
84 {
85   if (ptr) {
86   #if defined(RAJA_ENABLE_CUDA)
87     cudaErrchk(cudaFree(ptr));
88   #elif defined(RAJA_ENABLE_HIP)
89     hipErrchk(hipFree(ptr));
90   #elif defined(RAJA_ENABLE_SYCL)
91     sycl_res->deallocate(ptr);
92   #endif
93     ptr = nullptr;
94   }
95 }
96 #endif
97 
98 }; // namespace memoryManager
99