958 lines · cpp
1//===----- Workshare.cpp - OpenMP workshare implementation ------ C++ -*-===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8//9// This file contains the implementation of the KMPC interface10// for the loop construct plus other worksharing constructs that use the same11// interface as loops.12//13//===----------------------------------------------------------------------===//14 15#include "Workshare.h"16#include "Debug.h"17#include "DeviceTypes.h"18#include "DeviceUtils.h"19#include "Interface.h"20#include "Mapping.h"21#include "State.h"22#include "Synchronization.h"23 24using namespace ompx;25 26// TODO:27struct DynamicScheduleTracker {28 int64_t Chunk;29 int64_t LoopUpperBound;30 int64_t NextLowerBound;31 int64_t Stride;32 kmp_sched_t ScheduleType;33 DynamicScheduleTracker *NextDST;34};35 36#define ASSERT0(...)37 38// used by the library for the interface with the app39#define DISPATCH_FINISHED 040#define DISPATCH_NOTFINISHED 141 42// used by dynamic scheduling43#define FINISHED 044#define NOT_FINISHED 145#define LAST_CHUNK 246 47// TODO: This variable is a hack inherited from the old runtime.48[[clang::loader_uninitialized]] static Local<uint64_t> Cnt;49 50template <typename T, typename ST> struct omptarget_nvptx_LoopSupport {51 ////////////////////////////////////////////////////////////////////////////////52 // Loop with static scheduling with chunk53 54 // Generic implementation of OMP loop scheduling with static policy55 /*! \brief Calculate initial bounds for static loop and stride56 * @param[in] loc location in code of the call (not used here)57 * @param[in] global_tid global thread id58 * @param[in] schetype type of scheduling (see omptarget-nvptx.h)59 * @param[in] plastiter pointer to last iteration60 * @param[in,out] pointer to loop lower bound. it will contain value of61 * lower bound of first chunk62 * @param[in,out] pointer to loop upper bound. It will contain value of63 * upper bound of first chunk64 * @param[in,out] pointer to loop stride. It will contain value of stride65 * between two successive chunks executed by the same thread66 * @param[in] loop increment bump67 * @param[in] chunk size68 */69 70 // helper function for static chunk71 static void ForStaticChunk(int &last, T &lb, T &ub, ST &stride, ST chunk,72 T entityId, T numberOfEntities) {73 // each thread executes multiple chunks all of the same size, except74 // the last one75 // distance between two successive chunks76 stride = numberOfEntities * chunk;77 lb = lb + entityId * chunk;78 T inputUb = ub;79 ub = lb + chunk - 1; // Clang uses i <= ub80 // Say ub' is the beginning of the last chunk. Then who ever has a81 // lower bound plus a multiple of the increment equal to ub' is82 // the last one.83 T beginingLastChunk = inputUb - (inputUb % chunk);84 last = ((beginingLastChunk - lb) % stride) == 0;85 }86 87 ////////////////////////////////////////////////////////////////////////////////88 // Loop with static scheduling without chunk89 90 // helper function for static no chunk91 static void ForStaticNoChunk(int &last, T &lb, T &ub, ST &stride, ST &chunk,92 T entityId, T numberOfEntities) {93 // No chunk size specified. Each thread or warp gets at most one94 // chunk; chunks are all almost of equal size95 T loopSize = ub - lb + 1;96 97 chunk = loopSize / numberOfEntities;98 T leftOver = loopSize - chunk * numberOfEntities;99 100 if (entityId < leftOver) {101 chunk++;102 lb = lb + entityId * chunk;103 } else {104 lb = lb + entityId * chunk + leftOver;105 }106 107 T inputUb = ub;108 ub = lb + chunk - 1; // Clang uses i <= ub109 last = lb <= inputUb && inputUb <= ub;110 stride = loopSize; // make sure we only do 1 chunk per warp111 }112 113 ////////////////////////////////////////////////////////////////////////////////114 // Support for Static Init115 116 static void for_static_init(int32_t, int32_t schedtype, int32_t *plastiter,117 T *plower, T *pupper, ST *pstride, ST chunk,118 bool IsSPMDExecutionMode) {119 int32_t gtid = omp_get_thread_num();120 int numberOfActiveOMPThreads = omp_get_num_threads();121 122 // All warps that are in excess of the maximum requested, do123 // not execute the loop124 ASSERT0(LT_FUSSY, gtid < numberOfActiveOMPThreads,125 "current thread is not needed here; error");126 127 // copy128 int lastiter = 0;129 T lb = *plower;130 T ub = *pupper;131 ST stride = *pstride;132 133 // init134 switch (SCHEDULE_WITHOUT_MODIFIERS(schedtype)) {135 case kmp_sched_static_chunk: {136 if (chunk > 0) {137 ForStaticChunk(lastiter, lb, ub, stride, chunk, gtid,138 numberOfActiveOMPThreads);139 break;140 }141 [[fallthrough]];142 } // note: if chunk <=0, use nochunk143 case kmp_sched_static_balanced_chunk: {144 if (chunk > 0) {145 // round up to make sure the chunk is enough to cover all iterations146 T tripCount = ub - lb + 1; // +1 because ub is inclusive147 T span = (tripCount + numberOfActiveOMPThreads - 1) /148 numberOfActiveOMPThreads;149 // perform chunk adjustment150 chunk = (span + chunk - 1) & ~(chunk - 1);151 152 ASSERT0(LT_FUSSY, ub >= lb, "ub must be >= lb.");153 T oldUb = ub;154 ForStaticChunk(lastiter, lb, ub, stride, chunk, gtid,155 numberOfActiveOMPThreads);156 if (ub > oldUb)157 ub = oldUb;158 break;159 }160 [[fallthrough]];161 } // note: if chunk <=0, use nochunk162 case kmp_sched_static_nochunk: {163 ForStaticNoChunk(lastiter, lb, ub, stride, chunk, gtid,164 numberOfActiveOMPThreads);165 break;166 }167 case kmp_sched_distr_static_chunk: {168 if (chunk > 0) {169 ForStaticChunk(lastiter, lb, ub, stride, chunk, omp_get_team_num(),170 omp_get_num_teams());171 break;172 }173 [[fallthrough]];174 } // note: if chunk <=0, use nochunk175 case kmp_sched_distr_static_nochunk: {176 ForStaticNoChunk(lastiter, lb, ub, stride, chunk, omp_get_team_num(),177 omp_get_num_teams());178 break;179 }180 case kmp_sched_distr_static_chunk_sched_static_chunkone: {181 ForStaticChunk(lastiter, lb, ub, stride, chunk,182 numberOfActiveOMPThreads * omp_get_team_num() + gtid,183 omp_get_num_teams() * numberOfActiveOMPThreads);184 break;185 }186 default: {187 // ASSERT(LT_FUSSY, 0, "unknown schedtype %d", (int)schedtype);188 ForStaticChunk(lastiter, lb, ub, stride, chunk, gtid,189 numberOfActiveOMPThreads);190 break;191 }192 }193 // copy back194 *plastiter = lastiter;195 *plower = lb;196 *pupper = ub;197 *pstride = stride;198 }199 200 ////////////////////////////////////////////////////////////////////////////////201 // Support for dispatch Init202 203 static int OrderedSchedule(kmp_sched_t schedule) {204 return schedule >= kmp_sched_ordered_first &&205 schedule <= kmp_sched_ordered_last;206 }207 208 static void dispatch_init(IdentTy *loc, int32_t threadId,209 kmp_sched_t schedule, T lb, T ub, ST st, ST chunk,210 DynamicScheduleTracker *DST) {211 int tid = mapping::getThreadIdInBlock();212 T tnum = omp_get_num_threads();213 T tripCount = ub - lb + 1; // +1 because ub is inclusive214 ASSERT0(LT_FUSSY, threadId < tnum,215 "current thread is not needed here; error");216 217 /* Currently just ignore the monotonic and non-monotonic modifiers218 * (the compiler isn't producing them * yet anyway).219 * When it is we'll want to look at them somewhere here and use that220 * information to add to our schedule choice. We shouldn't need to pass221 * them on, they merely affect which schedule we can legally choose for222 * various dynamic cases. (In particular, whether or not a stealing scheme223 * is legal).224 */225 schedule = SCHEDULE_WITHOUT_MODIFIERS(schedule);226 227 // Process schedule.228 if (tnum == 1 || tripCount <= 1 || OrderedSchedule(schedule)) {229 if (OrderedSchedule(schedule))230 __kmpc_barrier(loc, threadId);231 schedule = kmp_sched_static_chunk;232 chunk = tripCount; // one thread gets the whole loop233 } else if (schedule == kmp_sched_runtime) {234 // process runtime235 omp_sched_t rtSched;236 int ChunkInt;237 omp_get_schedule(&rtSched, &ChunkInt);238 chunk = ChunkInt;239 switch (rtSched) {240 case omp_sched_static: {241 if (chunk > 0)242 schedule = kmp_sched_static_chunk;243 else244 schedule = kmp_sched_static_nochunk;245 break;246 }247 case omp_sched_auto: {248 schedule = kmp_sched_static_chunk;249 chunk = 1;250 break;251 }252 case omp_sched_dynamic:253 case omp_sched_guided: {254 schedule = kmp_sched_dynamic;255 break;256 }257 }258 } else if (schedule == kmp_sched_auto) {259 schedule = kmp_sched_static_chunk;260 chunk = 1;261 } else {262 // ASSERT(LT_FUSSY,263 // schedule == kmp_sched_dynamic || schedule == kmp_sched_guided,264 // "unknown schedule %d & chunk %lld\n", (int)schedule,265 // (long long)chunk);266 }267 268 // init schedules269 if (schedule == kmp_sched_static_chunk) {270 ASSERT0(LT_FUSSY, chunk > 0, "bad chunk value");271 // save sched state272 DST->ScheduleType = schedule;273 // save ub274 DST->LoopUpperBound = ub;275 // compute static chunk276 ST stride;277 int lastiter = 0;278 ForStaticChunk(lastiter, lb, ub, stride, chunk, threadId, tnum);279 // save computed params280 DST->Chunk = chunk;281 DST->NextLowerBound = lb;282 DST->Stride = stride;283 } else if (schedule == kmp_sched_static_balanced_chunk) {284 ASSERT0(LT_FUSSY, chunk > 0, "bad chunk value");285 // save sched state286 DST->ScheduleType = schedule;287 // save ub288 DST->LoopUpperBound = ub;289 // compute static chunk290 ST stride;291 int lastiter = 0;292 // round up to make sure the chunk is enough to cover all iterations293 T span = (tripCount + tnum - 1) / tnum;294 // perform chunk adjustment295 chunk = (span + chunk - 1) & ~(chunk - 1);296 297 T oldUb = ub;298 ForStaticChunk(lastiter, lb, ub, stride, chunk, threadId, tnum);299 ASSERT0(LT_FUSSY, ub >= lb, "ub must be >= lb.");300 if (ub > oldUb)301 ub = oldUb;302 // save computed params303 DST->Chunk = chunk;304 DST->NextLowerBound = lb;305 DST->Stride = stride;306 } else if (schedule == kmp_sched_static_nochunk) {307 ASSERT0(LT_FUSSY, chunk == 0, "bad chunk value");308 // save sched state309 DST->ScheduleType = schedule;310 // save ub311 DST->LoopUpperBound = ub;312 // compute static chunk313 ST stride;314 int lastiter = 0;315 ForStaticNoChunk(lastiter, lb, ub, stride, chunk, threadId, tnum);316 // save computed params317 DST->Chunk = chunk;318 DST->NextLowerBound = lb;319 DST->Stride = stride;320 } else if (schedule == kmp_sched_dynamic || schedule == kmp_sched_guided) {321 // save data322 DST->ScheduleType = schedule;323 if (chunk < 1)324 chunk = 1;325 DST->Chunk = chunk;326 DST->LoopUpperBound = ub;327 DST->NextLowerBound = lb;328 __kmpc_barrier(loc, threadId);329 if (tid == 0) {330 Cnt = 0;331 fence::team(atomic::seq_cst);332 }333 __kmpc_barrier(loc, threadId);334 }335 }336 337 ////////////////////////////////////////////////////////////////////////////////338 // Support for dispatch next339 340 static uint64_t NextIter() {341 __kmpc_impl_lanemask_t active = mapping::activemask();342 uint32_t leader = utils::ffs(active) - 1;343 uint32_t change = utils::popc(active);344 __kmpc_impl_lanemask_t lane_mask_lt = mapping::lanemaskLT();345 unsigned int rank = utils::popc(active & lane_mask_lt);346 uint64_t warp_res = 0;347 if (rank == 0) {348 warp_res = atomic::add(&Cnt, change, atomic::seq_cst);349 }350 warp_res = utils::shuffle(active, warp_res, leader, mapping::getWarpSize());351 return warp_res + rank;352 }353 354 static int DynamicNextChunk(T &lb, T &ub, T chunkSize, T loopLowerBound,355 T loopUpperBound) {356 T N = NextIter();357 lb = loopLowerBound + N * chunkSize;358 ub = lb + chunkSize - 1; // Clang uses i <= ub359 360 // 3 result cases:361 // a. lb and ub < loopUpperBound --> NOT_FINISHED362 // b. lb < loopUpperBound and ub >= loopUpperBound: last chunk -->363 // NOT_FINISHED364 // c. lb and ub >= loopUpperBound: empty chunk --> FINISHED365 // a.366 if (lb <= loopUpperBound && ub < loopUpperBound) {367 return NOT_FINISHED;368 }369 // b.370 if (lb <= loopUpperBound) {371 ub = loopUpperBound;372 return LAST_CHUNK;373 }374 // c. if we are here, we are in case 'c'375 lb = loopUpperBound + 2;376 ub = loopUpperBound + 1;377 return FINISHED;378 }379 380 static int dispatch_next(IdentTy *loc, int32_t gtid, int32_t *plast,381 T *plower, T *pupper, ST *pstride,382 DynamicScheduleTracker *DST) {383 // ID of a thread in its own warp384 385 // automatically selects thread or warp ID based on selected implementation386 ASSERT0(LT_FUSSY, gtid < omp_get_num_threads(),387 "current thread is not needed here; error");388 // retrieve schedule389 kmp_sched_t schedule = DST->ScheduleType;390 391 // xxx reduce to one392 if (schedule == kmp_sched_static_chunk ||393 schedule == kmp_sched_static_nochunk) {394 T myLb = DST->NextLowerBound;395 T ub = DST->LoopUpperBound;396 // finished?397 if (myLb > ub) {398 return DISPATCH_FINISHED;399 }400 // not finished, save current bounds401 ST chunk = DST->Chunk;402 *plower = myLb;403 T myUb = myLb + chunk - 1; // Clang uses i <= ub404 if (myUb > ub)405 myUb = ub;406 *pupper = myUb;407 *plast = (int32_t)(myUb == ub);408 409 // increment next lower bound by the stride410 ST stride = DST->Stride;411 DST->NextLowerBound = myLb + stride;412 return DISPATCH_NOTFINISHED;413 }414 ASSERT0(LT_FUSSY,415 schedule == kmp_sched_dynamic || schedule == kmp_sched_guided,416 "bad sched");417 T myLb, myUb;418 int finished = DynamicNextChunk(myLb, myUb, DST->Chunk, DST->NextLowerBound,419 DST->LoopUpperBound);420 421 if (finished == FINISHED)422 return DISPATCH_FINISHED;423 424 // not finished (either not finished or last chunk)425 *plast = (int32_t)(finished == LAST_CHUNK);426 *plower = myLb;427 *pupper = myUb;428 *pstride = 1;429 430 return DISPATCH_NOTFINISHED;431 }432 433 static void dispatch_fini() {434 // nothing435 }436 437 ////////////////////////////////////////////////////////////////////////////////438 // end of template class that encapsulate all the helper functions439 ////////////////////////////////////////////////////////////////////////////////440};441 442////////////////////////////////////////////////////////////////////////////////443// KMP interface implementation (dyn loops)444////////////////////////////////////////////////////////////////////////////////445 446// TODO: Expand the dispatch API to take a DST pointer which can then be447// allocated properly without malloc.448// For now, each team will contain an LDS pointer (ThreadDST) to a global array449// of references to the DST structs allocated (in global memory) for each thread450// in the team. The global memory array is allocated during the init phase if it451// was not allocated already and will be deallocated when the dispatch phase452// ends:453//454// __kmpc_dispatch_init455//456// ** Dispatch loop **457//458// __kmpc_dispatch_deinit459//460[[clang::loader_uninitialized]] static Local<DynamicScheduleTracker **>461 ThreadDST;462 463// Create a new DST, link the current one, and define the new as current.464static DynamicScheduleTracker *pushDST() {465 int32_t ThreadIndex = mapping::getThreadIdInBlock();466 // Each block will allocate an array of pointers to DST structs. The array is467 // equal in length to the number of threads in that block.468 if (!ThreadDST) {469 // Allocate global memory array of pointers to DST structs:470 if (mapping::isMainThreadInGenericMode() || ThreadIndex == 0)471 ThreadDST = static_cast<DynamicScheduleTracker **>(472 memory::allocGlobal(mapping::getNumberOfThreadsInBlock() *473 sizeof(DynamicScheduleTracker *),474 "new ThreadDST array"));475 synchronize::threads(atomic::seq_cst);476 477 // Initialize the array pointers:478 ThreadDST[ThreadIndex] = nullptr;479 }480 481 // Create a DST struct for the current thread:482 DynamicScheduleTracker *NewDST = static_cast<DynamicScheduleTracker *>(483 memory::allocGlobal(sizeof(DynamicScheduleTracker), "new DST"));484 *NewDST = DynamicScheduleTracker({0});485 486 // Add the new DST struct to the array of DST structs:487 NewDST->NextDST = ThreadDST[ThreadIndex];488 ThreadDST[ThreadIndex] = NewDST;489 return NewDST;490}491 492// Return the current DST.493static DynamicScheduleTracker *peekDST() {494 return ThreadDST[mapping::getThreadIdInBlock()];495}496 497// Pop the current DST and restore the last one.498static void popDST() {499 int32_t ThreadIndex = mapping::getThreadIdInBlock();500 DynamicScheduleTracker *CurrentDST = ThreadDST[ThreadIndex];501 DynamicScheduleTracker *OldDST = CurrentDST->NextDST;502 memory::freeGlobal(CurrentDST, "remove DST");503 ThreadDST[ThreadIndex] = OldDST;504 505 // Check if we need to deallocate the global array. Ensure all threads506 // in the block have finished deallocating the individual DSTs.507 synchronize::threads(atomic::seq_cst);508 if (!ThreadDST[ThreadIndex] && !ThreadIndex) {509 memory::freeGlobal(ThreadDST, "remove ThreadDST array");510 ThreadDST = nullptr;511 }512 synchronize::threads(atomic::seq_cst);513}514 515void workshare::init(bool IsSPMD) {516 if (mapping::isInitialThreadInLevel0(IsSPMD))517 ThreadDST = nullptr;518}519 520extern "C" {521 522// init523void __kmpc_dispatch_init_4(IdentTy *loc, int32_t tid, int32_t schedule,524 int32_t lb, int32_t ub, int32_t st, int32_t chunk) {525 DynamicScheduleTracker *DST = pushDST();526 omptarget_nvptx_LoopSupport<int32_t, int32_t>::dispatch_init(527 loc, tid, (kmp_sched_t)schedule, lb, ub, st, chunk, DST);528}529 530void __kmpc_dispatch_init_4u(IdentTy *loc, int32_t tid, int32_t schedule,531 uint32_t lb, uint32_t ub, int32_t st,532 int32_t chunk) {533 DynamicScheduleTracker *DST = pushDST();534 omptarget_nvptx_LoopSupport<uint32_t, int32_t>::dispatch_init(535 loc, tid, (kmp_sched_t)schedule, lb, ub, st, chunk, DST);536}537 538void __kmpc_dispatch_init_8(IdentTy *loc, int32_t tid, int32_t schedule,539 int64_t lb, int64_t ub, int64_t st, int64_t chunk) {540 DynamicScheduleTracker *DST = pushDST();541 omptarget_nvptx_LoopSupport<int64_t, int64_t>::dispatch_init(542 loc, tid, (kmp_sched_t)schedule, lb, ub, st, chunk, DST);543}544 545void __kmpc_dispatch_init_8u(IdentTy *loc, int32_t tid, int32_t schedule,546 uint64_t lb, uint64_t ub, int64_t st,547 int64_t chunk) {548 DynamicScheduleTracker *DST = pushDST();549 omptarget_nvptx_LoopSupport<uint64_t, int64_t>::dispatch_init(550 loc, tid, (kmp_sched_t)schedule, lb, ub, st, chunk, DST);551}552 553// next554int __kmpc_dispatch_next_4(IdentTy *loc, int32_t tid, int32_t *p_last,555 int32_t *p_lb, int32_t *p_ub, int32_t *p_st) {556 DynamicScheduleTracker *DST = peekDST();557 return omptarget_nvptx_LoopSupport<int32_t, int32_t>::dispatch_next(558 loc, tid, p_last, p_lb, p_ub, p_st, DST);559}560 561int __kmpc_dispatch_next_4u(IdentTy *loc, int32_t tid, int32_t *p_last,562 uint32_t *p_lb, uint32_t *p_ub, int32_t *p_st) {563 DynamicScheduleTracker *DST = peekDST();564 return omptarget_nvptx_LoopSupport<uint32_t, int32_t>::dispatch_next(565 loc, tid, p_last, p_lb, p_ub, p_st, DST);566}567 568int __kmpc_dispatch_next_8(IdentTy *loc, int32_t tid, int32_t *p_last,569 int64_t *p_lb, int64_t *p_ub, int64_t *p_st) {570 DynamicScheduleTracker *DST = peekDST();571 return omptarget_nvptx_LoopSupport<int64_t, int64_t>::dispatch_next(572 loc, tid, p_last, p_lb, p_ub, p_st, DST);573}574 575int __kmpc_dispatch_next_8u(IdentTy *loc, int32_t tid, int32_t *p_last,576 uint64_t *p_lb, uint64_t *p_ub, int64_t *p_st) {577 DynamicScheduleTracker *DST = peekDST();578 return omptarget_nvptx_LoopSupport<uint64_t, int64_t>::dispatch_next(579 loc, tid, p_last, p_lb, p_ub, p_st, DST);580}581 582// fini583void __kmpc_dispatch_fini_4(IdentTy *loc, int32_t tid) {584 omptarget_nvptx_LoopSupport<int32_t, int32_t>::dispatch_fini();585}586 587void __kmpc_dispatch_fini_4u(IdentTy *loc, int32_t tid) {588 omptarget_nvptx_LoopSupport<uint32_t, int32_t>::dispatch_fini();589}590 591void __kmpc_dispatch_fini_8(IdentTy *loc, int32_t tid) {592 omptarget_nvptx_LoopSupport<int64_t, int64_t>::dispatch_fini();593}594 595void __kmpc_dispatch_fini_8u(IdentTy *loc, int32_t tid) {596 omptarget_nvptx_LoopSupport<uint64_t, int64_t>::dispatch_fini();597}598 599// deinit600void __kmpc_dispatch_deinit(IdentTy *loc, int32_t tid) { popDST(); }601 602////////////////////////////////////////////////////////////////////////////////603// KMP interface implementation (static loops)604////////////////////////////////////////////////////////////////////////////////605 606void __kmpc_for_static_init_4(IdentTy *loc, int32_t global_tid,607 int32_t schedtype, int32_t *plastiter,608 int32_t *plower, int32_t *pupper,609 int32_t *pstride, int32_t incr, int32_t chunk) {610 omptarget_nvptx_LoopSupport<int32_t, int32_t>::for_static_init(611 global_tid, schedtype, plastiter, plower, pupper, pstride, chunk,612 mapping::isSPMDMode());613}614 615void __kmpc_for_static_init_4u(IdentTy *loc, int32_t global_tid,616 int32_t schedtype, int32_t *plastiter,617 uint32_t *plower, uint32_t *pupper,618 int32_t *pstride, int32_t incr, int32_t chunk) {619 omptarget_nvptx_LoopSupport<uint32_t, int32_t>::for_static_init(620 global_tid, schedtype, plastiter, plower, pupper, pstride, chunk,621 mapping::isSPMDMode());622}623 624void __kmpc_for_static_init_8(IdentTy *loc, int32_t global_tid,625 int32_t schedtype, int32_t *plastiter,626 int64_t *plower, int64_t *pupper,627 int64_t *pstride, int64_t incr, int64_t chunk) {628 omptarget_nvptx_LoopSupport<int64_t, int64_t>::for_static_init(629 global_tid, schedtype, plastiter, plower, pupper, pstride, chunk,630 mapping::isSPMDMode());631}632 633void __kmpc_for_static_init_8u(IdentTy *loc, int32_t global_tid,634 int32_t schedtype, int32_t *plastiter,635 uint64_t *plower, uint64_t *pupper,636 int64_t *pstride, int64_t incr, int64_t chunk) {637 omptarget_nvptx_LoopSupport<uint64_t, int64_t>::for_static_init(638 global_tid, schedtype, plastiter, plower, pupper, pstride, chunk,639 mapping::isSPMDMode());640}641 642void __kmpc_distribute_static_init_4(IdentTy *loc, int32_t global_tid,643 int32_t schedtype, int32_t *plastiter,644 int32_t *plower, int32_t *pupper,645 int32_t *pstride, int32_t incr,646 int32_t chunk) {647 omptarget_nvptx_LoopSupport<int32_t, int32_t>::for_static_init(648 global_tid, schedtype, plastiter, plower, pupper, pstride, chunk,649 mapping::isSPMDMode());650}651 652void __kmpc_distribute_static_init_4u(IdentTy *loc, int32_t global_tid,653 int32_t schedtype, int32_t *plastiter,654 uint32_t *plower, uint32_t *pupper,655 int32_t *pstride, int32_t incr,656 int32_t chunk) {657 omptarget_nvptx_LoopSupport<uint32_t, int32_t>::for_static_init(658 global_tid, schedtype, plastiter, plower, pupper, pstride, chunk,659 mapping::isSPMDMode());660}661 662void __kmpc_distribute_static_init_8(IdentTy *loc, int32_t global_tid,663 int32_t schedtype, int32_t *plastiter,664 int64_t *plower, int64_t *pupper,665 int64_t *pstride, int64_t incr,666 int64_t chunk) {667 omptarget_nvptx_LoopSupport<int64_t, int64_t>::for_static_init(668 global_tid, schedtype, plastiter, plower, pupper, pstride, chunk,669 mapping::isSPMDMode());670}671 672void __kmpc_distribute_static_init_8u(IdentTy *loc, int32_t global_tid,673 int32_t schedtype, int32_t *plastiter,674 uint64_t *plower, uint64_t *pupper,675 int64_t *pstride, int64_t incr,676 int64_t chunk) {677 omptarget_nvptx_LoopSupport<uint64_t, int64_t>::for_static_init(678 global_tid, schedtype, plastiter, plower, pupper, pstride, chunk,679 mapping::isSPMDMode());680}681 682void __kmpc_for_static_fini(IdentTy *loc, int32_t global_tid) {}683 684void __kmpc_distribute_static_fini(IdentTy *loc, int32_t global_tid) {}685}686 687namespace ompx {688 689/// Helper class to hide the generic loop nest and provide the template argument690/// throughout.691template <typename Ty> class StaticLoopChunker {692 693 /// Generic loop nest that handles block and/or thread distribution in the694 /// absence of user specified chunk sizes. This implicitly picks a block chunk695 /// size equal to the number of threads in the block and a thread chunk size696 /// equal to one. In contrast to the chunked version we can get away with a697 /// single loop in this case698 static void NormalizedLoopNestNoChunk(void (*LoopBody)(Ty, void *), void *Arg,699 Ty NumBlocks, Ty BId, Ty NumThreads,700 Ty TId, Ty NumIters,701 uint8_t OneIterationPerThread) {702 Ty KernelIteration = NumBlocks * NumThreads;703 704 // Start index in the normalized space.705 Ty IV = BId * NumThreads + TId;706 ASSERT(IV >= 0, "Bad index");707 708 // Cover the entire iteration space, assumptions in the caller might allow709 // to simplify this loop to a conditional.710 if (IV < NumIters) {711 do {712 713 // Execute the loop body.714 LoopBody(IV, Arg);715 716 // Every thread executed one block and thread chunk now.717 IV += KernelIteration;718 719 if (OneIterationPerThread)720 return;721 722 } while (IV < NumIters);723 }724 }725 726 /// Generic loop nest that handles block and/or thread distribution in the727 /// presence of user specified chunk sizes (for at least one of them).728 static void NormalizedLoopNestChunked(void (*LoopBody)(Ty, void *), void *Arg,729 Ty BlockChunk, Ty NumBlocks, Ty BId,730 Ty ThreadChunk, Ty NumThreads, Ty TId,731 Ty NumIters,732 uint8_t OneIterationPerThread) {733 Ty KernelIteration = NumBlocks * BlockChunk;734 735 // Start index in the chunked space.736 Ty IV = BId * BlockChunk + TId;737 ASSERT(IV >= 0, "Bad index");738 739 // Cover the entire iteration space, assumptions in the caller might allow740 // to simplify this loop to a conditional.741 do {742 743 Ty BlockChunkLeft =744 BlockChunk >= TId * ThreadChunk ? BlockChunk - TId * ThreadChunk : 0;745 Ty ThreadChunkLeft =746 ThreadChunk <= BlockChunkLeft ? ThreadChunk : BlockChunkLeft;747 748 while (ThreadChunkLeft--) {749 750 // Given the blocking it's hard to keep track of what to execute.751 if (IV >= NumIters)752 return;753 754 // Execute the loop body.755 LoopBody(IV, Arg);756 757 if (OneIterationPerThread)758 return;759 760 ++IV;761 }762 763 IV += KernelIteration;764 765 } while (IV < NumIters);766 }767 768public:769 /// Worksharing `for`-loop.770 /// \param[in] Loc Description of source location771 /// \param[in] LoopBody Function which corresponds to loop body772 /// \param[in] Arg Pointer to struct which contains loop body args773 /// \param[in] NumIters Number of loop iterations774 /// \param[in] NumThreads Number of GPU threads775 /// \param[in] ThreadChunk Size of thread chunk776 /// \param[in] OneIterationPerThread If true/nonzero, each thread executes777 /// only one loop iteration or one thread chunk. This avoids an outer loop778 /// over all loop iterations/chunks.779 static void For(IdentTy *Loc, void (*LoopBody)(Ty, void *), void *Arg,780 Ty NumIters, Ty NumThreads, Ty ThreadChunk,781 uint8_t OneIterationPerThread) {782 ASSERT(NumIters >= 0, "Bad iteration count");783 ASSERT(ThreadChunk >= 0, "Bad thread count");784 785 // All threads need to participate but we don't know if we are in a786 // parallel at all or if the user might have used a `num_threads` clause787 // on the parallel and reduced the number compared to the block size.788 // Since nested parallels are possible too we need to get the thread id789 // from the `omp` getter and not the mapping directly.790 Ty TId = omp_get_thread_num();791 792 // There are no blocks involved here.793 Ty BlockChunk = 0;794 Ty NumBlocks = 1;795 Ty BId = 0;796 797 // If the thread chunk is not specified we pick a default now.798 if (ThreadChunk == 0)799 ThreadChunk = 1;800 801 // If we know we have more threads than iterations we can indicate that to802 // avoid an outer loop.803 if (OneIterationPerThread)804 ASSERT(NumThreads >= NumIters, "Broken assumption");805 806 if (ThreadChunk != 1)807 NormalizedLoopNestChunked(LoopBody, Arg, BlockChunk, NumBlocks, BId,808 ThreadChunk, NumThreads, TId, NumIters,809 OneIterationPerThread);810 else811 NormalizedLoopNestNoChunk(LoopBody, Arg, NumBlocks, BId, NumThreads, TId,812 NumIters, OneIterationPerThread);813 }814 815 /// Worksharing `distribute`-loop.816 /// \param[in] Loc Description of source location817 /// \param[in] LoopBody Function which corresponds to loop body818 /// \param[in] Arg Pointer to struct which contains loop body args819 /// \param[in] NumIters Number of loop iterations820 /// \param[in] BlockChunk Size of block chunk821 /// \param[in] OneIterationPerThread If true/nonzero, each thread executes822 /// only one loop iteration or one thread chunk. This avoids an outer loop823 /// over all loop iterations/chunks.824 static void Distribute(IdentTy *Loc, void (*LoopBody)(Ty, void *), void *Arg,825 Ty NumIters, Ty BlockChunk,826 uint8_t OneIterationPerThread) {827 ASSERT(icv::Level == 0, "Bad distribute");828 ASSERT(icv::ActiveLevel == 0, "Bad distribute");829 ASSERT(state::ParallelRegionFn == nullptr, "Bad distribute");830 ASSERT(state::ParallelTeamSize == 1, "Bad distribute");831 832 ASSERT(NumIters >= 0, "Bad iteration count");833 ASSERT(BlockChunk >= 0, "Bad block count");834 835 // There are no threads involved here.836 Ty ThreadChunk = 0;837 Ty NumThreads = 1;838 Ty TId = 0;839 840 // All teams need to participate.841 Ty NumBlocks = mapping::getNumberOfBlocksInKernel();842 Ty BId = mapping::getBlockIdInKernel();843 844 // If the block chunk is not specified we pick a default now.845 if (BlockChunk == 0)846 BlockChunk = NumThreads;847 848 // If we know we have more blocks than iterations we can indicate that to849 // avoid an outer loop.850 if (OneIterationPerThread)851 ASSERT(NumBlocks >= NumIters, "Broken assumption");852 853 if (BlockChunk != NumThreads)854 NormalizedLoopNestChunked(LoopBody, Arg, BlockChunk, NumBlocks, BId,855 ThreadChunk, NumThreads, TId, NumIters,856 OneIterationPerThread);857 else858 NormalizedLoopNestNoChunk(LoopBody, Arg, NumBlocks, BId, NumThreads, TId,859 NumIters, OneIterationPerThread);860 861 ASSERT(icv::Level == 0, "Bad distribute");862 ASSERT(icv::ActiveLevel == 0, "Bad distribute");863 ASSERT(state::ParallelRegionFn == nullptr, "Bad distribute");864 ASSERT(state::ParallelTeamSize == 1, "Bad distribute");865 }866 867 /// Worksharing `distribute parallel for`-loop.868 /// \param[in] Loc Description of source location869 /// \param[in] LoopBody Function which corresponds to loop body870 /// \param[in] Arg Pointer to struct which contains loop body args871 /// \param[in] NumIters Number of loop iterations872 /// \param[in] NumThreads Number of GPU threads873 /// \param[in] BlockChunk Size of block chunk874 /// \param[in] ThreadChunk Size of thread chunk875 /// \param[in] OneIterationPerThread If true/nonzero, each thread executes876 /// only one loop iteration or one thread chunk. This avoids an outer loop877 /// over all loop iterations/chunks.878 static void DistributeFor(IdentTy *Loc, void (*LoopBody)(Ty, void *),879 void *Arg, Ty NumIters, Ty NumThreads,880 Ty BlockChunk, Ty ThreadChunk,881 uint8_t OneIterationPerThread) {882 ASSERT(icv::Level == 1, "Bad distribute");883 ASSERT(icv::ActiveLevel == 1, "Bad distribute");884 ASSERT(state::ParallelRegionFn == nullptr, "Bad distribute");885 886 ASSERT(NumIters >= 0, "Bad iteration count");887 ASSERT(BlockChunk >= 0, "Bad block count");888 ASSERT(ThreadChunk >= 0, "Bad thread count");889 890 // All threads need to participate but the user might have used a891 // `num_threads` clause on the parallel and reduced the number compared to892 // the block size.893 Ty TId = mapping::getThreadIdInBlock();894 895 // All teams need to participate.896 Ty NumBlocks = mapping::getNumberOfBlocksInKernel();897 Ty BId = mapping::getBlockIdInKernel();898 899 // If the block chunk is not specified we pick a default now.900 if (BlockChunk == 0)901 BlockChunk = NumThreads;902 903 // If the thread chunk is not specified we pick a default now.904 if (ThreadChunk == 0)905 ThreadChunk = 1;906 907 // If we know we have more threads (across all blocks) than iterations we908 // can indicate that to avoid an outer loop.909 if (OneIterationPerThread)910 ASSERT(NumBlocks * NumThreads >= NumIters, "Broken assumption");911 912 if (BlockChunk != NumThreads || ThreadChunk != 1)913 NormalizedLoopNestChunked(LoopBody, Arg, BlockChunk, NumBlocks, BId,914 ThreadChunk, NumThreads, TId, NumIters,915 OneIterationPerThread);916 else917 NormalizedLoopNestNoChunk(LoopBody, Arg, NumBlocks, BId, NumThreads, TId,918 NumIters, OneIterationPerThread);919 920 ASSERT(icv::Level == 1, "Bad distribute");921 ASSERT(icv::ActiveLevel == 1, "Bad distribute");922 ASSERT(state::ParallelRegionFn == nullptr, "Bad distribute");923 }924};925 926} // namespace ompx927 928#define OMP_LOOP_ENTRY(BW, TY) \929 [[gnu::flatten, clang::always_inline]] void \930 __kmpc_distribute_for_static_loop##BW( \931 IdentTy *loc, void (*fn)(TY, void *), void *arg, TY num_iters, \932 TY num_threads, TY block_chunk, TY thread_chunk, \933 uint8_t one_iteration_per_thread) { \934 ompx::StaticLoopChunker<TY>::DistributeFor( \935 loc, fn, arg, num_iters, num_threads, block_chunk, thread_chunk, \936 one_iteration_per_thread); \937 } \938 [[gnu::flatten, clang::always_inline]] void \939 __kmpc_distribute_static_loop##BW(IdentTy *loc, void (*fn)(TY, void *), \940 void *arg, TY num_iters, TY block_chunk, \941 uint8_t one_iteration_per_thread) { \942 ompx::StaticLoopChunker<TY>::Distribute( \943 loc, fn, arg, num_iters, block_chunk, one_iteration_per_thread); \944 } \945 [[gnu::flatten, clang::always_inline]] void __kmpc_for_static_loop##BW( \946 IdentTy *loc, void (*fn)(TY, void *), void *arg, TY num_iters, \947 TY num_threads, TY thread_chunk, uint8_t one_iteration_per_thread) { \948 ompx::StaticLoopChunker<TY>::For(loc, fn, arg, num_iters, num_threads, \949 thread_chunk, one_iteration_per_thread); \950 }951 952extern "C" {953OMP_LOOP_ENTRY(_4, int32_t)954OMP_LOOP_ENTRY(_4u, uint32_t)955OMP_LOOP_ENTRY(_8, int64_t)956OMP_LOOP_ENTRY(_8u, uint64_t)957}958