Use cuda_stream/device also for UVM and HostPinned

masterleinad · Jan 9, 2024 · 26f1bb2 · 26f1bb2
1 parent 67b7d68
commit 26f1bb2
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 5 deletions.
diff --git a/core/src/Cuda/Kokkos_CudaSpace.cpp b/core/src/Cuda/Kokkos_CudaSpace.cpp
@@ -513,11 +513,11 @@ SharedAllocationRecord<Kokkos::CudaSpace, void>::SharedAllocationRecord(
   this->base_t::_fill_host_accessible_header_info(header, arg_label);
 
   // Copy to device memory
-  cudaSetDevice(arg_space.device());
+  cudaSetDevice(arg_space.cuda_device());
   cudaMemcpyAsync(RecordBase::m_alloc_ptr, &header,
                   sizeof(SharedAllocationHeader), cudaMemcpyDefault,
-                  arg_space.stream());
-  cudaStreamSynchronize(arg_space.stream());
+                  arg_space.cuda_stream());
+  cudaStreamSynchronize(arg_space.cuda_stream());
 }
 
 SharedAllocationRecord<Kokkos::CudaSpace, void>::SharedAllocationRecord(

diff --git a/core/src/Cuda/Kokkos_CudaSpace.hpp b/core/src/Cuda/Kokkos_CudaSpace.hpp
@@ -118,9 +118,9 @@ class CudaSpace {
   /**\brief Return Name of the MemorySpace */
   static constexpr const char* name() { return m_name; }
 
-  int device() const { return m_device; }
+  int cuda_device() const { return m_device; }
 
-  cudaStream_t stream() const { return m_stream; }
+  cudaStream_t cuda_stream() const { return m_stream; }
 
  private:
   int m_device;
@@ -198,6 +198,10 @@ class CudaUVMSpace {
   /**\brief Return Name of the MemorySpace */
   static constexpr const char* name() { return m_name; }
 
+  int cuda_device() const { return m_device; }
+
+  cudaStream_t cuda_stream() const { return m_stream; }
+
 #ifdef KOKKOS_IMPL_DEBUG_CUDA_PIN_UVM_TO_HOST
   static bool cuda_pin_uvm_to_host();
   static void cuda_set_pin_uvm_to_host(bool val);
@@ -284,6 +288,10 @@ class CudaHostPinnedSpace {
   /**\brief Return Name of the MemorySpace */
   static constexpr const char* name() { return m_name; }
 
+  int cuda_device() const { return m_device; }
+
+  cudaStream_t cuda_stream() const { return m_stream; }
+
  private:
   int m_device;
   cudaStream_t m_stream;