[OpenMP] Add OpenMP 5.1 omp_realloc function implementation

[protze.joachim]

Can we have tests for ptr = NULL, size=0? Probably just replace alloc and free by these cases and check returned address after each call?

openmp/runtime/src/kmp_alloc.cpp

1585

Shouldn't this be after handling size==0?

[AndreyChurbanov]

In D90971#2382007, @protze.joachim wrote:

Can we have tests for ptr = NULL, size=0? Probably just replace alloc and free by these cases and check returned address after each call?

Nice suggestion, agree with it.

openmp/runtime/src/kmp_alloc.cpp

1585

Actually, the assertion does not make sense at all, because size_t is unsigned.

[Nawrin]

Added new tests and remove debug assert

[AndreyChurbanov]

Please also change the openmp/runtime/tools/generate-def.pl file in order to not break Windows build.
See https://reviews.llvm.org/D91478 as an example of the change made for omp_calloc function.

[Nawrin]

Added fix for windows build

[protze.joachim]

The tests have the same format type issue as in D90967. Feel free to include the fixes for the calloc tests in this review.

PS: this summary does not show the inline code suggestions. They are only inline below :(

openmp/runtime/test/api/omp_realloc_def_fb.c

20–21

openmp/runtime/test/api/omp_realloc_null_ptr.c

20–21 ↗

(On Diff #305540)

openmp/runtime/test/api/omp_realloc_size_0.c

20–21 ↗

(On Diff #305540)

[Nawrin]

Fix test errors for both realloc and calloc implementation.

[AndreyChurbanov]

LGTM

openmp/runtime/src/dllexports

kmp_set_disp_num_buffers                    890
     omp_free                                895
     omp_get_device_num                      896
     omp_init_allocator                      897
     omp_destroy_allocator                   898
     %ifndef stub
         __kmpc_set_default_allocator
         __kmpc_get_default_allocator
         __kmpc_alloc
+        __kmpc_realloc
         __kmpc_free
         __kmpc_init_allocator
         __kmpc_destroy_allocator
     %endif
     omp_set_affinity_format                 748
     omp_get_affinity_format                 749
     omp_display_affinity                    750
     omp_capture_affinity                    751
     ompc_set_affinity_format                752
     ompc_get_affinity_format                753
     ompc_display_affinity                   754
     ompc_capture_affinity                   755
     omp_pause_resource                      756
     omp_pause_resource_all                  757
     omp_get_supported_active_levels         758
     omp_fulfill_event                       759
     omp_display_env                         733
+    omp_realloc                             777
 
     omp_null_allocator                     DATA
     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

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

#   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);
     extern omp_allocator_handle_t __KAI_KMPC_CONVENTION omp_get_default_allocator(void);
 #   ifdef __cplusplus
     extern void *__KAI_KMPC_CONVENTION omp_alloc(size_t size, omp_allocator_handle_t a = omp_null_allocator);
+    extern void *__KAI_KMPC_CONVENTION omp_realloc(void *ptr, size_t size,
+                                                   omp_allocator_handle_t allocator = omp_null_allocator,
+                                                   omp_allocator_handle_t free_allocator = omp_null_allocator);
     extern void __KAI_KMPC_CONVENTION omp_free(void * ptr, omp_allocator_handle_t a = omp_null_allocator);
 #   else
     extern void *__KAI_KMPC_CONVENTION omp_alloc(size_t size, omp_allocator_handle_t a);
+    extern void *__KAI_KMPC_CONVENTION omp_realloc(void *ptr, size_t size, omp_allocator_handle_t allocator,
+                                                   omp_allocator_handle_t free_allocator);
     extern void __KAI_KMPC_CONVENTION omp_free(void *ptr, omp_allocator_handle_t a);
 #   endif
 
     /* OpenMP 5.0 Affinity Format */
     extern void __KAI_KMPC_CONVENTION omp_set_affinity_format(char const *);
     extern size_t __KAI_KMPC_CONVENTION omp_get_affinity_format(char *, size_t);
     extern void __KAI_KMPC_CONVENTION omp_display_affinity(char const *);
     extern size_t __KAI_KMPC_CONVENTION omp_capture_affinity(char *, size_t, char const *);

openmp/runtime/src/kmp.h

 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);
 extern omp_allocator_handle_t __kmpc_get_default_allocator(int gtid);
 extern void *__kmpc_alloc(int gtid, 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();
 
 /* ------------------------------------------------------------------------ */
 
 #define KMP_UINT64_MAX                                                         \

openmp/runtime/src/kmp_alloc.cpp

 
 omp_allocator_handle_t __kmpc_get_default_allocator(int gtid) {
   return __kmp_threads[gtid]->th.th_def_allocator;
 }
 
 typedef struct kmp_mem_desc { // Memory block descriptor
   void *ptr_alloc; // Pointer returned by allocator
   size_t size_a; // Size of allocated memory block (initial+descriptor+align)
+  size_t size_orig; // Original size requested
   void *ptr_align; // Pointer to aligned memory, returned
   kmp_allocator_t *allocator; // allocator
 } kmp_mem_desc_t;
 static int alignment = sizeof(void *); // let's align to pointer size
 
 void *__kmpc_alloc(int gtid, size_t size, omp_allocator_handle_t allocator) {
   void *ptr = NULL;
   kmp_allocator_t *al;
   KMP_DEBUG_ASSERT(__kmp_init_serial);
   if (allocator == omp_null_allocator)
     allocator = __kmp_threads[gtid]->th.th_def_allocator;
 
   KE_TRACE(25, ("__kmpc_alloc: T#%d (%d, %p)\n", gtid, (int)size, allocator));
   al = RCAST(kmp_allocator_t *, CCAST(omp_allocator_handle_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
   int align = alignment; // default alignment
   if (allocator > kmp_max_mem_alloc && al->alignment > 0) {
     align = al->alignment; // alignment requested by user
   }
+  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 {
void *__kmpc_alloc(int gtid, size_t size, omp_allocator_handle_t allocator) {
   desc.allocator = al;
   *((kmp_mem_desc_t *)addr_descr) = desc; // save descriptor contents
   KMP_MB();
 
   KE_TRACE(25, ("__kmpc_alloc returns %p, T#%d\n", desc.ptr_align, gtid));
   return desc.ptr_align;
 }
 
+void *__kmpc_realloc(int gtid, void *ptr, size_t size,
+                     omp_allocator_handle_t allocator,
+                     omp_allocator_handle_t free_allocator) {
+  void *nptr = NULL;
+  KMP_DEBUG_ASSERT(__kmp_init_serial);
+  KMP_DEBUG_ASSERT(size > 0);
+
+  if (size == 0) {
+    if (ptr != NULL)
+      __kmpc_free(gtid, ptr, free_allocator);
+    return nptr;
+  }
+
+  KE_TRACE(25, ("__kmpc_realloc: T#%d (%p, %d, %p, %p)\n", gtid, ptr, (int)size,
+                allocator, free_allocator));
+
+  nptr = __kmpc_alloc(gtid, size, allocator);
+
+  if (nptr != NULL && ptr != NULL) {
+    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
+
+    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);
+    KMP_DEBUG_ASSERT(desc.size_orig > 0);
+    KMP_DEBUG_ASSERT(desc.size_orig < desc.size_a);
+    KMP_MEMCPY((char *)nptr, (char *)ptr,
+               (size_t)((size < desc.size_orig) ? size : desc.size_orig));
+  }
+
+  if (nptr != NULL) {
+    __kmpc_free(gtid, ptr, free_allocator);
+  }
+
+  KE_TRACE(25, ("__kmpc_realloc returns %p, T#%d\n", nptr, gtid));
+  return nptr;
+}
+
 void __kmpc_free(int gtid, void *ptr, const omp_allocator_handle_t allocator) {
   KE_TRACE(25, ("__kmpc_free: T#%d free(%p,%p)\n", gtid, ptr, allocator));
   if (ptr == NULL)
     return;
 
   kmp_allocator_t *al;
   omp_allocator_handle_t oal;
   al = RCAST(kmp_allocator_t *, CCAST(omp_allocator_handle_t, allocator));

openmp/runtime/src/kmp_csupport.cpp

void __kmpc_doacross_fini(ident_t *loc, int gtid) {
   KA_TRACE(20, ("__kmpc_doacross_fini() exit: T#%d\n", gtid));
 }
 
 /* omp_alloc/omp_free only defined for C/C++, not for Fortran */
 void *omp_alloc(size_t size, omp_allocator_handle_t allocator) {
   return __kmpc_alloc(__kmp_entry_gtid(), size, allocator);
 }
 
+void *omp_realloc(void *ptr, size_t size, omp_allocator_handle_t allocator,
+                  omp_allocator_handle_t free_allocator) {
+  return __kmpc_realloc(__kmp_entry_gtid(), ptr, size, allocator,
+                        free_allocator);
+}
+
 void omp_free(void *ptr, omp_allocator_handle_t allocator) {
   __kmpc_free(__kmp_entry_gtid(), ptr, allocator);
 }
 
 int __kmpc_get_target_offload(void) {
   if (!__kmp_init_serial) {
     __kmp_serial_initialize();
   }

openmp/runtime/src/kmp_stub.cpp

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;
 #endif /* KMP_OS_WINDOWS */
 void *omp_alloc(size_t size, const omp_allocator_handle_t allocator) {
   i;
   return malloc(size);
 }
+void *omp_realloc(void *ptr, size_t size,
+                  const omp_allocator_handle_t allocator,
+                  const omp_allocator_handle_t free_allocator) {
+  i;
+  return realloc(ptr, size);
+}
 void omp_free(void *ptr, const omp_allocator_handle_t allocator) {
   i;
   free(ptr);
 }
 /* OpenMP 5.0 Affinity Format */
 void omp_set_affinity_format(char const *format) { i; }
 size_t omp_get_affinity_format(char *buffer, size_t size) {
   i;

openmp/runtime/test/api/omp_realloc_def_fb.c

+// RUN: %libomp-compile-and-run
+
+#include <stdio.h>
+#include <omp.h>
+
+int main() {
+  omp_alloctrait_t at[2];
+  omp_allocator_handle_t a;
+  omp_allocator_handle_t f_a;
+  void *ptr[2];
+  void *nptr[2];
+  at[0].key = omp_atk_pool_size;
+  at[0].value = 2 * 1024 * 1024;
+  at[1].key = omp_atk_fallback;
+  at[1].value = omp_atv_default_mem_fb;
+
+  a = omp_init_allocator(omp_large_cap_mem_space, 2, at);
+  f_a = omp_init_allocator(omp_default_mem_space, 2, at);
+  printf("allocator large created: %p\n", a);
+  printf("allocator default created: %p\n", f_a);
+
+  #pragma omp parallel num_threads(2)
+  {
+    int i = omp_get_thread_num();
+    ptr[i] = omp_alloc(1024 * 1024, f_a);
+    #pragma omp barrier
+    nptr[i] = omp_realloc(ptr[i], 1024 * 1024, a, f_a);
+    #pragma omp barrier
+    printf("th %d, nptr %p\n", i, nptr[i]);
+    omp_free(nptr[i], a);
+  }
+  // Both pointers should be non-NULL
+  if (nptr[0] != NULL && nptr[1] != NULL) {
+    printf("passed\n");
+    return 0;
+  } else {
+    printf("failed: pointers %p %p\n", nptr[0], nptr[1]);
+    return 1;
+  }
+}