[OpenMP] Use persistent memory for omp_large_cap_mem

[hbae]

Summary

This change enables volatile use of persistent memory for omp_large_cap_mem*
on supported systems. It depends on libmemkind's support for persistent memory,
and requirements/details can be found at the following url.

https://pmem.io/2020/01/20/memkind-dax-kmem.html

[AndreyChurbanov]

LGTM

openmp/runtime/src/kmp_alloc.cpp

 static void **mk_default;
 static void **mk_interleave;
 static void **mk_hbw;
 static void **mk_hbw_interleave;
 static void **mk_hbw_preferred;
 static void **mk_hugetlb;
 static void **mk_hbw_hugetlb;
 static void **mk_hbw_preferred_hugetlb;
+static void **mk_dax_kmem;
+static void **mk_dax_kmem_all;
+static void **mk_dax_kmem_preferred;
 
 #if KMP_OS_UNIX && KMP_DYNAMIC_LIB
 static inline void chk_kind(void ***pkind) {
   KMP_DEBUG_ASSERT(pkind);
   if (*pkind) // symbol found
     if (kmp_mk_check(**pkind)) // kind not available or error
       *pkind = NULL;
 }
if (kmp_mk_check && kmp_mk_alloc && kmp_mk_free && mk_default &&
       chk_kind(&mk_hbw_preferred);
       mk_hugetlb = (void **)dlsym(h_memkind, "MEMKIND_HUGETLB");
       chk_kind(&mk_hugetlb);
       mk_hbw_hugetlb = (void **)dlsym(h_memkind, "MEMKIND_HBW_HUGETLB");
       chk_kind(&mk_hbw_hugetlb);
       mk_hbw_preferred_hugetlb =
           (void **)dlsym(h_memkind, "MEMKIND_HBW_PREFERRED_HUGETLB");
       chk_kind(&mk_hbw_preferred_hugetlb);
+      mk_dax_kmem = (void **)dlsym(h_memkind, "MEMKIND_DAX_KMEM");
+      chk_kind(&mk_dax_kmem);
+      mk_dax_kmem_all = (void **)dlsym(h_memkind, "MEMKIND_DAX_KMEM_ALL");
+      chk_kind(&mk_dax_kmem_all);
+      mk_dax_kmem_preferred =
+          (void **)dlsym(h_memkind, "MEMKIND_DAX_KMEM_PREFERRED");
+      chk_kind(&mk_dax_kmem_preferred);
       KE_TRACE(25, ("__kmp_init_memkind: memkind library initialized\n"));
       return; // success
     }
     dlclose(h_memkind); // failure
-    h_memkind = NULL;
   }
-  kmp_mk_check = NULL;
-  kmp_mk_alloc = NULL;
-  kmp_mk_free = NULL;
-  mk_default = NULL;
-  mk_interleave = NULL;
-  mk_hbw = NULL;
-  mk_hbw_interleave = NULL;
-  mk_hbw_preferred = NULL;
-  mk_hugetlb = NULL;
-  mk_hbw_hugetlb = NULL;
-  mk_hbw_preferred_hugetlb = NULL;
-#else
+#else // !(KMP_OS_UNIX && KMP_DYNAMIC_LIB)
   kmp_mk_lib_name = "";
+#endif // !(KMP_OS_UNIX && KMP_DYNAMIC_LIB)
   h_memkind = NULL;
   kmp_mk_check = NULL;
   kmp_mk_alloc = NULL;
   kmp_mk_free = NULL;
   mk_default = NULL;
   mk_interleave = NULL;
   mk_hbw = NULL;
   mk_hbw_interleave = NULL;
   mk_hbw_preferred = NULL;
   mk_hugetlb = NULL;
   mk_hbw_hugetlb = NULL;
   mk_hbw_preferred_hugetlb = NULL;
-#endif
+  mk_dax_kmem = NULL;
+  mk_dax_kmem_all = NULL;
+  mk_dax_kmem_preferred = NULL;
 }
 
 void __kmp_fini_memkind() {
 #if KMP_OS_UNIX && KMP_DYNAMIC_LIB
   if (__kmp_memkind_available)
     KE_TRACE(25, ("__kmp_fini_memkind: finalize memkind library\n"));
   if (h_memkind) {
     dlclose(h_memkind);
     h_memkind = NULL;
   }
   kmp_mk_check = NULL;
   kmp_mk_alloc = NULL;
   kmp_mk_free = NULL;
   mk_default = NULL;
   mk_interleave = NULL;
   mk_hbw = NULL;
   mk_hbw_interleave = NULL;
   mk_hbw_preferred = NULL;
   mk_hugetlb = NULL;
   mk_hbw_hugetlb = NULL;
   mk_hbw_preferred_hugetlb = NULL;
+  mk_dax_kmem = NULL;
+  mk_dax_kmem_all = NULL;
+  mk_dax_kmem_preferred = NULL;
 #endif
 }
 
 omp_allocator_handle_t __kmpc_init_allocator(int gtid, omp_memspace_handle_t ms,
                                              int ntraits,
                                              omp_alloctrait_t traits[]) {
   // OpenMP 5.0 only allows predefined memspaces
   KMP_DEBUG_ASSERT(ms == omp_default_mem_space || ms == omp_low_lat_mem_space ||
if (ms == omp_high_bw_mem_space) {
         // manual says: "Using this function in production code may result in
         // serious performance penalty".
         al->memkind = mk_hbw_preferred;
       } else {
         // HBW is requested but not available --> return NULL allocator
         __kmp_free(al);
         return omp_null_allocator;
       }
+    } else if (ms == omp_large_cap_mem_space) {
+      if (mk_dax_kmem_all) {
+        // All pmem nodes are visited
+        al->memkind = mk_dax_kmem_all;
+      } else if (mk_dax_kmem) {
+        // Only closest pmem node is visited
+        al->memkind = mk_dax_kmem;
+      } else {
+        __kmp_free(al);
+        return omp_null_allocator;
+      }
     } else {
       if (al->memkind == (void *)omp_atv_interleaved && mk_interleave) {
         al->memkind = mk_interleave;
       } else {
         al->memkind = mk_default;
       }
     }
   } else {
void *__kmpc_alloc(int gtid, size_t size, omp_allocator_handle_t allocator) {
   desc.size_orig = size;
   desc.size_a = size + sz_desc + align;
 
   if (__kmp_memkind_available) {
     if (allocator < kmp_max_mem_alloc) {
       // pre-defined allocator
       if (allocator == omp_high_bw_mem_alloc && mk_hbw_preferred) {
         ptr = kmp_mk_alloc(*mk_hbw_preferred, desc.size_a);
+      } else if (allocator == omp_large_cap_mem_alloc && mk_dax_kmem_all) {
+        ptr = kmp_mk_alloc(*mk_dax_kmem_all, desc.size_a);
       } else {
         ptr = kmp_mk_alloc(*mk_default, desc.size_a);
       }
     } else if (al->pool_size > 0) {
       // custom allocator with pool size requested
       kmp_uint64 used =
           KMP_TEST_THEN_ADD64((kmp_int64 *)&al->pool_used, desc.size_a);
       if (used + desc.size_a > al->pool_size) {
if (allocator < kmp_max_mem_alloc) {
           return __kmpc_alloc(gtid, size, (omp_allocator_handle_t)al);
         }
       }
     }
   } else if (allocator < kmp_max_mem_alloc) {
     // pre-defined allocator
     if (allocator == omp_high_bw_mem_alloc) {
       // ptr = NULL;
+    } else if (allocator == omp_large_cap_mem_alloc) {
+      // warnings?
     } else {
       ptr = __kmp_thread_malloc(__kmp_thread_from_gtid(gtid), desc.size_a);
     }
   } else if (al->pool_size > 0) {
     // custom allocator with pool size requested
     kmp_uint64 used =
         KMP_TEST_THEN_ADD64((kmp_int64 *)&al->pool_used, desc.size_a);
     if (used + desc.size_a > al->pool_size) {
void __kmpc_free(int gtid, void *ptr, const omp_allocator_handle_t allocator) {
   oal = (omp_allocator_handle_t)al; // cast to void* for comparisons
   KMP_DEBUG_ASSERT(al);
 
   if (__kmp_memkind_available) {
     if (oal < kmp_max_mem_alloc) {
       // pre-defined allocator
       if (oal == omp_high_bw_mem_alloc && mk_hbw_preferred) {
         kmp_mk_free(*mk_hbw_preferred, desc.ptr_alloc);
+      } else if (oal == omp_large_cap_mem_alloc && mk_dax_kmem_all) {
+        kmp_mk_free(*mk_dax_kmem_all, desc.ptr_alloc);
       } else {
         kmp_mk_free(*mk_default, desc.ptr_alloc);
       }
     } else {
       if (al->pool_size > 0) { // custom allocator with pool size requested
         kmp_uint64 used =
             KMP_TEST_THEN_ADD64((kmp_int64 *)&al->pool_used, -desc.size_a);
         (void)used; // to suppress compiler warning