[OpenMP] Simplify variable sharing and increase shared memory size

[jdoerfert]

Summary

In order to avoid malloc/free, up to NUM_SHARED_VARIABLES_IN_SHARED_MEM
(=64) variables are communicated in dedicated shared memory instead. The
simplification does avoid the need for an "init" and requires "deinit"
only if we ever communicate more than NUM_SHARED_VARIABLES_IN_SHARED_MEM
variables.

[jdoerfert]

Fix minor problems

[tianshilei1992]

LGTM

[tianshilei1992]

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

158

but I don't like the name…Starting with _ is not a typical LLVM name.

[jdoerfert]

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

158

I will use static and remove the _

openmp/libomptarget/deviceRTLs/amdgcn/src/target_impl.h

 
 ////////////////////////////////////////////////////////////////////////////////
 // The following def must match the absolute limit hardwired in the host RTL
 // max number of threads per team
 #define MAX_THREADS_PER_TEAM 1024
 
 #define WARPSIZE 64
 
-// Maximum number of preallocated arguments to an outlined parallel/simd
-// function. Anything more requires dynamic memory allocation.
-#define MAX_SHARED_ARGS 20
-
 // Maximum number of omp state objects per SM allocated statically in global
 // memory.
 #define OMP_STATE_COUNT 32
 #define MAX_SM 64
 
 #define OMP_ACTIVE_PARALLEL_LEVEL 128
 
 // Data sharing related quantities, need to match what is used in the compiler.

openmp/libomptarget/deviceRTLs/common/omptarget.h

 // used by dynamic scheduling
 #define FINISHED 0
 #define NOT_FINISHED 1
 #define LAST_CHUNK 2
 
 #define BARRIER_COUNTER 0
 #define ORDERED_COUNTER 1
 
-// arguments needed for L0 parallelism only.
-class omptarget_nvptx_SharedArgs {
-public:
-  // All these methods must be called by the master thread only.
-  INLINE void Init() {
-    args = buffer;
-    nArgs = MAX_SHARED_ARGS;
-  }
-  INLINE void DeInit() {
-    // Free any memory allocated for outlined parallel function with a large
-    // number of arguments.
-    if (nArgs > MAX_SHARED_ARGS) {
-      SafeFree(args, "new extended args");
-      Init();
-    }
-  }
-  INLINE void EnsureSize(size_t size) {
-    if (size > nArgs) {
-      if (nArgs > MAX_SHARED_ARGS) {
-        SafeFree(args, "new extended args");
-      }
-      args = (void **)SafeMalloc(size * sizeof(void *), "new extended args");
-      nArgs = size;
-    }
-  }
-  // Called by all threads.
-  INLINE void **GetArgs() const { return args; };
-
-private:
-  // buffer of pre-allocated arguments.
-  void *buffer[MAX_SHARED_ARGS];
-  // pointer to arguments buffer.
-  // starts off as a pointer to 'buffer' but can be dynamically allocated.
-  void **args;
-  // starts off as MAX_SHARED_ARGS but can increase in size.
-  uint32_t nArgs;
-};
-
-extern omptarget_nvptx_SharedArgs EXTERN_SHARED(omptarget_nvptx_globalArgs);
-
 // Worker slot type which is initialized with the default worker slot
 // size of 4*32 bytes.
 struct __kmpc_data_sharing_slot {
   __kmpc_data_sharing_slot *Next;
   __kmpc_data_sharing_slot *Prev;
   void *PrevSlotStackPtr;
   void *DataEnd;
   char Data[DS_Worker_Warp_Slot_Size];

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

for (unsigned i = 0; i < MainSharedStack.NumWarps; ++i) {
     MainSharedStack.Usage[i] = 0;
   }
   for (unsigned i = 0; i < WorkerSharedStack.NumWarps; ++i) {
     WorkerSharedStack.SizeUsage[i] = 0;
     WorkerSharedStack.Usage[i] = 0;
   }
 }
 
+/// Allocate storage in shared memory to communicate arguments from the main
+/// thread to the workers in generic mode. If we exceed
+/// NUM_SHARED_VARIABLES_IN_SHARED_MEM we will malloc space for communication.
+#define NUM_SHARED_VARIABLES_IN_SHARED_MEM 64
+
+[[clang::loader_uninitialized]] static void
+    *SharedMemVariableSharingSpace[NUM_SHARED_VARIABLES_IN_SHARED_MEM];
+#pragma omp allocate(SharedMemVariableSharingSpace)                            \
+    allocator(omp_pteam_mem_alloc)
+[[clang::loader_uninitialized]] static void **SharedMemVariableSharingSpacePtr;
+#pragma omp allocate(SharedMemVariableSharingSpacePtr)                         \
+    allocator(omp_pteam_mem_alloc)
+
 // 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();
+  if (nArgs <= NUM_SHARED_VARIABLES_IN_SHARED_MEM) {
+    SharedMemVariableSharingSpacePtr = &SharedMemVariableSharingSpace[0];
+  } else {
+    SharedMemVariableSharingSpacePtr =
+        (void **)SafeMalloc(nArgs * sizeof(void *), "new extended args");
+  }
+  *GlobalArgs = SharedMemVariableSharingSpacePtr;
 }
 
 // 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 (SharedMemVariableSharingSpacePtr != &SharedMemVariableSharingSpace[0])
+    SafeFree(SharedMemVariableSharingSpacePtr, "new extended args");
 }
 
 // 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();
+  *GlobalArgs = SharedMemVariableSharingSpacePtr;
 }
 
 // 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/common/src/omp_data.cu

 ////////////////////////////////////////////////////////////////////////////////
 uint32_t SHARED(execution_param);
 
 ////////////////////////////////////////////////////////////////////////////////
 // Scratchpad for teams reduction.
 ////////////////////////////////////////////////////////////////////////////////
 void *SHARED(ReductionScratchpadPtr);
 
-////////////////////////////////////////////////////////////////////////////////
-// Data sharing related variables.
-////////////////////////////////////////////////////////////////////////////////
-omptarget_nvptx_SharedArgs SHARED(omptarget_nvptx_globalArgs);
-
 #pragma omp end declare target

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

omptarget_nvptx_threadPrivateContext->SetTopLevelTaskDescr(
       threadId, currTeamDescr.LevelZeroTaskDescr());
 
   // set number of threads and thread limit in team to started value
   omptarget_nvptx_TaskDescr *currTaskDescr =
       omptarget_nvptx_threadPrivateContext->GetTopLevelTaskDescr(threadId);
   nThreads = GetNumberOfWorkersInTeam();
   threadLimit = nThreads;
 
-  omptarget_nvptx_globalArgs.Init();
-
   __kmpc_data_sharing_init_stack();
   __kmpc_impl_target_init();
 }
 
 static void __kmpc_generic_kernel_deinit() {
   PRINT0(LD_IO, "call to __kmpc_kernel_deinit\n");
   // Enqueue omp state object for use by another team.
   int slot = usedSlotIdx;

openmp/libomptarget/deviceRTLs/nvptx/src/target_impl.h

 
 ////////////////////////////////////////////////////////////////////////////////
 // The following def must match the absolute limit hardwired in the host RTL
 // max number of threads per team
 #define MAX_THREADS_PER_TEAM 1024
 
 #define WARPSIZE 32
 
-// Maximum number of preallocated arguments to an outlined parallel/simd
-// function. Anything more requires dynamic memory allocation.
-#define MAX_SHARED_ARGS 20
-
 // Maximum number of omp state objects per SM allocated statically in global
 // memory.
 #if __CUDA_ARCH__ >= 600
 #define OMP_STATE_COUNT 32
 #else
 #define OMP_STATE_COUNT 16
 #endif