[OpenMP] Add check for target allocator regardless of the availability of libmemkind

[Nawrin]

Summary

Our current implementation only checks for target allocator when libmemkind is not available. This patch adds check for target allocator regardless of the presence of libmemkind.

[jlpeyton]

LGTM

openmp/runtime/src/kmp_alloc.cpp

void *__kmp_alloc(int gtid, size_t algn, size_t size,
   desc.size_a = size + sz_desc + align;
   bool is_pinned = false;
   if (allocator > kmp_max_mem_alloc)
     is_pinned = al->pinned;
 
   // Use default allocator if libmemkind is not available
   int use_default_allocator = (__kmp_memkind_available) ? false : true;
 
+  if (KMP_IS_TARGET_MEM_ALLOC(allocator)) {
+    // Use size input directly as the memory may not be accessible on host.
+    // Use default device for now.
+    if (__kmp_target_mem_available) {
+      kmp_int32 device =
+          __kmp_threads[gtid]->th.th_current_task->td_icvs.default_device;
+      if (allocator == llvm_omp_target_host_mem_alloc)
+        ptr = kmp_target_alloc_host(size, device);
+      else if (allocator == llvm_omp_target_shared_mem_alloc)
+        ptr = kmp_target_alloc_shared(size, device);
+      else // allocator == llvm_omp_target_device_mem_alloc
+        ptr = kmp_target_alloc_device(size, device);
+    }
+    return ptr;
+  }
+
+  if (allocator >= kmp_max_mem_alloc && KMP_IS_TARGET_MEM_SPACE(al->memspace)) {
+    if (__kmp_target_mem_available) {
+      kmp_int32 device =
+          __kmp_threads[gtid]->th.th_current_task->td_icvs.default_device;
+      if (al->memspace == llvm_omp_target_host_mem_space)
+        ptr = kmp_target_alloc_host(size, device);
+      else if (al->memspace == llvm_omp_target_shared_mem_space)
+        ptr = kmp_target_alloc_shared(size, device);
+      else // al->memspace == llvm_omp_target_device_mem_space
+        ptr = kmp_target_alloc_device(size, device);
+    }
+    return ptr;
+  }
+
   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 {
if (allocator < kmp_max_mem_alloc) {
           ptr = __kmp_alloc(gtid, algn, size, (omp_allocator_handle_t)al);
           if (is_pinned && kmp_target_lock_mem)
             kmp_target_lock_mem(ptr, size, default_device);
           return ptr;
         }
       }
     }
   } else if (allocator < kmp_max_mem_alloc) {
-    if (KMP_IS_TARGET_MEM_ALLOC(allocator)) {
-      // Use size input directly as the memory may not be accessible on host.
-      // Use default device for now.
-      if (__kmp_target_mem_available) {
-        kmp_int32 device =
-            __kmp_threads[gtid]->th.th_current_task->td_icvs.default_device;
-        if (allocator == llvm_omp_target_host_mem_alloc)
-          ptr = kmp_target_alloc_host(size, device);
-        else if (allocator == llvm_omp_target_shared_mem_alloc)
-          ptr = kmp_target_alloc_shared(size, device);
-        else // allocator == llvm_omp_target_device_mem_alloc
-          ptr = kmp_target_alloc_device(size, device);
-      }
-      return ptr;
-    }
-
     // pre-defined allocator
     if (allocator == omp_high_bw_mem_alloc) {
       KMP_WARNING(OmpNoAllocator, "omp_high_bw_mem_alloc");
     } else if (allocator == omp_large_cap_mem_alloc) {
       KMP_WARNING(OmpNoAllocator, "omp_large_cap_mem_alloc");
     } else if (allocator == omp_const_mem_alloc) {
       KMP_WARNING(OmpNoAllocator, "omp_const_mem_alloc");
     } else if (allocator == omp_low_lat_mem_alloc) {
       KMP_WARNING(OmpNoAllocator, "omp_low_lat_mem_alloc");
     } else if (allocator == omp_cgroup_mem_alloc) {
       KMP_WARNING(OmpNoAllocator, "omp_cgroup_mem_alloc");
     } else if (allocator == omp_pteam_mem_alloc) {
       KMP_WARNING(OmpNoAllocator, "omp_pteam_mem_alloc");
     } else if (allocator == omp_thread_mem_alloc) {
       KMP_WARNING(OmpNoAllocator, "omp_thread_mem_alloc");
     } else { // default allocator requested
       use_default_allocator = true;
     }
     if (use_default_allocator) {
       ptr = __kmp_thread_malloc(__kmp_thread_from_gtid(gtid), desc.size_a);
       use_default_allocator = false;
     }
-  } else if (KMP_IS_TARGET_MEM_SPACE(al->memspace)) {
-    if (__kmp_target_mem_available) {
-      kmp_int32 device =
-          __kmp_threads[gtid]->th.th_current_task->td_icvs.default_device;
-      if (al->memspace == llvm_omp_target_host_mem_space)
-        ptr = kmp_target_alloc_host(size, device);
-      else if (al->memspace == llvm_omp_target_shared_mem_space)
-        ptr = kmp_target_alloc_shared(size, device);
-      else // al->memspace == llvm_omp_target_device_mem_space
-        ptr = kmp_target_alloc_device(size, device);
-    }
-    return ptr;
   } 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) {
       // not enough space, need to go fallback path
       KMP_TEST_THEN_ADD64((kmp_int64 *)&al->pool_used, -desc.size_a);
       if (al->fb == omp_atv_default_mem_fb) {