Add unified memory arch support for AMD

This is intended for MI300A, but for now you have to set -DKokkos_ENABLE_IMPL_HIP_UNIFIED_MEMORY=ON This marks HIPSpace as host accessible, and thus will make create_mirror_view a no-op.
crtrott · Mar 5, 2024 · 9eff88f · 9eff88f
1 parent 8746494
commit 9eff88f
Show file tree

Hide file tree

Showing 7 changed files with 37 additions and 0 deletions.
diff --git a/algorithms/src/sorting/impl/Kokkos_SortImpl.hpp b/algorithms/src/sorting/impl/Kokkos_SortImpl.hpp
@@ -357,9 +357,14 @@ sort_device_view_with_comparator(
 
   using ViewType = Kokkos::View<DataType, Properties...>;
   using MemSpace = typename ViewType::memory_space;
+// Note with HIP unified memory this code path is still the right thing to do
+// We don't have thrust in yet, and the create_mirror_view_and_copy will do
+// the right thing.
+#ifndef KOKKOS_ENABLE_IMPL_HIP_UNIFIED_MEMORY
   static_assert(!SpaceAccessibility<HostSpace, MemSpace>::accessible,
                 "Impl::sort_device_view_with_comparator: should not be called "
                 "on a view that is already accessible on the host");
+#endif
 
   copy_to_host_run_stdsort_copy_back(exec, view, comparator);
 }

diff --git a/cmake/KokkosCore_config.h.in b/cmake/KokkosCore_config.h.in
@@ -41,6 +41,7 @@
 #cmakedefine KOKKOS_ENABLE_IMPL_CUDA_EMULATE_UNIFIED_MEMORY
 #cmakedefine KOKKOS_ENABLE_HIP_RELOCATABLE_DEVICE_CODE
 #cmakedefine KOKKOS_ENABLE_HIP_MULTIPLE_KERNEL_INSTANTIATIONS
+#cmakedefine KOKKOS_ENABLE_IMPL_HIP_UNIFIED_MEMORY
 #cmakedefine KOKKOS_ENABLE_IMPL_HPX_ASYNC_DISPATCH
 #cmakedefine KOKKOS_ENABLE_DEBUG
 #cmakedefine KOKKOS_ENABLE_DEBUG_DUALVIEW_MODIFY_CHECK

diff --git a/cmake/kokkos_enable_options.cmake b/cmake/kokkos_enable_options.cmake
@@ -73,6 +73,7 @@ KOKKOS_ENABLE_OPTION(TUNING               OFF "Whether to create bindings for tu
 KOKKOS_ENABLE_OPTION(AGGRESSIVE_VECTORIZATION OFF "Whether to aggressively vectorize loops")
 KOKKOS_ENABLE_OPTION(COMPILE_AS_CMAKE_LANGUAGE OFF "Whether to use native cmake language support")
 KOKKOS_ENABLE_OPTION(HIP_MULTIPLE_KERNEL_INSTANTIATIONS OFF "Whether multiple kernels are instantiated at compile time - improve performance but increase compile time")
+KOKKOS_ENABLE_OPTION(IMPL_HIP_UNIFIED_MEMORY OFF "Whether to leverage unified memory architectures for HIP")
 
 # This option will go away eventually, but allows fallback to old implementation when needed.
 KOKKOS_ENABLE_OPTION(DESUL_ATOMICS_EXTERNAL OFF "Whether to use an external desul installation")

diff --git a/core/src/HIP/Kokkos_HIP.cpp b/core/src/HIP/Kokkos_HIP.cpp
@@ -146,6 +146,10 @@ void HIP::print_configuration(std::ostream& os, bool /*verbose*/) const {
 #else
   os << "no\n";
 #endif
+#ifdef KOKKOS_ENABLE_IMPL_HIP_UNIFIED_MEMORY
+  os << "  KOKKOS_ENABLE_IMPL_HIP_UNIFIED_MEMORY: ";
+  os << "yes\n";
+#endif
 
   os << "\nRuntime Configuration:\n";
 

diff --git a/core/src/HIP/Kokkos_HIP_SharedAllocationRecord.cpp b/core/src/HIP/Kokkos_HIP_SharedAllocationRecord.cpp
@@ -23,8 +23,13 @@
 #include <HIP/Kokkos_HIP_SharedAllocationRecord.hpp>
 #include <impl/Kokkos_SharedAlloc_timpl.hpp>
 
+#ifndef KOKKOS_ENABLE_IMPL_HIP_UNIFIED_MEMORY
 KOKKOS_IMPL_HOST_INACCESSIBLE_SHARED_ALLOCATION_RECORD_EXPLICIT_INSTANTIATION(
     Kokkos::HIPSpace);
+#else
+KOKKOS_IMPL_SHARED_ALLOCATION_RECORD_EXPLICIT_INSTANTIATION(
+    Kokkos::HIPSpace);
+#endif
 KOKKOS_IMPL_SHARED_ALLOCATION_RECORD_EXPLICIT_INSTANTIATION(
     Kokkos::HIPHostPinnedSpace);
 KOKKOS_IMPL_SHARED_ALLOCATION_RECORD_EXPLICIT_INSTANTIATION(

diff --git a/core/src/HIP/Kokkos_HIP_SharedAllocationRecord.hpp b/core/src/HIP/Kokkos_HIP_SharedAllocationRecord.hpp
@@ -20,8 +20,12 @@
 #include <HIP/Kokkos_HIP_Space.hpp>
 #include <impl/Kokkos_SharedAlloc.hpp>
 
+#if !defined(KOKKOS_ENABLE_IMPL_HIP_UNIFIED_MEMORY)
 KOKKOS_IMPL_HOST_INACCESSIBLE_SHARED_ALLOCATION_SPECIALIZATION(
     Kokkos::HIPSpace);
+#else
+KOKKOS_IMPL_SHARED_ALLOCATION_SPECIALIZATION(Kokkos::HIPSpace);
+#endif
 KOKKOS_IMPL_SHARED_ALLOCATION_SPECIALIZATION(Kokkos::HIPHostPinnedSpace);
 KOKKOS_IMPL_SHARED_ALLOCATION_SPECIALIZATION(Kokkos::HIPManagedSpace);
 

diff --git a/core/src/HIP/Kokkos_HIP_Space.hpp b/core/src/HIP/Kokkos_HIP_Space.hpp
@@ -65,6 +65,18 @@ class HIPSpace {
   ~HIPSpace()                              = default;
 
   /**\brief  Allocate untracked memory in the hip space */
+#ifdef KOKKOS_ENABLE_IMPL_HIP_UNIFIED_MEMORY
+  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);
+  }
+#else
   // FIXME_HIP Use execution space instance
   void* allocate(const HIP&, const size_t arg_alloc_size) const {
     return allocate(arg_alloc_size);
@@ -74,6 +86,7 @@ class HIPSpace {
                  const size_t arg_logical_size = 0) const {
     return allocate(arg_label, arg_alloc_size, arg_logical_size);
   }
+#endif
   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;
@@ -267,7 +280,11 @@ static_assert(Kokkos::Impl::MemorySpaceAccess<HIPSpace, HIPSpace>::assignable);
 template <>
 struct MemorySpaceAccess<HostSpace, HIPSpace> {
   enum : bool { assignable = false };
+#if !defined(KOKKOS_ENABLE_IMPL_HIP_UNIFIED_MEMORY)
   enum : bool { accessible = false };
+#else
+  enum : bool { accessible = true };
+#endif
   enum : bool { deepcopy = true };
 };