[Libomptarget] Introduce new globalization runtime calls

[jhuber6]

Summary

This patch introduces the new globalization runtime to be used by D97680. These
runtime calls will replace the __kmpc_data_sharing_push_stack and
__kmpc_data_sharing_pop_stack functions.

[tianshilei1992]

LGTM

[jhuber6]

This interface doesn't use __kmpc_init_stack so the data sharing buffer needs to be initialized differently.

[jhuber6]

Changing implementation, had some bugs.

openmp/libomptarget/deviceRTLs/common/src/data_sharing.cu

INLINE static void data_sharing_init_stack_common() {
 
   for (int WID = 0; WID < DS_Max_Warp_Number; WID++) {
     __kmpc_data_sharing_slot *RootS = teamDescr->GetPreallocatedSlotAddr(WID);
     DataSharingState.SlotPtr[WID] = RootS;
     DataSharingState.StackPtr[WID] = (void *)&RootS->Data[0];
   }
 }
 
+EXTERN void *__kmpc_alloc_shared(size_t DataSize) {
+  return (void *)SafeMalloc(DataSize, "Alloc Shared");
+}
+
+EXTERN void __kmpc_free_shared(void *FrameStart) {
+  SafeFree(FrameStart, "Free Shared");
+}
+
 // Initialize data sharing data structure. This function needs to be called
 // once at the beginning of a data sharing context (coincides with the kernel
 // initialization). This function is called only by the MASTER thread of each
 // team in non-SPMD mode.
 EXTERN void __kmpc_data_sharing_init_stack() {
   ASSERT0(LT_FUSSY, isRuntimeInitialized(), "Runtime must be initialized.");
   // This function initializes the stack pointer with the pointer to the
   // statically allocated shared memory slots. The size of a shared memory
if (SlotEmpty && SlotP->Prev) {
       // Remove the slot.
       SlotP = SlotP->Prev;
       SafeFree(SlotP->Next, "Free slot.");
       SlotP->Next = 0;
     }
   }
 }
 
-// Begin a data sharing context. Maintain a list of references to shared
-// variables. This list of references to shared variables will be passed
-// to one or more threads.
-// In L0 data sharing this is called by master thread.
-// In L1 data sharing this is called by active warp master thread.
-EXTERN void __kmpc_begin_sharing_variables(void ***GlobalArgs, size_t nArgs) {
-  omptarget_nvptx_globalArgs.EnsureSize(nArgs);
-  *GlobalArgs = omptarget_nvptx_globalArgs.GetArgs();
+void** SHARED(GlobalArgs);
+void* globalArgsStorage[64];
+#pragma omp allocate(globalArgsStorage) allocator(omp_pteam_mem_alloc)
+EXTERN void __kmpc_begin_sharing_variables(void ***GlobalArgsPtr, size_t nArgs) {
+  if (nArgs > 64) {
+    *GlobalArgsPtr = GlobalArgs = reinterpret_cast<void**>(__kmpc_alloc_shared(nArgs * sizeof(void*)));
+  } else {
+    *GlobalArgsPtr = GlobalArgs = &globalArgsStorage[0];
   }
-
-// End a data sharing context. There is no need to have a list of refs
-// to shared variables because the context in which those variables were
-// shared has now ended. This should clean-up the list of references only
-// without affecting the actual global storage of the variables.
-// In L0 data sharing this is called by master thread.
-// In L1 data sharing this is called by active warp master thread.
+}
+
 EXTERN void __kmpc_end_sharing_variables() {
-  omptarget_nvptx_globalArgs.DeInit();
+  if (GlobalArgs != &globalArgsStorage[0])
+    __kmpc_free_shared(GlobalArgs);
 }
 
 // This function will return a list of references to global variables. This
 // is how the workers will get a reference to the globalized variable. The
 // members of this list will be passed to the outlined parallel function
 // preserving the order.
 // Called by all workers.
-EXTERN void __kmpc_get_shared_variables(void ***GlobalArgs) {
-  *GlobalArgs = omptarget_nvptx_globalArgs.GetArgs();
+
+EXTERN void __kmpc_get_shared_variables(void ***GlobalArgsPtr) {
+  *GlobalArgsPtr = GlobalArgs;
 }
 
 // This function is used to init static memory manager. This manager is used to
 // manage statically allocated global memory. This memory is allocated by the
 // compiler and used to correctly implement globalization of the variables in
 // target, teams and distribute regions.
 EXTERN void __kmpc_get_team_static_memory(int16_t isSPMDExecutionMode,
                                           const void *buf, size_t size,

openmp/libomptarget/deviceRTLs/interface.h

 typedef enum omp_proc_bind_t {
   omp_proc_bind_false = 0,
   omp_proc_bind_true = 1,
   omp_proc_bind_master = 2,
   omp_proc_bind_close = 3,
   omp_proc_bind_spread = 4
 } omp_proc_bind_t;
 
 EXTERN double omp_get_wtick(void);
 EXTERN double omp_get_wtime(void);
 
 EXTERN void omp_set_num_threads(int num);
 EXTERN int omp_get_num_threads(void);
 EXTERN int omp_get_max_threads(void);
 EXTERN int omp_get_thread_limit(void);
 EXTERN int omp_get_thread_num(void);
 EXTERN int omp_get_num_procs(void);
 EXTERN int omp_in_parallel(void);
 EXTERN int omp_in_final(void);
 EXTERN void omp_set_dynamic(int flag);
 EXTERN int omp_get_dynamic(void);
 EXTERN void omp_set_nested(int flag);
 EXTERN int omp_get_nested(void);
 EXTERN void omp_set_max_active_levels(int level);
 EXTERN int omp_get_max_active_levels(void);
 EXTERN int omp_get_level(void);
 EXTERN int omp_get_active_level(void);
 EXTERN int omp_get_ancestor_thread_num(int level);
 EXTERN int omp_get_team_size(int level);
 
 EXTERN void omp_init_lock(omp_lock_t *lock);
 EXTERN void omp_init_nest_lock(omp_nest_lock_t *lock);
 EXTERN void omp_destroy_lock(omp_lock_t *lock);
 EXTERN void omp_destroy_nest_lock(omp_nest_lock_t *lock);
 EXTERN void omp_set_lock(omp_lock_t *lock);
 EXTERN void omp_set_nest_lock(omp_nest_lock_t *lock);

 
 EXTERN void __kmpc_get_team_static_memory(int16_t isSPMDExecutionMode,
                                           const void *buf, size_t size,
                                           int16_t is_shared, const void **res);
 
 EXTERN void __kmpc_restore_team_static_memory(int16_t isSPMDExecutionMode,
                                               int16_t is_shared);
 
+/// Allocate \p Bytes in "shareable" memory and return the address. Needs to be
+/// called balanced with __kmpc_free_shared like a stack (push/pop). Can be
+/// called by any thread, allocation happens per-thread.
+EXTERN void *__kmpc_alloc_shared(uint64_t Bytes);
+
+/// Deallocate \p Ptr. Needs to be called balanced with __kmpc_alloc_shared like
+/// a stack (push/pop). Can be called by any thread. \p Ptr must be allocated by
+/// __kmpc_alloc_shared by the same thread.
+EXTERN void __kmpc_free_shared(void *Ptr);
+
 #endif