Fixup Cuda/HIP graph header includes and pointer casting

dalg24 · Jan 21, 2024 · 88bcf13 · 88bcf13
1 parent 237475e
commit 88bcf13
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/core/src/Cuda/Kokkos_Cuda_GraphNodeKernel.hpp b/core/src/Cuda/Kokkos_Cuda_GraphNodeKernel.hpp
@@ -77,7 +77,7 @@ class GraphNodeKernelImpl<Kokkos::Cuda, PolicyType, Functor, PatternTag,
 
   ~GraphNodeKernelImpl() {
     if (m_driver_storage) {
-      Kokkos::kokkos_free<Kokkos::CudaSpace>(m_driver_storage);
+      Kokkos::CudaSpace().deallocate(m_driver_storage, sizeof(base_t));
     }
   }
 
@@ -92,8 +92,8 @@ class GraphNodeKernelImpl<Kokkos::Cuda, PolicyType, Functor, PatternTag,
 
   Kokkos::ObservingRawPtr<base_t> allocate_driver_memory_buffer() const {
     KOKKOS_EXPECTS(m_driver_storage == nullptr)
-    m_driver_storage = Kokkos::kokkos_malloc<Kokkos::CudaSpace>(
-        "GraphNodeKernel global memory functor storage", sizeof(base_t));
+    m_driver_storage = static_cast<base_t*>(Kokkos::CudaSpace().allocate(
+        "GraphNodeKernel global memory functor storage", sizeof(base_t)));
     KOKKOS_ENSURES(m_driver_storage != nullptr)
     return m_driver_storage;
   }

diff --git a/core/src/HIP/Kokkos_HIP_GraphNodeKernel.hpp b/core/src/HIP/Kokkos_HIP_GraphNodeKernel.hpp
@@ -58,7 +58,7 @@ class GraphNodeKernelImpl<Kokkos::HIP, PolicyType, Functor, PatternTag, Args...>
 
   ~GraphNodeKernelImpl() {
     if (m_driver_storage) {
-      Kokkos::kokkos_free<Kokkos::HIPSpace>(m_driver_storage);
+      Kokkos::HIPSpace().deallocate(m_driver_storage, sizeof(base_t));
     }
   }
 
@@ -76,8 +76,8 @@ class GraphNodeKernelImpl<Kokkos::HIP, PolicyType, Functor, PatternTag, Args...>
 
   Kokkos::ObservingRawPtr<base_t> allocate_driver_memory_buffer() const {
     KOKKOS_EXPECTS(m_driver_storage == nullptr);
-    Kokkos::kokkos_malloc<Kokkos::HIPSpace>(
-        "GraphNodeKernel global memory functor storage", sizeof(base_t));
+    m_driver_storage = static_cast<base_t*>(Kokkos::HIPSpace().allocate(
+        "GraphNodeKernel global memory functor storage", sizeof(base_t)));
     KOKKOS_ENSURES(m_driver_storage != nullptr);
     return m_driver_storage;
   }