[OpenMP] Add allocator support for target memory

[hbae]

Summary

This is a preview of allocator support for target memory that depends on the
offload runtime API which allocates memory as described below.

llvm_omp_target_alloc_host(size_t size, int device_num);

• Returns non-migratable memory owned by host.
• Memory is accessible by host and device(s).

llvm_omp_target_alloc_shared(size_t size, int device_num);

• Returns migratable memory owned by host and device.
• Memory is accessible by host and device.

llvm_omp_target_alloc_device(size_t size, int device_num);

• Returns memory owned by device.
• Memory is only accessible by device.

New memory space and predefined allocator names are

• llvm_omp_target_host_mem_space
• llvm_omp_target_shared_mem_space
• llvm_omp_target_device_mem_space
• llvm_omp_target_host_mem_alloc
• llvm_omp_target_shared_mem_alloc
• llvm_omp_target_device_mem_alloc

[hbae]

Changed prefix from preomp to llvm_omp.

[jdoerfert]

Do you also have the API calls and libomptarget integration?

[hbae]

There will be a separate patch for libomptarget.

[jdoerfert]

LG, make sure to clang format the code.

openmp/runtime/src/dllexports

kmp_set_disp_num_buffers                    890
     omp_default_mem_alloc                  DATA
     omp_large_cap_mem_alloc                DATA
     omp_const_mem_alloc                    DATA
     omp_high_bw_mem_alloc                  DATA
     omp_low_lat_mem_alloc                  DATA
     omp_cgroup_mem_alloc                   DATA
     omp_pteam_mem_alloc                    DATA
     omp_thread_mem_alloc                   DATA
+    # Preview of target memory support
+    llvm_omp_target_host_mem_alloc         DATA
+    llvm_omp_target_shared_mem_alloc       DATA
+    llvm_omp_target_device_mem_alloc       DATA
 
     omp_default_mem_space                  DATA
     omp_large_cap_mem_space                DATA
     omp_const_mem_space                    DATA
     omp_high_bw_mem_space                  DATA
     omp_low_lat_mem_space                  DATA
+    # Preview of target memory support
+    llvm_omp_target_host_mem_space         DATA
+    llvm_omp_target_shared_mem_space       DATA
+    llvm_omp_target_device_mem_space       DATA
 
 %ifndef stub
     # Ordinals between 900 and 999 are reserved
 
     # Ordinals between 1000 and 1999 are reserved
     # for user-callable uppercase Fortran entries.
 
 

openmp/runtime/src/include/omp.h.var

#   if defined(_WIN32)
     extern __KMP_IMP omp_allocator_handle_t const omp_default_mem_alloc;
     extern __KMP_IMP omp_allocator_handle_t const omp_large_cap_mem_alloc;
     extern __KMP_IMP omp_allocator_handle_t const omp_const_mem_alloc;
     extern __KMP_IMP omp_allocator_handle_t const omp_high_bw_mem_alloc;
     extern __KMP_IMP omp_allocator_handle_t const omp_low_lat_mem_alloc;
     extern __KMP_IMP omp_allocator_handle_t const omp_cgroup_mem_alloc;
     extern __KMP_IMP omp_allocator_handle_t const omp_pteam_mem_alloc;
     extern __KMP_IMP omp_allocator_handle_t const omp_thread_mem_alloc;
+    /* Preview of target memory support */
+    extern __KMP_IMP omp_allocator_handle_t const llvm_omp_target_host_mem_alloc;
+    extern __KMP_IMP omp_allocator_handle_t const llvm_omp_target_shared_mem_alloc;
+    extern __KMP_IMP omp_allocator_handle_t const llvm_omp_target_device_mem_alloc;
+
     typedef omp_uintptr_t omp_memspace_handle_t;
     extern __KMP_IMP omp_memspace_handle_t const omp_default_mem_space;
     extern __KMP_IMP omp_memspace_handle_t const omp_large_cap_mem_space;
     extern __KMP_IMP omp_memspace_handle_t const omp_const_mem_space;
     extern __KMP_IMP omp_memspace_handle_t const omp_high_bw_mem_space;
     extern __KMP_IMP omp_memspace_handle_t const omp_low_lat_mem_space;
+    /* Preview of target memory support */
+    extern __KMP_IMP omp_memspace_handle_t const llvm_omp_target_host_mem_space;
+    extern __KMP_IMP omp_memspace_handle_t const llvm_omp_target_shared_mem_space;
+    extern __KMP_IMP omp_memspace_handle_t const llvm_omp_target_device_mem_space;
 #   else
 #       if __cplusplus >= 201103
     typedef enum omp_allocator_handle_t : omp_uintptr_t
 #       else
     typedef enum omp_allocator_handle_t
 #       endif
     {
       omp_null_allocator = 0,
       omp_default_mem_alloc = 1,
       omp_large_cap_mem_alloc = 2,
       omp_const_mem_alloc = 3,
       omp_high_bw_mem_alloc = 4,
       omp_low_lat_mem_alloc = 5,
       omp_cgroup_mem_alloc = 6,
       omp_pteam_mem_alloc = 7,
       omp_thread_mem_alloc = 8,
+      /* Preview of target memory support */
+      llvm_omp_target_host_mem_alloc = 100,
+      llvm_omp_target_shared_mem_alloc = 101,
+      llvm_omp_target_device_mem_alloc = 102,
       KMP_ALLOCATOR_MAX_HANDLE = UINTPTR_MAX
     } omp_allocator_handle_t;
 #       if __cplusplus >= 201103
     typedef enum omp_memspace_handle_t : omp_uintptr_t
 #       else
     typedef enum omp_memspace_handle_t
 #       endif
     {
       omp_default_mem_space = 0,
       omp_large_cap_mem_space = 1,
       omp_const_mem_space = 2,
       omp_high_bw_mem_space = 3,
       omp_low_lat_mem_space = 4,
+      /* Preview of target memory support */
+      llvm_omp_target_host_mem_space = 100,
+      llvm_omp_target_shared_mem_space = 101,
+      llvm_omp_target_device_mem_space = 102,
       KMP_MEMSPACE_MAX_HANDLE = UINTPTR_MAX
     } omp_memspace_handle_t;
 #   endif
     extern omp_allocator_handle_t __KAI_KMPC_CONVENTION omp_init_allocator(omp_memspace_handle_t m,
                                                        int ntraits, omp_alloctrait_t traits[]);
     extern void __KAI_KMPC_CONVENTION omp_destroy_allocator(omp_allocator_handle_t allocator);
 
     extern void __KAI_KMPC_CONVENTION omp_set_default_allocator(omp_allocator_handle_t a);

openmp/runtime/src/include/omp_lib.h.var

!
       integer(kind=omp_allocator_handle_kind)omp_low_lat_mem_alloc
       parameter(omp_low_lat_mem_alloc=5)
       integer(kind=omp_allocator_handle_kind)omp_cgroup_mem_alloc
       parameter(omp_cgroup_mem_alloc=6)
       integer(kind=omp_allocator_handle_kind)omp_pteam_mem_alloc
       parameter(omp_pteam_mem_alloc=7)
       integer(kind=omp_allocator_handle_kind)omp_thread_mem_alloc
       parameter(omp_thread_mem_alloc=8)
+      ! Preview of target memory support
+      integer(kind=omp_allocator_handle_kind)llvm_omp_target_host_mem_alloc
+      parameter(llvm_omp_target_host_mem_alloc=100)
+      integer(kind=omp_allocator_handle_kind)llvm_omp_target_shared_mem_alloc
+      parameter(llvm_omp_target_shared_mem_alloc=101)
+      integer(kind=omp_allocator_handle_kind)llvm_omp_target_device_mem_alloc
+      parameter(llvm_omp_target_device_mem_alloc=102)
 
       integer(kind=omp_memspace_handle_kind)omp_default_mem_space
       parameter(omp_default_mem_space=0)
       integer(kind=omp_memspace_handle_kind)omp_large_cap_mem_space
       parameter(omp_large_cap_mem_space=1)
       integer(kind=omp_memspace_handle_kind)omp_const_mem_space
       parameter(omp_const_mem_space=2)
       integer(kind=omp_memspace_handle_kind)omp_high_bw_mem_space
       parameter(omp_high_bw_mem_space=3)
       integer(kind=omp_memspace_handle_kind)omp_low_lat_mem_space
       parameter(omp_low_lat_mem_space=4)
+      ! Preview of target memory support
+      integer(kind=omp_memspace_handle_kind)llvm_omp_target_host_mem_space
+      parameter(llvm_omp_target_host_mem_space=100)
+      integer(kind=omp_memspace_handle_kind)llvm_omp_target_shared_mem_space
+      parameter(llvm_omp_target_shared_mem_space=101)
+      integer(kind=omp_memspace_handle_kind)llvm_omp_target_device_mem_space
+      parameter(llvm_omp_target_device_mem_space=102)
 
       integer(kind=omp_pause_resource_kind)omp_pause_resume
       parameter(omp_pause_resume=0)
       integer(kind=omp_pause_resource_kind)omp_pause_soft
       parameter(omp_pause_soft=1)
       integer(kind=omp_pause_resource_kind)omp_pause_hard
       parameter(omp_pause_hard=2)
 

openmp/runtime/src/include/omp_lib.f90.var

!
         integer (kind=omp_allocator_handle_kind), parameter :: omp_default_mem_alloc = 1
         integer (kind=omp_allocator_handle_kind), parameter :: omp_large_cap_mem_alloc = 2
         integer (kind=omp_allocator_handle_kind), parameter :: omp_const_mem_alloc = 3
         integer (kind=omp_allocator_handle_kind), parameter :: omp_high_bw_mem_alloc = 4
         integer (kind=omp_allocator_handle_kind), parameter :: omp_low_lat_mem_alloc = 5
         integer (kind=omp_allocator_handle_kind), parameter :: omp_cgroup_mem_alloc = 6
         integer (kind=omp_allocator_handle_kind), parameter :: omp_pteam_mem_alloc = 7
         integer (kind=omp_allocator_handle_kind), parameter :: omp_thread_mem_alloc = 8
+        ! Preview of target memory support
+        integer (kind=omp_allocator_handle_kind), parameter :: llvm_omp_target_host_mem_alloc = 100
+        integer (kind=omp_allocator_handle_kind), parameter :: llvm_omp_target_shared_mem_alloc = 101
+        integer (kind=omp_allocator_handle_kind), parameter :: llvm_omp_target_device_mem_alloc = 102
 
         integer (kind=omp_memspace_handle_kind), parameter :: omp_default_mem_space = 0
         integer (kind=omp_memspace_handle_kind), parameter :: omp_large_cap_mem_space = 1
         integer (kind=omp_memspace_handle_kind), parameter :: omp_const_mem_space = 2
         integer (kind=omp_memspace_handle_kind), parameter :: omp_high_bw_mem_space = 3
         integer (kind=omp_memspace_handle_kind), parameter :: omp_low_lat_mem_space = 4
+        ! Preview of target memory support
+        integer (kind=omp_memspace_handle_kind), parameter :: llvm_omp_target_host_mem_space = 100
+        integer (kind=omp_memspace_handle_kind), parameter :: llvm_omp_target_shared_mem_space = 101
+        integer (kind=omp_memspace_handle_kind), parameter :: llvm_omp_target_device_mem_space = 102
 
         integer (kind=omp_pause_resource_kind), parameter :: omp_pause_resume = 0
         integer (kind=omp_pause_resource_kind), parameter :: omp_pause_soft = 1
         integer (kind=omp_pause_resource_kind), parameter :: omp_pause_hard = 2
 
         interface
 
 !         ***

openmp/runtime/src/kmp.h

 #define omp_atv_default ((omp_uintptr_t)-1)
 
 typedef void *omp_memspace_handle_t;
 extern omp_memspace_handle_t const omp_default_mem_space;
 extern omp_memspace_handle_t const omp_large_cap_mem_space;
 extern omp_memspace_handle_t const omp_const_mem_space;
 extern omp_memspace_handle_t const omp_high_bw_mem_space;
 extern omp_memspace_handle_t const omp_low_lat_mem_space;
+// Preview of target memory support
+extern omp_memspace_handle_t const llvm_omp_target_host_mem_space;
+extern omp_memspace_handle_t const llvm_omp_target_shared_mem_space;
+extern omp_memspace_handle_t const llvm_omp_target_device_mem_space;
 
 typedef struct {
   omp_alloctrait_key_t key;
   omp_uintptr_t value;
 } omp_alloctrait_t;
 
 typedef void *omp_allocator_handle_t;
 extern omp_allocator_handle_t const omp_null_allocator;
 extern omp_allocator_handle_t const omp_default_mem_alloc;
 extern omp_allocator_handle_t const omp_large_cap_mem_alloc;
 extern omp_allocator_handle_t const omp_const_mem_alloc;
 extern omp_allocator_handle_t const omp_high_bw_mem_alloc;
 extern omp_allocator_handle_t const omp_low_lat_mem_alloc;
 extern omp_allocator_handle_t const omp_cgroup_mem_alloc;
 extern omp_allocator_handle_t const omp_pteam_mem_alloc;
 extern omp_allocator_handle_t const omp_thread_mem_alloc;
+// Preview of target memory support
+extern omp_allocator_handle_t const llvm_omp_target_host_mem_alloc;
+extern omp_allocator_handle_t const llvm_omp_target_shared_mem_alloc;
+extern omp_allocator_handle_t const llvm_omp_target_device_mem_alloc;
 extern omp_allocator_handle_t const kmp_max_mem_alloc;
 extern omp_allocator_handle_t __kmp_def_allocator;
 
 // end of duplicate type definitions from omp.h
 #endif
 
 extern int __kmp_memkind_available;
 
extern void *__kmpc_calloc(int gtid, size_t nmemb, size_t sz,
                            omp_allocator_handle_t al);
 extern void *__kmpc_realloc(int gtid, void *ptr, size_t sz,
                             omp_allocator_handle_t al,
                             omp_allocator_handle_t free_al);
 extern void __kmpc_free(int gtid, void *ptr, omp_allocator_handle_t al);
 
 extern void __kmp_init_memkind();
 extern void __kmp_fini_memkind();
+extern void __kmp_init_target_mem();
 
 /* ------------------------------------------------------------------------ */
 
 #define KMP_UINT64_MAX                                                         \
   (~((kmp_uint64)1 << ((sizeof(kmp_uint64) * (1 << 3)) - 1)))
 
 #define KMP_MIN_NTH 1
 

openmp/runtime/src/kmp_alloc.cpp

 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;
+// Preview of target memory support
+static void *(*kmp_target_alloc_host)(size_t size, int device);
+static void *(*kmp_target_alloc_shared)(size_t size, int device);
+static void *(*kmp_target_alloc_device)(size_t size, int device);
+static void *(*kmp_target_free)(void *ptr, int device);
+static bool __kmp_target_mem_available;
+#define KMP_IS_TARGET_MEM_SPACE(MS)                                            \
+  (MS == llvm_omp_target_host_mem_space ||                                     \
+   MS == llvm_omp_target_shared_mem_space ||                                   \
+   MS == llvm_omp_target_device_mem_space)
+#define KMP_IS_TARGET_MEM_ALLOC(MA)                                            \
+  (MA == llvm_omp_target_host_mem_alloc ||                                     \
+   MA == llvm_omp_target_shared_mem_alloc ||                                   \
+   MA == llvm_omp_target_device_mem_alloc)
 
 #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_OS_UNIX && KMP_DYNAMIC_LIB
   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
 }
+// Preview of target memory support
+void __kmp_init_target_mem() {
+  *(void **)(&kmp_target_alloc_host) = KMP_DLSYM("llvm_omp_target_alloc_host");
+  *(void **)(&kmp_target_alloc_shared) =
+      KMP_DLSYM("llvm_omp_target_alloc_shared");
+  *(void **)(&kmp_target_alloc_device) =
+      KMP_DLSYM("llvm_omp_target_alloc_device");
+  *(void **)(&kmp_target_free) = KMP_DLSYM("omp_target_free");
+  __kmp_target_mem_available = kmp_target_alloc_host &&
+                               kmp_target_alloc_shared &&
+                               kmp_target_alloc_device && kmp_target_free;
+}
 
 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 ||
                    ms == omp_large_cap_mem_space || ms == omp_const_mem_space ||
-                   ms == omp_high_bw_mem_space);
+                   ms == omp_high_bw_mem_space || KMP_IS_TARGET_MEM_SPACE(ms));
   kmp_allocator_t *al;
   int i;
   al = (kmp_allocator_t *)__kmp_allocate(sizeof(kmp_allocator_t)); // zeroed
   al->memspace = ms; // not used currently
   for (i = 0; i < ntraits; ++i) {
     switch (traits[i].key) {
     case omp_atk_sync_hint:
     case omp_atk_access:
if (ms == omp_high_bw_mem_space) {
       }
     } else {
       if (al->memkind == (void *)omp_atv_interleaved && mk_interleave) {
         al->memkind = mk_interleave;
       } else {
         al->memkind = mk_default;
       }
     }
+  } else if (KMP_IS_TARGET_MEM_SPACE(ms) && !__kmp_target_mem_available) {
+    __kmp_free(al);
+    return omp_null_allocator;
   } else {
     if (ms == omp_high_bw_mem_space) {
       // cannot detect HBW memory presence without memkind library
       __kmp_free(al);
       return omp_null_allocator;
     }
   }
   return (omp_allocator_handle_t)al;
if (allocator < kmp_max_mem_alloc) {
         } else if (al->fb == omp_atv_allocator_fb) {
           KMP_ASSERT(al != al->fb_data);
           al = al->fb_data;
           return __kmpc_alloc(gtid, size, (omp_allocator_handle_t)al);
         }
       }
     }
   } 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) {
       // 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 (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) {
if (ptr == NULL)
     return;
 
   kmp_allocator_t *al;
   omp_allocator_handle_t oal;
   al = RCAST(kmp_allocator_t *, CCAST(omp_allocator_handle_t, allocator));
   kmp_mem_desc_t desc;
   kmp_uintptr_t addr_align; // address to return to caller
   kmp_uintptr_t addr_descr; // address of memory block descriptor
+  if (KMP_IS_TARGET_MEM_ALLOC(allocator) ||
+      (allocator > kmp_max_mem_alloc &&
+       KMP_IS_TARGET_MEM_SPACE(al->memspace))) {
+    KMP_DEBUG_ASSERT(kmp_target_free);
+    kmp_int32 device =
+        __kmp_threads[gtid]->th.th_current_task->td_icvs.default_device;
+    kmp_target_free(ptr, device);
+    return;
+  }
 
   addr_align = (kmp_uintptr_t)ptr;
   addr_descr = addr_align - sizeof(kmp_mem_desc_t);
   desc = *((kmp_mem_desc_t *)addr_descr); // read descriptor
 
   KMP_DEBUG_ASSERT(desc.ptr_align == ptr);
   if (allocator) {
     KMP_DEBUG_ASSERT(desc.allocator == al || desc.allocator == al->fb_data);

openmp/runtime/src/kmp_global.cpp

 omp_allocator_handle_t const omp_low_lat_mem_alloc =
     (omp_allocator_handle_t const)5;
 omp_allocator_handle_t const omp_cgroup_mem_alloc =
     (omp_allocator_handle_t const)6;
 omp_allocator_handle_t const omp_pteam_mem_alloc =
     (omp_allocator_handle_t const)7;
 omp_allocator_handle_t const omp_thread_mem_alloc =
     (omp_allocator_handle_t const)8;
+// Preview of target memory support
+omp_allocator_handle_t const llvm_omp_target_host_mem_alloc =
+    (omp_allocator_handle_t const)100;
+omp_allocator_handle_t const llvm_omp_target_shared_mem_alloc =
+    (omp_allocator_handle_t const)101;
+omp_allocator_handle_t const llvm_omp_target_device_mem_alloc =
+    (omp_allocator_handle_t const)102;
 omp_allocator_handle_t const kmp_max_mem_alloc =
     (omp_allocator_handle_t const)1024;
 omp_allocator_handle_t __kmp_def_allocator = omp_default_mem_alloc;
 
 omp_memspace_handle_t const omp_default_mem_space =
     (omp_memspace_handle_t const)0;
 omp_memspace_handle_t const omp_large_cap_mem_space =
     (omp_memspace_handle_t const)1;
 omp_memspace_handle_t const omp_const_mem_space =
     (omp_memspace_handle_t const)2;
 omp_memspace_handle_t const omp_high_bw_mem_space =
     (omp_memspace_handle_t const)3;
 omp_memspace_handle_t const omp_low_lat_mem_space =
     (omp_memspace_handle_t const)4;
+// Preview of target memory support
+omp_memspace_handle_t const llvm_omp_target_host_mem_space =
+    (omp_memspace_handle_t const)100;
+omp_memspace_handle_t const llvm_omp_target_shared_mem_space =
+    (omp_memspace_handle_t const)101;
+omp_memspace_handle_t const llvm_omp_target_device_mem_space =
+    (omp_memspace_handle_t const)102;
 
 /* This check ensures that the compiler is passing the correct data type for the
    flags formal parameter of the function kmpc_omp_task_alloc(). If the type is
    not a 4-byte type, then give an error message about a non-positive length
    array pointing here.  If that happens, the kmp_tasking_flags_t structure must
    be redefined to have exactly 32 bits. */
 KMP_BUILD_ASSERT(sizeof(kmp_tasking_flags_t) == 4);
 

openmp/runtime/src/kmp_runtime.cpp

__kmp_print_storage_map_gtid(
       sizeof(kmp_info_t *) * num_thr, "%s_%d.t_threads", header, team_id);
 
   __kmp_print_storage_map_gtid(-1, &team->t.t_disp_buffer[0],
                                &team->t.t_disp_buffer[num_disp_buff],
                                sizeof(dispatch_shared_info_t) * num_disp_buff,
                                "%s_%d.t_disp_buffer", header, team_id);
 }
 
-static void __kmp_init_allocator() { __kmp_init_memkind(); }
+static void __kmp_init_allocator() {
+  __kmp_init_memkind();
+  __kmp_init_target_mem();
+}
 static void __kmp_fini_allocator() { __kmp_fini_memkind(); }
 
 /* ------------------------------------------------------------------------ */
 
 #if KMP_DYNAMIC_LIB
 #if KMP_OS_WINDOWS
 
 BOOL WINAPI DllMain(HINSTANCE hInstDLL, DWORD fdwReason, LPVOID lpReserved) {

openmp/runtime/src/kmp_stub.cpp

 omp_allocator_handle_t const omp_low_lat_mem_alloc =
     (omp_allocator_handle_t const)5;
 omp_allocator_handle_t const omp_cgroup_mem_alloc =
     (omp_allocator_handle_t const)6;
 omp_allocator_handle_t const omp_pteam_mem_alloc =
     (omp_allocator_handle_t const)7;
 omp_allocator_handle_t const omp_thread_mem_alloc =
     (omp_allocator_handle_t const)8;
+// Preview of target memory support
+omp_allocator_handle_t const llvm_omp_target_host_mem_alloc =
+    (omp_allocator_handle_t const)100;
+omp_allocator_handle_t const llvm_omp_target_shared_mem_alloc =
+    (omp_allocator_handle_t const)101;
+omp_allocator_handle_t const llvm_omp_target_device_mem_alloc =
+    (omp_allocator_handle_t const)102;
 
 omp_memspace_handle_t const omp_default_mem_space =
     (omp_memspace_handle_t const)0;
 omp_memspace_handle_t const omp_large_cap_mem_space =
     (omp_memspace_handle_t const)1;
 omp_memspace_handle_t const omp_const_mem_space =
     (omp_memspace_handle_t const)2;
 omp_memspace_handle_t const omp_high_bw_mem_space =
     (omp_memspace_handle_t const)3;
 omp_memspace_handle_t const omp_low_lat_mem_space =
     (omp_memspace_handle_t const)4;
+// Preview of target memory support
+omp_memspace_handle_t const llvm_omp_target_host_mem_space =
+    (omp_memspace_handle_t const)100;
+omp_memspace_handle_t const llvm_omp_target_shared_mem_space =
+    (omp_memspace_handle_t const)101;
+omp_memspace_handle_t const llvm_omp_target_device_mem_space =
+    (omp_memspace_handle_t const)102;
 #endif /* KMP_OS_WINDOWS */
 void *omp_alloc(size_t size, const omp_allocator_handle_t allocator) {
   i;
   return malloc(size);
 }
 void *omp_calloc(size_t nmemb, size_t size,
                  const omp_allocator_handle_t allocator) {
   i;