Merge pull request #4478 from masterleinad/fix_impl_shared_alloc

Fix impl_shared_alloc for GPU+threads
kokkos · Oct 28, 2021 · 8b17f8e · 8b17f8e
2 parents 78f0bd5 + 43ada2b
commit 8b17f8e
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 17 deletions.
diff --git a/core/src/Cuda/Kokkos_CudaSpace.cpp b/core/src/Cuda/Kokkos_CudaSpace.cpp
@@ -492,9 +492,10 @@ SharedAllocationRecord<Kokkos::CudaSpace, void>::~SharedAllocationRecord() {
   const char *label = nullptr;
   if (Kokkos::Profiling::profileLibraryLoaded()) {
     SharedAllocationHeader header;
+    Kokkos::Cuda exec;
     Kokkos::Impl::DeepCopy<Kokkos::CudaSpace, HostSpace>(
-        &header, RecordBase::m_alloc_ptr, sizeof(SharedAllocationHeader));
-    Kokkos::fence(
+        exec, &header, RecordBase::m_alloc_ptr, sizeof(SharedAllocationHeader));
+    exec.fence(
         "SharedAllocationRecord<Kokkos::CudaSpace, "
         "void>::~SharedAllocationRecord(): fence after copying header from "
         "HostSpace");
@@ -552,9 +553,10 @@ SharedAllocationRecord<Kokkos::CudaSpace, void>::SharedAllocationRecord(
   this->base_t::_fill_host_accessible_header_info(header, arg_label);
 
   // Copy to device memory
-  Kokkos::Impl::DeepCopy<CudaSpace, HostSpace>(RecordBase::m_alloc_ptr, &header,
-                                               sizeof(SharedAllocationHeader));
-  Kokkos::fence(
+  Kokkos::Cuda exec;
+  Kokkos::Impl::DeepCopy<CudaSpace, HostSpace>(
+      exec, RecordBase::m_alloc_ptr, &header, sizeof(SharedAllocationHeader));
+  exec.fence(
       "SharedAllocationRecord<Kokkos::CudaSpace, "
       "void>::SharedAllocationRecord(): fence after copying header from "
       "HostSpace");

diff --git a/core/src/HIP/Kokkos_HIP_Space.cpp b/core/src/HIP/Kokkos_HIP_Space.cpp
@@ -265,9 +265,10 @@ SharedAllocationRecord<Kokkos::Experimental::HIPSpace,
   const char* label = nullptr;
   if (Kokkos::Profiling::profileLibraryLoaded()) {
     SharedAllocationHeader header;
+    Kokkos::Experimental::HIP exec;
     Kokkos::Impl::DeepCopy<Kokkos::Experimental::HIPSpace, HostSpace>(
-        &header, RecordBase::m_alloc_ptr, sizeof(SharedAllocationHeader));
-    Kokkos::fence(
+        exec, &header, RecordBase::m_alloc_ptr, sizeof(SharedAllocationHeader));
+    exec.fence(
         "SharedAllocationRecord<Kokkos::Experimental::HIPSpace, "
         "void>::~SharedAllocationRecord(): fence after copying header from "
         "HostSpace");
@@ -307,9 +308,10 @@ SharedAllocationRecord<Kokkos::Experimental::HIPSpace, void>::
   this->base_t::_fill_host_accessible_header_info(header, arg_label);
 
   // Copy to device memory
+  Kokkos::Experimental::HIP exec;
   Kokkos::Impl::DeepCopy<Kokkos::Experimental::HIPSpace, HostSpace>(
-      RecordBase::m_alloc_ptr, &header, sizeof(SharedAllocationHeader));
-  Kokkos::fence(
+      exec, RecordBase::m_alloc_ptr, &header, sizeof(SharedAllocationHeader));
+  exec.fence(
       "SharedAllocationRecord<Kokkos::Experimental::HIPSpace, "
       "void>::SharedAllocationRecord(): fence after copying header from "
       "HostSpace");

diff --git a/core/src/SYCL/Kokkos_SYCL_Space.cpp b/core/src/SYCL/Kokkos_SYCL_Space.cpp
@@ -252,9 +252,10 @@ SharedAllocationRecord<Kokkos::Experimental::SYCLDeviceUSMSpace, void>::
   this->base_t::_fill_host_accessible_header_info(header, label);
 
   // Copy to device memory
+  Kokkos::Experimental::SYCL exec;
   Kokkos::Impl::DeepCopy<Kokkos::Experimental::SYCLDeviceUSMSpace, HostSpace>(
-      RecordBase::m_alloc_ptr, &header, sizeof(SharedAllocationHeader));
-  Kokkos::fence(
+      exec, RecordBase::m_alloc_ptr, &header, sizeof(SharedAllocationHeader));
+  exec.fence(
       "SharedAllocationRecord<Kokkos::Experimental::SYCLDeviceUSMSpace, "
       "void>::SharedAllocationRecord(): fence after copying header from "
       "HostSpace");
@@ -316,10 +317,11 @@ SharedAllocationRecord<Kokkos::Experimental::SYCLDeviceUSMSpace,
   const char* label = nullptr;
   if (Kokkos::Profiling::profileLibraryLoaded()) {
     SharedAllocationHeader header;
+    Kokkos::Experimental::SYCL exec;
     Kokkos::Impl::DeepCopy<Kokkos::Experimental::SYCLDeviceUSMSpace,
-                           Kokkos::HostSpace>(&header, RecordBase::m_alloc_ptr,
-                                              sizeof(SharedAllocationHeader));
-    Kokkos::fence(
+                           Kokkos::HostSpace>(
+        exec, &header, RecordBase::m_alloc_ptr, sizeof(SharedAllocationHeader));
+    exec.fence(
         "SharedAllocationRecord<Kokkos::Experimental::SYCLDeviceUSMSpace, "
         "void>::~SharedAllocationRecord(): fence after copying header from "
         "HostSpace");

diff --git a/core/src/impl/Kokkos_SharedAlloc_timpl.hpp b/core/src/impl/Kokkos_SharedAlloc_timpl.hpp
@@ -262,9 +262,10 @@ auto HostInaccessibleSharedAllocationRecordCommon<MemorySpace>::get_record(
       alloc_ptr ? SharedAllocationHeader::get_header(alloc_ptr) : nullptr;
 
   if (alloc_ptr) {
-    Kokkos::Impl::DeepCopy<HostSpace, MemorySpace>(
-        &head, head_cuda, sizeof(SharedAllocationHeader));
-    Kokkos::fence(
+    typename MemorySpace::execution_space exec_space;
+    Kokkos::Impl::DeepCopy<HostSpace, MemorySpace, decltype(exec_space)>(
+        exec_space, &head, head_cuda, sizeof(SharedAllocationHeader));
+    exec_space.fence(
         "HostInaccessibleSharedAllocationRecordCommon::get_record(): fence "
         "after copying header to HostSpace");
   }