13de4cd12SGeorge Rokos //===----RTLs/cuda/src/rtl.cpp - Target RTLs Implementation ------- C++ -*-===//
23de4cd12SGeorge Rokos //
357b08b09SChandler Carruth // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
457b08b09SChandler Carruth // See https://llvm.org/LICENSE.txt for license information.
557b08b09SChandler Carruth // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
63de4cd12SGeorge Rokos //
73de4cd12SGeorge Rokos //===----------------------------------------------------------------------===//
83de4cd12SGeorge Rokos //
93de4cd12SGeorge Rokos // RTL for CUDA machine
103de4cd12SGeorge Rokos //
113de4cd12SGeorge Rokos //===----------------------------------------------------------------------===//
123de4cd12SGeorge Rokos 
133de4cd12SGeorge Rokos #include <cassert>
143de4cd12SGeorge Rokos #include <cstddef>
153de4cd12SGeorge Rokos #include <cuda.h>
163de4cd12SGeorge Rokos #include <list>
17a5153dbcSJohannes Doerfert #include <memory>
18feed674dSShilei Tian #include <mutex>
193de4cd12SGeorge Rokos #include <string>
203de4cd12SGeorge Rokos #include <vector>
213de4cd12SGeorge Rokos 
22ae95ceebSJoseph Huber #include "Debug.h"
231546d319SGeorge Rokos #include "omptargetplugin.h"
243de4cd12SGeorge Rokos 
253de4cd12SGeorge Rokos #define TARGET_NAME CUDA
26ae95ceebSJoseph Huber #define DEBUG_PREFIX "Target " GETNAME(TARGET_NAME) " RTL"
2778769ec4SMichael Kruse 
280871d6d5SShilei Tian #include "MemoryManager.h"
290871d6d5SShilei Tian 
3078769ec4SMichael Kruse // Utility for retrieving and printing CUDA error string.
31ae95ceebSJoseph Huber #ifdef OMPTARGET_DEBUG
3278769ec4SMichael Kruse #define CUDA_ERR_STRING(err)                                                   \
3378769ec4SMichael Kruse   do {                                                                         \
34ae95ceebSJoseph Huber     if (getDebugLevel() > 0) {                                                 \
3503111e5eSYe Luo       const char *errStr = nullptr;                                            \
3603111e5eSYe Luo       CUresult errStr_status = cuGetErrorString(err, &errStr);                 \
3703111e5eSYe Luo       if (errStr_status == CUDA_ERROR_INVALID_VALUE)                           \
382ce16810SJoseph Huber         REPORT("Unrecognized CUDA error code: %d\n", err);                     \
3903111e5eSYe Luo       else if (errStr_status == CUDA_SUCCESS)                                  \
402ce16810SJoseph Huber         REPORT("CUDA error is: %s\n", errStr);                                 \
4103111e5eSYe Luo       else {                                                                   \
422ce16810SJoseph Huber         REPORT("Unresolved CUDA error code: %d\n", err);                       \
432ce16810SJoseph Huber         REPORT("Unsuccessful cuGetErrorString return status: %d\n",            \
4403111e5eSYe Luo                errStr_status);                                                 \
4503111e5eSYe Luo       }                                                                        \
462ce16810SJoseph Huber     } else {                                                                   \
472ce16810SJoseph Huber       const char *errStr = nullptr;                                            \
482ce16810SJoseph Huber       CUresult errStr_status = cuGetErrorString(err, &errStr);                 \
492ce16810SJoseph Huber       if (errStr_status == CUDA_SUCCESS)                                       \
502ce16810SJoseph Huber         REPORT("%s \n", errStr);                                               \
5178769ec4SMichael Kruse     }                                                                          \
5278769ec4SMichael Kruse   } while (false)
53b305d26bSSergey Dmitriev #else // OMPTARGET_DEBUG
542ce16810SJoseph Huber #define CUDA_ERR_STRING(err)                                                   \
552ce16810SJoseph Huber   do {                                                                         \
562ce16810SJoseph Huber     const char *errStr = nullptr;                                              \
572ce16810SJoseph Huber     CUresult errStr_status = cuGetErrorString(err, &errStr);                   \
582ce16810SJoseph Huber     if (errStr_status == CUDA_SUCCESS)                                         \
592ce16810SJoseph Huber       REPORT("%s \n", errStr);                                                 \
602ce16810SJoseph Huber   } while (false)
61b305d26bSSergey Dmitriev #endif // OMPTARGET_DEBUG
623de4cd12SGeorge Rokos 
63a81c68aeSShilei Tian #include "elf_common.h"
643de4cd12SGeorge Rokos 
653de4cd12SGeorge Rokos /// Keep entries table per device.
663de4cd12SGeorge Rokos struct FuncOrGblEntryTy {
673de4cd12SGeorge Rokos   __tgt_target_table Table;
683de4cd12SGeorge Rokos   std::vector<__tgt_offload_entry> Entries;
693de4cd12SGeorge Rokos };
703de4cd12SGeorge Rokos 
713de4cd12SGeorge Rokos enum ExecutionModeType {
72c13df8e5SGeorge Rokos   SPMD, // constructors, destructors,
73c13df8e5SGeorge Rokos   // combined constructs (`teams distribute parallel for [simd]`)
74c13df8e5SGeorge Rokos   GENERIC, // everything else
753de4cd12SGeorge Rokos   NONE
763de4cd12SGeorge Rokos };
773de4cd12SGeorge Rokos 
78dc34b1c9SRon Lieberman /// Use a single entity to encode a kernel and a set of flags.
793de4cd12SGeorge Rokos struct KernelTy {
803de4cd12SGeorge Rokos   CUfunction Func;
813de4cd12SGeorge Rokos 
823de4cd12SGeorge Rokos   // execution mode of kernel
833de4cd12SGeorge Rokos   // 0 - SPMD mode (without master warp)
843de4cd12SGeorge Rokos   // 1 - Generic mode (with master warp)
853de4cd12SGeorge Rokos   int8_t ExecutionMode;
863de4cd12SGeorge Rokos 
87aa27cfc1SJohannes Doerfert   /// Maximal number of threads per block for this kernel.
88aa27cfc1SJohannes Doerfert   int MaxThreadsPerBlock = 0;
89aa27cfc1SJohannes Doerfert 
903de4cd12SGeorge Rokos   KernelTy(CUfunction _Func, int8_t _ExecutionMode)
913de4cd12SGeorge Rokos       : Func(_Func), ExecutionMode(_ExecutionMode) {}
923de4cd12SGeorge Rokos };
933de4cd12SGeorge Rokos 
944c6a098aSKazuaki Ishizaki /// Device environment data
95feed674dSShilei Tian /// Manually sync with the deviceRTL side for now, move to a dedicated header
96feed674dSShilei Tian /// file later.
97e1c7a46dSGuansong Zhang struct omptarget_device_environmentTy {
98e1c7a46dSGuansong Zhang   int32_t debug_level;
99e1c7a46dSGuansong Zhang };
100e1c7a46dSGuansong Zhang 
101feed674dSShilei Tian namespace {
102feed674dSShilei Tian bool checkResult(CUresult Err, const char *ErrMsg) {
103feed674dSShilei Tian   if (Err == CUDA_SUCCESS)
104feed674dSShilei Tian     return true;
105feed674dSShilei Tian 
1062ce16810SJoseph Huber   REPORT("%s", ErrMsg);
107feed674dSShilei Tian   CUDA_ERR_STRING(Err);
108feed674dSShilei Tian   return false;
109feed674dSShilei Tian }
1104031bb98SShilei Tian 
111a014fbbcSShilei Tian int memcpyDtoD(const void *SrcPtr, void *DstPtr, int64_t Size,
112a014fbbcSShilei Tian                CUstream Stream) {
11395f0d1edSJon Chesterfield   CUresult Err =
11495f0d1edSJon Chesterfield       cuMemcpyDtoDAsync((CUdeviceptr)DstPtr, (CUdeviceptr)SrcPtr, Size, Stream);
115a014fbbcSShilei Tian 
116a014fbbcSShilei Tian   if (Err != CUDA_SUCCESS) {
1172ce16810SJoseph Huber     REPORT("Error when copying data from device to device. Pointers: src "
118a014fbbcSShilei Tian            "= " DPxMOD ", dst = " DPxMOD ", size = %" PRId64 "\n",
119a014fbbcSShilei Tian            DPxPTR(SrcPtr), DPxPTR(DstPtr), Size);
120a014fbbcSShilei Tian     CUDA_ERR_STRING(Err);
121a014fbbcSShilei Tian     return OFFLOAD_FAIL;
122a014fbbcSShilei Tian   }
123a014fbbcSShilei Tian 
124a014fbbcSShilei Tian   return OFFLOAD_SUCCESS;
125a014fbbcSShilei Tian }
126a014fbbcSShilei Tian 
1274031bb98SShilei Tian // Structure contains per-device data
1284031bb98SShilei Tian struct DeviceDataTy {
1295272d29eSJohannes Doerfert   /// List that contains all the kernels.
1305272d29eSJohannes Doerfert   std::list<KernelTy> KernelsList;
1315272d29eSJohannes Doerfert 
1324031bb98SShilei Tian   std::list<FuncOrGblEntryTy> FuncGblEntries;
1335272d29eSJohannes Doerfert 
1344031bb98SShilei Tian   CUcontext Context = nullptr;
1354031bb98SShilei Tian   // Device properties
1364031bb98SShilei Tian   int ThreadsPerBlock = 0;
1374031bb98SShilei Tian   int BlocksPerGrid = 0;
1384031bb98SShilei Tian   int WarpSize = 0;
1394031bb98SShilei Tian   // OpenMP properties
1404031bb98SShilei Tian   int NumTeams = 0;
1414031bb98SShilei Tian   int NumThreads = 0;
1424031bb98SShilei Tian };
143feed674dSShilei Tian 
144feed674dSShilei Tian class StreamManagerTy {
145feed674dSShilei Tian   int NumberOfDevices;
1464031bb98SShilei Tian   // The initial size of stream pool
1474031bb98SShilei Tian   int EnvNumInitialStreams;
148feed674dSShilei Tian   // Per-device stream mutex
149feed674dSShilei Tian   std::vector<std::unique_ptr<std::mutex>> StreamMtx;
150feed674dSShilei Tian   // Per-device stream Id indicates the next available stream in the pool
151feed674dSShilei Tian   std::vector<int> NextStreamId;
152feed674dSShilei Tian   // Per-device stream pool
153feed674dSShilei Tian   std::vector<std::vector<CUstream>> StreamPool;
1544031bb98SShilei Tian   // Reference to per-device data
1554031bb98SShilei Tian   std::vector<DeviceDataTy> &DeviceData;
156feed674dSShilei Tian 
157feed674dSShilei Tian   // If there is no CUstream left in the pool, we will resize the pool to
158feed674dSShilei Tian   // allocate more CUstream. This function should be called with device mutex,
159feed674dSShilei Tian   // and we do not resize to smaller one.
160feed674dSShilei Tian   void resizeStreamPool(const int DeviceId, const size_t NewSize) {
161feed674dSShilei Tian     std::vector<CUstream> &Pool = StreamPool[DeviceId];
162feed674dSShilei Tian     const size_t CurrentSize = Pool.size();
163feed674dSShilei Tian     assert(NewSize > CurrentSize && "new size is not larger than current size");
164feed674dSShilei Tian 
1654031bb98SShilei Tian     CUresult Err = cuCtxSetCurrent(DeviceData[DeviceId].Context);
1664031bb98SShilei Tian     if (!checkResult(Err, "Error returned from cuCtxSetCurrent\n")) {
167feed674dSShilei Tian       // We will return if cannot switch to the right context in case of
168feed674dSShilei Tian       // creating bunch of streams that are not corresponding to the right
169feed674dSShilei Tian       // device. The offloading will fail later because selected CUstream is
170feed674dSShilei Tian       // nullptr.
171feed674dSShilei Tian       return;
172feed674dSShilei Tian     }
173feed674dSShilei Tian 
1744031bb98SShilei Tian     Pool.resize(NewSize, nullptr);
1754031bb98SShilei Tian 
176feed674dSShilei Tian     for (size_t I = CurrentSize; I < NewSize; ++I) {
1774031bb98SShilei Tian       checkResult(cuStreamCreate(&Pool[I], CU_STREAM_NON_BLOCKING),
1784031bb98SShilei Tian                   "Error returned from cuStreamCreate\n");
179feed674dSShilei Tian     }
180feed674dSShilei Tian   }
181feed674dSShilei Tian 
182feed674dSShilei Tian public:
1834031bb98SShilei Tian   StreamManagerTy(const int NumberOfDevices,
1844031bb98SShilei Tian                   std::vector<DeviceDataTy> &DeviceData)
1854031bb98SShilei Tian       : NumberOfDevices(NumberOfDevices), EnvNumInitialStreams(32),
1864031bb98SShilei Tian         DeviceData(DeviceData) {
187feed674dSShilei Tian     StreamPool.resize(NumberOfDevices);
188feed674dSShilei Tian     NextStreamId.resize(NumberOfDevices);
189feed674dSShilei Tian     StreamMtx.resize(NumberOfDevices);
190feed674dSShilei Tian 
1914031bb98SShilei Tian     if (const char *EnvStr = getenv("LIBOMPTARGET_NUM_INITIAL_STREAMS"))
1924031bb98SShilei Tian       EnvNumInitialStreams = std::stoi(EnvStr);
193feed674dSShilei Tian 
194feed674dSShilei Tian     // Initialize the next stream id
195feed674dSShilei Tian     std::fill(NextStreamId.begin(), NextStreamId.end(), 0);
196feed674dSShilei Tian 
197feed674dSShilei Tian     // Initialize stream mutex
198feed674dSShilei Tian     for (std::unique_ptr<std::mutex> &Ptr : StreamMtx)
199feed674dSShilei Tian       Ptr = std::make_unique<std::mutex>();
200feed674dSShilei Tian   }
201feed674dSShilei Tian 
202feed674dSShilei Tian   ~StreamManagerTy() {
203feed674dSShilei Tian     // Destroy streams
204feed674dSShilei Tian     for (int I = 0; I < NumberOfDevices; ++I) {
2054031bb98SShilei Tian       checkResult(cuCtxSetCurrent(DeviceData[I].Context),
2064031bb98SShilei Tian                   "Error returned from cuCtxSetCurrent\n");
207feed674dSShilei Tian 
208feed674dSShilei Tian       for (CUstream &S : StreamPool[I]) {
2094031bb98SShilei Tian         if (S)
21095f0d1edSJon Chesterfield           checkResult(cuStreamDestroy(S),
21195f0d1edSJon Chesterfield                       "Error returned from cuStreamDestroy\n");
212feed674dSShilei Tian       }
213feed674dSShilei Tian     }
214feed674dSShilei Tian   }
215feed674dSShilei Tian 
216feed674dSShilei Tian   // Get a CUstream from pool. Per-device next stream id always points to the
217feed674dSShilei Tian   // next available CUstream. That means, CUstreams [0, id-1] have been
218feed674dSShilei Tian   // assigned, and [id,] are still available. If there is no CUstream left, we
219feed674dSShilei Tian   // will ask more CUstreams from CUDA RT. Each time a CUstream is assigned,
220feed674dSShilei Tian   // the id will increase one.
221feed674dSShilei Tian   // xxxxxs+++++++++
222feed674dSShilei Tian   //      ^
223feed674dSShilei Tian   //      id
224feed674dSShilei Tian   // After assignment, the pool becomes the following and s is assigned.
225feed674dSShilei Tian   // xxxxxs+++++++++
226feed674dSShilei Tian   //       ^
227feed674dSShilei Tian   //       id
228feed674dSShilei Tian   CUstream getStream(const int DeviceId) {
229feed674dSShilei Tian     const std::lock_guard<std::mutex> Lock(*StreamMtx[DeviceId]);
230feed674dSShilei Tian     int &Id = NextStreamId[DeviceId];
231feed674dSShilei Tian     // No CUstream left in the pool, we need to request from CUDA RT
232ad94fce8SJohannes Doerfert     if (Id == static_cast<int>(StreamPool[DeviceId].size())) {
233feed674dSShilei Tian       // By default we double the stream pool every time
234feed674dSShilei Tian       resizeStreamPool(DeviceId, Id * 2);
235feed674dSShilei Tian     }
236feed674dSShilei Tian     return StreamPool[DeviceId][Id++];
237feed674dSShilei Tian   }
238feed674dSShilei Tian 
239feed674dSShilei Tian   // Return a CUstream back to pool. As mentioned above, per-device next
240feed674dSShilei Tian   // stream is always points to the next available CUstream, so when we return
241feed674dSShilei Tian   // a CUstream, we need to first decrease the id, and then copy the CUstream
242feed674dSShilei Tian   // back.
243feed674dSShilei Tian   // It is worth noting that, the order of streams return might be different
244feed674dSShilei Tian   // from that they're assigned, that saying, at some point, there might be
245feed674dSShilei Tian   // two identical CUstreams.
246feed674dSShilei Tian   // xxax+a+++++
247feed674dSShilei Tian   //     ^
248feed674dSShilei Tian   //     id
249feed674dSShilei Tian   // However, it doesn't matter, because they're always on the two sides of
250feed674dSShilei Tian   // id. The left one will in the end be overwritten by another CUstream.
251feed674dSShilei Tian   // Therefore, after several execution, the order of pool might be different
252feed674dSShilei Tian   // from its initial state.
253feed674dSShilei Tian   void returnStream(const int DeviceId, CUstream Stream) {
254feed674dSShilei Tian     const std::lock_guard<std::mutex> Lock(*StreamMtx[DeviceId]);
255feed674dSShilei Tian     int &Id = NextStreamId[DeviceId];
256feed674dSShilei Tian     assert(Id > 0 && "Wrong stream ID");
257feed674dSShilei Tian     StreamPool[DeviceId][--Id] = Stream;
258feed674dSShilei Tian   }
259feed674dSShilei Tian 
2604031bb98SShilei Tian   bool initializeDeviceStreamPool(const int DeviceId) {
2614031bb98SShilei Tian     assert(StreamPool[DeviceId].empty() && "stream pool has been initialized");
2624031bb98SShilei Tian 
2634031bb98SShilei Tian     resizeStreamPool(DeviceId, EnvNumInitialStreams);
2644031bb98SShilei Tian 
2654031bb98SShilei Tian     // Check the size of stream pool
266ad94fce8SJohannes Doerfert     if (static_cast<int>(StreamPool[DeviceId].size()) != EnvNumInitialStreams)
2674031bb98SShilei Tian       return false;
2684031bb98SShilei Tian 
2694031bb98SShilei Tian     // Check whether each stream is valid
2704031bb98SShilei Tian     for (CUstream &S : StreamPool[DeviceId])
2714031bb98SShilei Tian       if (!S)
2724031bb98SShilei Tian         return false;
2734031bb98SShilei Tian 
2744031bb98SShilei Tian     return true;
275feed674dSShilei Tian   }
276feed674dSShilei Tian };
277feed674dSShilei Tian 
2784031bb98SShilei Tian class DeviceRTLTy {
2793de4cd12SGeorge Rokos   int NumberOfDevices;
2804031bb98SShilei Tian   // OpenMP environment properties
2813de4cd12SGeorge Rokos   int EnvNumTeams;
2823de4cd12SGeorge Rokos   int EnvTeamLimit;
2834031bb98SShilei Tian   // OpenMP requires flags
284c5fe030cSGheorghe-Teodor Bercea   int64_t RequiresFlags;
285c5fe030cSGheorghe-Teodor Bercea 
2864031bb98SShilei Tian   static constexpr const int HardTeamLimit = 1U << 16U; // 64k
2874031bb98SShilei Tian   static constexpr const int HardThreadLimit = 1024;
2884031bb98SShilei Tian   static constexpr const int DefaultNumTeams = 128;
2894031bb98SShilei Tian   static constexpr const int DefaultNumThreads = 128;
2903de4cd12SGeorge Rokos 
2914031bb98SShilei Tian   std::unique_ptr<StreamManagerTy> StreamManager;
2924031bb98SShilei Tian   std::vector<DeviceDataTy> DeviceData;
2934031bb98SShilei Tian   std::vector<CUmodule> Modules;
294feed674dSShilei Tian 
2950871d6d5SShilei Tian   /// A class responsible for interacting with device native runtime library to
2960871d6d5SShilei Tian   /// allocate and free memory.
2970871d6d5SShilei Tian   class CUDADeviceAllocatorTy : public DeviceAllocatorTy {
2980871d6d5SShilei Tian     const int DeviceId;
2990871d6d5SShilei Tian     const std::vector<DeviceDataTy> &DeviceData;
3000871d6d5SShilei Tian 
3010871d6d5SShilei Tian   public:
3020871d6d5SShilei Tian     CUDADeviceAllocatorTy(int DeviceId, std::vector<DeviceDataTy> &DeviceData)
3030871d6d5SShilei Tian         : DeviceId(DeviceId), DeviceData(DeviceData) {}
3040871d6d5SShilei Tian 
3050871d6d5SShilei Tian     void *allocate(size_t Size, void *) override {
3060871d6d5SShilei Tian       if (Size == 0)
3070871d6d5SShilei Tian         return nullptr;
3080871d6d5SShilei Tian 
3090871d6d5SShilei Tian       CUresult Err = cuCtxSetCurrent(DeviceData[DeviceId].Context);
3100871d6d5SShilei Tian       if (!checkResult(Err, "Error returned from cuCtxSetCurrent\n"))
3110871d6d5SShilei Tian         return nullptr;
3120871d6d5SShilei Tian 
3130871d6d5SShilei Tian       CUdeviceptr DevicePtr;
31495f0d1edSJon Chesterfield       Err = cuMemAlloc(&DevicePtr, Size);
31595f0d1edSJon Chesterfield       if (!checkResult(Err, "Error returned from cuMemAlloc\n"))
3160871d6d5SShilei Tian         return nullptr;
3170871d6d5SShilei Tian 
3180871d6d5SShilei Tian       return (void *)DevicePtr;
3190871d6d5SShilei Tian     }
3200871d6d5SShilei Tian 
3210871d6d5SShilei Tian     int free(void *TgtPtr) override {
3220871d6d5SShilei Tian       CUresult Err = cuCtxSetCurrent(DeviceData[DeviceId].Context);
3230871d6d5SShilei Tian       if (!checkResult(Err, "Error returned from cuCtxSetCurrent\n"))
3240871d6d5SShilei Tian         return OFFLOAD_FAIL;
3250871d6d5SShilei Tian 
32695f0d1edSJon Chesterfield       Err = cuMemFree((CUdeviceptr)TgtPtr);
32795f0d1edSJon Chesterfield       if (!checkResult(Err, "Error returned from cuMemFree\n"))
3280871d6d5SShilei Tian         return OFFLOAD_FAIL;
3290871d6d5SShilei Tian 
3300871d6d5SShilei Tian       return OFFLOAD_SUCCESS;
3310871d6d5SShilei Tian     }
3320871d6d5SShilei Tian   };
3330871d6d5SShilei Tian 
3340871d6d5SShilei Tian   /// A vector of device allocators
3350871d6d5SShilei Tian   std::vector<CUDADeviceAllocatorTy> DeviceAllocators;
3360871d6d5SShilei Tian 
3370871d6d5SShilei Tian   /// A vector of memory managers. Since the memory manager is non-copyable and
3380871d6d5SShilei Tian   // non-removable, we wrap them into std::unique_ptr.
3390871d6d5SShilei Tian   std::vector<std::unique_ptr<MemoryManagerTy>> MemoryManagers;
3400871d6d5SShilei Tian 
3410871d6d5SShilei Tian   /// Whether use memory manager
3420871d6d5SShilei Tian   bool UseMemoryManager = true;
3430871d6d5SShilei Tian 
3443de4cd12SGeorge Rokos   // Record entry point associated with device
3454031bb98SShilei Tian   void addOffloadEntry(const int DeviceId, const __tgt_offload_entry entry) {
3464031bb98SShilei Tian     FuncOrGblEntryTy &E = DeviceData[DeviceId].FuncGblEntries.back();
3473de4cd12SGeorge Rokos     E.Entries.push_back(entry);
3483de4cd12SGeorge Rokos   }
3493de4cd12SGeorge Rokos 
350ae209397SJoseph Huber   // Return a pointer to the entry associated with the pointer
351ae209397SJoseph Huber   const __tgt_offload_entry *getOffloadEntry(const int DeviceId,
352ae209397SJoseph Huber                                              const void *Addr) const {
3534031bb98SShilei Tian     for (const __tgt_offload_entry &Itr :
3544031bb98SShilei Tian          DeviceData[DeviceId].FuncGblEntries.back().Entries)
3554031bb98SShilei Tian       if (Itr.addr == Addr)
356ae209397SJoseph Huber         return &Itr;
3573de4cd12SGeorge Rokos 
358ae209397SJoseph Huber     return nullptr;
3593de4cd12SGeorge Rokos   }
3603de4cd12SGeorge Rokos 
3613de4cd12SGeorge Rokos   // Return the pointer to the target entries table
3624031bb98SShilei Tian   __tgt_target_table *getOffloadEntriesTable(const int DeviceId) {
3634031bb98SShilei Tian     FuncOrGblEntryTy &E = DeviceData[DeviceId].FuncGblEntries.back();
3643de4cd12SGeorge Rokos 
3654031bb98SShilei Tian     if (E.Entries.empty())
3664031bb98SShilei Tian       return nullptr;
3673de4cd12SGeorge Rokos 
3683de4cd12SGeorge Rokos     // Update table info according to the entries and return the pointer
3694031bb98SShilei Tian     E.Table.EntriesBegin = E.Entries.data();
3704031bb98SShilei Tian     E.Table.EntriesEnd = E.Entries.data() + E.Entries.size();
3713de4cd12SGeorge Rokos 
3723de4cd12SGeorge Rokos     return &E.Table;
3733de4cd12SGeorge Rokos   }
3743de4cd12SGeorge Rokos 
3753de4cd12SGeorge Rokos   // Clear entries table for a device
3764031bb98SShilei Tian   void clearOffloadEntriesTable(const int DeviceId) {
3774031bb98SShilei Tian     DeviceData[DeviceId].FuncGblEntries.emplace_back();
3784031bb98SShilei Tian     FuncOrGblEntryTy &E = DeviceData[DeviceId].FuncGblEntries.back();
3793de4cd12SGeorge Rokos     E.Entries.clear();
3804031bb98SShilei Tian     E.Table.EntriesBegin = E.Table.EntriesEnd = nullptr;
3813de4cd12SGeorge Rokos   }
3823de4cd12SGeorge Rokos 
3835449fbb5SJohannes Doerfert   CUstream getStream(const int DeviceId, __tgt_async_info *AsyncInfo) const {
3845449fbb5SJohannes Doerfert     assert(AsyncInfo && "AsyncInfo is nullptr");
3854031bb98SShilei Tian 
3865449fbb5SJohannes Doerfert     if (!AsyncInfo->Queue)
3875449fbb5SJohannes Doerfert       AsyncInfo->Queue = StreamManager->getStream(DeviceId);
3884031bb98SShilei Tian 
3895449fbb5SJohannes Doerfert     return reinterpret_cast<CUstream>(AsyncInfo->Queue);
390b305d26bSSergey Dmitriev   }
3914031bb98SShilei Tian 
3924031bb98SShilei Tian public:
3934031bb98SShilei Tian   // This class should not be copied
3944031bb98SShilei Tian   DeviceRTLTy(const DeviceRTLTy &) = delete;
3954031bb98SShilei Tian   DeviceRTLTy(DeviceRTLTy &&) = delete;
3964031bb98SShilei Tian 
3974031bb98SShilei Tian   DeviceRTLTy()
3984031bb98SShilei Tian       : NumberOfDevices(0), EnvNumTeams(-1), EnvTeamLimit(-1),
3994031bb98SShilei Tian         RequiresFlags(OMP_REQ_UNDEFINED) {
400b305d26bSSergey Dmitriev 
4013de4cd12SGeorge Rokos     DP("Start initializing CUDA\n");
4023de4cd12SGeorge Rokos 
4034031bb98SShilei Tian     CUresult Err = cuInit(0);
4047baff00eSJon Chesterfield     if (Err == CUDA_ERROR_INVALID_HANDLE) {
4057baff00eSJon Chesterfield       // Can't call cuGetErrorString if dlsym failed
4067baff00eSJon Chesterfield       DP("Failed to load CUDA shared library\n");
4077baff00eSJon Chesterfield       return;
4087baff00eSJon Chesterfield     }
4094031bb98SShilei Tian     if (!checkResult(Err, "Error returned from cuInit\n")) {
4103de4cd12SGeorge Rokos       return;
4113de4cd12SGeorge Rokos     }
4123de4cd12SGeorge Rokos 
4134031bb98SShilei Tian     Err = cuDeviceGetCount(&NumberOfDevices);
4144031bb98SShilei Tian     if (!checkResult(Err, "Error returned from cuDeviceGetCount\n"))
4153de4cd12SGeorge Rokos       return;
4163de4cd12SGeorge Rokos 
4173de4cd12SGeorge Rokos     if (NumberOfDevices == 0) {
4183de4cd12SGeorge Rokos       DP("There are no devices supporting CUDA.\n");
4193de4cd12SGeorge Rokos       return;
4203de4cd12SGeorge Rokos     }
4213de4cd12SGeorge Rokos 
4224031bb98SShilei Tian     DeviceData.resize(NumberOfDevices);
4233de4cd12SGeorge Rokos 
4243de4cd12SGeorge Rokos     // Get environment variables regarding teams
4254031bb98SShilei Tian     if (const char *EnvStr = getenv("OMP_TEAM_LIMIT")) {
4263de4cd12SGeorge Rokos       // OMP_TEAM_LIMIT has been set
4274031bb98SShilei Tian       EnvTeamLimit = std::stoi(EnvStr);
4283de4cd12SGeorge Rokos       DP("Parsed OMP_TEAM_LIMIT=%d\n", EnvTeamLimit);
4293de4cd12SGeorge Rokos     }
4304031bb98SShilei Tian     if (const char *EnvStr = getenv("OMP_NUM_TEAMS")) {
4313de4cd12SGeorge Rokos       // OMP_NUM_TEAMS has been set
4324031bb98SShilei Tian       EnvNumTeams = std::stoi(EnvStr);
4333de4cd12SGeorge Rokos       DP("Parsed OMP_NUM_TEAMS=%d\n", EnvNumTeams);
4343de4cd12SGeorge Rokos     }
435c5fe030cSGheorghe-Teodor Bercea 
436feed674dSShilei Tian     StreamManager =
4374031bb98SShilei Tian         std::make_unique<StreamManagerTy>(NumberOfDevices, DeviceData);
4380871d6d5SShilei Tian 
4390871d6d5SShilei Tian     for (int I = 0; I < NumberOfDevices; ++I)
4400871d6d5SShilei Tian       DeviceAllocators.emplace_back(I, DeviceData);
4410871d6d5SShilei Tian 
4420871d6d5SShilei Tian     // Get the size threshold from environment variable
4430871d6d5SShilei Tian     std::pair<size_t, bool> Res = MemoryManagerTy::getSizeThresholdFromEnv();
4440871d6d5SShilei Tian     UseMemoryManager = Res.second;
4450871d6d5SShilei Tian     size_t MemoryManagerThreshold = Res.first;
4460871d6d5SShilei Tian 
4470871d6d5SShilei Tian     if (UseMemoryManager)
4480871d6d5SShilei Tian       for (int I = 0; I < NumberOfDevices; ++I)
4490871d6d5SShilei Tian         MemoryManagers.emplace_back(std::make_unique<MemoryManagerTy>(
4500871d6d5SShilei Tian             DeviceAllocators[I], MemoryManagerThreshold));
4513de4cd12SGeorge Rokos   }
4523de4cd12SGeorge Rokos 
4534031bb98SShilei Tian   ~DeviceRTLTy() {
4540871d6d5SShilei Tian     // We first destruct memory managers in case that its dependent data are
4550871d6d5SShilei Tian     // destroyed before it.
4560871d6d5SShilei Tian     for (auto &M : MemoryManagers)
4570871d6d5SShilei Tian       M.release();
4580871d6d5SShilei Tian 
459feed674dSShilei Tian     StreamManager = nullptr;
460feed674dSShilei Tian 
4614031bb98SShilei Tian     for (CUmodule &M : Modules)
4624031bb98SShilei Tian       // Close module
4634031bb98SShilei Tian       if (M)
4644031bb98SShilei Tian         checkResult(cuModuleUnload(M), "Error returned from cuModuleUnload\n");
4654031bb98SShilei Tian 
4664031bb98SShilei Tian     for (DeviceDataTy &D : DeviceData) {
4674031bb98SShilei Tian       // Destroy context
468c5348aecSYe Luo       if (D.Context) {
469c5348aecSYe Luo         checkResult(cuCtxSetCurrent(D.Context),
470c5348aecSYe Luo                     "Error returned from cuCtxSetCurrent\n");
471c5348aecSYe Luo         CUdevice Device;
472c5348aecSYe Luo         checkResult(cuCtxGetDevice(&Device),
473c5348aecSYe Luo                     "Error returned from cuCtxGetDevice\n");
47495f0d1edSJon Chesterfield         checkResult(cuDevicePrimaryCtxRelease(Device),
47595f0d1edSJon Chesterfield                     "Error returned from cuDevicePrimaryCtxRelease\n");
476c5348aecSYe Luo       }
4773de4cd12SGeorge Rokos     }
4783de4cd12SGeorge Rokos   }
4793de4cd12SGeorge Rokos 
4804031bb98SShilei Tian   // Check whether a given DeviceId is valid
4814031bb98SShilei Tian   bool isValidDeviceId(const int DeviceId) const {
4824031bb98SShilei Tian     return DeviceId >= 0 && DeviceId < NumberOfDevices;
48332ed2927SShilei Tian   }
48432ed2927SShilei Tian 
485cb038927SShilei Tian   int getNumOfDevices() const { return NumberOfDevices; }
4864031bb98SShilei Tian 
4874031bb98SShilei Tian   void setRequiresFlag(const int64_t Flags) { this->RequiresFlags = Flags; }
4884031bb98SShilei Tian 
4894031bb98SShilei Tian   int initDevice(const int DeviceId) {
4904031bb98SShilei Tian     CUdevice Device;
4914031bb98SShilei Tian 
4924031bb98SShilei Tian     DP("Getting device %d\n", DeviceId);
4934031bb98SShilei Tian     CUresult Err = cuDeviceGet(&Device, DeviceId);
4944031bb98SShilei Tian     if (!checkResult(Err, "Error returned from cuDeviceGet\n"))
49532ed2927SShilei Tian       return OFFLOAD_FAIL;
4963de4cd12SGeorge Rokos 
497c5348aecSYe Luo     // Query the current flags of the primary context and set its flags if
498c5348aecSYe Luo     // it is inactive
499c5348aecSYe Luo     unsigned int FormerPrimaryCtxFlags = 0;
500c5348aecSYe Luo     int FormerPrimaryCtxIsActive = 0;
501c5348aecSYe Luo     Err = cuDevicePrimaryCtxGetState(Device, &FormerPrimaryCtxFlags,
502c5348aecSYe Luo                                      &FormerPrimaryCtxIsActive);
503c5348aecSYe Luo     if (!checkResult(Err, "Error returned from cuDevicePrimaryCtxGetState\n"))
504c5348aecSYe Luo       return OFFLOAD_FAIL;
505c5348aecSYe Luo 
506c5348aecSYe Luo     if (FormerPrimaryCtxIsActive) {
507c5348aecSYe Luo       DP("The primary context is active, no change to its flags\n");
508c5348aecSYe Luo       if ((FormerPrimaryCtxFlags & CU_CTX_SCHED_MASK) !=
509c5348aecSYe Luo           CU_CTX_SCHED_BLOCKING_SYNC)
510c5348aecSYe Luo         DP("Warning the current flags are not CU_CTX_SCHED_BLOCKING_SYNC\n");
511c5348aecSYe Luo     } else {
512c5348aecSYe Luo       DP("The primary context is inactive, set its flags to "
513c5348aecSYe Luo          "CU_CTX_SCHED_BLOCKING_SYNC\n");
51495f0d1edSJon Chesterfield       Err = cuDevicePrimaryCtxSetFlags(Device, CU_CTX_SCHED_BLOCKING_SYNC);
51595f0d1edSJon Chesterfield       if (!checkResult(Err, "Error returned from cuDevicePrimaryCtxSetFlags\n"))
516c5348aecSYe Luo         return OFFLOAD_FAIL;
517c5348aecSYe Luo     }
518c5348aecSYe Luo 
519c5348aecSYe Luo     // Retain the per device primary context and save it to use whenever this
520c5348aecSYe Luo     // device is selected.
521c5348aecSYe Luo     Err = cuDevicePrimaryCtxRetain(&DeviceData[DeviceId].Context, Device);
522c5348aecSYe Luo     if (!checkResult(Err, "Error returned from cuDevicePrimaryCtxRetain\n"))
5233de4cd12SGeorge Rokos       return OFFLOAD_FAIL;
5243de4cd12SGeorge Rokos 
5254031bb98SShilei Tian     Err = cuCtxSetCurrent(DeviceData[DeviceId].Context);
5264031bb98SShilei Tian     if (!checkResult(Err, "Error returned from cuCtxSetCurrent\n"))
5274031bb98SShilei Tian       return OFFLOAD_FAIL;
528a5153dbcSJohannes Doerfert 
529feed674dSShilei Tian     // Initialize stream pool
5304031bb98SShilei Tian     if (!StreamManager->initializeDeviceStreamPool(DeviceId))
5314031bb98SShilei Tian       return OFFLOAD_FAIL;
532a5153dbcSJohannes Doerfert 
533bb51d398SJonas Hahnfeld     // Query attributes to determine number of threads/block and blocks/grid.
5344031bb98SShilei Tian     int MaxGridDimX;
5354031bb98SShilei Tian     Err = cuDeviceGetAttribute(&MaxGridDimX, CU_DEVICE_ATTRIBUTE_MAX_GRID_DIM_X,
5364031bb98SShilei Tian                                Device);
5374031bb98SShilei Tian     if (Err != CUDA_SUCCESS) {
5384031bb98SShilei Tian       DP("Error getting max grid dimension, use default value %d\n",
5394031bb98SShilei Tian          DeviceRTLTy::DefaultNumTeams);
5404031bb98SShilei Tian       DeviceData[DeviceId].BlocksPerGrid = DeviceRTLTy::DefaultNumTeams;
5414031bb98SShilei Tian     } else if (MaxGridDimX <= DeviceRTLTy::HardTeamLimit) {
5424031bb98SShilei Tian       DP("Using %d CUDA blocks per grid\n", MaxGridDimX);
5434031bb98SShilei Tian       DeviceData[DeviceId].BlocksPerGrid = MaxGridDimX;
5443de4cd12SGeorge Rokos     } else {
5453de4cd12SGeorge Rokos       DP("Max CUDA blocks per grid %d exceeds the hard team limit %d, capping "
546bb51d398SJonas Hahnfeld          "at the hard limit\n",
5474031bb98SShilei Tian          MaxGridDimX, DeviceRTLTy::HardTeamLimit);
5484031bb98SShilei Tian       DeviceData[DeviceId].BlocksPerGrid = DeviceRTLTy::HardTeamLimit;
5493de4cd12SGeorge Rokos     }
5503de4cd12SGeorge Rokos 
551bb51d398SJonas Hahnfeld     // We are only exploiting threads along the x axis.
5524031bb98SShilei Tian     int MaxBlockDimX;
5534031bb98SShilei Tian     Err = cuDeviceGetAttribute(&MaxBlockDimX,
5544031bb98SShilei Tian                                CU_DEVICE_ATTRIBUTE_MAX_BLOCK_DIM_X, Device);
5554031bb98SShilei Tian     if (Err != CUDA_SUCCESS) {
5564031bb98SShilei Tian       DP("Error getting max block dimension, use default value %d\n",
5574031bb98SShilei Tian          DeviceRTLTy::DefaultNumThreads);
5584031bb98SShilei Tian       DeviceData[DeviceId].ThreadsPerBlock = DeviceRTLTy::DefaultNumThreads;
5594031bb98SShilei Tian     } else if (MaxBlockDimX <= DeviceRTLTy::HardThreadLimit) {
5604031bb98SShilei Tian       DP("Using %d CUDA threads per block\n", MaxBlockDimX);
5614031bb98SShilei Tian       DeviceData[DeviceId].ThreadsPerBlock = MaxBlockDimX;
5623de4cd12SGeorge Rokos     } else {
5634031bb98SShilei Tian       DP("Max CUDA threads per block %d exceeds the hard thread limit %d, "
5644031bb98SShilei Tian          "capping at the hard limit\n",
5654031bb98SShilei Tian          MaxBlockDimX, DeviceRTLTy::HardThreadLimit);
5664031bb98SShilei Tian       DeviceData[DeviceId].ThreadsPerBlock = DeviceRTLTy::HardThreadLimit;
5673de4cd12SGeorge Rokos     }
5683de4cd12SGeorge Rokos 
5694031bb98SShilei Tian     // Get and set warp size
5704031bb98SShilei Tian     int WarpSize;
5714031bb98SShilei Tian     Err =
5724031bb98SShilei Tian         cuDeviceGetAttribute(&WarpSize, CU_DEVICE_ATTRIBUTE_WARP_SIZE, Device);
5734031bb98SShilei Tian     if (Err != CUDA_SUCCESS) {
5744031bb98SShilei Tian       DP("Error getting warp size, assume default value 32\n");
5754031bb98SShilei Tian       DeviceData[DeviceId].WarpSize = 32;
576bb51d398SJonas Hahnfeld     } else {
5774031bb98SShilei Tian       DP("Using warp size %d\n", WarpSize);
5784031bb98SShilei Tian       DeviceData[DeviceId].WarpSize = WarpSize;
5793de4cd12SGeorge Rokos     }
5803de4cd12SGeorge Rokos 
5813de4cd12SGeorge Rokos     // Adjust teams to the env variables
5824031bb98SShilei Tian     if (EnvTeamLimit > 0 && DeviceData[DeviceId].BlocksPerGrid > EnvTeamLimit) {
5833de4cd12SGeorge Rokos       DP("Capping max CUDA blocks per grid to OMP_TEAM_LIMIT=%d\n",
5844031bb98SShilei Tian          EnvTeamLimit);
5854031bb98SShilei Tian       DeviceData[DeviceId].BlocksPerGrid = EnvTeamLimit;
5863de4cd12SGeorge Rokos     }
5873de4cd12SGeorge Rokos 
588119a9ea1SJoseph Huber     INFO(OMP_INFOTYPE_PLUGIN_KERNEL, DeviceId,
589ae209397SJoseph Huber          "Device supports up to %d CUDA blocks and %d threads with a "
590ae209397SJoseph Huber          "warp size of %d\n",
591ae209397SJoseph Huber          DeviceData[DeviceId].BlocksPerGrid,
592ae209397SJoseph Huber          DeviceData[DeviceId].ThreadsPerBlock, DeviceData[DeviceId].WarpSize);
5933de4cd12SGeorge Rokos 
5943de4cd12SGeorge Rokos     // Set default number of teams
5954031bb98SShilei Tian     if (EnvNumTeams > 0) {
5963de4cd12SGeorge Rokos       DP("Default number of teams set according to environment %d\n",
5974031bb98SShilei Tian          EnvNumTeams);
5984031bb98SShilei Tian       DeviceData[DeviceId].NumTeams = EnvNumTeams;
5993de4cd12SGeorge Rokos     } else {
6004031bb98SShilei Tian       DeviceData[DeviceId].NumTeams = DeviceRTLTy::DefaultNumTeams;
6013de4cd12SGeorge Rokos       DP("Default number of teams set according to library's default %d\n",
6024031bb98SShilei Tian          DeviceRTLTy::DefaultNumTeams);
6033de4cd12SGeorge Rokos     }
6044031bb98SShilei Tian 
6054031bb98SShilei Tian     if (DeviceData[DeviceId].NumTeams > DeviceData[DeviceId].BlocksPerGrid) {
6063de4cd12SGeorge Rokos       DP("Default number of teams exceeds device limit, capping at %d\n",
6074031bb98SShilei Tian          DeviceData[DeviceId].BlocksPerGrid);
6084031bb98SShilei Tian       DeviceData[DeviceId].NumTeams = DeviceData[DeviceId].BlocksPerGrid;
6093de4cd12SGeorge Rokos     }
6103de4cd12SGeorge Rokos 
6113de4cd12SGeorge Rokos     // Set default number of threads
6124031bb98SShilei Tian     DeviceData[DeviceId].NumThreads = DeviceRTLTy::DefaultNumThreads;
6133de4cd12SGeorge Rokos     DP("Default number of threads set according to library's default %d\n",
6144031bb98SShilei Tian        DeviceRTLTy::DefaultNumThreads);
6154031bb98SShilei Tian     if (DeviceData[DeviceId].NumThreads >
6164031bb98SShilei Tian         DeviceData[DeviceId].ThreadsPerBlock) {
6173de4cd12SGeorge Rokos       DP("Default number of threads exceeds device limit, capping at %d\n",
6184031bb98SShilei Tian          DeviceData[DeviceId].ThreadsPerBlock);
6194031bb98SShilei Tian       DeviceData[DeviceId].NumTeams = DeviceData[DeviceId].ThreadsPerBlock;
6203de4cd12SGeorge Rokos     }
6213de4cd12SGeorge Rokos 
6223de4cd12SGeorge Rokos     return OFFLOAD_SUCCESS;
6233de4cd12SGeorge Rokos   }
6243de4cd12SGeorge Rokos 
6254031bb98SShilei Tian   __tgt_target_table *loadBinary(const int DeviceId,
6264031bb98SShilei Tian                                  const __tgt_device_image *Image) {
6274031bb98SShilei Tian     // Set the context we are using
6284031bb98SShilei Tian     CUresult Err = cuCtxSetCurrent(DeviceData[DeviceId].Context);
6294031bb98SShilei Tian     if (!checkResult(Err, "Error returned from cuCtxSetCurrent\n"))
6304031bb98SShilei Tian       return nullptr;
6313de4cd12SGeorge Rokos 
6323de4cd12SGeorge Rokos     // Clear the offload table as we are going to create a new one.
6334031bb98SShilei Tian     clearOffloadEntriesTable(DeviceId);
6343de4cd12SGeorge Rokos 
6353de4cd12SGeorge Rokos     // Create the module and extract the function pointers.
6364031bb98SShilei Tian     CUmodule Module;
6374031bb98SShilei Tian     DP("Load data from image " DPxMOD "\n", DPxPTR(Image->ImageStart));
6384031bb98SShilei Tian     Err = cuModuleLoadDataEx(&Module, Image->ImageStart, 0, nullptr, nullptr);
6394031bb98SShilei Tian     if (!checkResult(Err, "Error returned from cuModuleLoadDataEx\n"))
6404031bb98SShilei Tian       return nullptr;
6413de4cd12SGeorge Rokos 
6423de4cd12SGeorge Rokos     DP("CUDA module successfully loaded!\n");
6434031bb98SShilei Tian 
6444031bb98SShilei Tian     Modules.push_back(Module);
6453de4cd12SGeorge Rokos 
6463de4cd12SGeorge Rokos     // Find the symbols in the module by name.
6474031bb98SShilei Tian     const __tgt_offload_entry *HostBegin = Image->EntriesBegin;
6484031bb98SShilei Tian     const __tgt_offload_entry *HostEnd = Image->EntriesEnd;
6493de4cd12SGeorge Rokos 
6505272d29eSJohannes Doerfert     std::list<KernelTy> &KernelsList = DeviceData[DeviceId].KernelsList;
6514031bb98SShilei Tian     for (const __tgt_offload_entry *E = HostBegin; E != HostEnd; ++E) {
6524031bb98SShilei Tian       if (!E->addr) {
6534031bb98SShilei Tian         // We return nullptr when something like this happens, the host should
6544031bb98SShilei Tian         // have always something in the address to uniquely identify the target
6554031bb98SShilei Tian         // region.
6564031bb98SShilei Tian         DP("Invalid binary: host entry '<null>' (size = %zd)...\n", E->size);
6574031bb98SShilei Tian         return nullptr;
6583de4cd12SGeorge Rokos       }
6593de4cd12SGeorge Rokos 
6604031bb98SShilei Tian       if (E->size) {
6614031bb98SShilei Tian         __tgt_offload_entry Entry = *E;
6624031bb98SShilei Tian         CUdeviceptr CUPtr;
6634031bb98SShilei Tian         size_t CUSize;
66495f0d1edSJon Chesterfield         Err = cuModuleGetGlobal(&CUPtr, &CUSize, Module, E->name);
6654031bb98SShilei Tian         // We keep this style here because we need the name
6664031bb98SShilei Tian         if (Err != CUDA_SUCCESS) {
6672ce16810SJoseph Huber           REPORT("Loading global '%s' Failed\n", E->name);
6684031bb98SShilei Tian           CUDA_ERR_STRING(Err);
6694031bb98SShilei Tian           return nullptr;
6703de4cd12SGeorge Rokos         }
6713de4cd12SGeorge Rokos 
6724031bb98SShilei Tian         if (CUSize != E->size) {
6734031bb98SShilei Tian           DP("Loading global '%s' - size mismatch (%zd != %zd)\n", E->name,
6744031bb98SShilei Tian              CUSize, E->size);
6754031bb98SShilei Tian           return nullptr;
6763de4cd12SGeorge Rokos         }
6773de4cd12SGeorge Rokos 
6783de4cd12SGeorge Rokos         DP("Entry point " DPxMOD " maps to global %s (" DPxMOD ")\n",
6794031bb98SShilei Tian            DPxPTR(E - HostBegin), E->name, DPxPTR(CUPtr));
6804031bb98SShilei Tian 
6814031bb98SShilei Tian         Entry.addr = (void *)(CUPtr);
6823de4cd12SGeorge Rokos 
683aace6d28SGheorghe-Teodor Bercea         // Note: In the current implementation declare target variables
684aace6d28SGheorghe-Teodor Bercea         // can either be link or to. This means that once unified
685aace6d28SGheorghe-Teodor Bercea         // memory is activated via the requires directive, the variable
686aace6d28SGheorghe-Teodor Bercea         // can be used directly from the host in both cases.
687aace6d28SGheorghe-Teodor Bercea         // TODO: when variables types other than to or link are added,
6884c6a098aSKazuaki Ishizaki         // the below condition should be changed to explicitly
689aace6d28SGheorghe-Teodor Bercea         // check for to and link variables types:
6904031bb98SShilei Tian         // (RequiresFlags & OMP_REQ_UNIFIED_SHARED_MEMORY && (e->flags &
6914031bb98SShilei Tian         // OMP_DECLARE_TARGET_LINK || e->flags == OMP_DECLARE_TARGET_TO))
6924031bb98SShilei Tian         if (RequiresFlags & OMP_REQ_UNIFIED_SHARED_MEMORY) {
693aace6d28SGheorghe-Teodor Bercea           // If unified memory is present any target link or to variables
694c5fe030cSGheorghe-Teodor Bercea           // can access host addresses directly. There is no longer a
695c5fe030cSGheorghe-Teodor Bercea           // need for device copies.
69695f0d1edSJon Chesterfield           cuMemcpyHtoD(CUPtr, E->addr, sizeof(void *));
6974031bb98SShilei Tian           DP("Copy linked variable host address (" DPxMOD
6984031bb98SShilei Tian              ") to device address (" DPxMOD ")\n",
6994031bb98SShilei Tian              DPxPTR(*((void **)E->addr)), DPxPTR(CUPtr));
700c5fe030cSGheorghe-Teodor Bercea         }
701c5fe030cSGheorghe-Teodor Bercea 
7024031bb98SShilei Tian         addOffloadEntry(DeviceId, Entry);
7033de4cd12SGeorge Rokos 
7043de4cd12SGeorge Rokos         continue;
7053de4cd12SGeorge Rokos       }
7063de4cd12SGeorge Rokos 
7074031bb98SShilei Tian       CUfunction Func;
7084031bb98SShilei Tian       Err = cuModuleGetFunction(&Func, Module, E->name);
7094031bb98SShilei Tian       // We keep this style here because we need the name
7104031bb98SShilei Tian       if (Err != CUDA_SUCCESS) {
7112ce16810SJoseph Huber         REPORT("Loading '%s' Failed\n", E->name);
7124031bb98SShilei Tian         CUDA_ERR_STRING(Err);
7134031bb98SShilei Tian         return nullptr;
7143de4cd12SGeorge Rokos       }
7153de4cd12SGeorge Rokos 
7163de4cd12SGeorge Rokos       DP("Entry point " DPxMOD " maps to %s (" DPxMOD ")\n",
7174031bb98SShilei Tian          DPxPTR(E - HostBegin), E->name, DPxPTR(Func));
7183de4cd12SGeorge Rokos 
7193de4cd12SGeorge Rokos       // default value GENERIC (in case symbol is missing from cubin file)
7203de4cd12SGeorge Rokos       int8_t ExecModeVal = ExecutionModeType::GENERIC;
7214031bb98SShilei Tian       std::string ExecModeNameStr(E->name);
7223de4cd12SGeorge Rokos       ExecModeNameStr += "_exec_mode";
7233de4cd12SGeorge Rokos       const char *ExecModeName = ExecModeNameStr.c_str();
7243de4cd12SGeorge Rokos 
7253de4cd12SGeorge Rokos       CUdeviceptr ExecModePtr;
7264031bb98SShilei Tian       size_t CUSize;
72795f0d1edSJon Chesterfield       Err = cuModuleGetGlobal(&ExecModePtr, &CUSize, Module, ExecModeName);
7284031bb98SShilei Tian       if (Err == CUDA_SUCCESS) {
7294031bb98SShilei Tian         if (CUSize != sizeof(int8_t)) {
7303de4cd12SGeorge Rokos           DP("Loading global exec_mode '%s' - size mismatch (%zd != %zd)\n",
7314031bb98SShilei Tian              ExecModeName, CUSize, sizeof(int8_t));
7324031bb98SShilei Tian           return nullptr;
7333de4cd12SGeorge Rokos         }
7343de4cd12SGeorge Rokos 
73595f0d1edSJon Chesterfield         Err = cuMemcpyDtoH(&ExecModeVal, ExecModePtr, CUSize);
7364031bb98SShilei Tian         if (Err != CUDA_SUCCESS) {
7372ce16810SJoseph Huber           REPORT("Error when copying data from device to host. Pointers: "
7383de4cd12SGeorge Rokos                  "host = " DPxMOD ", device = " DPxMOD ", size = %zd\n",
7394031bb98SShilei Tian                  DPxPTR(&ExecModeVal), DPxPTR(ExecModePtr), CUSize);
7404031bb98SShilei Tian           CUDA_ERR_STRING(Err);
7414031bb98SShilei Tian           return nullptr;
7423de4cd12SGeorge Rokos         }
7433de4cd12SGeorge Rokos 
7443de4cd12SGeorge Rokos         if (ExecModeVal < 0 || ExecModeVal > 1) {
7453de4cd12SGeorge Rokos           DP("Error wrong exec_mode value specified in cubin file: %d\n",
7463de4cd12SGeorge Rokos              ExecModeVal);
7474031bb98SShilei Tian           return nullptr;
7483de4cd12SGeorge Rokos         }
7493de4cd12SGeorge Rokos       } else {
7502ce16810SJoseph Huber         REPORT("Loading global exec_mode '%s' - symbol missing, using default "
7514031bb98SShilei Tian                "value GENERIC (1)\n",
7524031bb98SShilei Tian                ExecModeName);
7534031bb98SShilei Tian         CUDA_ERR_STRING(Err);
7543de4cd12SGeorge Rokos       }
7553de4cd12SGeorge Rokos 
7564031bb98SShilei Tian       KernelsList.emplace_back(Func, ExecModeVal);
7573de4cd12SGeorge Rokos 
7584031bb98SShilei Tian       __tgt_offload_entry Entry = *E;
7594031bb98SShilei Tian       Entry.addr = &KernelsList.back();
7604031bb98SShilei Tian       addOffloadEntry(DeviceId, Entry);
7613de4cd12SGeorge Rokos     }
7623de4cd12SGeorge Rokos 
763e1c7a46dSGuansong Zhang     // send device environment data to the device
764e1c7a46dSGuansong Zhang     {
7654031bb98SShilei Tian       omptarget_device_environmentTy DeviceEnv{0};
766e1c7a46dSGuansong Zhang 
767e1c7a46dSGuansong Zhang #ifdef OMPTARGET_DEBUG
7684031bb98SShilei Tian       if (const char *EnvStr = getenv("LIBOMPTARGET_DEVICE_RTL_DEBUG"))
7694031bb98SShilei Tian         DeviceEnv.debug_level = std::stoi(EnvStr);
770e1c7a46dSGuansong Zhang #endif
771e1c7a46dSGuansong Zhang 
7724031bb98SShilei Tian       const char *DeviceEnvName = "omptarget_device_environment";
7734031bb98SShilei Tian       CUdeviceptr DeviceEnvPtr;
7744031bb98SShilei Tian       size_t CUSize;
775e1c7a46dSGuansong Zhang 
77695f0d1edSJon Chesterfield       Err = cuModuleGetGlobal(&DeviceEnvPtr, &CUSize, Module, DeviceEnvName);
7774031bb98SShilei Tian       if (Err == CUDA_SUCCESS) {
7784031bb98SShilei Tian         if (CUSize != sizeof(DeviceEnv)) {
7792ce16810SJoseph Huber           REPORT(
7802ce16810SJoseph Huber               "Global device_environment '%s' - size mismatch (%zu != %zu)\n",
7814031bb98SShilei Tian               DeviceEnvName, CUSize, sizeof(int32_t));
7824031bb98SShilei Tian           CUDA_ERR_STRING(Err);
7834031bb98SShilei Tian           return nullptr;
784e1c7a46dSGuansong Zhang         }
785e1c7a46dSGuansong Zhang 
78695f0d1edSJon Chesterfield         Err = cuMemcpyHtoD(DeviceEnvPtr, &DeviceEnv, CUSize);
7874031bb98SShilei Tian         if (Err != CUDA_SUCCESS) {
7882ce16810SJoseph Huber           REPORT("Error when copying data from host to device. Pointers: "
789e1c7a46dSGuansong Zhang                  "host = " DPxMOD ", device = " DPxMOD ", size = %zu\n",
7904031bb98SShilei Tian                  DPxPTR(&DeviceEnv), DPxPTR(DeviceEnvPtr), CUSize);
7914031bb98SShilei Tian           CUDA_ERR_STRING(Err);
7924031bb98SShilei Tian           return nullptr;
793e1c7a46dSGuansong Zhang         }
794e1c7a46dSGuansong Zhang 
7954031bb98SShilei Tian         DP("Sending global device environment data %zu bytes\n", CUSize);
796e1c7a46dSGuansong Zhang       } else {
7974031bb98SShilei Tian         DP("Finding global device environment '%s' - symbol missing.\n",
7984031bb98SShilei Tian            DeviceEnvName);
7994031bb98SShilei Tian         DP("Continue, considering this is a device RTL which does not accept "
8004031bb98SShilei Tian            "environment setting.\n");
801e1c7a46dSGuansong Zhang       }
802e1c7a46dSGuansong Zhang     }
803e1c7a46dSGuansong Zhang 
8044031bb98SShilei Tian     return getOffloadEntriesTable(DeviceId);
8053de4cd12SGeorge Rokos   }
8063de4cd12SGeorge Rokos 
8070871d6d5SShilei Tian   void *dataAlloc(const int DeviceId, const int64_t Size) {
8080871d6d5SShilei Tian     if (UseMemoryManager)
8090871d6d5SShilei Tian       return MemoryManagers[DeviceId]->allocate(Size, nullptr);
8104031bb98SShilei Tian 
8110871d6d5SShilei Tian     return DeviceAllocators[DeviceId].allocate(Size, nullptr);
8123de4cd12SGeorge Rokos   }
8133de4cd12SGeorge Rokos 
8144031bb98SShilei Tian   int dataSubmit(const int DeviceId, const void *TgtPtr, const void *HstPtr,
8155449fbb5SJohannes Doerfert                  const int64_t Size, __tgt_async_info *AsyncInfo) const {
8165449fbb5SJohannes Doerfert     assert(AsyncInfo && "AsyncInfo is nullptr");
8173de4cd12SGeorge Rokos 
8184031bb98SShilei Tian     CUresult Err = cuCtxSetCurrent(DeviceData[DeviceId].Context);
8194031bb98SShilei Tian     if (!checkResult(Err, "Error returned from cuCtxSetCurrent\n"))
8203de4cd12SGeorge Rokos       return OFFLOAD_FAIL;
8213de4cd12SGeorge Rokos 
8225449fbb5SJohannes Doerfert     CUstream Stream = getStream(DeviceId, AsyncInfo);
8233de4cd12SGeorge Rokos 
82495f0d1edSJon Chesterfield     Err = cuMemcpyHtoDAsync((CUdeviceptr)TgtPtr, HstPtr, Size, Stream);
8254031bb98SShilei Tian     if (Err != CUDA_SUCCESS) {
8262ce16810SJoseph Huber       REPORT("Error when copying data from host to device. Pointers: host "
8272ce16810SJoseph Huber              "= " DPxMOD ", device = " DPxMOD ", size = %" PRId64 "\n",
8284031bb98SShilei Tian              DPxPTR(HstPtr), DPxPTR(TgtPtr), Size);
8294031bb98SShilei Tian       CUDA_ERR_STRING(Err);
8303de4cd12SGeorge Rokos       return OFFLOAD_FAIL;
8313de4cd12SGeorge Rokos     }
8323de4cd12SGeorge Rokos 
8333de4cd12SGeorge Rokos     return OFFLOAD_SUCCESS;
8343de4cd12SGeorge Rokos   }
8353de4cd12SGeorge Rokos 
8364031bb98SShilei Tian   int dataRetrieve(const int DeviceId, void *HstPtr, const void *TgtPtr,
8375449fbb5SJohannes Doerfert                    const int64_t Size, __tgt_async_info *AsyncInfo) const {
8385449fbb5SJohannes Doerfert     assert(AsyncInfo && "AsyncInfo is nullptr");
8394031bb98SShilei Tian 
8404031bb98SShilei Tian     CUresult Err = cuCtxSetCurrent(DeviceData[DeviceId].Context);
8414031bb98SShilei Tian     if (!checkResult(Err, "Error returned from cuCtxSetCurrent\n"))
84203ff643dSShilei Tian       return OFFLOAD_FAIL;
84303ff643dSShilei Tian 
8445449fbb5SJohannes Doerfert     CUstream Stream = getStream(DeviceId, AsyncInfo);
84503ff643dSShilei Tian 
84695f0d1edSJon Chesterfield     Err = cuMemcpyDtoHAsync(HstPtr, (CUdeviceptr)TgtPtr, Size, Stream);
8474031bb98SShilei Tian     if (Err != CUDA_SUCCESS) {
8482ce16810SJoseph Huber       REPORT("Error when copying data from device to host. Pointers: host "
8492ce16810SJoseph Huber              "= " DPxMOD ", device = " DPxMOD ", size = %" PRId64 "\n",
8504031bb98SShilei Tian              DPxPTR(HstPtr), DPxPTR(TgtPtr), Size);
8514031bb98SShilei Tian       CUDA_ERR_STRING(Err);
8523de4cd12SGeorge Rokos       return OFFLOAD_FAIL;
8533de4cd12SGeorge Rokos     }
8543de4cd12SGeorge Rokos 
8554031bb98SShilei Tian     return OFFLOAD_SUCCESS;
8564031bb98SShilei Tian   }
8574031bb98SShilei Tian 
858a014fbbcSShilei Tian   int dataExchange(int SrcDevId, const void *SrcPtr, int DstDevId, void *DstPtr,
8595449fbb5SJohannes Doerfert                    int64_t Size, __tgt_async_info *AsyncInfo) const {
8605449fbb5SJohannes Doerfert     assert(AsyncInfo && "AsyncInfo is nullptr");
861a014fbbcSShilei Tian 
862a014fbbcSShilei Tian     CUresult Err = cuCtxSetCurrent(DeviceData[SrcDevId].Context);
863a014fbbcSShilei Tian     if (!checkResult(Err, "Error returned from cuCtxSetCurrent\n"))
864a014fbbcSShilei Tian       return OFFLOAD_FAIL;
865a014fbbcSShilei Tian 
8665449fbb5SJohannes Doerfert     CUstream Stream = getStream(SrcDevId, AsyncInfo);
867a014fbbcSShilei Tian 
868a014fbbcSShilei Tian     // If they are two devices, we try peer to peer copy first
869a014fbbcSShilei Tian     if (SrcDevId != DstDevId) {
870a014fbbcSShilei Tian       int CanAccessPeer = 0;
871a014fbbcSShilei Tian       Err = cuDeviceCanAccessPeer(&CanAccessPeer, SrcDevId, DstDevId);
872a014fbbcSShilei Tian       if (Err != CUDA_SUCCESS) {
8732ce16810SJoseph Huber         REPORT("Error returned from cuDeviceCanAccessPeer. src = %" PRId32
874a014fbbcSShilei Tian                ", dst = %" PRId32 "\n",
875a014fbbcSShilei Tian                SrcDevId, DstDevId);
876a014fbbcSShilei Tian         CUDA_ERR_STRING(Err);
877a014fbbcSShilei Tian         return memcpyDtoD(SrcPtr, DstPtr, Size, Stream);
878a014fbbcSShilei Tian       }
879a014fbbcSShilei Tian 
880a014fbbcSShilei Tian       if (!CanAccessPeer) {
881a014fbbcSShilei Tian         DP("P2P memcpy not supported so fall back to D2D memcpy");
882a014fbbcSShilei Tian         return memcpyDtoD(SrcPtr, DstPtr, Size, Stream);
883a014fbbcSShilei Tian       }
884a014fbbcSShilei Tian 
885a014fbbcSShilei Tian       Err = cuCtxEnablePeerAccess(DeviceData[DstDevId].Context, 0);
886a014fbbcSShilei Tian       if (Err != CUDA_SUCCESS) {
8872ce16810SJoseph Huber         REPORT("Error returned from cuCtxEnablePeerAccess. src = %" PRId32
888a014fbbcSShilei Tian                ", dst = %" PRId32 "\n",
889a014fbbcSShilei Tian                SrcDevId, DstDevId);
890a014fbbcSShilei Tian         CUDA_ERR_STRING(Err);
891a014fbbcSShilei Tian         return memcpyDtoD(SrcPtr, DstPtr, Size, Stream);
892a014fbbcSShilei Tian       }
893a014fbbcSShilei Tian 
894a014fbbcSShilei Tian       Err = cuMemcpyPeerAsync((CUdeviceptr)DstPtr, DeviceData[DstDevId].Context,
895a014fbbcSShilei Tian                               (CUdeviceptr)SrcPtr, DeviceData[SrcDevId].Context,
896a014fbbcSShilei Tian                               Size, Stream);
897a014fbbcSShilei Tian       if (Err == CUDA_SUCCESS)
898a014fbbcSShilei Tian         return OFFLOAD_SUCCESS;
899a014fbbcSShilei Tian 
9002ce16810SJoseph Huber       REPORT("Error returned from cuMemcpyPeerAsync. src_ptr = " DPxMOD
9012ce16810SJoseph Huber              ", src_id =%" PRId32 ", dst_ptr = " DPxMOD ", dst_id =%" PRId32
9022ce16810SJoseph Huber              "\n",
90345bb073dSYe Luo              DPxPTR(SrcPtr), SrcDevId, DPxPTR(DstPtr), DstDevId);
904a014fbbcSShilei Tian       CUDA_ERR_STRING(Err);
905a014fbbcSShilei Tian     }
906a014fbbcSShilei Tian 
907a014fbbcSShilei Tian     return memcpyDtoD(SrcPtr, DstPtr, Size, Stream);
908a014fbbcSShilei Tian   }
909a014fbbcSShilei Tian 
9100871d6d5SShilei Tian   int dataDelete(const int DeviceId, void *TgtPtr) {
9110871d6d5SShilei Tian     if (UseMemoryManager)
9120871d6d5SShilei Tian       return MemoryManagers[DeviceId]->free(TgtPtr);
9134031bb98SShilei Tian 
9140871d6d5SShilei Tian     return DeviceAllocators[DeviceId].free(TgtPtr);
9154031bb98SShilei Tian   }
9164031bb98SShilei Tian 
917aa27cfc1SJohannes Doerfert   int runTargetTeamRegion(const int DeviceId, void *TgtEntryPtr, void **TgtArgs,
918aa27cfc1SJohannes Doerfert                           ptrdiff_t *TgtOffsets, const int ArgNum,
919aa27cfc1SJohannes Doerfert                           const int TeamNum, const int ThreadLimit,
9204031bb98SShilei Tian                           const unsigned int LoopTripCount,
9214031bb98SShilei Tian                           __tgt_async_info *AsyncInfo) const {
9224031bb98SShilei Tian     CUresult Err = cuCtxSetCurrent(DeviceData[DeviceId].Context);
9234031bb98SShilei Tian     if (!checkResult(Err, "Error returned from cuCtxSetCurrent\n"))
9244031bb98SShilei Tian       return OFFLOAD_FAIL;
9254031bb98SShilei Tian 
9263de4cd12SGeorge Rokos     // All args are references.
9274031bb98SShilei Tian     std::vector<void *> Args(ArgNum);
9284031bb98SShilei Tian     std::vector<void *> Ptrs(ArgNum);
9293de4cd12SGeorge Rokos 
9304031bb98SShilei Tian     for (int I = 0; I < ArgNum; ++I) {
9314031bb98SShilei Tian       Ptrs[I] = (void *)((intptr_t)TgtArgs[I] + TgtOffsets[I]);
9324031bb98SShilei Tian       Args[I] = &Ptrs[I];
9331546d319SGeorge Rokos     }
9343de4cd12SGeorge Rokos 
935aa27cfc1SJohannes Doerfert     KernelTy *KernelInfo = reinterpret_cast<KernelTy *>(TgtEntryPtr);
9363de4cd12SGeorge Rokos 
937aa27cfc1SJohannes Doerfert     int CudaThreadsPerBlock;
9384031bb98SShilei Tian     if (ThreadLimit > 0) {
9394031bb98SShilei Tian       DP("Setting CUDA threads per block to requested %d\n", ThreadLimit);
9404031bb98SShilei Tian       CudaThreadsPerBlock = ThreadLimit;
9413de4cd12SGeorge Rokos       // Add master warp if necessary
9423de4cd12SGeorge Rokos       if (KernelInfo->ExecutionMode == GENERIC) {
9434031bb98SShilei Tian         DP("Adding master warp: +%d threads\n", DeviceData[DeviceId].WarpSize);
9444031bb98SShilei Tian         CudaThreadsPerBlock += DeviceData[DeviceId].WarpSize;
9453de4cd12SGeorge Rokos       }
946c13df8e5SGeorge Rokos     } else {
947c13df8e5SGeorge Rokos       DP("Setting CUDA threads per block to default %d\n",
9484031bb98SShilei Tian          DeviceData[DeviceId].NumThreads);
9494031bb98SShilei Tian       CudaThreadsPerBlock = DeviceData[DeviceId].NumThreads;
950c13df8e5SGeorge Rokos     }
9513de4cd12SGeorge Rokos 
9524031bb98SShilei Tian     if (CudaThreadsPerBlock > DeviceData[DeviceId].ThreadsPerBlock) {
9533de4cd12SGeorge Rokos       DP("Threads per block capped at device limit %d\n",
9544031bb98SShilei Tian          DeviceData[DeviceId].ThreadsPerBlock);
9554031bb98SShilei Tian       CudaThreadsPerBlock = DeviceData[DeviceId].ThreadsPerBlock;
9563de4cd12SGeorge Rokos     }
9573de4cd12SGeorge Rokos 
958aa27cfc1SJohannes Doerfert     if (!KernelInfo->MaxThreadsPerBlock) {
959aa27cfc1SJohannes Doerfert       Err = cuFuncGetAttribute(&KernelInfo->MaxThreadsPerBlock,
9604031bb98SShilei Tian                                CU_FUNC_ATTRIBUTE_MAX_THREADS_PER_BLOCK,
9614031bb98SShilei Tian                                KernelInfo->Func);
962aa27cfc1SJohannes Doerfert       if (!checkResult(Err, "Error returned from cuFuncGetAttribute\n"))
963aa27cfc1SJohannes Doerfert         return OFFLOAD_FAIL;
964aa27cfc1SJohannes Doerfert     }
965aa27cfc1SJohannes Doerfert 
966aa27cfc1SJohannes Doerfert     if (KernelInfo->MaxThreadsPerBlock < CudaThreadsPerBlock) {
967aa27cfc1SJohannes Doerfert       DP("Threads per block capped at kernel limit %d\n",
968aa27cfc1SJohannes Doerfert          KernelInfo->MaxThreadsPerBlock);
969aa27cfc1SJohannes Doerfert       CudaThreadsPerBlock = KernelInfo->MaxThreadsPerBlock;
9703de4cd12SGeorge Rokos     }
9713de4cd12SGeorge Rokos 
9724031bb98SShilei Tian     unsigned int CudaBlocksPerGrid;
9734031bb98SShilei Tian     if (TeamNum <= 0) {
9744031bb98SShilei Tian       if (LoopTripCount > 0 && EnvNumTeams < 0) {
975c13df8e5SGeorge Rokos         if (KernelInfo->ExecutionMode == SPMD) {
9764031bb98SShilei Tian           // We have a combined construct, i.e. `target teams distribute
9774031bb98SShilei Tian           // parallel for [simd]`. We launch so many teams so that each thread
9784031bb98SShilei Tian           // will execute one iteration of the loop. round up to the nearest
9794031bb98SShilei Tian           // integer
9804031bb98SShilei Tian           CudaBlocksPerGrid = ((LoopTripCount - 1) / CudaThreadsPerBlock) + 1;
981c13df8e5SGeorge Rokos         } else {
982c13df8e5SGeorge Rokos           // If we reach this point, then we have a non-combined construct, i.e.
983c13df8e5SGeorge Rokos           // `teams distribute` with a nested `parallel for` and each team is
984c13df8e5SGeorge Rokos           // assigned one iteration of the `distribute` loop. E.g.:
985c13df8e5SGeorge Rokos           //
986c13df8e5SGeorge Rokos           // #pragma omp target teams distribute
987c13df8e5SGeorge Rokos           // for(...loop_tripcount...) {
988c13df8e5SGeorge Rokos           //   #pragma omp parallel for
989c13df8e5SGeorge Rokos           //   for(...) {}
990c13df8e5SGeorge Rokos           // }
991c13df8e5SGeorge Rokos           //
992c13df8e5SGeorge Rokos           // Threads within a team will execute the iterations of the `parallel`
993c13df8e5SGeorge Rokos           // loop.
9944031bb98SShilei Tian           CudaBlocksPerGrid = LoopTripCount;
995c13df8e5SGeorge Rokos         }
99645bb073dSYe Luo         DP("Using %d teams due to loop trip count %" PRIu32
9974031bb98SShilei Tian            " and number of threads per block %d\n",
9984031bb98SShilei Tian            CudaBlocksPerGrid, LoopTripCount, CudaThreadsPerBlock);
9993de4cd12SGeorge Rokos       } else {
10004031bb98SShilei Tian         DP("Using default number of teams %d\n", DeviceData[DeviceId].NumTeams);
10014031bb98SShilei Tian         CudaBlocksPerGrid = DeviceData[DeviceId].NumTeams;
10023de4cd12SGeorge Rokos       }
10034031bb98SShilei Tian     } else if (TeamNum > DeviceData[DeviceId].BlocksPerGrid) {
10043de4cd12SGeorge Rokos       DP("Capping number of teams to team limit %d\n",
10054031bb98SShilei Tian          DeviceData[DeviceId].BlocksPerGrid);
10064031bb98SShilei Tian       CudaBlocksPerGrid = DeviceData[DeviceId].BlocksPerGrid;
10073de4cd12SGeorge Rokos     } else {
10084031bb98SShilei Tian       DP("Using requested number of teams %d\n", TeamNum);
10094031bb98SShilei Tian       CudaBlocksPerGrid = TeamNum;
10103de4cd12SGeorge Rokos     }
10113de4cd12SGeorge Rokos 
1012119a9ea1SJoseph Huber     INFO(OMP_INFOTYPE_PLUGIN_KERNEL, DeviceId,
1013ae209397SJoseph Huber          "Launching kernel %s with %d blocks and %d threads in %s "
1014ae209397SJoseph Huber          "mode\n",
1015ae209397SJoseph Huber          (getOffloadEntry(DeviceId, TgtEntryPtr))
1016ae209397SJoseph Huber              ? getOffloadEntry(DeviceId, TgtEntryPtr)->name
1017ae209397SJoseph Huber              : "(null)",
1018ae209397SJoseph Huber          CudaBlocksPerGrid, CudaThreadsPerBlock,
1019ae209397SJoseph Huber          (KernelInfo->ExecutionMode == SPMD) ? "SPMD" : "Generic");
10203de4cd12SGeorge Rokos 
10214031bb98SShilei Tian     CUstream Stream = getStream(DeviceId, AsyncInfo);
10224031bb98SShilei Tian     Err = cuLaunchKernel(KernelInfo->Func, CudaBlocksPerGrid, /* gridDimY */ 1,
10234031bb98SShilei Tian                          /* gridDimZ */ 1, CudaThreadsPerBlock,
10244031bb98SShilei Tian                          /* blockDimY */ 1, /* blockDimZ */ 1,
10254031bb98SShilei Tian                          /* sharedMemBytes */ 0, Stream, &Args[0], nullptr);
10264031bb98SShilei Tian     if (!checkResult(Err, "Error returned from cuLaunchKernel\n"))
10274031bb98SShilei Tian       return OFFLOAD_FAIL;
10284031bb98SShilei Tian 
10294031bb98SShilei Tian     DP("Launch of entry point at " DPxMOD " successful!\n",
10304031bb98SShilei Tian        DPxPTR(TgtEntryPtr));
10314031bb98SShilei Tian 
10324031bb98SShilei Tian     return OFFLOAD_SUCCESS;
10334031bb98SShilei Tian   }
10344031bb98SShilei Tian 
10355449fbb5SJohannes Doerfert   int synchronize(const int DeviceId, __tgt_async_info *AsyncInfo) const {
10365449fbb5SJohannes Doerfert     CUstream Stream = reinterpret_cast<CUstream>(AsyncInfo->Queue);
10374031bb98SShilei Tian     CUresult Err = cuStreamSynchronize(Stream);
10384031bb98SShilei Tian     if (Err != CUDA_SUCCESS) {
10392ce16810SJoseph Huber       REPORT("Error when synchronizing stream. stream = " DPxMOD
10404031bb98SShilei Tian              ", async info ptr = " DPxMOD "\n",
10415449fbb5SJohannes Doerfert              DPxPTR(Stream), DPxPTR(AsyncInfo));
10424031bb98SShilei Tian       CUDA_ERR_STRING(Err);
10433de4cd12SGeorge Rokos       return OFFLOAD_FAIL;
10443de4cd12SGeorge Rokos     }
10453de4cd12SGeorge Rokos 
10464031bb98SShilei Tian     // Once the stream is synchronized, return it to stream pool and reset
10475449fbb5SJohannes Doerfert     // AsyncInfo. This is to make sure the synchronization only works for its
10484031bb98SShilei Tian     // own tasks.
10495449fbb5SJohannes Doerfert     StreamManager->returnStream(DeviceId,
10505449fbb5SJohannes Doerfert                                 reinterpret_cast<CUstream>(AsyncInfo->Queue));
10515449fbb5SJohannes Doerfert     AsyncInfo->Queue = nullptr;
10523de4cd12SGeorge Rokos 
10533de4cd12SGeorge Rokos     return OFFLOAD_SUCCESS;
10543de4cd12SGeorge Rokos   }
10554031bb98SShilei Tian };
10564031bb98SShilei Tian 
10574031bb98SShilei Tian DeviceRTLTy DeviceRTL;
10584031bb98SShilei Tian } // namespace
10594031bb98SShilei Tian 
10604031bb98SShilei Tian // Exposed library API function
10614031bb98SShilei Tian #ifdef __cplusplus
10624031bb98SShilei Tian extern "C" {
10634031bb98SShilei Tian #endif
10644031bb98SShilei Tian 
10654031bb98SShilei Tian int32_t __tgt_rtl_is_valid_binary(__tgt_device_image *image) {
10664031bb98SShilei Tian   return elf_check_machine(image, /* EM_CUDA */ 190);
10674031bb98SShilei Tian }
10684031bb98SShilei Tian 
10694031bb98SShilei Tian int32_t __tgt_rtl_number_of_devices() { return DeviceRTL.getNumOfDevices(); }
10704031bb98SShilei Tian 
10714031bb98SShilei Tian int64_t __tgt_rtl_init_requires(int64_t RequiresFlags) {
107240470eb2SGeorge Rokos   DP("Init requires flags to %" PRId64 "\n", RequiresFlags);
10734031bb98SShilei Tian   DeviceRTL.setRequiresFlag(RequiresFlags);
10744031bb98SShilei Tian   return RequiresFlags;
10754031bb98SShilei Tian }
10764031bb98SShilei Tian 
1077a014fbbcSShilei Tian int32_t __tgt_rtl_is_data_exchangable(int32_t src_dev_id, int dst_dev_id) {
1078a014fbbcSShilei Tian   if (DeviceRTL.isValidDeviceId(src_dev_id) &&
1079a014fbbcSShilei Tian       DeviceRTL.isValidDeviceId(dst_dev_id))
1080a014fbbcSShilei Tian     return 1;
1081a014fbbcSShilei Tian 
1082a014fbbcSShilei Tian   return 0;
1083a014fbbcSShilei Tian }
1084a014fbbcSShilei Tian 
10854031bb98SShilei Tian int32_t __tgt_rtl_init_device(int32_t device_id) {
10864031bb98SShilei Tian   assert(DeviceRTL.isValidDeviceId(device_id) && "device_id is invalid");
10874031bb98SShilei Tian 
10884031bb98SShilei Tian   return DeviceRTL.initDevice(device_id);
10894031bb98SShilei Tian }
10904031bb98SShilei Tian 
10914031bb98SShilei Tian __tgt_target_table *__tgt_rtl_load_binary(int32_t device_id,
10924031bb98SShilei Tian                                           __tgt_device_image *image) {
10934031bb98SShilei Tian   assert(DeviceRTL.isValidDeviceId(device_id) && "device_id is invalid");
10944031bb98SShilei Tian 
10954031bb98SShilei Tian   return DeviceRTL.loadBinary(device_id, image);
10964031bb98SShilei Tian }
10974031bb98SShilei Tian 
1098*2468fdd9SGeorge Rokos void *__tgt_rtl_data_alloc(int32_t device_id, int64_t size, void *,
1099*2468fdd9SGeorge Rokos                            int32_t kind) {
11004031bb98SShilei Tian   assert(DeviceRTL.isValidDeviceId(device_id) && "device_id is invalid");
11014031bb98SShilei Tian 
1102*2468fdd9SGeorge Rokos   if (kind != TARGET_ALLOC_DEFAULT) {
1103*2468fdd9SGeorge Rokos     REPORT("Invalid target data allocation kind or requested allocator not "
1104*2468fdd9SGeorge Rokos            "implemented yet\n");
1105*2468fdd9SGeorge Rokos     return NULL;
1106*2468fdd9SGeorge Rokos   }
1107*2468fdd9SGeorge Rokos 
11084031bb98SShilei Tian   return DeviceRTL.dataAlloc(device_id, size);
11094031bb98SShilei Tian }
11104031bb98SShilei Tian 
11114031bb98SShilei Tian int32_t __tgt_rtl_data_submit(int32_t device_id, void *tgt_ptr, void *hst_ptr,
11124031bb98SShilei Tian                               int64_t size) {
11134031bb98SShilei Tian   assert(DeviceRTL.isValidDeviceId(device_id) && "device_id is invalid");
11144031bb98SShilei Tian 
11155449fbb5SJohannes Doerfert   __tgt_async_info AsyncInfo;
11164031bb98SShilei Tian   const int32_t rc = __tgt_rtl_data_submit_async(device_id, tgt_ptr, hst_ptr,
11175449fbb5SJohannes Doerfert                                                  size, &AsyncInfo);
11184031bb98SShilei Tian   if (rc != OFFLOAD_SUCCESS)
11194031bb98SShilei Tian     return OFFLOAD_FAIL;
11204031bb98SShilei Tian 
11215449fbb5SJohannes Doerfert   return __tgt_rtl_synchronize(device_id, &AsyncInfo);
11224031bb98SShilei Tian }
11234031bb98SShilei Tian 
11244031bb98SShilei Tian int32_t __tgt_rtl_data_submit_async(int32_t device_id, void *tgt_ptr,
11254031bb98SShilei Tian                                     void *hst_ptr, int64_t size,
11264031bb98SShilei Tian                                     __tgt_async_info *async_info_ptr) {
11274031bb98SShilei Tian   assert(DeviceRTL.isValidDeviceId(device_id) && "device_id is invalid");
11284031bb98SShilei Tian   assert(async_info_ptr && "async_info_ptr is nullptr");
11294031bb98SShilei Tian 
11304031bb98SShilei Tian   return DeviceRTL.dataSubmit(device_id, tgt_ptr, hst_ptr, size,
11314031bb98SShilei Tian                               async_info_ptr);
11324031bb98SShilei Tian }
11334031bb98SShilei Tian 
11344031bb98SShilei Tian int32_t __tgt_rtl_data_retrieve(int32_t device_id, void *hst_ptr, void *tgt_ptr,
11354031bb98SShilei Tian                                 int64_t size) {
11364031bb98SShilei Tian   assert(DeviceRTL.isValidDeviceId(device_id) && "device_id is invalid");
11374031bb98SShilei Tian 
11385449fbb5SJohannes Doerfert   __tgt_async_info AsyncInfo;
11394031bb98SShilei Tian   const int32_t rc = __tgt_rtl_data_retrieve_async(device_id, hst_ptr, tgt_ptr,
11405449fbb5SJohannes Doerfert                                                    size, &AsyncInfo);
11414031bb98SShilei Tian   if (rc != OFFLOAD_SUCCESS)
11424031bb98SShilei Tian     return OFFLOAD_FAIL;
11434031bb98SShilei Tian 
11445449fbb5SJohannes Doerfert   return __tgt_rtl_synchronize(device_id, &AsyncInfo);
11454031bb98SShilei Tian }
11464031bb98SShilei Tian 
11474031bb98SShilei Tian int32_t __tgt_rtl_data_retrieve_async(int32_t device_id, void *hst_ptr,
11484031bb98SShilei Tian                                       void *tgt_ptr, int64_t size,
11494031bb98SShilei Tian                                       __tgt_async_info *async_info_ptr) {
11504031bb98SShilei Tian   assert(DeviceRTL.isValidDeviceId(device_id) && "device_id is invalid");
11514031bb98SShilei Tian   assert(async_info_ptr && "async_info_ptr is nullptr");
11524031bb98SShilei Tian 
11534031bb98SShilei Tian   return DeviceRTL.dataRetrieve(device_id, hst_ptr, tgt_ptr, size,
11544031bb98SShilei Tian                                 async_info_ptr);
11554031bb98SShilei Tian }
11564031bb98SShilei Tian 
1157a014fbbcSShilei Tian int32_t __tgt_rtl_data_exchange_async(int32_t src_dev_id, void *src_ptr,
1158a014fbbcSShilei Tian                                       int dst_dev_id, void *dst_ptr,
1159a014fbbcSShilei Tian                                       int64_t size,
11605449fbb5SJohannes Doerfert                                       __tgt_async_info *AsyncInfo) {
1161a014fbbcSShilei Tian   assert(DeviceRTL.isValidDeviceId(src_dev_id) && "src_dev_id is invalid");
1162a014fbbcSShilei Tian   assert(DeviceRTL.isValidDeviceId(dst_dev_id) && "dst_dev_id is invalid");
11635449fbb5SJohannes Doerfert   assert(AsyncInfo && "AsyncInfo is nullptr");
1164a014fbbcSShilei Tian 
1165a014fbbcSShilei Tian   return DeviceRTL.dataExchange(src_dev_id, src_ptr, dst_dev_id, dst_ptr, size,
11665449fbb5SJohannes Doerfert                                 AsyncInfo);
1167a014fbbcSShilei Tian }
1168a014fbbcSShilei Tian 
1169a014fbbcSShilei Tian int32_t __tgt_rtl_data_exchange(int32_t src_dev_id, void *src_ptr,
1170a014fbbcSShilei Tian                                 int32_t dst_dev_id, void *dst_ptr,
1171a014fbbcSShilei Tian                                 int64_t size) {
1172a014fbbcSShilei Tian   assert(DeviceRTL.isValidDeviceId(src_dev_id) && "src_dev_id is invalid");
1173a014fbbcSShilei Tian   assert(DeviceRTL.isValidDeviceId(dst_dev_id) && "dst_dev_id is invalid");
1174a014fbbcSShilei Tian 
11755449fbb5SJohannes Doerfert   __tgt_async_info AsyncInfo;
1176a014fbbcSShilei Tian   const int32_t rc = __tgt_rtl_data_exchange_async(
11775449fbb5SJohannes Doerfert       src_dev_id, src_ptr, dst_dev_id, dst_ptr, size, &AsyncInfo);
1178a014fbbcSShilei Tian   if (rc != OFFLOAD_SUCCESS)
1179a014fbbcSShilei Tian     return OFFLOAD_FAIL;
1180a014fbbcSShilei Tian 
11815449fbb5SJohannes Doerfert   return __tgt_rtl_synchronize(src_dev_id, &AsyncInfo);
1182a014fbbcSShilei Tian }
1183a014fbbcSShilei Tian 
11844031bb98SShilei Tian int32_t __tgt_rtl_data_delete(int32_t device_id, void *tgt_ptr) {
11854031bb98SShilei Tian   assert(DeviceRTL.isValidDeviceId(device_id) && "device_id is invalid");
11864031bb98SShilei Tian 
11874031bb98SShilei Tian   return DeviceRTL.dataDelete(device_id, tgt_ptr);
11884031bb98SShilei Tian }
11894031bb98SShilei Tian 
11904031bb98SShilei Tian int32_t __tgt_rtl_run_target_team_region(int32_t device_id, void *tgt_entry_ptr,
11914031bb98SShilei Tian                                          void **tgt_args,
11924031bb98SShilei Tian                                          ptrdiff_t *tgt_offsets,
11934031bb98SShilei Tian                                          int32_t arg_num, int32_t team_num,
11944031bb98SShilei Tian                                          int32_t thread_limit,
11954031bb98SShilei Tian                                          uint64_t loop_tripcount) {
11964031bb98SShilei Tian   assert(DeviceRTL.isValidDeviceId(device_id) && "device_id is invalid");
11974031bb98SShilei Tian 
11985449fbb5SJohannes Doerfert   __tgt_async_info AsyncInfo;
11994031bb98SShilei Tian   const int32_t rc = __tgt_rtl_run_target_team_region_async(
12004031bb98SShilei Tian       device_id, tgt_entry_ptr, tgt_args, tgt_offsets, arg_num, team_num,
12015449fbb5SJohannes Doerfert       thread_limit, loop_tripcount, &AsyncInfo);
12024031bb98SShilei Tian   if (rc != OFFLOAD_SUCCESS)
12034031bb98SShilei Tian     return OFFLOAD_FAIL;
12044031bb98SShilei Tian 
12055449fbb5SJohannes Doerfert   return __tgt_rtl_synchronize(device_id, &AsyncInfo);
12064031bb98SShilei Tian }
12074031bb98SShilei Tian 
12084031bb98SShilei Tian int32_t __tgt_rtl_run_target_team_region_async(
12094031bb98SShilei Tian     int32_t device_id, void *tgt_entry_ptr, void **tgt_args,
12104031bb98SShilei Tian     ptrdiff_t *tgt_offsets, int32_t arg_num, int32_t team_num,
12114031bb98SShilei Tian     int32_t thread_limit, uint64_t loop_tripcount,
12124031bb98SShilei Tian     __tgt_async_info *async_info_ptr) {
12134031bb98SShilei Tian   assert(DeviceRTL.isValidDeviceId(device_id) && "device_id is invalid");
12144031bb98SShilei Tian 
12154031bb98SShilei Tian   return DeviceRTL.runTargetTeamRegion(
12164031bb98SShilei Tian       device_id, tgt_entry_ptr, tgt_args, tgt_offsets, arg_num, team_num,
12174031bb98SShilei Tian       thread_limit, loop_tripcount, async_info_ptr);
12184031bb98SShilei Tian }
12193de4cd12SGeorge Rokos 
12203de4cd12SGeorge Rokos int32_t __tgt_rtl_run_target_region(int32_t device_id, void *tgt_entry_ptr,
122132ed2927SShilei Tian                                     void **tgt_args, ptrdiff_t *tgt_offsets,
122203ff643dSShilei Tian                                     int32_t arg_num) {
12234031bb98SShilei Tian   assert(DeviceRTL.isValidDeviceId(device_id) && "device_id is invalid");
12244031bb98SShilei Tian 
12255449fbb5SJohannes Doerfert   __tgt_async_info AsyncInfo;
12264031bb98SShilei Tian   const int32_t rc = __tgt_rtl_run_target_region_async(
12275449fbb5SJohannes Doerfert       device_id, tgt_entry_ptr, tgt_args, tgt_offsets, arg_num, &AsyncInfo);
122803ff643dSShilei Tian   if (rc != OFFLOAD_SUCCESS)
122903ff643dSShilei Tian     return OFFLOAD_FAIL;
123003ff643dSShilei Tian 
12315449fbb5SJohannes Doerfert   return __tgt_rtl_synchronize(device_id, &AsyncInfo);
123203ff643dSShilei Tian }
123303ff643dSShilei Tian 
123403ff643dSShilei Tian int32_t __tgt_rtl_run_target_region_async(int32_t device_id,
123503ff643dSShilei Tian                                           void *tgt_entry_ptr, void **tgt_args,
123603ff643dSShilei Tian                                           ptrdiff_t *tgt_offsets,
123732ed2927SShilei Tian                                           int32_t arg_num,
12384031bb98SShilei Tian                                           __tgt_async_info *async_info_ptr) {
12394031bb98SShilei Tian   assert(DeviceRTL.isValidDeviceId(device_id) && "device_id is invalid");
12404031bb98SShilei Tian 
124103ff643dSShilei Tian   return __tgt_rtl_run_target_team_region_async(
12424031bb98SShilei Tian       device_id, tgt_entry_ptr, tgt_args, tgt_offsets, arg_num,
12434031bb98SShilei Tian       /* team num*/ 1, /* thread_limit */ 1, /* loop_tripcount */ 0,
12444031bb98SShilei Tian       async_info_ptr);
124532ed2927SShilei Tian }
124632ed2927SShilei Tian 
12474031bb98SShilei Tian int32_t __tgt_rtl_synchronize(int32_t device_id,
12484031bb98SShilei Tian                               __tgt_async_info *async_info_ptr) {
12494031bb98SShilei Tian   assert(DeviceRTL.isValidDeviceId(device_id) && "device_id is invalid");
12504031bb98SShilei Tian   assert(async_info_ptr && "async_info_ptr is nullptr");
12514031bb98SShilei Tian   assert(async_info_ptr->Queue && "async_info_ptr->Queue is nullptr");
125232ed2927SShilei Tian 
12534031bb98SShilei Tian   return DeviceRTL.synchronize(device_id, async_info_ptr);
12543de4cd12SGeorge Rokos }
12553de4cd12SGeorge Rokos 
12563de4cd12SGeorge Rokos #ifdef __cplusplus
12573de4cd12SGeorge Rokos }
12583de4cd12SGeorge Rokos #endif
1259