Fix SharedAllocationRecord to allocate using the correct execution sp…

…ace instance (kokkos#6789) * Fix SharedAllocationRecord * Fix SYCL
diehlpk · Feb 7, 2024 · 391f2d1 · 391f2d1
1 parent 442e4d4
commit 391f2d1
Show file tree

Hide file tree

Showing 6 changed files with 93 additions and 4 deletions.
diff --git a/core/src/Cuda/Kokkos_CudaSpace.hpp b/core/src/Cuda/Kokkos_CudaSpace.hpp
@@ -168,6 +168,16 @@ class CudaUVMSpace {
   ~CudaUVMSpace()                                  = default;
 
   /**\brief  Allocate untracked memory in the cuda space */
+  template <typename ExecutionSpace>
+  void* allocate(const ExecutionSpace&, const size_t arg_alloc_size) const {
+    return allocate(arg_alloc_size);
+  }
+  template <typename ExecutionSpace>
+  void* allocate(const ExecutionSpace&, const char* arg_label,
+                 const size_t arg_alloc_size,
+                 const size_t arg_logical_size = 0) const {
+    return allocate(arg_label, arg_alloc_size, arg_logical_size);
+  }
   void* allocate(const size_t arg_alloc_size) const;
   void* allocate(const char* arg_label, const size_t arg_alloc_size,
                  const size_t arg_logical_size = 0) const;
@@ -250,6 +260,16 @@ class CudaHostPinnedSpace {
   ~CudaHostPinnedSpace()                                         = default;
 
   /**\brief  Allocate untracked memory in the space */
+  template <typename ExecutionSpace>
+  void* allocate(const ExecutionSpace&, const size_t arg_alloc_size) const {
+    return allocate(arg_alloc_size);
+  }
+  template <typename ExecutionSpace>
+  void* allocate(const ExecutionSpace&, const char* arg_label,
+                 const size_t arg_alloc_size,
+                 const size_t arg_logical_size = 0) const {
+    return allocate(arg_label, arg_alloc_size, arg_logical_size);
+  }
   void* allocate(const size_t arg_alloc_size) const;
   void* allocate(const char* arg_label, const size_t arg_alloc_size,
                  const size_t arg_logical_size = 0) const;

diff --git a/core/src/HIP/Kokkos_HIP_Space.hpp b/core/src/HIP/Kokkos_HIP_Space.hpp
@@ -65,6 +65,15 @@ class HIPSpace {
   ~HIPSpace()                              = default;
 
   /**\brief  Allocate untracked memory in the hip space */
+  // FIXME_HIP Use execution space instance
+  void* allocate(const HIP&, const size_t arg_alloc_size) const {
+    return allocate(arg_alloc_size);
+  }
+  // FIXME_HIP Use execution space instance
+  void* allocate(const HIP&, const char* arg_label, const size_t arg_alloc_size,
+                 const size_t arg_logical_size = 0) const {
+    return allocate(arg_label, arg_alloc_size, arg_logical_size);
+  }
   void* allocate(const size_t arg_alloc_size) const;
   void* allocate(const char* arg_label, const size_t arg_alloc_size,
                  const size_t arg_logical_size = 0) const;
@@ -125,6 +134,16 @@ class HIPHostPinnedSpace {
   ~HIPHostPinnedSpace()                                        = default;
 
   /**\brief  Allocate untracked memory in the space */
+  template <typename ExecutionSpace>
+  void* allocate(const ExecutionSpace&, const size_t arg_alloc_size) const {
+    return allocate(arg_alloc_size);
+  }
+  template <typename ExecutionSpace>
+  void* allocate(const ExecutionSpace&, const char* arg_label,
+                 const size_t arg_alloc_size,
+                 const size_t arg_logical_size = 0) const {
+    return allocate(arg_label, arg_alloc_size, arg_logical_size);
+  }
   void* allocate(const size_t arg_alloc_size) const;
   void* allocate(const char* arg_label, const size_t arg_alloc_size,
                  const size_t arg_logical_size = 0) const;
@@ -188,6 +207,16 @@ class HIPManagedSpace {
   ~HIPManagedSpace()                                     = default;
 
   /**\brief  Allocate untracked memory in the space */
+  template <typename ExecutionSpace>
+  void* allocate(const ExecutionSpace&, const size_t arg_alloc_size) const {
+    return allocate(arg_alloc_size);
+  }
+  template <typename ExecutionSpace>
+  void* allocate(const ExecutionSpace&, const char* arg_label,
+                 const size_t arg_alloc_size,
+                 const size_t arg_logical_size = 0) const {
+    return allocate(arg_label, arg_alloc_size, arg_logical_size);
+  }
   void* allocate(const size_t arg_alloc_size) const;
   void* allocate(const char* arg_label, const size_t arg_alloc_size,
                  const size_t arg_logical_size = 0) const;

diff --git a/core/src/Kokkos_HostSpace.hpp b/core/src/Kokkos_HostSpace.hpp
@@ -93,6 +93,16 @@ class HostSpace {
 #endif
 
   /**\brief  Allocate untracked memory in the space */
+  template <typename ExecutionSpace>
+  void* allocate(const ExecutionSpace&, const size_t arg_alloc_size) const {
+    return allocate(arg_alloc_size);
+  }
+  template <typename ExecutionSpace>
+  void* allocate(const ExecutionSpace&, const char* arg_label,
+                 const size_t arg_alloc_size,
+                 const size_t arg_logical_size = 0) const {
+    return allocate(arg_label, arg_alloc_size, arg_logical_size);
+  }
   void* allocate(const size_t arg_alloc_size) const;
   void* allocate(const char* arg_label, const size_t arg_alloc_size,
                  const size_t arg_logical_size = 0) const;

diff --git a/core/src/OpenMPTarget/Kokkos_OpenMPTargetSpace.hpp b/core/src/OpenMPTarget/Kokkos_OpenMPTargetSpace.hpp
@@ -98,6 +98,16 @@ class OpenMPTargetSpace {
   ~OpenMPTargetSpace()                                   = default;
 
   /**\brief  Allocate untracked memory in the space */
+  // FIXME_OPENMPTARGET Use execution space instance
+  void* allocate(const OpenMPTarget&, const size_t arg_alloc_size) const {
+    return allocate(arg_alloc_size);
+  }
+  // FIXME_OPENMPTARGET Use execution space instance
+  void* allocate(const OpenMPTarget&, const char* arg_label,
+                 const size_t arg_alloc_size,
+                 const size_t arg_logical_size = 0) const {
+    return allocate(arg_label, arg_alloc_size, arg_logical_size);
+  }
   void* allocate(const size_t arg_alloc_size) const;
   void* allocate(const char* arg_label, const size_t arg_alloc_size,
                  const size_t arg_logical_size = 0) const;

diff --git a/core/src/SYCL/Kokkos_SYCL_Space.hpp b/core/src/SYCL/Kokkos_SYCL_Space.hpp
@@ -82,6 +82,16 @@ class SYCLSharedUSMSpace {
   SYCLSharedUSMSpace();
   explicit SYCLSharedUSMSpace(sycl::queue queue);
 
+  template <typename ExecutionSpace>
+  void* allocate(const ExecutionSpace&, const size_t arg_alloc_size) const {
+    return allocate(arg_alloc_size);
+  }
+  template <typename ExecutionSpace>
+  void* allocate(const ExecutionSpace&, const char* arg_label,
+                 const size_t arg_alloc_size,
+                 const size_t arg_logical_size = 0) const {
+    return allocate(arg_label, arg_alloc_size, arg_logical_size);
+  }
   void* allocate(const SYCL& exec_space,
                  const std::size_t arg_alloc_size) const;
   void* allocate(const SYCL& exec_space, const char* arg_label,
@@ -113,6 +123,16 @@ class SYCLHostUSMSpace {
   SYCLHostUSMSpace();
   explicit SYCLHostUSMSpace(sycl::queue queue);
 
+  template <typename ExecutionSpace>
+  void* allocate(const ExecutionSpace&, const size_t arg_alloc_size) const {
+    return allocate(arg_alloc_size);
+  }
+  template <typename ExecutionSpace>
+  void* allocate(const ExecutionSpace&, const char* arg_label,
+                 const size_t arg_alloc_size,
+                 const size_t arg_logical_size = 0) const {
+    return allocate(arg_label, arg_alloc_size, arg_logical_size);
+  }
   void* allocate(const SYCL& exec_space,
                  const std::size_t arg_alloc_size) const;
   void* allocate(const SYCL& exec_space, const char* arg_label,

diff --git a/core/src/impl/Kokkos_SharedAlloc.hpp b/core/src/impl/Kokkos_SharedAlloc.hpp
@@ -252,14 +252,14 @@ class SharedAllocationRecordCommon : public SharedAllocationRecord<void, void> {
   ~SharedAllocationRecordCommon();
   template <class ExecutionSpace>
   SharedAllocationRecordCommon(
-      ExecutionSpace const&, MemorySpace const& space, std::string const& label,
-      std::size_t alloc_size,
+      ExecutionSpace const& exec, MemorySpace const& space,
+      std::string const& label, std::size_t alloc_size,
       record_base_t::function_type dealloc = &deallocate)
       : SharedAllocationRecord<void, void>(
 #ifdef KOKKOS_ENABLE_DEBUG
             &s_root_record,
 #endif
-            checked_allocation_with_header(space, label, alloc_size),
+            checked_allocation_with_header(exec, space, label, alloc_size),
             sizeof(SharedAllocationHeader) + alloc_size, dealloc, label),
         m_space(space) {
     auto& header = *SharedAllocationRecord<void, void>::m_alloc_ptr;
@@ -315,7 +315,7 @@ class HostInaccessibleSharedAllocationRecordCommon
 #ifdef KOKKOS_ENABLE_DEBUG
             &s_root_record,
 #endif
-            checked_allocation_with_header(space, label, alloc_size),
+            checked_allocation_with_header(exec, space, label, alloc_size),
             sizeof(SharedAllocationHeader) + alloc_size, dealloc, label),
         m_space(space) {
     SharedAllocationHeader header;