[OpenMP] Create and use `__kmpc_is_generic_main_thread`

[jdoerfert]

Summary

In order to fold calls based on high-level knowledge and control flow
tracking it helps to expose the information as a runtime call. The
logic: !SPMD && getTID() == getMasterTID() was used in various places
and is now encapsulated in __kmpc_is_generic_main_thread. As part of
this rewrite we replaced eager computation of arguments with on-demand
computation, especially helpful if the calls can be folded and arguments
don't need to be computed consequently.

[tianshilei1992]

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

458

why do we still check location to tell if SPMD mode? Don't we want to completely remove it?

[jdoerfert]

Address comment

[tianshilei1992]

LGTM.

openmp/libomptarget/deviceRTLs/common/omptargeti.h

 }
 
 INLINE omptarget_nvptx_TaskDescr *getMyTopTaskDescriptor(int threadId) {
   return omptarget_nvptx_threadPrivateContext->GetTopLevelTaskDescr(threadId);
 }
 
 INLINE omptarget_nvptx_TaskDescr *
 getMyTopTaskDescriptor(bool isSPMDExecutionMode) {
-  return getMyTopTaskDescriptor(GetLogicalThreadIdInBlock(isSPMDExecutionMode));
+  return getMyTopTaskDescriptor(GetLogicalThreadIdInBlock());
 }
 
 ////////////////////////////////////////////////////////////////////////////////
 // Memory management runtime functions.
 ////////////////////////////////////////////////////////////////////////////////
 
 INLINE void omptarget_nvptx_SimpleMemoryManager::Release() {
   ASSERT0(LT_FUSSY, usedSlotIdx < MAX_SM,

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

 //
 //===----------------------------------------------------------------------===//
 #pragma omp declare target
 
 #include "common/omptarget.h"
 #include "target/shuffle.h"
 #include "target_impl.h"
 
-// Return true if this is the master thread.
-INLINE static bool IsMasterThread(bool isSPMDExecutionMode) {
-  return !isSPMDExecutionMode && GetMasterThreadID() == GetThreadIdInBlock();
-}
-
 ////////////////////////////////////////////////////////////////////////////////
 // Runtime functions for trunk data sharing scheme.
 ////////////////////////////////////////////////////////////////////////////////
 
 static constexpr unsigned MinBytes = 8;
 
 template <unsigned BytesPerThread, unsigned NThreads = MAX_THREADS_PER_TEAM>
 struct alignas(32) ThreadStackTy {
static void *__kmpc_alloc_for_warp(AllocTy Alloc, unsigned Bytes,
   FP[0] = __kmpc_impl_shfl_sync(CurActive, FP[0], LeaderID);
   if (sizeof(Ptr) == 8)
     FP[1] = __kmpc_impl_shfl_sync(CurActive, FP[1], LeaderID);
   return (void *)&((char *)(Ptr))[(GetLaneId() - LeaderID) * Bytes];
 }
 
 EXTERN void *__kmpc_alloc_shared(size_t Bytes) {
   Bytes = Bytes + (Bytes % MinBytes);
-  if (IsMasterThread(__kmpc_is_spmd_exec_mode())) {
+  int TID = GetThreadIdInBlock();
+  if (__kmpc_is_generic_main_thread(TID)) {
     // Main thread alone, use shared memory if space is available.
     if (MainSharedStack.Usage[0] + Bytes <= MainSharedStack.MaxSize) {
       void *Ptr = &MainSharedStack.Data[MainSharedStack.Usage[0]];
       MainSharedStack.Usage[0] += Bytes;
       MainSharedStack.Sizes[MainSharedStack.SizeUsage[0]++] = Bytes;
       return Ptr;
     }
   } else {
-    int TID = GetThreadIdInBlock();
     int WID = GetWarpId();
     unsigned WarpBytes = Bytes * WARPSIZE;
     auto AllocSharedStack = [&]() {
       unsigned WarpOffset = WID * WorkerSharedStack.MaxSizePerWarp;
       void *Ptr =
           &WorkerSharedStack.Data[WarpOffset + WorkerSharedStack.Usage[WID]];
       WorkerSharedStack.Usage[WID] += WarpBytes;
       WorkerSharedStack.Sizes[WorkerSharedStack.SizeUsage[WID]++] = WarpBytes;
       return Ptr;
     };
     if (TID < WorkerSharedStack.NumThreads &&
         WorkerSharedStack.Usage[WID] + WarpBytes <=
             WorkerSharedStack.MaxSizePerWarp)
       return __kmpc_alloc_for_warp(AllocSharedStack, Bytes, WarpBytes);
   }
   // Fallback to malloc
-  int TID = GetThreadIdInBlock();
   unsigned WarpBytes = Bytes * WARPSIZE;
   auto AllocGlobal = [&] {
     return SafeMalloc(WarpBytes, "AllocGlobalFallback");
   };
   return __kmpc_alloc_for_warp(AllocGlobal, Bytes, WarpBytes);
 }
 
 EXTERN void __kmpc_free_shared(void *Ptr) {

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

EXTERN int omp_get_thread_limit(void) {
   if (__kmpc_is_spmd_exec_mode())
     return GetNumberOfThreadsInBlock();
   int rc = threadLimit;
   PRINT(LD_IO, "call omp_get_thread_limit() return %d\n", rc);
   return rc;
 }
 
 EXTERN int omp_get_thread_num() {
-  bool isSPMDExecutionMode = __kmpc_is_spmd_exec_mode();
-  int tid = GetLogicalThreadIdInBlock(isSPMDExecutionMode);
-  int rc = GetOmpThreadId(tid, isSPMDExecutionMode);
+  int rc = GetOmpThreadId();
   PRINT(LD_IO, "call omp_get_thread_num() returns %d\n", rc);
   return rc;
 }
 
 EXTERN int omp_get_num_procs(void) {
   int rc = GetNumberOfProcsInDevice(__kmpc_is_spmd_exec_mode());
   PRINT(LD_IO, "call omp_get_num_procs() returns %d\n", rc);
   return rc;

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

INLINE static void dispatch_init(kmp_Ident *loc, int32_t threadId,
                                    kmp_sched_t schedule, T lb, T ub, ST st,
                                    ST chunk) {
     if (isRuntimeUninitialized()) {
       // In SPMD mode no need to check parallelism level - dynamic scheduling
       // may appear only in L2 parallel regions with lightweight runtime.
       ASSERT0(LT_FUSSY, __kmpc_is_spmd_exec_mode(), "Expected non-SPMD mode.");
       return;
     }
-    int tid = GetLogicalThreadIdInBlock(__kmpc_is_spmd_exec_mode());
+    int tid = GetLogicalThreadIdInBlock();
     omptarget_nvptx_TaskDescr *currTaskDescr = getMyTopTaskDescriptor(tid);
     T tnum = GetNumberOfOmpThreads(__kmpc_is_spmd_exec_mode());
     T tripCount = ub - lb + 1; // +1 because ub is inclusive
     ASSERT0(LT_FUSSY, threadId < tnum,
             "current thread is not needed here; error");
 
     /* Currently just ignore the monotonic and non-monotonic modifiers
      * (the compiler isn't producing them * yet anyway).
if (isRuntimeUninitialized()) {
       if (*plast)
         return DISPATCH_FINISHED;
       *plast = 1;
       return DISPATCH_NOTFINISHED;
     }
     // ID of a thread in its own warp
 
     // automatically selects thread or warp ID based on selected implementation
-    int tid = GetLogicalThreadIdInBlock(__kmpc_is_spmd_exec_mode());
+    int tid = GetLogicalThreadIdInBlock();
     ASSERT0(LT_FUSSY, gtid < GetNumberOfOmpThreads(__kmpc_is_spmd_exec_mode()),
             "current thread is not needed here; error");
     // retrieve schedule
     kmp_sched_t schedule =
         omptarget_nvptx_threadPrivateContext->ScheduleType(tid);
 
     // xxx reduce to one
     if (schedule == kmp_sched_static_chunk ||
         schedule == kmp_sched_static_nochunk) {
       T myLb = omptarget_nvptx_threadPrivateContext->NextLowerBound(tid);

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

static void __kmpc_spmd_kernel_deinit(bool RequiresFullRuntime) {
   }
 }
 
 // Return true if the current target region is executed in SPMD mode.
 EXTERN int8_t __kmpc_is_spmd_exec_mode() {
   return (execution_param & ModeMask) == Spmd;
 }
 
+EXTERN int8_t __kmpc_is_generic_main_thread(kmp_int32 Tid) {
+  return !__kmpc_is_spmd_exec_mode() && GetMasterThreadID() == Tid;
+}
+
 EXTERN bool __kmpc_kernel_parallel(void**WorkFn);
 
 static void __kmpc_target_region_state_machine(ident_t *Ident) {
 
   int TId = GetThreadIdInBlock();
   do {
     void* WorkFn = 0;
 

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

EXTERN void __kmpc_serialized_parallel(kmp_Ident *loc, uint32_t global_tid) {
 
   if (isRuntimeUninitialized()) {
     ASSERT0(LT_FUSSY, __kmpc_is_spmd_exec_mode(),
             "Expected SPMD mode with uninitialized runtime.");
     return;
   }
 
   // assume this is only called for nested parallel
-  int threadId = GetLogicalThreadIdInBlock(__kmpc_is_spmd_exec_mode());
+  int threadId = GetLogicalThreadIdInBlock();
 
   // unlike actual parallel, threads in the same team do not share
   // the workTaskDescr in this case and num threads is fixed to 1
 
   // get current task
   omptarget_nvptx_TaskDescr *currTaskDescr = getMyTopTaskDescriptor(threadId);
   currTaskDescr->SaveLoopData();
 
EXTERN void __kmpc_end_serialized_parallel(kmp_Ident *loc,
 
   if (isRuntimeUninitialized()) {
     ASSERT0(LT_FUSSY, __kmpc_is_spmd_exec_mode(),
             "Expected SPMD mode with uninitialized runtime.");
     return;
   }
 
   // pop stack
-  int threadId = GetLogicalThreadIdInBlock(__kmpc_is_spmd_exec_mode());
+  int threadId = GetLogicalThreadIdInBlock();
   omptarget_nvptx_TaskDescr *currTaskDescr = getMyTopTaskDescriptor(threadId);
   // set new top
   omptarget_nvptx_threadPrivateContext->SetTopLevelTaskDescr(
       threadId, currTaskDescr->GetPrevTaskDescr());
   // free
   SafeFree(currTaskDescr, "new seq parallel task");
   currTaskDescr = getMyTopTaskDescriptor(threadId);
   currTaskDescr->RestoreLoopData();
 }
 
 EXTERN uint16_t __kmpc_parallel_level(kmp_Ident *loc, uint32_t global_tid) {
   PRINT0(LD_IO, "call to __kmpc_parallel_level\n");
 
   return parallelLevel[GetWarpId()] & (OMP_ACTIVE_PARALLEL_LEVEL - 1);
 }
 
 // This kmpc call returns the thread id across all teams. It's value is
 // cached by the compiler and used when calling the runtime. On nvptx
 // it's cheap to recalculate this value so we never use the result
 // of this call.
 EXTERN int32_t __kmpc_global_thread_num(kmp_Ident *loc) {
-  int tid = GetLogicalThreadIdInBlock(__kmpc_is_spmd_exec_mode());
-  return GetOmpThreadId(tid, __kmpc_is_spmd_exec_mode());
+  return GetOmpThreadId();
 }
 
 ////////////////////////////////////////////////////////////////////////////////
 // push params
 ////////////////////////////////////////////////////////////////////////////////
 
 EXTERN void __kmpc_push_num_threads(kmp_Ident *loc, int32_t tid,
                                     int32_t num_threads) {
   PRINT(LD_IO, "call kmpc_push_num_threads %d\n", num_threads);
   ASSERT0(LT_FUSSY, isRuntimeInitialized(),
           "Runtime must be initialized.");
-  tid = GetLogicalThreadIdInBlock(__kmpc_is_spmd_exec_mode());
+  tid = GetLogicalThreadIdInBlock();
   omptarget_nvptx_threadPrivateContext->NumThreadsForNextParallel(tid) =
       num_threads;
 }
 
 // Do nothing. The host guarantees we started the requested number of
 // teams and we only need inspection of gridDim.
 
 EXTERN void __kmpc_push_num_teams(kmp_Ident *loc, int32_t tid,

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

 }
 #endif
 
 INLINE
 static int32_t nvptx_parallel_reduce_nowait(
     int32_t global_tid, int32_t num_vars, size_t reduce_size, void *reduce_data,
     kmp_ShuffleReductFctPtr shflFct, kmp_InterWarpCopyFctPtr cpyFct,
     bool isSPMDExecutionMode, bool isRuntimeUninitialized) {
-  uint32_t BlockThreadId = GetLogicalThreadIdInBlock(isSPMDExecutionMode);
+  uint32_t BlockThreadId = GetLogicalThreadIdInBlock();
   uint32_t NumThreads = GetNumberOfOmpThreads(isSPMDExecutionMode);
   if (NumThreads == 1)
     return 1;
     /*
      * This reduce function handles reduction within a team. It handles
      * parallel regions in both L1 and L2 parallelism levels. It also
      * supports Generic, SPMD, and NoOMP modes.
      *

 EXTERN int32_t __kmpc_nvptx_teams_reduce_nowait_v2(
     kmp_Ident *loc, int32_t global_tid, void *global_buffer,
     int32_t num_of_records, void *reduce_data, kmp_ShuffleReductFctPtr shflFct,
     kmp_InterWarpCopyFctPtr cpyFct, kmp_ListGlobalFctPtr lgcpyFct,
     kmp_ListGlobalFctPtr lgredFct, kmp_ListGlobalFctPtr glcpyFct,
     kmp_ListGlobalFctPtr glredFct) {
 
   // Terminate all threads in non-SPMD mode except for the master thread.
-  if (!__kmpc_is_spmd_exec_mode() && GetThreadIdInBlock() != GetMasterThreadID())
+  if (!__kmpc_is_spmd_exec_mode() &&
+      !__kmpc_is_generic_main_thread(GetThreadIdInBlock()))
     return 0;
 
-  uint32_t ThreadId = GetLogicalThreadIdInBlock(__kmpc_is_spmd_exec_mode());
+  uint32_t ThreadId = GetLogicalThreadIdInBlock();
 
   // In non-generic mode all workers participate in the teams reduction.
   // In generic mode only the team master participates in the teams
   // reduction because the workers are waiting for parallel work.
   uint32_t NumThreads =
       __kmpc_is_spmd_exec_mode() ? GetNumberOfOmpThreads(/*isSPMDExecutionMode=*/true)
                          : /*Master thread only*/ 1;
   uint32_t TeamId = GetBlockIdInKernel();

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

 
 ////////////////////////////////////////////////////////////////////////////////
 // get thread id in team
 
 // This function may be called in a parallel region by the workers
 // or a serial region by the master.  If the master (whose CUDA thread
 // id is GetMasterThreadID()) calls this routine, we return 0 because
 // it is a shadow for the first worker.
-int GetLogicalThreadIdInBlock(bool isSPMDExecutionMode) {
+int GetLogicalThreadIdInBlock() {
   // Implemented using control flow (predication) instead of with a modulo
   // operation.
   int tid = GetThreadIdInBlock();
-  if (!isSPMDExecutionMode && tid >= GetMasterThreadID())
+  if (__kmpc_is_generic_main_thread(tid))
     return 0;
   else
     return tid;
 }
 
 ////////////////////////////////////////////////////////////////////////////////
 //
 // OpenMP Thread Support Layer
 //
 ////////////////////////////////////////////////////////////////////////////////
 
-int GetOmpThreadId(int threadId, bool isSPMDExecutionMode) {
+int GetOmpThreadId() {
+  int tid = GetThreadIdInBlock();
+  if (__kmpc_is_generic_main_thread(tid))
+    return 0;
   // omp_thread_num
   int rc;
   if ((parallelLevel[GetWarpId()] & (OMP_ACTIVE_PARALLEL_LEVEL - 1)) > 1) {
     rc = 0;
-  } else if (isSPMDExecutionMode) {
-    rc = GetThreadIdInBlock();
+  } else if (__kmpc_is_spmd_exec_mode()) {
+    rc = tid;
   } else {
     omptarget_nvptx_TaskDescr *currTaskDescr =
-        omptarget_nvptx_threadPrivateContext->GetTopLevelTaskDescr(threadId);
+        omptarget_nvptx_threadPrivateContext->GetTopLevelTaskDescr(tid);
     ASSERT0(LT_FUSSY, currTaskDescr, "expected a top task descr");
     rc = currTaskDescr->ThreadId();
   }
   return rc;
 }
 
 int GetNumberOfOmpThreads(bool isSPMDExecutionMode) {
   // omp_num_threads

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

 }
 
 EXTERN void __kmpc_barrier(kmp_Ident *loc_ref, int32_t tid) {
   if (isRuntimeUninitialized()) {
     ASSERT0(LT_FUSSY, __kmpc_is_spmd_exec_mode(),
             "Expected SPMD mode with uninitialized runtime.");
     __kmpc_barrier_simple_spmd(loc_ref, tid);
   } else {
-    tid = GetLogicalThreadIdInBlock(__kmpc_is_spmd_exec_mode());
+    tid = GetLogicalThreadIdInBlock();
     int numberOfActiveOMPThreads =
         GetNumberOfOmpThreads(__kmpc_is_spmd_exec_mode());
     if (numberOfActiveOMPThreads > 1) {
       if (__kmpc_is_spmd_exec_mode()) {
         __kmpc_barrier_simple_spmd(loc_ref, tid);
       } else {
         // The #threads parameter must be rounded up to the WARPSIZE.
         int threads =

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

omptarget_nvptx_ExplicitTaskDescr *newExplicitTaskDescr =
           newKmpTaskDescr, sizeof(omptarget_nvptx_TaskDescr));
   ASSERT0(LT_FUSSY, &newExplicitTaskDescr->kmpTaskDescr == newKmpTaskDescr,
           "bad assumptions");
   omptarget_nvptx_TaskDescr *newTaskDescr = &newExplicitTaskDescr->taskDescr;
   ASSERT0(LT_FUSSY, (uint64_t)newTaskDescr == (uint64_t)newExplicitTaskDescr,
           "bad assumptions");
 
   // 2. push new context: update new task descriptor
-  int tid = GetLogicalThreadIdInBlock(__kmpc_is_spmd_exec_mode());
+  int tid = GetLogicalThreadIdInBlock();
   omptarget_nvptx_TaskDescr *parentTaskDescr = getMyTopTaskDescriptor(tid);
   newTaskDescr->CopyForExplicitTask(parentTaskDescr);
   // set new task descriptor as top
   omptarget_nvptx_threadPrivateContext->SetTopLevelTaskDescr(tid, newTaskDescr);
 
   // 3. call sub
   PRINT(LD_TASK, "call task sub 0x%llx(task descr 0x%llx)\n",
         (unsigned long long)newKmpTaskDescr->sub,
omptarget_nvptx_ExplicitTaskDescr *newExplicitTaskDescr =
           newKmpTaskDescr, sizeof(omptarget_nvptx_TaskDescr));
   ASSERT0(LT_FUSSY, &newExplicitTaskDescr->kmpTaskDescr == newKmpTaskDescr,
           "bad assumptions");
   omptarget_nvptx_TaskDescr *newTaskDescr = &newExplicitTaskDescr->taskDescr;
   ASSERT0(LT_FUSSY, (uint64_t)newTaskDescr == (uint64_t)newExplicitTaskDescr,
           "bad assumptions");
 
   // 2. push new context: update new task descriptor
-  int tid = GetLogicalThreadIdInBlock(__kmpc_is_spmd_exec_mode());
+  int tid = GetLogicalThreadIdInBlock();
   omptarget_nvptx_TaskDescr *parentTaskDescr = getMyTopTaskDescriptor(tid);
   newTaskDescr->CopyForExplicitTask(parentTaskDescr);
   // set new task descriptor as top
   omptarget_nvptx_threadPrivateContext->SetTopLevelTaskDescr(tid, newTaskDescr);
   // 3... noting to call... is inline
   // 4 & 5 ... done in complete
 }
 
ASSERT0(LT_FUSSY, &newExplicitTaskDescr->kmpTaskDescr == newKmpTaskDescr,
           "bad assumptions");
   omptarget_nvptx_TaskDescr *newTaskDescr = &newExplicitTaskDescr->taskDescr;
   ASSERT0(LT_FUSSY, (uint64_t)newTaskDescr == (uint64_t)newExplicitTaskDescr,
           "bad assumptions");
   // 2. get parent
   omptarget_nvptx_TaskDescr *parentTaskDescr = newTaskDescr->GetPrevTaskDescr();
   // 3... noting to call... is inline
   // 4. pop context
-  int tid = GetLogicalThreadIdInBlock(__kmpc_is_spmd_exec_mode());
+  int tid = GetLogicalThreadIdInBlock();
   omptarget_nvptx_threadPrivateContext->SetTopLevelTaskDescr(tid,
                                                              parentTaskDescr);
   // 5. free
   SafeFree(newExplicitTaskDescr, "explicit task descriptor");
 }
 
 EXTERN void __kmpc_omp_wait_deps(kmp_Ident *loc, uint32_t global_tid,
                                  int32_t depNum, void *depList,

openmp/libomptarget/deviceRTLs/common/support.h

 bool isRuntimeUninitialized();
 bool isRuntimeInitialized();
 
 ////////////////////////////////////////////////////////////////////////////////
 // get info from machine
 ////////////////////////////////////////////////////////////////////////////////
 
 // get global ids to locate tread/team info (constant regardless of OMP)
-int GetLogicalThreadIdInBlock(bool isSPMDExecutionMode);
+int GetLogicalThreadIdInBlock();
 int GetMasterThreadID();
 int GetNumberOfWorkersInTeam();
 
 // get OpenMP thread and team ids
-int GetOmpThreadId(int threadId,
-                   bool isSPMDExecutionMode); // omp_thread_num
+int GetOmpThreadId();                         // omp_thread_num
 int GetOmpTeamId();                           // omp_team_num
 
 // get OpenMP number of threads and team
 int GetNumberOfOmpThreads(bool isSPMDExecutionMode); // omp_num_threads
 int GetNumberOfOmpTeams();                           // omp_num_teams
 
 // get OpenMP number of procs
 int GetNumberOfProcsInTeam(bool isSPMDExecutionMode);

openmp/libomptarget/deviceRTLs/interface.h

 EXTERN void __kmpc_parallel_51(ident_t *ident, kmp_int32 global_tid,
                                kmp_int32 if_expr, kmp_int32 num_threads,
                                int proc_bind, void *fn, void *wrapper_fn,
                                void **args, size_t nargs);
 
 // SPMD execution mode interrogation function.
 EXTERN int8_t __kmpc_is_spmd_exec_mode();
 
+/// Return true if the hardware thread id \p Tid represents the OpenMP main
+/// thread in generic mode outside of a parallel region.
+EXTERN int8_t __kmpc_is_generic_main_thread(kmp_int32 Tid);
+
 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