| //===----- Workshare.cpp - OpenMP workshare implementation ------ 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 | |
| // | |
| //===----------------------------------------------------------------------===// | |
| // | |
| // This file contains the implementation of the KMPC interface | |
| // for the loop construct plus other worksharing constructs that use the same | |
| // interface as loops. | |
| // | |
| //===----------------------------------------------------------------------===// | |
| #include "Workshare.h" | |
| #include "Debug.h" | |
| #include "Interface.h" | |
| #include "Mapping.h" | |
| #include "State.h" | |
| #include "Synchronization.h" | |
| #include "Types.h" | |
| #include "Utils.h" | |
| using namespace ompx; | |
| // TODO: | |
| struct DynamicScheduleTracker { | |
| int64_t Chunk; | |
| int64_t LoopUpperBound; | |
| int64_t NextLowerBound; | |
| int64_t Stride; | |
| kmp_sched_t ScheduleType; | |
| DynamicScheduleTracker *NextDST; | |
| }; | |
| #define ASSERT0(...) | |
| // used by the library for the interface with the app | |
| #define DISPATCH_FINISHED 0 | |
| #define DISPATCH_NOTFINISHED 1 | |
| // used by dynamic scheduling | |
| #define FINISHED 0 | |
| #define NOT_FINISHED 1 | |
| #define LAST_CHUNK 2 | |
| #pragma omp begin declare target device_type(nohost) | |
| // TODO: This variable is a hack inherited from the old runtime. | |
| static uint64_t SHARED(Cnt); | |
| template <typename T, typename ST> struct omptarget_nvptx_LoopSupport { | |
| //////////////////////////////////////////////////////////////////////////////// | |
| // Loop with static scheduling with chunk | |
| // Generic implementation of OMP loop scheduling with static policy | |
| /*! \brief Calculate initial bounds for static loop and stride | |
| * @param[in] loc location in code of the call (not used here) | |
| * @param[in] global_tid global thread id | |
| * @param[in] schetype type of scheduling (see omptarget-nvptx.h) | |
| * @param[in] plastiter pointer to last iteration | |
| * @param[in,out] pointer to loop lower bound. it will contain value of | |
| * lower bound of first chunk | |
| * @param[in,out] pointer to loop upper bound. It will contain value of | |
| * upper bound of first chunk | |
| * @param[in,out] pointer to loop stride. It will contain value of stride | |
| * between two successive chunks executed by the same thread | |
| * @param[in] loop increment bump | |
| * @param[in] chunk size | |
| */ | |
| // helper function for static chunk | |
| static void ForStaticChunk(int &last, T &lb, T &ub, ST &stride, ST chunk, | |
| T entityId, T numberOfEntities) { | |
| // each thread executes multiple chunks all of the same size, except | |
| // the last one | |
| // distance between two successive chunks | |
| stride = numberOfEntities * chunk; | |
| lb = lb + entityId * chunk; | |
| T inputUb = ub; | |
| ub = lb + chunk - 1; // Clang uses i <= ub | |
| // Say ub' is the begining of the last chunk. Then who ever has a | |
| // lower bound plus a multiple of the increment equal to ub' is | |
| // the last one. | |
| T beginingLastChunk = inputUb - (inputUb % chunk); | |
| last = ((beginingLastChunk - lb) % stride) == 0; | |
| } | |
| //////////////////////////////////////////////////////////////////////////////// | |
| // Loop with static scheduling without chunk | |
| // helper function for static no chunk | |
| static void ForStaticNoChunk(int &last, T &lb, T &ub, ST &stride, ST &chunk, | |
| T entityId, T numberOfEntities) { | |
| // No chunk size specified. Each thread or warp gets at most one | |
| // chunk; chunks are all almost of equal size | |
| T loopSize = ub - lb + 1; | |
| chunk = loopSize / numberOfEntities; | |
| T leftOver = loopSize - chunk * numberOfEntities; | |
| if (entityId < leftOver) { | |
| chunk++; | |
| lb = lb + entityId * chunk; | |
| } else { | |
| lb = lb + entityId * chunk + leftOver; | |
| } | |
| T inputUb = ub; | |
| ub = lb + chunk - 1; // Clang uses i <= ub | |
| last = lb <= inputUb && inputUb <= ub; | |
| stride = loopSize; // make sure we only do 1 chunk per warp | |
| } | |
| //////////////////////////////////////////////////////////////////////////////// | |
| // Support for Static Init | |
| static void for_static_init(int32_t, int32_t schedtype, int32_t *plastiter, | |
| T *plower, T *pupper, ST *pstride, ST chunk, | |
| bool IsSPMDExecutionMode) { | |
| int32_t gtid = omp_get_thread_num(); | |
| int numberOfActiveOMPThreads = omp_get_num_threads(); | |
| // All warps that are in excess of the maximum requested, do | |
| // not execute the loop | |
| ASSERT0(LT_FUSSY, gtid < numberOfActiveOMPThreads, | |
| "current thread is not needed here; error"); | |
| // copy | |
| int lastiter = 0; | |
| T lb = *plower; | |
| T ub = *pupper; | |
| ST stride = *pstride; | |
| // init | |
| switch (SCHEDULE_WITHOUT_MODIFIERS(schedtype)) { | |
| case kmp_sched_static_chunk: { | |
| if (chunk > 0) { | |
| ForStaticChunk(lastiter, lb, ub, stride, chunk, gtid, | |
| numberOfActiveOMPThreads); | |
| break; | |
| } | |
| [[fallthrough]]; | |
| } // note: if chunk <=0, use nochunk | |
| case kmp_sched_static_balanced_chunk: { | |
| if (chunk > 0) { | |
| // round up to make sure the chunk is enough to cover all iterations | |
| T tripCount = ub - lb + 1; // +1 because ub is inclusive | |
| T span = (tripCount + numberOfActiveOMPThreads - 1) / | |
| numberOfActiveOMPThreads; | |
| // perform chunk adjustment | |
| chunk = (span + chunk - 1) & ~(chunk - 1); | |
| ASSERT0(LT_FUSSY, ub >= lb, "ub must be >= lb."); | |
| T oldUb = ub; | |
| ForStaticChunk(lastiter, lb, ub, stride, chunk, gtid, | |
| numberOfActiveOMPThreads); | |
| if (ub > oldUb) | |
| ub = oldUb; | |
| break; | |
| } | |
| [[fallthrough]]; | |
| } // note: if chunk <=0, use nochunk | |
| case kmp_sched_static_nochunk: { | |
| ForStaticNoChunk(lastiter, lb, ub, stride, chunk, gtid, | |
| numberOfActiveOMPThreads); | |
| break; | |
| } | |
| case kmp_sched_distr_static_chunk: { | |
| if (chunk > 0) { | |
| ForStaticChunk(lastiter, lb, ub, stride, chunk, omp_get_team_num(), | |
| omp_get_num_teams()); | |
| break; | |
| } | |
| [[fallthrough]]; | |
| } // note: if chunk <=0, use nochunk | |
| case kmp_sched_distr_static_nochunk: { | |
| ForStaticNoChunk(lastiter, lb, ub, stride, chunk, omp_get_team_num(), | |
| omp_get_num_teams()); | |
| break; | |
| } | |
| case kmp_sched_distr_static_chunk_sched_static_chunkone: { | |
| ForStaticChunk(lastiter, lb, ub, stride, chunk, | |
| numberOfActiveOMPThreads * omp_get_team_num() + gtid, | |
| omp_get_num_teams() * numberOfActiveOMPThreads); | |
| break; | |
| } | |
| default: { | |
| // ASSERT(LT_FUSSY, 0, "unknown schedtype %d", (int)schedtype); | |
| ForStaticChunk(lastiter, lb, ub, stride, chunk, gtid, | |
| numberOfActiveOMPThreads); | |
| break; | |
| } | |
| } | |
| // copy back | |
| *plastiter = lastiter; | |
| *plower = lb; | |
| *pupper = ub; | |
| *pstride = stride; | |
| } | |
| //////////////////////////////////////////////////////////////////////////////// | |
| // Support for dispatch Init | |
| static int OrderedSchedule(kmp_sched_t schedule) { | |
| return schedule >= kmp_sched_ordered_first && | |
| schedule <= kmp_sched_ordered_last; | |
| } | |
| static void dispatch_init(IdentTy *loc, int32_t threadId, | |
| kmp_sched_t schedule, T lb, T ub, ST st, ST chunk, | |
| DynamicScheduleTracker *DST) { | |
| int tid = mapping::getThreadIdInBlock(); | |
| T tnum = omp_get_num_threads(); | |
| 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). | |
| * When it is we'll want to look at them somewhere here and use that | |
| * information to add to our schedule choice. We shouldn't need to pass | |
| * them on, they merely affect which schedule we can legally choose for | |
| * various dynamic cases. (In particular, whether or not a stealing scheme | |
| * is legal). | |
| */ | |
| schedule = SCHEDULE_WITHOUT_MODIFIERS(schedule); | |
| // Process schedule. | |
| if (tnum == 1 || tripCount <= 1 || OrderedSchedule(schedule)) { | |
| if (OrderedSchedule(schedule)) | |
| __kmpc_barrier(loc, threadId); | |
| 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; | |
| 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; | |
| } | |
| case omp_sched_auto: { | |
| schedule = kmp_sched_static_chunk; | |
| chunk = 1; | |
| break; | |
| } | |
| case omp_sched_dynamic: | |
| case omp_sched_guided: { | |
| schedule = kmp_sched_dynamic; | |
| break; | |
| } | |
| } | |
| } else if (schedule == kmp_sched_auto) { | |
| schedule = kmp_sched_static_chunk; | |
| chunk = 1; | |
| } else { | |
| // ASSERT(LT_FUSSY, | |
| // schedule == kmp_sched_dynamic || schedule == kmp_sched_guided, | |
| // "unknown schedule %d & chunk %lld\n", (int)schedule, | |
| // (long long)chunk); | |
| } | |
| // init schedules | |
| if (schedule == kmp_sched_static_chunk) { | |
| ASSERT0(LT_FUSSY, chunk > 0, "bad chunk value"); | |
| // save sched state | |
| DST->ScheduleType = schedule; | |
| // save ub | |
| DST->LoopUpperBound = ub; | |
| // compute static chunk | |
| ST stride; | |
| int lastiter = 0; | |
| ForStaticChunk(lastiter, lb, ub, stride, chunk, threadId, tnum); | |
| // save computed params | |
| DST->Chunk = chunk; | |
| DST->NextLowerBound = lb; | |
| DST->Stride = stride; | |
| } else if (schedule == kmp_sched_static_balanced_chunk) { | |
| ASSERT0(LT_FUSSY, chunk > 0, "bad chunk value"); | |
| // save sched state | |
| DST->ScheduleType = schedule; | |
| // save ub | |
| DST->LoopUpperBound = ub; | |
| // compute static chunk | |
| ST stride; | |
| int lastiter = 0; | |
| // round up to make sure the chunk is enough to cover all iterations | |
| T span = (tripCount + tnum - 1) / tnum; | |
| // perform chunk adjustment | |
| chunk = (span + chunk - 1) & ~(chunk - 1); | |
| T oldUb = ub; | |
| ForStaticChunk(lastiter, lb, ub, stride, chunk, threadId, tnum); | |
| ASSERT0(LT_FUSSY, ub >= lb, "ub must be >= lb."); | |
| if (ub > oldUb) | |
| ub = oldUb; | |
| // save computed params | |
| DST->Chunk = chunk; | |
| DST->NextLowerBound = lb; | |
| DST->Stride = stride; | |
| } else if (schedule == kmp_sched_static_nochunk) { | |
| ASSERT0(LT_FUSSY, chunk == 0, "bad chunk value"); | |
| // save sched state | |
| DST->ScheduleType = schedule; | |
| // save ub | |
| DST->LoopUpperBound = ub; | |
| // compute static chunk | |
| ST stride; | |
| int lastiter = 0; | |
| ForStaticNoChunk(lastiter, lb, ub, stride, chunk, threadId, tnum); | |
| // save computed params | |
| DST->Chunk = chunk; | |
| DST->NextLowerBound = lb; | |
| DST->Stride = stride; | |
| } else if (schedule == kmp_sched_dynamic || schedule == kmp_sched_guided) { | |
| // save data | |
| DST->ScheduleType = schedule; | |
| if (chunk < 1) | |
| chunk = 1; | |
| DST->Chunk = chunk; | |
| DST->LoopUpperBound = ub; | |
| DST->NextLowerBound = lb; | |
| __kmpc_barrier(loc, threadId); | |
| if (tid == 0) { | |
| Cnt = 0; | |
| fence::team(atomic::seq_cst); | |
| } | |
| __kmpc_barrier(loc, threadId); | |
| } | |
| } | |
| //////////////////////////////////////////////////////////////////////////////// | |
| // Support for dispatch next | |
| static uint64_t NextIter() { | |
| __kmpc_impl_lanemask_t active = mapping::activemask(); | |
| uint32_t leader = utils::ffs(active) - 1; | |
| uint32_t change = utils::popc(active); | |
| __kmpc_impl_lanemask_t lane_mask_lt = mapping::lanemaskLT(); | |
| unsigned int rank = utils::popc(active & lane_mask_lt); | |
| uint64_t warp_res = 0; | |
| if (rank == 0) { | |
| warp_res = atomic::add(&Cnt, change, atomic::seq_cst); | |
| } | |
| warp_res = utils::shuffle(active, warp_res, leader); | |
| return warp_res + rank; | |
| } | |
| static int DynamicNextChunk(T &lb, T &ub, T chunkSize, T loopLowerBound, | |
| T loopUpperBound) { | |
| T N = NextIter(); | |
| lb = loopLowerBound + N * chunkSize; | |
| ub = lb + chunkSize - 1; // Clang uses i <= ub | |
| // 3 result cases: | |
| // a. lb and ub < loopUpperBound --> NOT_FINISHED | |
| // b. lb < loopUpperBound and ub >= loopUpperBound: last chunk --> | |
| // NOT_FINISHED | |
| // c. lb and ub >= loopUpperBound: empty chunk --> FINISHED | |
| // a. | |
| if (lb <= loopUpperBound && ub < loopUpperBound) { | |
| return NOT_FINISHED; | |
| } | |
| // b. | |
| if (lb <= loopUpperBound) { | |
| ub = loopUpperBound; | |
| return LAST_CHUNK; | |
| } | |
| // c. if we are here, we are in case 'c' | |
| lb = loopUpperBound + 2; | |
| ub = loopUpperBound + 1; | |
| return FINISHED; | |
| } | |
| static int dispatch_next(IdentTy *loc, int32_t gtid, int32_t *plast, | |
| T *plower, T *pupper, ST *pstride, | |
| DynamicScheduleTracker *DST) { | |
| // ID of a thread in its own warp | |
| // automatically selects thread or warp ID based on selected implementation | |
| ASSERT0(LT_FUSSY, gtid < omp_get_num_threads(), | |
| "current thread is not needed here; error"); | |
| // retrieve schedule | |
| kmp_sched_t schedule = DST->ScheduleType; | |
| // xxx reduce to one | |
| if (schedule == kmp_sched_static_chunk || | |
| schedule == kmp_sched_static_nochunk) { | |
| T myLb = DST->NextLowerBound; | |
| T ub = DST->LoopUpperBound; | |
| // finished? | |
| if (myLb > ub) { | |
| return DISPATCH_FINISHED; | |
| } | |
| // not finished, save current bounds | |
| ST chunk = DST->Chunk; | |
| *plower = myLb; | |
| T myUb = myLb + chunk - 1; // Clang uses i <= ub | |
| if (myUb > ub) | |
| myUb = ub; | |
| *pupper = myUb; | |
| *plast = (int32_t)(myUb == ub); | |
| // increment next lower bound by the stride | |
| ST stride = DST->Stride; | |
| DST->NextLowerBound = myLb + stride; | |
| return DISPATCH_NOTFINISHED; | |
| } | |
| ASSERT0(LT_FUSSY, | |
| schedule == kmp_sched_dynamic || schedule == kmp_sched_guided, | |
| "bad sched"); | |
| T myLb, myUb; | |
| int finished = DynamicNextChunk(myLb, myUb, DST->Chunk, DST->NextLowerBound, | |
| DST->LoopUpperBound); | |
| if (finished == FINISHED) | |
| return DISPATCH_FINISHED; | |
| // not finished (either not finished or last chunk) | |
| *plast = (int32_t)(finished == LAST_CHUNK); | |
| *plower = myLb; | |
| *pupper = myUb; | |
| *pstride = 1; | |
| return DISPATCH_NOTFINISHED; | |
| } | |
| static void dispatch_fini() { | |
| // nothing | |
| } | |
| //////////////////////////////////////////////////////////////////////////////// | |
| // end of template class that encapsulate all the helper functions | |
| //////////////////////////////////////////////////////////////////////////////// | |
| }; | |
| //////////////////////////////////////////////////////////////////////////////// | |
| // KMP interface implementation (dyn loops) | |
| //////////////////////////////////////////////////////////////////////////////// | |
| // TODO: Expand the dispatch API to take a DST pointer which can then be | |
| // allocated properly without malloc. | |
| // For now, each team will contain an LDS pointer (ThreadDST) to a global array | |
| // of references to the DST structs allocated (in global memory) for each thread | |
| // in the team. The global memory array is allocated during the init phase if it | |
| // was not allocated already and will be deallocated when the dispatch phase | |
| // ends: | |
| // | |
| // __kmpc_dispatch_init | |
| // | |
| // ** Dispatch loop ** | |
| // | |
| // __kmpc_dispatch_deinit | |
| // | |
| static DynamicScheduleTracker **SHARED(ThreadDST); | |
| // Create a new DST, link the current one, and define the new as current. | |
| static DynamicScheduleTracker *pushDST() { | |
| int32_t ThreadIndex = mapping::getThreadIdInBlock(); | |
| // Each block will allocate an array of pointers to DST structs. The array is | |
| // equal in length to the number of threads in that block. | |
| if (!ThreadDST) { | |
| // Allocate global memory array of pointers to DST structs: | |
| if (mapping::isMainThreadInGenericMode() || ThreadIndex == 0) | |
| ThreadDST = static_cast<DynamicScheduleTracker **>( | |
| memory::allocGlobal(mapping::getNumberOfThreadsInBlock() * | |
| sizeof(DynamicScheduleTracker *), | |
| "new ThreadDST array")); | |
| synchronize::threads(atomic::seq_cst); | |
| // Initialize the array pointers: | |
| ThreadDST[ThreadIndex] = nullptr; | |
| } | |
| // Create a DST struct for the current thread: | |
| DynamicScheduleTracker *NewDST = static_cast<DynamicScheduleTracker *>( | |
| memory::allocGlobal(sizeof(DynamicScheduleTracker), "new DST")); | |
| *NewDST = DynamicScheduleTracker({0}); | |
| // Add the new DST struct to the array of DST structs: | |
| NewDST->NextDST = ThreadDST[ThreadIndex]; | |
| ThreadDST[ThreadIndex] = NewDST; | |
| return NewDST; | |
| } | |
| // Return the current DST. | |
| static DynamicScheduleTracker *peekDST() { | |
| return ThreadDST[mapping::getThreadIdInBlock()]; | |
| } | |
| // Pop the current DST and restore the last one. | |
| static void popDST() { | |
| int32_t ThreadIndex = mapping::getThreadIdInBlock(); | |
| DynamicScheduleTracker *CurrentDST = ThreadDST[ThreadIndex]; | |
| DynamicScheduleTracker *OldDST = CurrentDST->NextDST; | |
| memory::freeGlobal(CurrentDST, "remove DST"); | |
| ThreadDST[ThreadIndex] = OldDST; | |
| // Check if we need to deallocate the global array. Ensure all threads | |
| // in the block have finished deallocating the individual DSTs. | |
| synchronize::threads(atomic::seq_cst); | |
| if (!ThreadDST[ThreadIndex] && !ThreadIndex) { | |
| memory::freeGlobal(ThreadDST, "remove ThreadDST array"); | |
| ThreadDST = nullptr; | |
| } | |
| synchronize::threads(atomic::seq_cst); | |
| } | |
| void workshare::init(bool IsSPMD) { | |
| if (mapping::isInitialThreadInLevel0(IsSPMD)) | |
| ThreadDST = nullptr; | |
| } | |
| extern "C" { | |
| // init | |
| void __kmpc_dispatch_init_4(IdentTy *loc, int32_t tid, int32_t schedule, | |
| int32_t lb, int32_t ub, int32_t st, int32_t chunk) { | |
| DynamicScheduleTracker *DST = pushDST(); | |
| omptarget_nvptx_LoopSupport<int32_t, int32_t>::dispatch_init( | |
| loc, tid, (kmp_sched_t)schedule, lb, ub, st, chunk, DST); | |
| } | |
| void __kmpc_dispatch_init_4u(IdentTy *loc, int32_t tid, int32_t schedule, | |
| uint32_t lb, uint32_t ub, int32_t st, | |
| int32_t chunk) { | |
| DynamicScheduleTracker *DST = pushDST(); | |
| omptarget_nvptx_LoopSupport<uint32_t, int32_t>::dispatch_init( | |
| loc, tid, (kmp_sched_t)schedule, lb, ub, st, chunk, DST); | |
| } | |
| void __kmpc_dispatch_init_8(IdentTy *loc, int32_t tid, int32_t schedule, | |
| int64_t lb, int64_t ub, int64_t st, int64_t chunk) { | |
| DynamicScheduleTracker *DST = pushDST(); | |
| omptarget_nvptx_LoopSupport<int64_t, int64_t>::dispatch_init( | |
| loc, tid, (kmp_sched_t)schedule, lb, ub, st, chunk, DST); | |
| } | |
| void __kmpc_dispatch_init_8u(IdentTy *loc, int32_t tid, int32_t schedule, | |
| uint64_t lb, uint64_t ub, int64_t st, | |
| int64_t chunk) { | |
| DynamicScheduleTracker *DST = pushDST(); | |
| omptarget_nvptx_LoopSupport<uint64_t, int64_t>::dispatch_init( | |
| loc, tid, (kmp_sched_t)schedule, lb, ub, st, chunk, DST); | |
| } | |
| // next | |
| int __kmpc_dispatch_next_4(IdentTy *loc, int32_t tid, int32_t *p_last, | |
| int32_t *p_lb, int32_t *p_ub, int32_t *p_st) { | |
| DynamicScheduleTracker *DST = peekDST(); | |
| return omptarget_nvptx_LoopSupport<int32_t, int32_t>::dispatch_next( | |
| loc, tid, p_last, p_lb, p_ub, p_st, DST); | |
| } | |
| int __kmpc_dispatch_next_4u(IdentTy *loc, int32_t tid, int32_t *p_last, | |
| uint32_t *p_lb, uint32_t *p_ub, int32_t *p_st) { | |
| DynamicScheduleTracker *DST = peekDST(); | |
| return omptarget_nvptx_LoopSupport<uint32_t, int32_t>::dispatch_next( | |
| loc, tid, p_last, p_lb, p_ub, p_st, DST); | |
| } | |
| int __kmpc_dispatch_next_8(IdentTy *loc, int32_t tid, int32_t *p_last, | |
| int64_t *p_lb, int64_t *p_ub, int64_t *p_st) { | |
| DynamicScheduleTracker *DST = peekDST(); | |
| return omptarget_nvptx_LoopSupport<int64_t, int64_t>::dispatch_next( | |
| loc, tid, p_last, p_lb, p_ub, p_st, DST); | |
| } | |
| int __kmpc_dispatch_next_8u(IdentTy *loc, int32_t tid, int32_t *p_last, | |
| uint64_t *p_lb, uint64_t *p_ub, int64_t *p_st) { | |
| DynamicScheduleTracker *DST = peekDST(); | |
| return omptarget_nvptx_LoopSupport<uint64_t, int64_t>::dispatch_next( | |
| loc, tid, p_last, p_lb, p_ub, p_st, DST); | |
| } | |
| // fini | |
| void __kmpc_dispatch_fini_4(IdentTy *loc, int32_t tid) { | |
| omptarget_nvptx_LoopSupport<int32_t, int32_t>::dispatch_fini(); | |
| } | |
| void __kmpc_dispatch_fini_4u(IdentTy *loc, int32_t tid) { | |
| omptarget_nvptx_LoopSupport<uint32_t, int32_t>::dispatch_fini(); | |
| } | |
| void __kmpc_dispatch_fini_8(IdentTy *loc, int32_t tid) { | |
| omptarget_nvptx_LoopSupport<int64_t, int64_t>::dispatch_fini(); | |
| } | |
| void __kmpc_dispatch_fini_8u(IdentTy *loc, int32_t tid) { | |
| omptarget_nvptx_LoopSupport<uint64_t, int64_t>::dispatch_fini(); | |
| } | |
| // deinit | |
| void __kmpc_dispatch_deinit(IdentTy *loc, int32_t tid) { popDST(); } | |
| //////////////////////////////////////////////////////////////////////////////// | |
| // KMP interface implementation (static loops) | |
| //////////////////////////////////////////////////////////////////////////////// | |
| void __kmpc_for_static_init_4(IdentTy *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) { | |
| omptarget_nvptx_LoopSupport<int32_t, int32_t>::for_static_init( | |
| global_tid, schedtype, plastiter, plower, pupper, pstride, chunk, | |
| mapping::isSPMDMode()); | |
| } | |
| void __kmpc_for_static_init_4u(IdentTy *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) { | |
| omptarget_nvptx_LoopSupport<uint32_t, int32_t>::for_static_init( | |
| global_tid, schedtype, plastiter, plower, pupper, pstride, chunk, | |
| mapping::isSPMDMode()); | |
| } | |
| void __kmpc_for_static_init_8(IdentTy *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) { | |
| omptarget_nvptx_LoopSupport<int64_t, int64_t>::for_static_init( | |
| global_tid, schedtype, plastiter, plower, pupper, pstride, chunk, | |
| mapping::isSPMDMode()); | |
| } | |
| void __kmpc_for_static_init_8u(IdentTy *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) { | |
| omptarget_nvptx_LoopSupport<uint64_t, int64_t>::for_static_init( | |
| global_tid, schedtype, plastiter, plower, pupper, pstride, chunk, | |
| mapping::isSPMDMode()); | |
| } | |
| void __kmpc_distribute_static_init_4(IdentTy *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) { | |
| omptarget_nvptx_LoopSupport<int32_t, int32_t>::for_static_init( | |
| global_tid, schedtype, plastiter, plower, pupper, pstride, chunk, | |
| mapping::isSPMDMode()); | |
| } | |
| void __kmpc_distribute_static_init_4u(IdentTy *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) { | |
| omptarget_nvptx_LoopSupport<uint32_t, int32_t>::for_static_init( | |
| global_tid, schedtype, plastiter, plower, pupper, pstride, chunk, | |
| mapping::isSPMDMode()); | |
| } | |
| void __kmpc_distribute_static_init_8(IdentTy *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) { | |
| omptarget_nvptx_LoopSupport<int64_t, int64_t>::for_static_init( | |
| global_tid, schedtype, plastiter, plower, pupper, pstride, chunk, | |
| mapping::isSPMDMode()); | |
| } | |
| void __kmpc_distribute_static_init_8u(IdentTy *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) { | |
| omptarget_nvptx_LoopSupport<uint64_t, int64_t>::for_static_init( | |
| global_tid, schedtype, plastiter, plower, pupper, pstride, chunk, | |
| mapping::isSPMDMode()); | |
| } | |
| void __kmpc_for_static_fini(IdentTy *loc, int32_t global_tid) {} | |
| void __kmpc_distribute_static_fini(IdentTy *loc, int32_t global_tid) {} | |
| } | |
| namespace ompx { | |
| /// Helper class to hide the generic loop nest and provide the template argument | |
| /// throughout. | |
| template <typename Ty> class StaticLoopChunker { | |
| /// Generic loop nest that handles block and/or thread distribution in the | |
| /// absence of user specified chunk sizes. This implicitly picks a block chunk | |
| /// size equal to the number of threads in the block and a thread chunk size | |
| /// equal to one. In contrast to the chunked version we can get away with a | |
| /// single loop in this case | |
| static void NormalizedLoopNestNoChunk(void (*LoopBody)(Ty, void *), void *Arg, | |
| Ty NumBlocks, Ty BId, Ty NumThreads, | |
| Ty TId, Ty NumIters, | |
| bool OneIterationPerThread) { | |
| Ty KernelIteration = NumBlocks * NumThreads; | |
| // Start index in the normalized space. | |
| Ty IV = BId * NumThreads + TId; | |
| ASSERT(IV >= 0, "Bad index"); | |
| // Cover the entire iteration space, assumptions in the caller might allow | |
| // to simplify this loop to a conditional. | |
| if (IV < NumIters) { | |
| do { | |
| // Execute the loop body. | |
| LoopBody(IV, Arg); | |
| // Every thread executed one block and thread chunk now. | |
| IV += KernelIteration; | |
| if (OneIterationPerThread) | |
| return; | |
| } while (IV < NumIters); | |
| } | |
| } | |
| /// Generic loop nest that handles block and/or thread distribution in the | |
| /// presence of user specified chunk sizes (for at least one of them). | |
| static void NormalizedLoopNestChunked(void (*LoopBody)(Ty, void *), void *Arg, | |
| Ty BlockChunk, Ty NumBlocks, Ty BId, | |
| Ty ThreadChunk, Ty NumThreads, Ty TId, | |
| Ty NumIters, | |
| bool OneIterationPerThread) { | |
| Ty KernelIteration = NumBlocks * BlockChunk; | |
| // Start index in the chunked space. | |
| Ty IV = BId * BlockChunk + TId; | |
| ASSERT(IV >= 0, "Bad index"); | |
| // Cover the entire iteration space, assumptions in the caller might allow | |
| // to simplify this loop to a conditional. | |
| do { | |
| Ty BlockChunkLeft = | |
| BlockChunk >= TId * ThreadChunk ? BlockChunk - TId * ThreadChunk : 0; | |
| Ty ThreadChunkLeft = | |
| ThreadChunk <= BlockChunkLeft ? ThreadChunk : BlockChunkLeft; | |
| while (ThreadChunkLeft--) { | |
| // Given the blocking it's hard to keep track of what to execute. | |
| if (IV >= NumIters) | |
| return; | |
| // Execute the loop body. | |
| LoopBody(IV, Arg); | |
| if (OneIterationPerThread) | |
| return; | |
| ++IV; | |
| } | |
| IV += KernelIteration; | |
| } while (IV < NumIters); | |
| } | |
| public: | |
| /// Worksharing `for`-loop. | |
| static void For(IdentTy *Loc, void (*LoopBody)(Ty, void *), void *Arg, | |
| Ty NumIters, Ty NumThreads, Ty ThreadChunk) { | |
| ASSERT(NumIters >= 0, "Bad iteration count"); | |
| ASSERT(ThreadChunk >= 0, "Bad thread count"); | |
| // All threads need to participate but we don't know if we are in a | |
| // parallel at all or if the user might have used a `num_threads` clause | |
| // on the parallel and reduced the number compared to the block size. | |
| // Since nested parallels are possible too we need to get the thread id | |
| // from the `omp` getter and not the mapping directly. | |
| Ty TId = omp_get_thread_num(); | |
| // There are no blocks involved here. | |
| Ty BlockChunk = 0; | |
| Ty NumBlocks = 1; | |
| Ty BId = 0; | |
| // If the thread chunk is not specified we pick a default now. | |
| if (ThreadChunk == 0) | |
| ThreadChunk = 1; | |
| // If we know we have more threads than iterations we can indicate that to | |
| // avoid an outer loop. | |
| bool OneIterationPerThread = false; | |
| if (config::getAssumeThreadsOversubscription()) { | |
| ASSERT(NumThreads >= NumIters, "Broken assumption"); | |
| OneIterationPerThread = true; | |
| } | |
| if (ThreadChunk != 1) | |
| NormalizedLoopNestChunked(LoopBody, Arg, BlockChunk, NumBlocks, BId, | |
| ThreadChunk, NumThreads, TId, NumIters, | |
| OneIterationPerThread); | |
| else | |
| NormalizedLoopNestNoChunk(LoopBody, Arg, NumBlocks, BId, NumThreads, TId, | |
| NumIters, OneIterationPerThread); | |
| } | |
| /// Worksharing `distrbute`-loop. | |
| static void Distribute(IdentTy *Loc, void (*LoopBody)(Ty, void *), void *Arg, | |
| Ty NumIters, Ty BlockChunk) { | |
| ASSERT(icv::Level == 0, "Bad distribute"); | |
| ASSERT(icv::ActiveLevel == 0, "Bad distribute"); | |
| ASSERT(state::ParallelRegionFn == nullptr, "Bad distribute"); | |
| ASSERT(state::ParallelTeamSize == 1, "Bad distribute"); | |
| ASSERT(NumIters >= 0, "Bad iteration count"); | |
| ASSERT(BlockChunk >= 0, "Bad block count"); | |
| // There are no threads involved here. | |
| Ty ThreadChunk = 0; | |
| Ty NumThreads = 1; | |
| Ty TId = 0; | |
| ASSERT(TId == mapping::getThreadIdInBlock(), "Bad thread id"); | |
| // All teams need to participate. | |
| Ty NumBlocks = mapping::getNumberOfBlocksInKernel(); | |
| Ty BId = mapping::getBlockIdInKernel(); | |
| // If the block chunk is not specified we pick a default now. | |
| if (BlockChunk == 0) | |
| BlockChunk = NumThreads; | |
| // If we know we have more blocks than iterations we can indicate that to | |
| // avoid an outer loop. | |
| bool OneIterationPerThread = false; | |
| if (config::getAssumeTeamsOversubscription()) { | |
| ASSERT(NumBlocks >= NumIters, "Broken assumption"); | |
| OneIterationPerThread = true; | |
| } | |
| if (BlockChunk != NumThreads) | |
| NormalizedLoopNestChunked(LoopBody, Arg, BlockChunk, NumBlocks, BId, | |
| ThreadChunk, NumThreads, TId, NumIters, | |
| OneIterationPerThread); | |
| else | |
| NormalizedLoopNestNoChunk(LoopBody, Arg, NumBlocks, BId, NumThreads, TId, | |
| NumIters, OneIterationPerThread); | |
| ASSERT(icv::Level == 0, "Bad distribute"); | |
| ASSERT(icv::ActiveLevel == 0, "Bad distribute"); | |
| ASSERT(state::ParallelRegionFn == nullptr, "Bad distribute"); | |
| ASSERT(state::ParallelTeamSize == 1, "Bad distribute"); | |
| } | |
| /// Worksharing `distrbute parallel for`-loop. | |
| static void DistributeFor(IdentTy *Loc, void (*LoopBody)(Ty, void *), | |
| void *Arg, Ty NumIters, Ty NumThreads, | |
| Ty BlockChunk, Ty ThreadChunk) { | |
| ASSERT(icv::Level == 1, "Bad distribute"); | |
| ASSERT(icv::ActiveLevel == 1, "Bad distribute"); | |
| ASSERT(state::ParallelRegionFn == nullptr, "Bad distribute"); | |
| ASSERT(NumIters >= 0, "Bad iteration count"); | |
| ASSERT(BlockChunk >= 0, "Bad block count"); | |
| ASSERT(ThreadChunk >= 0, "Bad thread count"); | |
| // All threads need to participate but the user might have used a | |
| // `num_threads` clause on the parallel and reduced the number compared to | |
| // the block size. | |
| Ty TId = mapping::getThreadIdInBlock(); | |
| // All teams need to participate. | |
| Ty NumBlocks = mapping::getNumberOfBlocksInKernel(); | |
| Ty BId = mapping::getBlockIdInKernel(); | |
| // If the block chunk is not specified we pick a default now. | |
| if (BlockChunk == 0) | |
| BlockChunk = NumThreads; | |
| // If the thread chunk is not specified we pick a default now. | |
| if (ThreadChunk == 0) | |
| ThreadChunk = 1; | |
| // If we know we have more threads (across all blocks) than iterations we | |
| // can indicate that to avoid an outer loop. | |
| bool OneIterationPerThread = false; | |
| if (config::getAssumeTeamsOversubscription() & | |
| config::getAssumeThreadsOversubscription()) { | |
| OneIterationPerThread = true; | |
| ASSERT(NumBlocks * NumThreads >= NumIters, "Broken assumption"); | |
| } | |
| if (BlockChunk != NumThreads || ThreadChunk != 1) | |
| NormalizedLoopNestChunked(LoopBody, Arg, BlockChunk, NumBlocks, BId, | |
| ThreadChunk, NumThreads, TId, NumIters, | |
| OneIterationPerThread); | |
| else | |
| NormalizedLoopNestNoChunk(LoopBody, Arg, NumBlocks, BId, NumThreads, TId, | |
| NumIters, OneIterationPerThread); | |
| ASSERT(icv::Level == 1, "Bad distribute"); | |
| ASSERT(icv::ActiveLevel == 1, "Bad distribute"); | |
| ASSERT(state::ParallelRegionFn == nullptr, "Bad distribute"); | |
| } | |
| }; | |
| } // namespace ompx | |
| #define OMP_LOOP_ENTRY(BW, TY) \ | |
| [[gnu::flatten, clang::always_inline]] void \ | |
| __kmpc_distribute_for_static_loop##BW( \ | |
| IdentTy *loc, void (*fn)(TY, void *), void *arg, TY num_iters, \ | |
| TY num_threads, TY block_chunk, TY thread_chunk) { \ | |
| ompx::StaticLoopChunker<TY>::DistributeFor( \ | |
| loc, fn, arg, num_iters + 1, num_threads, block_chunk, thread_chunk); \ | |
| } \ | |
| [[gnu::flatten, clang::always_inline]] void \ | |
| __kmpc_distribute_static_loop##BW(IdentTy *loc, void (*fn)(TY, void *), \ | |
| void *arg, TY num_iters, \ | |
| TY block_chunk) { \ | |
| ompx::StaticLoopChunker<TY>::Distribute(loc, fn, arg, num_iters + 1, \ | |
| block_chunk); \ | |
| } \ | |
| [[gnu::flatten, clang::always_inline]] void __kmpc_for_static_loop##BW( \ | |
| IdentTy *loc, void (*fn)(TY, void *), void *arg, TY num_iters, \ | |
| TY num_threads, TY thread_chunk) { \ | |
| ompx::StaticLoopChunker<TY>::For(loc, fn, arg, num_iters + 1, num_threads, \ | |
| thread_chunk); \ | |
| } | |
| extern "C" { | |
| OMP_LOOP_ENTRY(_4, int32_t) | |
| OMP_LOOP_ENTRY(_4u, uint32_t) | |
| OMP_LOOP_ENTRY(_8, int64_t) | |
| OMP_LOOP_ENTRY(_8u, uint64_t) | |
| } | |
| #pragma omp end declare target |