Re: [PATCH, OpenMP, Fortran] requires unified_shared_memory 2/2: insert USM allocators into libgfortran

[Chung-Lin Tang <cltang@codesourcery.com>]

[-- Attachment #1: Type: text/plain, Size: 2960 bytes --]



On 2022/8/15 7:15 PM, Chung-Lin Tang wrote:
> On 2022/8/15 7:06 PM, Chung-Lin Tang wrote:
>>
>> I know this is a big pile of yarn wrt how the main program/libgomp/libgfortran interacts, but it's
>> finally working. Again tested without regressions. Preparing to commit to devel/omp/gcc-12, and seeking
>> approval for mainline when the requires patches are in.
> 
> Just realized that I don't have the new testcases added in this patch.
> Will supplement them later :P

Here's the USM allocator/libgfortran patch, with a libgomp.fortran testcase added.

Thanks,
Chung-Lin

2022-09-05  Chung-Lin Tang  <cltang@codesourcery.com>

libgcc/
     * Makefile.in (crtoffloadend$(objext)): Add $(PICFLAG) to compile rule.
     * offloadstuff.c (GOMP_offload_register_ver): Add declaration of weak
     symbol.
     (__OFFLOAD_TABLE__): Likewise.
     (init_non_offload): New function.

libgfortran/

     * gfortran.map (GFORTRAN_13): New namespace.
     (_gfortran_mem_allocators_init): New name inside GFORTRAN_13.
     * libgfortran.h (mem_allocators_init): New exported declaration.
     * runtime/main.c (do_init): Rename from init, add run-once guard code.
     (cleanup): Add run-once guard code.
     (GOMP_post_offload_register_callback): Declare weak symbol.
     (GOMP_pre_gomp_target_fini_callback): Likewise.
     (init): New constructor to register offload callbacks, or call do_init
     when not OpenMP.
     * runtime/memory.c (gfortran_malloc): New pointer variable.
     (gfortran_calloc): Likewise.
     (gfortran_realloc): Likewise.
     (gfortran_free): Likewise.
     (mem_allocators_init): New function.
     (xmalloc): Use gfortran_malloc.
     (xmallocarray): Use gfortran_malloc.
     (xcalloc): Use gfortran_calloc.
     (xrealloc): Use gfortran_realloc.
     (xfree): Use gfortran_free.

libgomp/

     * libgomp.map (GOMP_5.1.2): New version namespace.
     (GOMP_post_offload_register_callback): New name inside GOMP_5.1.2.
     (GOMP_pre_gomp_target_fini_callback): Likewise.
     (GOMP_DEFINE_CALLBACK_SET): Macro to define callback set.
     (post_offload_register): Define callback set for after offload image
     register.
     (pre_gomp_target_fini): Define callback set for before gomp_target_fini
     is called.
     (libgfortran_malloc_usm): New function.
     (libgfortran_calloc_usm): Likewise
     (libgfortran_realloc_usm): Likewise
     (libgfortran_free_usm): Likewise.
     (_gfortran_mem_allocators_init): Declare weak symbol.
     (gomp_libgfortran_omp_allocators_init): New function.
     (GOMP_offload_register_ver): Add handling of host_table == NULL, calling
     into libgfortran to set unified_shared_memory allocators, and execution
     of post_offload_register callbacks.
     (gomp_target_init): Register all pre_gomp_target_fini callbacks to run
     at end of main using atexit().

     * testsuite/libgomp.fortran/target-unified_shared_memory-1.f90: New test.








[-- Attachment #2: fortran-usm-v2.patch --]
[-- Type: text/plain, Size: 10636 bytes --]

diff --git a/libgcc/Makefile.in b/libgcc/Makefile.in
index 09b3ec8bc2e..70720cc910c 100644
--- a/libgcc/Makefile.in
+++ b/libgcc/Makefile.in
@@ -1045,8 +1045,9 @@ crtbeginT$(objext): $(srcdir)/crtstuff.c
 crtoffloadbegin$(objext): $(srcdir)/offloadstuff.c
 	$(crt_compile) $(CRTSTUFF_T_CFLAGS) -c $< -DCRT_BEGIN
 
+# crtoffloadend contains a constructor with calls to libgomp, so build as PIC.
 crtoffloadend$(objext): $(srcdir)/offloadstuff.c
-	$(crt_compile) $(CRTSTUFF_T_CFLAGS) -c $< -DCRT_END
+	$(crt_compile) $(CRTSTUFF_T_CFLAGS) $(PICFLAG) -c $< -DCRT_END
 
 crtoffloadtable$(objext): $(srcdir)/offloadstuff.c
 	$(crt_compile) $(CRTSTUFF_T_CFLAGS) -c $< -DCRT_TABLE
diff --git a/libgcc/offloadstuff.c b/libgcc/offloadstuff.c
index 10e1fe19c8e..2edb6810021 100644
--- a/libgcc/offloadstuff.c
+++ b/libgcc/offloadstuff.c
@@ -63,6 +63,19 @@ const void *const __offload_vars_end[0]
   __attribute__ ((__used__, visibility ("hidden"),
 		  section (OFFLOAD_VAR_TABLE_SECTION_NAME))) = { };
 
+extern void GOMP_offload_register_ver (unsigned, const void *, int,
+				       const void *);
+extern const void *const __OFFLOAD_TABLE__[0] __attribute__ ((weak));
+static void __attribute__((constructor))
+init_non_offload (void)
+{
+  /* If an OpenMP program has no offloading, post-offload_register callbacks
+     that need to run will require a call to GOMP_offload_register_ver, in
+     order to properly trigger those callbacks during init.  */
+  if (__OFFLOAD_TABLE__ == NULL)
+    GOMP_offload_register_ver (0, NULL, 0, NULL);
+}
+
 #elif defined CRT_TABLE
 
 extern const void *const __offload_func_table[];
diff --git a/libgfortran/gfortran.map b/libgfortran/gfortran.map
index e0e795c3d48..55d2a529acd 100644
--- a/libgfortran/gfortran.map
+++ b/libgfortran/gfortran.map
@@ -1759,3 +1759,8 @@ GFORTRAN_12 {
   _gfortran_transfer_real128_write;
 #endif
 } GFORTRAN_10.2;
+
+GFORTRAN_13 {
+  global:
+  _gfortran_mem_allocators_init;
+} GFORTRAN_12;
diff --git a/libgfortran/libgfortran.h b/libgfortran/libgfortran.h
index 0b893a51851..e518b3989cf 100644
--- a/libgfortran/libgfortran.h
+++ b/libgfortran/libgfortran.h
@@ -874,6 +874,11 @@ internal_proto(xrealloc);
 extern void xfree (void *);
 internal_proto(xfree);
 
+#ifndef LIBGFOR_MINIMAL
+extern void mem_allocators_init (void *, void *, void *, void *);
+export_proto(mem_allocators_init);
+#endif
+
 /* environ.c */
 
 extern void init_variables (void);
diff --git a/libgfortran/runtime/main.c b/libgfortran/runtime/main.c
index 5162a8fecb0..8aa688e1d0f 100644
--- a/libgfortran/runtime/main.c
+++ b/libgfortran/runtime/main.c
@@ -61,9 +61,16 @@ get_args (int *argc, char ***argv)
 
 /* Initialize the runtime library.  */
 
-static void __attribute__((constructor))
-init (void)
+static void
+do_init (void)
 {
+#ifndef LIBGFOR_MINIMAL
+  static bool do_init_ran = false;
+  if (do_init_ran)
+    return;
+  do_init_ran = true;
+#endif
+
   /* Must be first */
   init_variables ();
 
@@ -82,5 +89,37 @@ init (void)
 static void __attribute__((destructor))
 cleanup (void)
 {
+#ifndef LIBGFOR_MINIMAL
+  static bool cleanup_ran = false;
+  if (cleanup_ran)
+    return;
+  cleanup_ran = true;
+#endif
+
   close_units ();
 }
+
+#ifndef LIBGFOR_MINIMAL
+extern void __attribute__((weak))
+GOMP_post_offload_register_callback (void (*func)(void));
+
+extern void __attribute__((weak))
+GOMP_pre_gomp_target_fini_callback (void (*func)(void));
+#endif
+
+static void __attribute__((constructor))
+init (void)
+{
+#ifndef LIBGFOR_MINIMAL
+  if (GOMP_post_offload_register_callback)
+    {
+      GOMP_post_offload_register_callback (do_init);
+      GOMP_pre_gomp_target_fini_callback (cleanup);
+      return;
+    }
+#endif
+
+  /* If libgomp is not present, then we can go ahead and call do_init
+     directly.  */
+  do_init ();
+}
diff --git a/libgfortran/runtime/memory.c b/libgfortran/runtime/memory.c
index cbcec7c9281..8bf5b605f86 100644
--- a/libgfortran/runtime/memory.c
+++ b/libgfortran/runtime/memory.c
@@ -26,6 +26,28 @@ see the files COPYING3 and COPYING.RUNTIME respectively.  If not, see
 #include "libgfortran.h"
 #include <errno.h>
 
+#ifndef LIBGFOR_MINIMAL
+static void * (*gfortran_malloc)(size_t) = malloc;
+static void * (*gfortran_calloc)(size_t, size_t) = calloc;
+static void * (*gfortran_realloc)(void *, size_t) = realloc;
+static void (*gfortran_free)(void *) = free;
+
+void
+mem_allocators_init (void *malloc_ptr, void *calloc_ptr,
+		     void *realloc_ptr, void *free_ptr)
+{
+  gfortran_malloc = malloc_ptr;
+  gfortran_calloc = calloc_ptr;
+  gfortran_realloc = realloc_ptr;
+  gfortran_free = free_ptr;
+}
+
+#else
+#define gfortran_malloc malloc
+#define gfortran_calloc calloc
+#define gfortran_realloc realloc
+#define gfortran_free free
+#endif
 
 void *
 xmalloc (size_t n)
@@ -35,7 +57,7 @@ xmalloc (size_t n)
   if (n == 0)
     n = 1;
 
-  p = malloc (n);
+  p = gfortran_malloc (n);
 
   if (p == NULL)
     os_error ("Memory allocation failed");
@@ -57,7 +79,7 @@ xmallocarray (size_t nmemb, size_t size)
       os_error ("Integer overflow in xmallocarray");
     }
 
-  p = malloc (prod);
+  p = gfortran_malloc (prod);
 
   if (!p)
     os_error ("Memory allocation failed in xmallocarray");
@@ -73,7 +95,7 @@ xcalloc (size_t nmemb, size_t size)
   if (!nmemb || !size)
     nmemb = size = 1;
 
-  void *p = calloc (nmemb, size);
+  void *p = gfortran_calloc (nmemb, size);
   if (!p)
     os_error ("Allocating cleared memory failed");
 
@@ -86,7 +108,7 @@ xrealloc (void *ptr, size_t size)
   if (size == 0)
     size = 1;
 
-  void *newp = realloc (ptr, size);
+  void *newp = gfortran_realloc (ptr, size);
   if (!newp)
     os_error ("Memory allocation failure in xrealloc");
 
@@ -96,5 +118,5 @@ xrealloc (void *ptr, size_t size)
 void
 xfree (void *ptr)
 {
-  free (ptr);
+  gfortran_free (ptr);
 }
diff --git a/libgomp/libgomp.map b/libgomp/libgomp.map
index 5af5c2d8929..c3af75cc800 100644
--- a/libgomp/libgomp.map
+++ b/libgomp/libgomp.map
@@ -624,3 +624,9 @@ GOMP_PLUGIN_1.3 {
 	GOMP_PLUGIN_goacc_profiling_dispatch;
 	GOMP_PLUGIN_goacc_thread;
 } GOMP_PLUGIN_1.2;
+
+GOMP_5.1.2 {
+  global:
+	GOMP_post_offload_register_callback;
+	GOMP_pre_gomp_target_fini_callback;
+} GOMP_5.1.1;
diff --git a/libgomp/target.c b/libgomp/target.c
index 997b2aa2f80..6a5c0bb1b36 100644
--- a/libgomp/target.c
+++ b/libgomp/target.c
@@ -2522,6 +2522,70 @@ gomp_requires_to_name (char *buf, size_t size, int requires_mask)
 		   (p == buf ? "" : ", "));
 }
 
+/* Macro to define a callback set with a name, and routine to register
+   a callback function into set.  */
+#define GOMP_DEFINE_CALLBACK_SET(name)				\
+  static unsigned int num_ ## name ## _callbacks = 0;		\
+  static void (*name ## _callbacks[4])(void);			\
+  void GOMP_ ## name ## _callback (void (*fn)(void))		\
+  {								\
+    if (num_ ## name ## _callbacks				\
+	< (sizeof (name ## _callbacks)				\
+	   / sizeof (name ## _callbacks[0])))			\
+      {								\
+	name ## _callbacks[num_ ## name ## _callbacks] = fn;	\
+	num_ ## name ## _callbacks += 1;			\
+      }								\
+  }
+GOMP_DEFINE_CALLBACK_SET(post_offload_register)
+GOMP_DEFINE_CALLBACK_SET(pre_gomp_target_fini)
+#undef GOMP_DEFINE_CALLBACK_SET
+
+/* Routines to insert into libgfortran, under unified_shared_memory.  */
+static void *
+libgfortran_malloc_usm (size_t size)
+{
+  return omp_alloc (size, ompx_unified_shared_mem_alloc);
+}
+
+static void *
+libgfortran_calloc_usm (size_t n, size_t size)
+{
+  return omp_calloc (n, size, ompx_unified_shared_mem_alloc);
+}
+
+static void *
+libgfortran_realloc_usm (void *ptr, size_t size)
+{
+  return omp_realloc (ptr, size, ompx_unified_shared_mem_alloc,
+		      ompx_unified_shared_mem_alloc);
+}
+
+static void
+libgfortran_free_usm (void *ptr)
+{
+  omp_free (ptr, ompx_unified_shared_mem_alloc);
+}
+
+extern void __attribute__((weak))
+_gfortran_mem_allocators_init (void *, void *, void *, void *);
+
+static void
+gomp_libgfortran_omp_allocators_init (int omp_requires_mask)
+{
+  static bool init = false;
+  if (init)
+    return;
+  init = true;
+
+  if ((omp_requires_mask & GOMP_REQUIRES_UNIFIED_SHARED_MEMORY)
+      && _gfortran_mem_allocators_init != NULL)
+    _gfortran_mem_allocators_init (libgfortran_malloc_usm,
+				   libgfortran_calloc_usm,
+				   libgfortran_realloc_usm,
+				   libgfortran_free_usm);
+}
+
 /* This function should be called from every offload image while loading.
    It gets the descriptor of the host func and var tables HOST_TABLE, TYPE of
    the target, and DATA.  */
@@ -2532,6 +2596,9 @@ GOMP_offload_register_ver (unsigned version, const void *host_table,
 {
   int i;
 
+  if (host_table == NULL)
+    goto end;
+
   if (GOMP_VERSION_LIB (version) > GOMP_VERSION)
     gomp_fatal ("Library too old for offload (version %u < %u)",
 		GOMP_VERSION, GOMP_VERSION_LIB (version));
@@ -2598,6 +2665,14 @@ GOMP_offload_register_ver (unsigned version, const void *host_table,
 
   num_offload_images++;
   gomp_mutex_unlock (&register_lock);
+
+  /* Call into libgfortran to initialize OpenMP memory allocators.  */
+  gomp_libgfortran_omp_allocators_init (omp_requires_mask);
+
+ end:
+  for (int i = 0; i < num_post_offload_register_callbacks; i++)
+    post_offload_register_callbacks[i] ();
+  num_post_offload_register_callbacks = 0;
 }
 
 /* Legacy entry point.  */
@@ -2710,7 +2785,7 @@ gomp_unload_device (struct gomp_device_descr *devicep)
   if (devicep->state == GOMP_DEVICE_INITIALIZED)
     {
       unsigned i;
-      
+
       /* Unload from device all images registered at the moment.  */
       for (i = 0; i < num_offload_images; i++)
 	{
@@ -4570,6 +4645,13 @@ gomp_target_init (void)
   devices = devs;
   if (atexit (gomp_target_fini) != 0)
     gomp_fatal ("atexit failed");
+
+  /* Register 'pre_gomp_target_fini' callbacks to run before gomp_target_fini
+     during finalization.  */
+  for (int i = 0; i < num_pre_gomp_target_fini_callbacks; i++)
+    if (atexit (pre_gomp_target_fini_callbacks[i]) != 0)
+      gomp_fatal ("atexit failed");
+  num_pre_gomp_target_fini_callbacks = 0;
 }
 
 #else /* PLUGIN_SUPPORT */
diff --git a/libgomp/testsuite/libgomp.fortran/target-unified_shared_memory-1.f90 b/libgomp/testsuite/libgomp.fortran/target-unified_shared_memory-1.f90
new file mode 100644
index 00000000000..f82b1c7527c
--- /dev/null
+++ b/libgomp/testsuite/libgomp.fortran/target-unified_shared_memory-1.f90
@@ -0,0 +1,13 @@
+! { dg-do run }
+
+program requires_unified_shared_memory
+  character(32) :: str
+  !$omp requires unified_shared_memory
+
+  str = trim (str)
+
+  !$omp target
+  block
+  end block
+
+end program requires_unified_shared_memory