diff --git a/openmp/runtime/src/kmp_alloc.cpp b/openmp/runtime/src/kmp_alloc.cpp
--- a/openmp/runtime/src/kmp_alloc.cpp
+++ b/openmp/runtime/src/kmp_alloc.cpp
@@ -1576,6 +1576,36 @@
   // 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
@@ -1645,22 +1675,6 @@
       }
     }
   } 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");
@@ -1683,18 +1697,6 @@
       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 =