[OpenMP][WIP] Move run-sched-var ICV to ICVStateTy to make TaskDescr obsolete

[jdoerfert]

Summary

Note: This is not for review.

As follow up on D98349 and D98678, this patch shows that we can
eliminate the members of TaskDescr by introducing a run-sched-var in the
ICVStateTy. Other functionality is already covered by the ICVStateTy.

The task.cu uses need to be converted to the new system, e.g., by
calling enterDataEnvironment instead of the task copy stuff. Other
than that, this looks like we only need to remove all the old plumbing.

Unclear how to merge/split D98349 and this one and the cleanup follow
up.

openmp/libomptarget/deviceRTLs/common/include/ICVs.h

 //===--------- ICVs.h - OpenMP ICV handling ----------------------- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 //
 //===----------------------------------------------------------------------===//
+
 #ifndef OMPTARGET_ICVS_H
 #define OMPTARGET_ICVS_H
 
+#include "interface.h"
+#include <bits/stdint-uintn.h>
+
 struct ICVStateTy {
   int nthreads_var;
 
   int levels_var;
 
   /// The `active-level` describes which of the parallel levels counted with the
   /// `levels-var` is active. There can only be one.
   ///
   /// active-levels-var is 1, if active_level is not 0, otherweise it is 0.
   int active_level;
 
+  /// Don't use UB type punning!
+  ///
+  ///{
+  struct RunSchedVarEncodingTy {
+    omp_sched_t ScheduleKind;
+    int ChunkSize;
+  };
+
+  uint64_t run_sched_var;
+
+  static_assert(
+      sizeof(run_sched_var) == sizeof(RunSchedVarEncodingTy),
+      "Schedule encoding is supposed to cover the entire run-sched-var ICV!");
+  ///}
+
   static bool ensureICVStateForThread(unsigned TId);
 
   static int &getICVForThread(int ICVStateTy::*Var);
   static int incICVForThread(int ICVStateTy::*Var, int UpdateVal);
   static int setICVForThread(int ICVStateTy::*Var, int UpdateVal);
+
+  static uint64_t &getICVForThread(uint64_t ICVStateTy::*Var);
+  static uint64_t incICVForThread(uint64_t ICVStateTy::*Var,
+                                  uint64_t UpdateVal);
+  static uint64_t setICVForThread(uint64_t ICVStateTy::*Var,
+                                  uint64_t UpdateVal);
 };
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
 #pragma omp declare target
 

openmp/libomptarget/deviceRTLs/common/omptarget.h

 
 extern DataSharingStateTy EXTERN_SHARED(DataSharingState);
 
 ////////////////////////////////////////////////////////////////////////////////
 // task ICV and (implicit & explicit) task state
 
 class omptarget_nvptx_TaskDescr {
 public:
-  // methods for flags
-  INLINE omp_sched_t GetRuntimeSched() const;
-  INLINE void SetRuntimeSched(omp_sched_t sched);
-  INLINE int InParallelRegion() const { return items.flags & TaskDescr_InPar; }
-  INLINE int InL2OrHigherParallelRegion() const {
-    return items.flags & TaskDescr_InParL2P;
-  }
-  INLINE int IsParallelConstruct() const {
-    return items.flags & TaskDescr_IsParConstr;
-  }
-  INLINE int IsTaskConstruct() const { return !IsParallelConstruct(); }
   // methods for other fields
-  INLINE uint16_t &ThreadId() { return items.threadId; }
-  INLINE uint64_t &RuntimeChunkSize() { return items.runtimeChunkSize; }
   INLINE omptarget_nvptx_TaskDescr *GetPrevTaskDescr() const { return prev; }
   INLINE void SetPrevTaskDescr(omptarget_nvptx_TaskDescr *taskDescr) {
     prev = taskDescr;
   }
   // init & copy
   INLINE void InitLevelZeroTaskDescr();
   INLINE void InitLevelOneTaskDescr(omptarget_nvptx_TaskDescr *parentTaskDescr);
   INLINE void Copy(omptarget_nvptx_TaskDescr *sourceTaskDescr);
   INLINE void CopyData(omptarget_nvptx_TaskDescr *sourceTaskDescr);
   INLINE void CopyParent(omptarget_nvptx_TaskDescr *parentTaskDescr);
   INLINE void CopyForExplicitTask(omptarget_nvptx_TaskDescr *parentTaskDescr);
   INLINE void CopyToWorkDescr(omptarget_nvptx_TaskDescr *masterTaskDescr);
   INLINE void CopyFromWorkDescr(omptarget_nvptx_TaskDescr *workTaskDescr);
   INLINE void CopyConvergentParent(omptarget_nvptx_TaskDescr *parentTaskDescr,
                                    uint16_t tid, uint16_t tnum);
 
 private:
-  // bits for flags: (6 used, 2 free)
-  //   3 bits (SchedMask) for runtime schedule
-  //   1 bit (InPar) if this thread has encountered one or more parallel region
-  //   1 bit (IsParConstr) if ICV for a parallel region (false = explicit task)
-  //   1 bit (InParL2+) if this thread has encountered L2 or higher parallel
-  //   region
-  static const uint8_t TaskDescr_SchedMask = (0x1 | 0x2 | 0x4);
-  static const uint8_t TaskDescr_InPar = 0x10;
-  static const uint8_t TaskDescr_IsParConstr = 0x20;
-  static const uint8_t TaskDescr_InParL2P = 0x40;
-
-  struct TaskDescr_items {
-    uint8_t flags; // 6 bit used (see flag above)
-    uint8_t unused;
-    uint16_t threadId;         // thread id
-    uint64_t runtimeChunkSize; // runtime chunk size
-  } items;
   omptarget_nvptx_TaskDescr *prev;
 };
 
 // build on kmp
 typedef struct omptarget_nvptx_ExplicitTaskDescr {
   omptarget_nvptx_TaskDescr
       taskDescr; // omptarget_nvptx task description (must be first)
   kmp_TaskDescr kmpTaskDescr; // kmp task description (must be last)

openmp/libomptarget/deviceRTLs/common/omptargeti.h

 // and functions.
 //
 //===----------------------------------------------------------------------===//
 
 ////////////////////////////////////////////////////////////////////////////////
 // Task Descriptor
 ////////////////////////////////////////////////////////////////////////////////
 
-INLINE omp_sched_t omptarget_nvptx_TaskDescr::GetRuntimeSched() const {
-  // sched starts from 1..4; encode it as 0..3; so add 1 here
-  uint8_t rc = (items.flags & TaskDescr_SchedMask) + 1;
-  return (omp_sched_t)rc;
-}
-
-INLINE void omptarget_nvptx_TaskDescr::SetRuntimeSched(omp_sched_t sched) {
-  // sched starts from 1..4; encode it as 0..3; so sub 1 here
-  uint8_t val = ((uint8_t)sched) - 1;
-  // clear current sched
-  items.flags &= ~TaskDescr_SchedMask;
-  // set new sched
-  items.flags |= val;
-}
-
 INLINE void omptarget_nvptx_TaskDescr::InitLevelZeroTaskDescr() {
-  // slow method
-  // flag:
-  //   default sched is static,
-  //   dyn is off (unused now anyway, but may need to sample from host ?)
-  //   not in parallel
-
-  items.flags = 0;
-  items.threadId = 0;         // is master
-  items.runtimeChunkSize = 1; // preferred chunking statik with chunk 1
 }
 
 // This is called when all threads are started together in SPMD mode.
 // OMP directives include target parallel, target distribute parallel for, etc.
 INLINE void omptarget_nvptx_TaskDescr::InitLevelOneTaskDescr(
     omptarget_nvptx_TaskDescr *parentTaskDescr) {
-  // slow method
-  // flag:
-  //   default sched is static,
-  //   dyn is off (unused now anyway, but may need to sample from host ?)
-  //   in L1 parallel
-
-  items.flags = TaskDescr_InPar | TaskDescr_IsParConstr; // set flag to parallel
-  items.threadId =
-      GetThreadIdInBlock();   // get ids from cuda (only called for 1st level)
-  items.runtimeChunkSize = 1; // preferred chunking statik with chunk 1
   prev = parentTaskDescr;
 }
 
 INLINE void omptarget_nvptx_TaskDescr::CopyData(
     omptarget_nvptx_TaskDescr *sourceTaskDescr) {
-  items = sourceTaskDescr->items;
 }
 
 INLINE void
 omptarget_nvptx_TaskDescr::Copy(omptarget_nvptx_TaskDescr *sourceTaskDescr) {
   CopyData(sourceTaskDescr);
   prev = sourceTaskDescr->prev;
 }
 
 INLINE void omptarget_nvptx_TaskDescr::CopyParent(
     omptarget_nvptx_TaskDescr *parentTaskDescr) {
   CopyData(parentTaskDescr);
   prev = parentTaskDescr;
 }
 
 INLINE void omptarget_nvptx_TaskDescr::CopyForExplicitTask(
     omptarget_nvptx_TaskDescr *parentTaskDescr) {
   CopyParent(parentTaskDescr);
-  items.flags = items.flags & ~TaskDescr_IsParConstr;
-  ASSERT0(LT_FUSSY, IsTaskConstruct(), "expected task");
 }
 
 INLINE void omptarget_nvptx_TaskDescr::CopyToWorkDescr(
     omptarget_nvptx_TaskDescr *masterTaskDescr) {
   CopyParent(masterTaskDescr);
-  // overwrite specific items;
-  items.flags |=
-      TaskDescr_InPar | TaskDescr_IsParConstr; // set flag to parallel
 }
 
 INLINE void omptarget_nvptx_TaskDescr::CopyFromWorkDescr(
     omptarget_nvptx_TaskDescr *workTaskDescr) {
   Copy(workTaskDescr);
-  //
-  // overwrite specific items;
-  //
-  // The threadID should be GetThreadIdInBlock() % GetMasterThreadID().
-  // This is so that the serial master (first lane in the master warp)
-  // gets a threadId of 0.
-  // However, we know that this function is always called in a parallel
-  // region where only workers are active.  The serial master thread
-  // never enters this region.  When a parallel region is executed serially,
-  // the threadId is set to 0 elsewhere and the kmpc_serialized_* functions
-  // are called, which never activate this region.
-  items.threadId =
-      GetThreadIdInBlock(); // get ids from cuda (only called for 1st level)
 }
 
 INLINE void omptarget_nvptx_TaskDescr::CopyConvergentParent(
     omptarget_nvptx_TaskDescr *parentTaskDescr, uint16_t tid, uint16_t tnum) {
   CopyParent(parentTaskDescr);
-  items.flags |= TaskDescr_InParL2P; // In L2+ parallelism
-  items.threadId = tid;
 }
 
 ////////////////////////////////////////////////////////////////////////////////
 // Thread Private Context
 ////////////////////////////////////////////////////////////////////////////////
 
 INLINE omptarget_nvptx_TaskDescr *
 omptarget_nvptx_ThreadPrivateContext::GetTopLevelTaskDescr(int tid) const {

openmp/libomptarget/deviceRTLs/common/src/ICVs.cpp

 
 #include "ICVs.h"
 #include "Mapping.h"
 #include "TeamState.h"
 #include "ThreadState.h"
 #include "omptarget.h"
 #include "support.h"
 #include "target_interface.h"
+#include <bits/stdint-uintn.h>
 #include <string.h>
 
 #define ICV_DEBUG(...)
 
 using namespace omp;
 
 #pragma omp declare target
 

 int omp_get_active_level(void) {
   return !!ICVStateTy::getICVForThread(&ICVStateTy::active_level);
 }
 
 int omp_in_parallel(void) {
   return !!ICVStateTy::getICVForThread(&ICVStateTy::active_level);
 }
 
+void omp_get_schedule(omp_sched_t *ScheduleKind, int *ChunkSize) {
+  uint64_t RunSchedVar =
+      ICVStateTy::getICVForThread(&ICVStateTy::run_sched_var);
+  ICVStateTy::RunSchedVarEncodingTy RunSchedVarEncoding;
+  memcpy(&RunSchedVarEncoding, &RunSchedVar, sizeof(RunSchedVarEncoding));
+  *ScheduleKind = RunSchedVarEncoding.ScheduleKind;
+  *ChunkSize = RunSchedVarEncoding.ChunkSize;
+}
+
+void omp_set_schedule(omp_sched_t ScheduleKind, int ChunkSize) {
+  ICVStateTy::RunSchedVarEncodingTy RunSchedVarEncoding;
+  RunSchedVarEncoding.ScheduleKind = ScheduleKind;
+  RunSchedVarEncoding.ChunkSize = ChunkSize;
+  uint64_t RunSchedVar;
+  memcpy(&RunSchedVar, &RunSchedVarEncoding, sizeof(RunSchedVarEncoding));
+  ICVStateTy::setICVForThread(&ICVStateTy::run_sched_var, RunSchedVar);
+}
+
 static int returnValIfLevelIsActive(int Level, int Val, int DefaultVal,
                                     int OutOfBoundsVal = -1) {
   if (Level == 0)
     return DefaultVal;
   int LevelsVar = omp_get_level();
   if (Level < 0 || Level > LevelsVar)
     return OutOfBoundsVal;
   int ActiveLevel = ICVStateTy::getICVForThread(&ICVStateTy::active_level);
void ThreadStateTy::dropForThread(unsigned TId) {
   if (!ThreadStates[TId])
     return;
 
   // assert(!ThreadStates[TId]->PreviousThreadState && "leftover thread state");
   free(ThreadStates[TId]);
   ThreadStates[TId] = nullptr;
 }
 
-DEVICE TeamStateTy SHARED(omp::TeamState);
+TeamStateTy SHARED(omp::TeamState);
 
-[[clang::loader_uninitialized]] DEVICE ThreadStateTy
+[[clang::loader_uninitialized]] ThreadStateTy
     *omp::ThreadStates[MAX_THREADS_PER_TEAM];
 #pragma omp allocate(omp::ThreadStates) allocator(omp_pteam_mem_alloc)
 
 #pragma omp end declare target

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

 }
 
 EXTERN int omp_get_max_active_levels(void) {
   int rc = 1;
   PRINT(LD_IO, "call omp_get_max_active_levels() returns %d\n", rc);
   return rc;
 }
 
-EXTERN void omp_get_schedule(omp_sched_t *kind, int *modifier) {
-  if (isRuntimeUninitialized()) {
-    ASSERT0(LT_FUSSY, isSPMDMode(),
-            "Expected SPMD mode only with uninitialized runtime.");
-    *kind = omp_sched_static;
-    *modifier = 1;
-  } else {
-    omptarget_nvptx_TaskDescr *currTaskDescr =
-        getMyTopTaskDescriptor(isSPMDMode());
-    *kind = currTaskDescr->GetRuntimeSched();
-    *modifier = currTaskDescr->RuntimeChunkSize();
-  }
-  PRINT(LD_IO, "call omp_get_schedule returns sched %d and modif %d\n",
-        (int)*kind, *modifier);
-}
-
-EXTERN void omp_set_schedule(omp_sched_t kind, int modifier) {
-  PRINT(LD_IO, "call omp_set_schedule(sched %d, modif %d)\n", (int)kind,
-        modifier);
-  if (isRuntimeUninitialized()) {
-    ASSERT0(LT_FUSSY, isSPMDMode(),
-            "Expected SPMD mode only with uninitialized runtime.");
-    return;
-  }
-  if (kind >= omp_sched_static && kind < omp_sched_auto) {
-    omptarget_nvptx_TaskDescr *currTaskDescr =
-        getMyTopTaskDescriptor(isSPMDMode());
-    currTaskDescr->SetRuntimeSched(kind);
-    currTaskDescr->RuntimeChunkSize() = modifier;
-    PRINT(LD_IOD, "omp_set_schedule did set sched %d & modif %" PRIu64 "\n",
-          (int)currTaskDescr->GetRuntimeSched(),
-          currTaskDescr->RuntimeChunkSize());
-  }
-}
-
 EXTERN omp_proc_bind_t omp_get_proc_bind(void) {
   PRINT0(LD_IO, "call omp_get_proc_bin() is true, regardless on state\n");
   return omp_proc_bind_true;
 }
 
 EXTERN int omp_get_num_places(void) {
   PRINT0(LD_IO, "call omp_get_num_places() returns 0\n");
   return 0;

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

if (tnum == 1 || tripCount <= 1 || OrderedSchedule(schedule)) {
         __kmpc_barrier(loc, threadId);
       PRINT(LD_LOOP,
             "go sequential as tnum=%ld, trip count %lld, ordered sched=%d\n",
             (long)tnum, (long long)tripCount, (int)schedule);
       schedule = kmp_sched_static_chunk;
       chunk = tripCount; // one thread gets the whole loop
     } else if (schedule == kmp_sched_runtime) {
       // process runtime
-      omp_sched_t rtSched = currTaskDescr->GetRuntimeSched();
-      chunk = currTaskDescr->RuntimeChunkSize();
+      omp_sched_t rtSched;
+      int ChunkInt;
+      omp_get_schedule(&rtSched, &ChunkInt);
+      chunk = ChunkInt;
       switch (rtSched) {
       case omp_sched_static: {
         if (chunk > 0)
           schedule = kmp_sched_static_chunk;
         else
           schedule = kmp_sched_static_nochunk;
         break;
       }

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

EXTERN void __kmpc_spmd_kernel_init(int, int16_t RequiresOMPRuntime) {
   // install new top descriptor
   omptarget_nvptx_threadPrivateContext->SetTopLevelTaskDescr(threadId,
                                                              newTaskDescr);
 
   // init thread private from init value
   PRINT(LD_PAR,
         "thread will execute parallel region with id %d in a team of "
         "%d threads\n",
-        (int)newTaskDescr->ThreadId(), (int)ThreadLimit);
+        (int)threadId, (int)ThreadLimit);
 }
 
 EXTERN void __kmpc_spmd_kernel_deinit_v2(int16_t RequiresOMPRuntime) {
   // We're not going to pop the task descr stack of each thread since
   // there are no more parallel regions in SPMD mode.
   if (!RequiresOMPRuntime)
     return;
 

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

EXTERN void __kmpc_kernel_prepare_parallel(void *WorkFn) {
 
   // This routine is only called by the team master.  The team master is
   // the first thread of the last warp.  It always has the logical thread
   // id of 0 (since it is a shadow for the first worker thread).
   const int threadId = 0;
   omptarget_nvptx_TaskDescr *currTaskDescr =
       omptarget_nvptx_threadPrivateContext->GetTopLevelTaskDescr(threadId);
   ASSERT0(LT_FUSSY, currTaskDescr, "expected a top task descr");
-  ASSERT0(LT_FUSSY, !currTaskDescr->InParallelRegion(),
-          "cannot be called in a parallel region.");
-  if (currTaskDescr->InParallelRegion()) {
-    PRINT0(LD_PAR, "already in parallel: go seq\n");
-    return;
-  }
 
   uint16_t NumThreads = determineNumberOfThreads();
   TeamState.ParallelTeamSize = NumThreads;
 
   ASSERT0(LT_FUSSY, GetThreadIdInBlock() == GetMasterThreadID(),
           "only team master can create parallel");
 
   // Set number of threads on work descriptor.
if (ThreadIsActive) {
     newTaskDescr->CopyFromWorkDescr(workDescr.WorkTaskDescr());
     // install new top descriptor
     omptarget_nvptx_threadPrivateContext->SetTopLevelTaskDescr(threadId,
                                                                newTaskDescr);
     // init private from int value
     PRINT(LD_PAR,
           "thread will execute parallel region with id %d in a team of "
           "%d threads\n",
-          (int)newTaskDescr->ThreadId(), (int)nThreads);
+          (int)threadId, (int)nThreads);
   }
 
   return ThreadIsActive;
 }
 
 EXTERN void __kmpc_kernel_end_parallel() {
   // pop stack
   PRINT0(LD_IO | LD_PAR, "call to __kmpc_kernel_end_parallel\n");
EXTERN void __kmpc_serialized_parallel(kmp_Ident *loc, uint32_t global_tid) {
 
   // allocate new task descriptor and copy value from current one, set prev to
   // it
   omptarget_nvptx_TaskDescr *newTaskDescr =
       (omptarget_nvptx_TaskDescr *)SafeMalloc(sizeof(omptarget_nvptx_TaskDescr),
                                               "new seq parallel task");
   newTaskDescr->CopyParent(currTaskDescr);
 
-  // tweak values for serialized parallel case:
-  // - each thread becomes ID 0 in its serialized parallel, and
-  // - there is only one thread per team
-  newTaskDescr->ThreadId() = 0;
-
   // set new task descriptor as top
   omptarget_nvptx_threadPrivateContext->SetTopLevelTaskDescr(threadId,
                                                              newTaskDescr);
 }
 
 EXTERN void __kmpc_end_serialized_parallel(kmp_Ident *loc,
                                            uint32_t global_tid) {
   PRINT0(LD_IO, "call to __kmpc_end_serialized_parallel\n");