[OpenMP] Add API for pinned memory

[carlo.bertolli]

Summary

This patch adds API support for the atk_pinned trait for omp_alloc. It does not implement kmp_target_lock_mem and kmp_target_unlock_mem in libomptarget, but prepares libomp for it. Patches to libomptarget to implement lock/unlock coming after this one.

[jdoerfert]

Generally looks good. Great we are gaining support upstream. I have one question though and I hope some host runtime ppl can take a look too.

openmp/runtime/src/kmp_alloc.cpp

1658

Calling lock_mem on alloc_device memory is probably not going to work well, right?
I am unsure we need to lock this here at all, it's either "_host" (=pinned), "_shared" (=managed), or "_device" (=GPU) memory.
When does locking make sense in those three cases and can we guarantee it won't cause problems?

1694

Same as above.

[carlo.bertolli]

Pinning does not apply to device memory on AMDGPU target, only to host memory, as explained here:
https://github.com/RadeonOpenCompute/ROCR-Runtime/blob/master/src/inc/hsa_ext_amd.h#L1526

CUDA API seems to agree with that:
https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__MEMORY.html#group__CUDART__MEMORY_1gb65da58f444e7230d3322b6126bb4902

openmp/runtime/src/kmp_alloc.cpp

1658

Good catch. For now, I moved the pinning to the llvm_omp_target_host_memo_alloc. However, given that this is a predefined memory allocator, I do not see any way the OpenMP API would allow users to set a special trait for it. Setting traits is something you can do when initializing an allocator, which needs to be seeded with a predefined memory space - not with a predefined allocator.
So, unless there is something in OpenMP 6.0 TR1 that would enable users to customize predefined allocators, this is dead code.
If you agree, I can remove support from these two cases you indicated. If not, I'll leave it up.

[jdoerfert]

openmp/runtime/src/kmp_alloc.cpp

1658

I think we should remove it. "_host" allocator should already pin memory, plus your point.

[carlo.bertolli]

Removed llvm_omp_target_host_mem_alloc as a case for pinning.

[carlo.bertolli]

ping

[jlpeyton]

openmp/runtime/src/kmp_alloc.cpp

2043

1. Should there be a check for is_pinned or similar?
2. Should the default device always be 0?

openmp/runtime/test/api/omp_pinned.c

1 ↗

(On Diff #478942)

Can we add another RUN line with env OMP_ALLOCATOR=omp_default_mem_space:pinned=true %libomp-run. This would have to be used in tandem with something like pinned_alloc2 = omp_get_default_allocator() below doing a second allocation and free.

[carlo.bertolli]

[OpenMP] Add test for default allocator and move unlock call to appropriate site.

[jlpeyton]

openmp/runtime/test/api/omp_pinned.c

2 ↗

(On Diff #482565)

%libomp-compile-and-run -> %libomp-run

As written now with %libomp-compile-and-run, LIT will set the environment variable for the re-compilation only. It translates to:

$ env OMP_ALLOCATOR=... <compilation_command> && ./executable

We can avoid the recompilation and have the environment variable apply to the run command by using %libomp-run

[carlo.bertolli]

[OpenMP] Only rerun lit test for pinned memory, no need to rebuild.

[jlpeyton]

LGTM

[jdoerfert]

Thanks Carlo, this is great :)

[xtian]

With the updated patch. LGTM

openmp/runtime/src/kmp.h

 typedef struct kmp_allocator_t {
   omp_memspace_handle_t memspace;
   void **memkind; // pointer to memkind
   size_t alignment;
   omp_alloctrait_value_t fb;
   kmp_allocator_t *fb_data;
   kmp_uint64 pool_size;
   kmp_uint64 pool_used;
+  bool pinned;
 } kmp_allocator_t;
 
 extern omp_allocator_handle_t __kmpc_init_allocator(int gtid,
                                                     omp_memspace_handle_t,
                                                     int ntraits,
                                                     omp_alloctrait_t traits[]);
 extern void __kmpc_destroy_allocator(int gtid, omp_allocator_handle_t al);
 extern void __kmpc_set_default_allocator(int gtid, omp_allocator_handle_t al);

openmp/runtime/src/kmp_alloc.cpp

 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;
 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_lock_mem)(void *ptr, size_t size, int device);
+static void *(*kmp_target_unlock_mem)(void *ptr, int device);
 static void *(*kmp_target_free_host)(void *ptr, int device);
 static void *(*kmp_target_free_shared)(void *ptr, int device);
 static void *(*kmp_target_free_device)(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)
omp_allocator_handle_t __kmpc_init_allocator(int gtid, omp_memspace_handle_t 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:
+      break;
     case omp_atk_pinned:
+      al->pinned = true;
       break;
     case omp_atk_alignment:
       __kmp_type_convert(traits[i].value, &(al->alignment));
       KMP_ASSERT(IS_POWER_OF_TWO(al->alignment));
       break;
     case omp_atk_pool_size:
       al->pool_size = traits[i].value;
       break;
void *__kmp_alloc(int gtid, size_t algn, size_t size,
                   omp_allocator_handle_t allocator) {
   void *ptr = NULL;
   kmp_allocator_t *al;
   KMP_DEBUG_ASSERT(__kmp_init_serial);
   if (size == 0)
     return NULL;
   if (allocator == omp_null_allocator)
     allocator = __kmp_threads[gtid]->th.th_def_allocator;
+  kmp_int32 default_device =
+      __kmp_threads[gtid]->th.th_current_task->td_icvs.default_device;
 
   al = RCAST(kmp_allocator_t *, allocator);
 
   int sz_desc = sizeof(kmp_mem_desc_t);
   kmp_mem_desc_t desc;
   kmp_uintptr_t addr; // address returned by allocator
   kmp_uintptr_t addr_align; // address to return to caller
   kmp_uintptr_t addr_descr; // address of memory block descriptor
   size_t align = alignment; // default alignment
   if (allocator > kmp_max_mem_alloc && al->alignment > align)
     align = al->alignment; // alignment required by allocator trait
   if (align < algn)
     align = algn; // max of allocator trait, parameter and sizeof(void*)
   desc.size_orig = 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_memkind_available) {
     if (allocator < kmp_max_mem_alloc) {
       // pre-defined allocator
       if (allocator == omp_high_bw_mem_alloc && mk_hbw_preferred) {
if (allocator < kmp_max_mem_alloc) {
         if (al->fb == omp_atv_default_mem_fb) {
           al = (kmp_allocator_t *)omp_default_mem_alloc;
           ptr = kmp_mk_alloc(*mk_default, desc.size_a);
         } else if (al->fb == omp_atv_abort_fb) {
           KMP_ASSERT(0); // abort fallback requested
         } else if (al->fb == omp_atv_allocator_fb) {
           KMP_ASSERT(al != al->fb_data);
           al = al->fb_data;
-          return __kmp_alloc(gtid, algn, size, (omp_allocator_handle_t)al);
+          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 ptr == NULL;
       } else {
         // pool has enough space
         ptr = kmp_mk_alloc(*al->memkind, desc.size_a);
         if (ptr == NULL) {
           if (al->fb == omp_atv_default_mem_fb) {
             al = (kmp_allocator_t *)omp_default_mem_alloc;
             ptr = kmp_mk_alloc(*mk_default, desc.size_a);
           } else if (al->fb == omp_atv_abort_fb) {
             KMP_ASSERT(0); // abort fallback requested
           } else if (al->fb == omp_atv_allocator_fb) {
             KMP_ASSERT(al != al->fb_data);
             al = al->fb_data;
-            return __kmp_alloc(gtid, algn, size, (omp_allocator_handle_t)al);
+            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 {
       // custom allocator, pool size not requested
       ptr = kmp_mk_alloc(*al->memkind, desc.size_a);
       if (ptr == NULL) {
         if (al->fb == omp_atv_default_mem_fb) {
           al = (kmp_allocator_t *)omp_default_mem_alloc;
           ptr = kmp_mk_alloc(*mk_default, desc.size_a);
         } else if (al->fb == omp_atv_abort_fb) {
           KMP_ASSERT(0); // abort fallback requested
         } else if (al->fb == omp_atv_allocator_fb) {
           KMP_ASSERT(al != al->fb_data);
           al = al->fb_data;
-          return __kmp_alloc(gtid, algn, size, (omp_allocator_handle_t)al);
+          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) {
if (__kmp_target_mem_available) {
           __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) {
         al = (kmp_allocator_t *)omp_default_mem_alloc;
         ptr = __kmp_thread_malloc(__kmp_thread_from_gtid(gtid), desc.size_a);
       } else if (al->fb == omp_atv_abort_fb) {
         KMP_ASSERT(0); // abort fallback requested
       } else if (al->fb == omp_atv_allocator_fb) {
         KMP_ASSERT(al != al->fb_data);
         al = al->fb_data;
-        return __kmp_alloc(gtid, algn, size, (omp_allocator_handle_t)al);
+        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 ptr == NULL;
     } else {
       // pool has enough space
       ptr = __kmp_thread_malloc(__kmp_thread_from_gtid(gtid), desc.size_a);
       if (ptr == NULL && al->fb == omp_atv_abort_fb) {
         KMP_ASSERT(0); // abort fallback requested
       } // no sense to look for another fallback because of same internal alloc
     }
   } else {
     // custom allocator, pool size not requested
     ptr = __kmp_thread_malloc(__kmp_thread_from_gtid(gtid), desc.size_a);
     if (ptr == NULL && al->fb == omp_atv_abort_fb) {
       KMP_ASSERT(0); // abort fallback requested
     } // no sense to look for another fallback because of same internal alloc
   }
   KE_TRACE(10, ("__kmp_alloc: T#%d %p=alloc(%d)\n", gtid, ptr, desc.size_a));
   if (ptr == NULL)
     return NULL;
 
+  if (is_pinned && kmp_target_lock_mem)
+    kmp_target_lock_mem(ptr, desc.size_a, default_device);
+
   addr = (kmp_uintptr_t)ptr;
   addr_align = (addr + sz_desc + align - 1) & ~(align - 1);
   addr_descr = addr_align - sz_desc;
 
   desc.ptr_alloc = ptr;
   desc.ptr_align = (void *)addr_align;
   desc.allocator = al;
   *((kmp_mem_desc_t *)addr_descr) = desc; // save descriptor contents
void ___kmpc_free(int gtid, void *ptr, omp_allocator_handle_t allocator) {
   KMP_DEBUG_ASSERT(desc.ptr_align == ptr);
   if (allocator) {
     KMP_DEBUG_ASSERT(desc.allocator == al || desc.allocator == al->fb_data);
   }
   al = desc.allocator;
   oal = (omp_allocator_handle_t)al; // cast to void* for comparisons
   KMP_DEBUG_ASSERT(al);
 
+  if (allocator > kmp_max_mem_alloc && kmp_target_unlock_mem && al->pinned) {
+    kmp_int32 device =
+        __kmp_threads[gtid]->th.th_current_task->td_icvs.default_device;
+    kmp_target_unlock_mem(desc.ptr_alloc, device);
+  }
+
   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 {
#if KMP_DEBUG
   KMP_DEBUG_ASSERT(descr.ptr_aligned == ptr);
   KMP_DEBUG_ASSERT(addr_allocated + sizeof(kmp_mem_descr_t) <= addr_aligned);
   KMP_DEBUG_ASSERT(descr.size_aligned < descr.size_allocated);
   KMP_DEBUG_ASSERT(addr_aligned + descr.size_aligned <=
                    addr_allocated + descr.size_allocated);
   memset(descr.ptr_allocated, 0xEF, descr.size_allocated);
 // Fill memory block with 0xEF, it helps catch using freed memory.
 #endif
 
 #ifndef LEAK_MEMORY
   KE_TRACE(10, ("   free( %p )\n", descr.ptr_allocated));
 #ifdef KMP_DEBUG
   _free_src_loc(descr.ptr_allocated, _file_, _line_);
 #else
   free_src_loc(descr.ptr_allocated KMP_SRC_LOC_PARM);
 #endif
 #endif