[OpenMP][libomptarget] Add runtime functions for default schedule for distribute

[gtbercea]

Summary

When using the default schedule for distribute, the default must ensure that at most one iteration is associated with every thread.

[ABataev]

libomptarget/deviceRTLs/nvptx/src/loop.cu

258

Seems to me, you need to add plastiter and update it for the last iteration. Otherwise it might break lastprivates

[ABataev]

libomptarget/deviceRTLs/nvptx/src/loop.cu

243

You don't use the initial value of the stride variable, you can declare on the first use.
The same for ub

[Hahnfeld]

Do we really need new entry points for this? I think we could avoid code duplication by letting the compiler generated code pass the correct chunk to __kmpc_for_static_init_??. This could either come from a (single) new query function (__kmpc_nvptx_distribute_default_chunk?) or can be hard coded to threadsPerBlock because this is only relevant for SPMD.

[gtbercea]

Due to most recent proposed changes to Clang in D52434, changes to the runtime are no longer required.

libomptarget/deviceRTLs/nvptx/src/interface.h

EXTERN void __kmpc_for_static_init_8(kmp_Indent *loc, int32_t global_tid,
                                      int64_t *pstride, int64_t incr,
                                      int64_t chunk);
 EXTERN void __kmpc_for_static_init_8u(kmp_Indent *loc, int32_t global_tid,
                                       int32_t sched, int32_t *plastiter1,
                                       uint64_t *plower, uint64_t *pupper,
                                       int64_t *pstride, int64_t incr,
                                       int64_t chunk);
 EXTERN
+
+void __kmpc_distribute_default_init_4(kmp_Indent *loc, int32_t global_tid,
+                                     int32_t sched, int32_t *plastiter,
+                                     int32_t *plower, int32_t *pupper,
+                                     int32_t *pstride, int32_t incr,
+                                     int32_t chunk);
+EXTERN
+void __kmpc_distribute_default_init_4u(kmp_Indent *loc, int32_t global_tid,
+                                      int32_t sched, int32_t *plastiter,
+                                      uint32_t *plower, uint32_t *pupper,
+                                      int32_t *pstride, int32_t incr,
+                                      int32_t chunk);
+EXTERN
+void __kmpc_distribute_default_init_8(kmp_Indent *loc, int32_t global_tid,
+                                     int32_t sched, int32_t *plastiter,
+                                     int64_t *plower, int64_t *pupper,
+                                     int64_t *pstride, int64_t incr,
+                                     int64_t chunk);
+EXTERN
+void __kmpc_distribute_default_init_8u(kmp_Indent *loc, int32_t global_tid,
+                                      int32_t sched, int32_t *plastiter1,
+                                      uint64_t *plower, uint64_t *pupper,
+                                      int64_t *pstride, int64_t incr,
+                                      int64_t chunk);
+EXTERN
 void __kmpc_for_static_init_4_simple_spmd(kmp_Indent *loc, int32_t global_tid,
                                           int32_t sched, int32_t *plastiter,
                                           int32_t *plower, int32_t *pupper,
                                           int32_t *pstride, int32_t incr,
                                           int32_t chunk);
 EXTERN
 void __kmpc_for_static_init_4u_simple_spmd(kmp_Indent *loc, int32_t global_tid,
                                            int32_t sched, int32_t *plastiter,

libomptarget/deviceRTLs/nvptx/src/loop.cu

PRINT(
         "Got sched: Active %d, total %d: lb %lld, ub %lld, stride %lld, last "
         "%d\n",
         GetNumberOfOmpThreads(tid, IsSPMDExecutionMode, IsRuntimeUninitialized),
         GetNumberOfWorkersInTeam(), P64(*plower), P64(*pupper), P64(*pstride),
         lastiter);
   }
 
   ////////////////////////////////////////////////////////////////////////////////
+  // Support for Static Init
+
+  INLINE static void distribute_default_init(T *plower, T *pupper, ST *pstride,
+                                             bool IsRuntimeUninitialized) {
+    int tid = GetLogicalThreadIdInBlock();
+
+    // copy
+    T lb = *plower;
+    T ub = *pupper;
+    ST stride = *pstride;
+    T entityId, numberOfEntities, chunkSize;
+
+    // init - the default schedule is static with chunk size equal to team size.
+    entityId = GetOmpTeamId();
+    numberOfEntities = GetNumberOfOmpTeams();
+    chunkSize = GetNumberOfOmpThreads(tid, true, IsRuntimeUninitialized);
+
+    stride = numberOfEntities * chunkSize;
+    lb = lb + entityId * chunkSize;
+    ub = lb + chunkSize - 1; // Clang uses i <= ub
+
+    // copy back
+    *plower = lb;
+    *pupper = ub;
+    *pstride = stride;
+  }
+
+  ////////////////////////////////////////////////////////////////////////////////
   // Support for dispatch Init
 
   INLINE static int OrderedSchedule(kmp_sched_t schedule) {
     return schedule >= kmp_sched_ordered_first &&
            schedule <= kmp_sched_ordered_last;
   }
 
   INLINE static void dispatch_init(kmp_Indent *loc, int32_t threadId,
EXTERN void __kmpc_for_static_init_8u(kmp_Indent *loc, int32_t global_tid,
                                       int64_t *pstride, int64_t incr,
                                       int64_t chunk) {
   PRINT0(LD_IO, "call kmpc_for_static_init_8u\n");
   omptarget_nvptx_LoopSupport<uint64_t, int64_t>::for_static_init(
       schedtype, plastiter, plower, pupper, pstride, chunk, isSPMDMode(),
       isRuntimeUninitialized());
 }
 
+////////////////////////////////////////////////////////////////////////////////
+// Static loops for distribute default schedule
+////////////////////////////////////////////////////////////////////////////////
+
+EXTERN
+void __kmpc_distribute_default_init_4(kmp_Indent *loc, int32_t global_tid,
+                                     int32_t schedtype, int32_t *plastiter,
+                                     int32_t *plower, int32_t *pupper,
+                                     int32_t *pstride, int32_t incr,
+                                     int32_t chunk) {
+  PRINT0(LD_IO, "call kmpc_distribute_default_init_4\n");
+  omptarget_nvptx_LoopSupport<int32_t, int32_t>::distribute_default_init(
+      plower, pupper, pstride, isRuntimeUninitialized());
+}
+
+EXTERN
+void __kmpc_distribute_default_init_4u(kmp_Indent *loc, int32_t global_tid,
+                                      int32_t schedtype, int32_t *plastiter,
+                                      uint32_t *plower, uint32_t *pupper,
+                                      int32_t *pstride, int32_t incr,
+                                      int32_t chunk) {
+  PRINT0(LD_IO, "call kmpc_distribute_default_init_4u\n");
+  omptarget_nvptx_LoopSupport<uint32_t, int32_t>::distribute_default_init(
+      plower, pupper, pstride, isRuntimeUninitialized());
+}
+
+EXTERN
+void __kmpc_distribute_default_init_8(kmp_Indent *loc, int32_t global_tid,
+                                     int32_t schedtype, int32_t *plastiter,
+                                     int64_t *plower, int64_t *pupper,
+                                     int64_t *pstride, int64_t incr,
+                                     int64_t chunk) {
+  PRINT0(LD_IO, "call kmpc_distribute_default_init_8\n");
+  omptarget_nvptx_LoopSupport<int64_t, int64_t>::distribute_default_init(
+      plower, pupper, pstride, isRuntimeUninitialized());
+}
+
+EXTERN
+void __kmpc_distribute_default_init_8u(kmp_Indent *loc, int32_t global_tid,
+                                      int32_t schedtype, int32_t *plastiter,
+                                      uint64_t *plower, uint64_t *pupper,
+                                      int64_t *pstride, int64_t incr,
+                                      int64_t chunk) {
+  PRINT0(LD_IO, "call kmpc_distribute_default_init_8u\n");
+  omptarget_nvptx_LoopSupport<uint64_t, int64_t>::distribute_default_init(
+      plower, pupper, pstride, isRuntimeUninitialized());
+}
+
 EXTERN
 void __kmpc_for_static_init_4_simple_spmd(kmp_Indent *loc, int32_t global_tid,
                                           int32_t schedtype, int32_t *plastiter,
                                           int32_t *plower, int32_t *pupper,
                                           int32_t *pstride, int32_t incr,
                                           int32_t chunk) {
   PRINT0(LD_IO, "call kmpc_for_static_init_4_simple_spmd\n");
   omptarget_nvptx_LoopSupport<int32_t, int32_t>::for_static_init(