Use std::aligned_alloc for allocations (kokkos#6341)

* Use std::aligned_alloc for allocations

* Deprecate unsupported AllocationMechanism enum values

* Fix typo

* Remove unused variable

* size_alignment_multiple -> size

* Remove unused RawMemoryAllocationFailure::AllocationMechanism enums

* Deprecate AllocationMechanism enums instead and handle unsupported cases in switch statement

core/src/Kokkos_HostSpace.hpp

@@ -64,25 +64,27 @@ class HostSpace {
   //! This memory space preferred device_type
   using device_type = Kokkos::Device<execution_space, memory_space>;
 
-  /**\brief  Default memory space instance */
-  HostSpace();
+  HostSpace()                     = default;
   HostSpace(HostSpace&& rhs)      = default;
   HostSpace(const HostSpace& rhs) = default;
   HostSpace& operator=(HostSpace&&) = default;
   HostSpace& operator=(const HostSpace&) = default;
   ~HostSpace()                           = default;
 
+#ifdef KOKKOS_ENABLE_DEPRECATED_CODE_4
   /**\brief  Non-default memory space instance to choose allocation mechansim,
    * if available */
 
-  enum AllocationMechanism {
+  enum KOKKOS_DEPRECATED AllocationMechanism {
     STD_MALLOC,
     POSIX_MEMALIGN,
     POSIX_MMAP,
     INTEL_MM_ALLOC
   };
 
+  KOKKOS_DEPRECATED
   explicit HostSpace(const AllocationMechanism&);
+#endif
 
   /**\brief  Allocate untracked memory in the space */
   void* allocate(const size_t arg_alloc_size) const;
@@ -114,7 +116,6 @@ class HostSpace {
   static constexpr const char* name() { return m_name; }
 
  private:
-  AllocationMechanism m_alloc_mech;
   static constexpr const char* m_name = "Host";
   friend class Kokkos::Impl::SharedAllocationRecord<Kokkos::HostSpace, void>;
 };
@@ -187,7 +188,7 @@ class SharedAllocationRecord<Kokkos::HostSpace, void>
   static RecordBase s_root_record;
 #endif
 
-  const Kokkos::HostSpace m_space;
+  Kokkos::HostSpace m_space;
 
  protected:
   ~SharedAllocationRecord();

core/src/OpenMP/Kokkos_OpenMP_Instance.cpp

@@ -256,8 +256,6 @@ void OpenMPInternal::initialize(int thread_count) {
       }
     }
 
-    OpenMP::memory_space space;
-
     // Before any other call to OMP query the maximum number of threads
     // and save the value for re-initialization unit testing.
 

core/src/impl/Kokkos_Error.cpp

@@ -102,11 +102,6 @@ void Experimental::RawMemoryAllocationFailure::print_error_message(
   o << "  (The allocation mechanism was ";
   switch (m_mechanism) {
     case AllocationMechanism::StdMalloc: o << "standard malloc()."; break;
-    case AllocationMechanism::PosixMemAlign: o << "posix_memalign()."; break;
-    case AllocationMechanism::PosixMMap: o << "POSIX mmap()."; break;
-    case AllocationMechanism::IntelMMAlloc:
-      o << "the Intel _mm_malloc() intrinsic.";
-      break;
     case AllocationMechanism::CudaMalloc: o << "cudaMalloc()."; break;
     case AllocationMechanism::CudaMallocManaged:
       o << "cudaMallocManaged().";
@@ -126,6 +121,7 @@ void Experimental::RawMemoryAllocationFailure::print_error_message(
     case AllocationMechanism::SYCLMallocHost:
       o << "sycl::malloc_host().";
       break;
+    default: o << "unsupported.";
   }
   append_additional_error_information(o);
   o << ")" << std::endl;

core/src/impl/Kokkos_Error.hpp

@@ -113,9 +113,11 @@ class RawMemoryAllocationFailure : public std::bad_alloc {
   };
   enum class AllocationMechanism {
     StdMalloc,
-    PosixMemAlign,
-    PosixMMap,
-    IntelMMAlloc,
+#ifdef KOKKOS_ENABLE_DEPRECATED_CODE_4
+    PosixMemAlign KOKKOS_DEPRECATED,
+    PosixMMap KOKKOS_DEPRECATED,
+    IntelMMAlloc KOKKOS_DEPRECATED,
+#endif
     CudaMalloc,
     CudaMallocManaged,
     CudaHostAlloc,

core/src/impl/Kokkos_HostSpace.cpp

@@ -55,41 +55,9 @@
 
 namespace Kokkos {
 
-/* Default allocation mechanism */
-HostSpace::HostSpace()
-    : m_alloc_mech(
-#if defined(KOKKOS_ENABLE_INTEL_MM_ALLOC)
-          HostSpace::INTEL_MM_ALLOC
-#else
-          HostSpace::STD_MALLOC
-#endif
-      ) {
-}
-
-/* Default allocation mechanism */
-HostSpace::HostSpace(const HostSpace::AllocationMechanism &arg_alloc_mech)
-    : m_alloc_mech(HostSpace::STD_MALLOC) {
-  if (arg_alloc_mech == STD_MALLOC) {
-    m_alloc_mech = HostSpace::STD_MALLOC;
-  }
-#if defined(KOKKOS_ENABLE_INTEL_MM_ALLOC)
-  else if (arg_alloc_mech == HostSpace::INTEL_MM_ALLOC) {
-    m_alloc_mech = HostSpace::INTEL_MM_ALLOC;
-  }
+#ifdef KOKKOS_ENABLE_DEPRECATED_CODE_4
+HostSpace::HostSpace(const HostSpace::AllocationMechanism &) : HostSpace() {}
 #endif
-  else {
-    const char *const mech =
-        (arg_alloc_mech == HostSpace::INTEL_MM_ALLOC)
-            ? "INTEL_MM_ALLOC"
-            : ((arg_alloc_mech == HostSpace::POSIX_MMAP) ? "POSIX_MMAP" : "");
-
-    std::string msg;
-    msg.append("Kokkos::HostSpace ");
-    msg.append(mech);
-    msg.append(" is not available");
-    Kokkos::Impl::throw_runtime_exception(msg);
-  }
-}
 
 void *HostSpace::allocate(const size_t arg_alloc_size) const {
   return allocate("[unlabeled]", arg_alloc_size);
@@ -119,30 +87,12 @@ void *HostSpace::impl_allocate(
   void *ptr = nullptr;
 
   if (arg_alloc_size) {
-    if (m_alloc_mech == STD_MALLOC) {
-      // Over-allocate to and round up to guarantee proper alignment.
-      size_t size_padded = arg_alloc_size + sizeof(void *) + alignment;
-
-      void *alloc_ptr = malloc(size_padded);
-
-      if (alloc_ptr) {
-        auto address = reinterpret_cast<uintptr_t>(alloc_ptr);
-
-        // offset enough to record the alloc_ptr
-        address += sizeof(void *);
-        uintptr_t rem    = address % alignment;
-        uintptr_t offset = rem ? (alignment - rem) : 0u;
-        address += offset;
-        ptr = reinterpret_cast<void *>(address);
-        // record the alloc'd pointer
-        address -= sizeof(void *);
-        *reinterpret_cast<void **>(address) = alloc_ptr;
-      }
-    }
-#if defined(KOKKOS_ENABLE_INTEL_MM_ALLOC)
-    else if (m_alloc_mech == INTEL_MM_ALLOC) {
-      ptr = _mm_malloc(arg_alloc_size, alignment);
-    }
+    // Over-allocate to multiple of alignment
+    size_t size = (arg_alloc_size + alignment - 1) / alignment * alignment;
+#ifdef KOKKOS_COMPILER_MSVC
+    ptr = _aligned_malloc(size, alignment);
+#else
+    ptr = std::aligned_alloc(alignment, size);
 #endif
   }
 
@@ -159,21 +109,6 @@ void *HostSpace::impl_allocate(
     Experimental::RawMemoryAllocationFailure::AllocationMechanism alloc_mec =
         Experimental::RawMemoryAllocationFailure::AllocationMechanism::
             StdMalloc;
-    switch (m_alloc_mech) {
-      case STD_MALLOC: break;  // default
-      case POSIX_MEMALIGN:
-        alloc_mec = Experimental::RawMemoryAllocationFailure::
-            AllocationMechanism::PosixMemAlign;
-        break;
-      case POSIX_MMAP:
-        alloc_mec = Experimental::RawMemoryAllocationFailure::
-            AllocationMechanism::PosixMMap;
-        break;
-      case INTEL_MM_ALLOC:
-        alloc_mec = Experimental::RawMemoryAllocationFailure::
-            AllocationMechanism::IntelMMAlloc;
-        break;
-    }
 
     throw Kokkos::Experimental::RawMemoryAllocationFailure(
         arg_alloc_size, alignment, failure_mode, alloc_mec);
@@ -208,14 +143,10 @@ void HostSpace::impl_deallocate(
       Kokkos::Profiling::deallocateData(arg_handle, arg_label, arg_alloc_ptr,
                                         reported_size);
     }
-    if (m_alloc_mech == STD_MALLOC) {
-      void *alloc_ptr = *(reinterpret_cast<void **>(arg_alloc_ptr) - 1);
-      free(alloc_ptr);
-    }
-#if defined(KOKKOS_ENABLE_INTEL_MM_ALLOC)
-    else if (m_alloc_mech == INTEL_MM_ALLOC) {
-      _mm_free(arg_alloc_ptr);
-    }
+#ifdef KOKKOS_COMPILER_MSVC
+    _aligned_free(arg_alloc_ptr);
+#else
+    std::free(arg_alloc_ptr);
 #endif
   }
 }