[OpenMP] Add omp_calloc implementation

This patch adds omp_calloc implementation according to OpenMP 5.1 specification. Differential Revision: https://reviews.llvm.org/D90967
llvm · Nov 13, 2020 · 938f1b8 · 938f1b8
1 parent bf3861b
commit 938f1b8
Show file tree

Hide file tree

Showing 8 changed files with 118 additions and 1 deletion.
diff --git a/openmp/runtime/src/dllexports b/openmp/runtime/src/dllexports
@@ -517,6 +517,7 @@ kmp_set_disp_num_buffers                    890
         __kmpc_set_default_allocator
         __kmpc_get_default_allocator
         __kmpc_alloc
+        __kmpc_calloc
         __kmpc_free
         __kmpc_init_allocator
         __kmpc_destroy_allocator
@@ -534,6 +535,7 @@ kmp_set_disp_num_buffers                    890
     omp_get_supported_active_levels         758
     omp_fulfill_event                       759
     omp_display_env                         733
+    omp_calloc                              776
 
     omp_null_allocator                     DATA
     omp_default_mem_alloc                  DATA

diff --git a/openmp/runtime/src/include/omp.h.var b/openmp/runtime/src/include/omp.h.var
@@ -323,9 +323,11 @@
     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_calloc(size_t nmemb, size_t size, omp_allocator_handle_t a = 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_calloc(size_t nmemb, size_t size, omp_allocator_handle_t a);
     extern void __KAI_KMPC_CONVENTION omp_free(void *ptr, omp_allocator_handle_t a);
 #   endif
 

diff --git a/openmp/runtime/src/kmp.h b/openmp/runtime/src/kmp.h
@@ -963,6 +963,8 @@ 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_calloc(int gtid, size_t nmemb, size_t sz,
+                           omp_allocator_handle_t al);
 extern void __kmpc_free(int gtid, void *ptr, omp_allocator_handle_t al);
 
 extern void __kmp_init_memkind();

diff --git a/openmp/runtime/src/kmp_alloc.cpp b/openmp/runtime/src/kmp_alloc.cpp
@@ -1445,6 +1445,10 @@ 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 (size == 0)
+    return NULL;
+
   if (allocator == omp_null_allocator)
     allocator = __kmp_threads[gtid]->th.th_def_allocator;
 
@@ -1575,6 +1579,39 @@ void *__kmpc_alloc(int gtid, size_t size, omp_allocator_handle_t allocator) {
   return desc.ptr_align;
 }
 
+void *__kmpc_calloc(int gtid, size_t nmemb, 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_calloc: T#%d (%d, %d, %p)\n", gtid, (int)nmemb,
+                (int)size, allocator));
+
+  al = RCAST(kmp_allocator_t *, CCAST(omp_allocator_handle_t, allocator));
+
+  if (nmemb == 0 || size == 0)
+    return ptr;
+
+  if ((SIZE_MAX - sizeof(kmp_mem_desc_t)) / size < nmemb) {
+    if (al->fb == omp_atv_abort_fb) {
+      KMP_ASSERT(0);
+    }
+    return ptr;
+  }
+
+  ptr = __kmpc_alloc(gtid, nmemb * size, allocator);
+
+  if (ptr) {
+    memset(ptr, 0x00, nmemb * size);
+  }
+  KE_TRACE(25, ("__kmpc_calloc returns %p, T#%d\n", ptr, gtid));
+  return ptr;
+}
+
 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)

diff --git a/openmp/runtime/src/kmp_csupport.cpp b/openmp/runtime/src/kmp_csupport.cpp
@@ -4202,11 +4202,15 @@ 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 */
+/* omp_alloc/omp_calloc/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_calloc(size_t nmemb, size_t size, omp_allocator_handle_t allocator) {
+  return __kmpc_calloc(__kmp_entry_gtid(), nmemb, size, allocator);
+}
+
 void omp_free(void *ptr, omp_allocator_handle_t allocator) {
   __kmpc_free(__kmp_entry_gtid(), ptr, allocator);
 }

diff --git a/openmp/runtime/src/kmp_stub.cpp b/openmp/runtime/src/kmp_stub.cpp
@@ -366,6 +366,11 @@ 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;
+  return calloc(nmemb, size);
+}
 void omp_free(void *ptr, const omp_allocator_handle_t allocator) {
   i;
   free(ptr);

diff --git a/openmp/runtime/test/api/omp_calloc_def_fb.c b/openmp/runtime/test/api/omp_calloc_def_fb.c
@@ -0,0 +1,32 @@
+// RUN: %libomp-compile-and-run
+
+#include <stdio.h>
+#include <omp.h>
+
+int main() {
+  omp_alloctrait_t at[2];
+  omp_allocator_handle_t a;
+  void *p[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);
+  printf("allocator large created: %p\n", a);
+  #pragma omp parallel num_threads(2)
+  {
+    int i = omp_get_thread_num();
+    p[i] = omp_calloc(1024, 1024, a);
+    #pragma omp barrier
+    printf("th %d, ptr %p\n", i, p[i]);
+    omp_free(p[i], a);
+  }
+  // Both pointers should be non-NULL
+  if (p[0] != NULL && p[1] != NULL) {
+    printf("passed\n");
+    return 0;
+  } else {
+    printf("failed: pointers %p %p\n", p[0], p[1]);
+    return 1;
+  }
+}
diff --git a/openmp/runtime/test/api/omp_calloc_size_0.c b/openmp/runtime/test/api/omp_calloc_size_0.c
@@ -0,0 +1,33 @@
+// RUN: %libomp-compile-and-run
+
+#include <stdio.h>
+#include <omp.h>
+
+int main()
+{
+  omp_alloctrait_t at[2];
+  omp_allocator_handle_t a;
+  void *p[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);
+  printf("allocator large created: %p\n", a);
+  #pragma omp parallel num_threads(2)
+  {
+    int i = omp_get_thread_num();
+    p[i] = omp_calloc(1024, 0, a);
+    #pragma omp barrier
+    printf("th %d, ptr %p\n", i, p[i]);
+    omp_free(p[i], a);
+  }
+  // Both pointers should be NULL
+  if (p[0] == NULL && p[1] == NULL) {
+    printf("passed\n");
+    return 0;
+  } else {
+    printf("failed: pointers %p %p\n", p[0], p[1]);
+    return 1;
+  }
+}