[OpenMP][Offloading] Added support for multiple streams so that multiple kernels can be executed concurrently

[tianshilei1992]

Summary

It will initialize a number of streams for each device at first. The number can be configured via environment variable LIBOMPTARGET_NUM_STREAMS. For each kernel submission, a stream will be selected in a round-robin manner.

[jdoerfert]

Thanks! Two comments below.

@ye-luo once the memory transfers are attached to a stream you should be able to offload synchronously from multiple threads at the same time. Could you pull the patch and test it?

openmp/libomptarget/plugins/cuda/src/rtl.cpp

95

Make it uint please.

529

We need the async versions at the HtoD and at the DtoH sides to use the streams. After the async call we directly have to wait for the stream to make it synchronous but on as specific stream.

[ye-luo]

@jdoerfert I can try it on a test program. miniQMC is choked by the linker at the moment. Is the "map" thread-safe now?

openmp/libomptarget/plugins/cuda/src/rtl.cpp

529

In this direction, the H2D, kernel and D2H optimally can be scheduled as a whole entity in the tasking runtime and use the same stream if they are on the same OpenMP pragma line.

[ye-luo]

Caught another issue.

openmp/libomptarget/plugins/cuda/src/rtl.cpp

862

This synchronization should be replaced with stream wait.

[jdoerfert]

In D74145#1863379, @ye-luo wrote:

@jdoerfert I can try it on a test program. miniQMC is choked by the linker at the moment. Is the "map" thread-safe now?

Map should be thread safe, yes.

[tianshilei1992]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

95

Right, in case of integer overflow, my bad...

[JonChesterfield]

Adding Ron to the list as he's maintaining the amdgcn equivalent to this

[JonChesterfield]

Are cuda streams available on all versions of cuda that the rest of openmp works from? I'm not sure when they were introduced.

A couple of minor comments inline. This seems to be a fairly straightforward wrapper over the cuda functionality.

openmp/libomptarget/plugins/cuda/src/rtl.cpp

95

vector of pointers to atomic_int is interesting. What's the advantage over vector<atomic_int>?

It might be worth putting a few asserts in the code to the effect that resizing the vector after the initial construction will break access from other threads.

259

If we do need the pointer wrapper, this should be make_unique

[tianshilei1992]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

95

atomic_int is not copyable. And the initialization of all these pointers are after the resize of vector, so we might not need to consider that.

259

make_unique only works since C++14.

[tianshilei1992]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

862

Are you referring to cudaStreamWaitEvent?

[jdoerfert]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

259

Do we have llvm::make_unique? But maybe not necessarily good to use it here anyway. @jon ok to stick with this for now?

[ye-luo]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

862

I mean cuStreamSynchronize

[JonChesterfield]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

95

Sure, but you're not copying the element anywhere.

Sadly I think we would need to provide the size of the vector up front. reserve() calls copy constructors (which I didn't expect) and they're deleted for atomic_int. I'm not sure the cuda api will permit that.

Which leads to the suggestion:

std::unique_ptr<std::atomic_int[]>> NextStreamId;
// ...
NextStreamId = std::make_unique<std::atomic_int[]>(NumberOfDevices);

This elides the NumberOfDevices heap allocations and the associated indirection on every access and makes it somewhat more obvious that we can't call various vector api functions.

It has the disadvantage that the integers will now definitely be in the same cache line, whereas previously there was a chance that the allocator would put them on different cache lines.

Overall I'm fine with either structure.

259

llvm::make_unique was removed by D66259, as we're now assuming C++14. They're semantically identical in this context so it doesn't matter much.

[tianshilei1992]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

259

Do you mean that we can assume -std=c++14 is always true?

862

Oh, I got you. Good one, in case of blocking other threads, although the offloading have finished.

[JonChesterfield]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

259

Other files in LLVM won't build with c++11 any more so >=14 seems a safe bet.

[tianshilei1992]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

259

That is cool! Thanks for the information. Will update this part correspondingly.

[tianshilei1992]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

259

Well, I just tried with make_unique but it turns out we're still using C++11 actually.

FAILED: libomptarget/plugins/cuda/CMakeFiles/omptarget.rtl.cuda.dir/src/rtl.cpp.o
/home/shiltian/.local/bin/clang++  -DOMPTARGET_DEBUG -DTARGET_NAME=CUDA -Domptarget_rtl_cuda_EXPORTS -I/home/shiltian/Documents/clion/llvm-project/openmp/libomptarget/include -I/opt/cuda/10.1/include -Wall -Wcast-qual -Wformat-pedantic -Wimplicit-fallthrough -Wsign-compare -Wno-extra -Wno-pedantic -std=gnu++11 -g -fPIC -MD -MT libomptarget/plugins/cuda/CMakeFiles/omptarget.rtl.cuda.dir/src/rtl.cpp.o -MF libomptarget/plugins/cuda/CMakeFiles/omptarget.rtl.cuda.dir/src/rtl.cpp.o.d -o libomptarget/plugins/cuda/CMakeFiles/omptarget.rtl.cuda.dir/src/rtl.cpp.o -c /home/shiltian/Documents/clion/llvm-project/openmp/libomptarget/plugins/cuda/src/rtl.cpp
/home/shiltian/Documents/clion/llvm-project/openmp/libomptarget/plugins/cuda/src/rtl.cpp:259:18: error: no member named 'make_unique' in namespace 'std'
      Ptr = std::make_unique<std::atomic_uint>(0);
            ~~~~~^
/home/shiltian/Documents/clion/llvm-project/openmp/libomptarget/plugins/cuda/src/rtl.cpp:259:46: error: expected '(' for function-style cast or type construction
      Ptr = std::make_unique<std::atomic_uint>(0);
                             ~~~~~~~~~~~~~~~~^
2 errors generated.
ninja: build stopped: subcommand failed.

[jdoerfert]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

259

Change the cmake in a separate commit. Llvm is on 14.

[tianshilei1992]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

259

So OpenMP will also switch to C++ 14 in a near future?

[JonChesterfield]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

259

Sounds good to me. Yep, let's change the cmake now.

[tianshilei1992]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

259

https://reviews.llvm.org/D74258

[jdoerfert]

We will probably need "version 2" functions soon which take additional information, e.g., the stream to be used. I would suggest to test this as is and merge it before we go there. It should already allow overlap between threads that offload. The "version 2" will only shrink the overhead per thread. That said, we are working on the nowait support so there will be other changes soon anyway.

@ye-luo Do you have a way to test this or do we need to fix the linker issue first?

openmp/libomptarget/plugins/cuda/src/rtl.cpp

246

The hardware will cap the number internally anyway so we should go higher here. Maybe 256?

[tianshilei1992]

I did a little experiment to show the performance improvement. Here is the micro benchmark:

#include <math.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <time.h>

void kernel() {
  const int num_threads = 64;

#pragma omp parallel for
  for (int i = 0; i < num_threads; ++i) {
    const size_t N = 1UL << 10;

#pragma omp target teams distribute parallel for
    for (size_t i = 0; i < N; ++i) {
      for (size_t j = 0; j < N / 2; ++j) {
        float x = sqrt(pow(3.14159, j));
      }
    }
  }
}

int main(int argc, char *argv[]) {
  const int N = 1000;

  const clock_t start = clock();

  for (int i = 0; i < N; ++i) {
    kernel();
  }

  const clock_t duration = (clock() - start) * 1000 / CLOCKS_PER_SEC / N;

  printf("Avg time: %ld ms\n", duration);

  return 0;
}

The execution result with multiple stream is:

$ /usr/local/cuda/bin/nvprof --output-profile parallel_offloading_ms.prof -f ./parallel_offloading
==32397== NVPROF is profiling process 32397, command: ./parallel_offloading
Avg time: 1081 ms
==32397== Generated result file: /home/shiltian/Documents/project/multiple_streams/tests/multistreams/parallel_offloading_ms.prof

And the result w/o multiple stream is:

$ /usr/local/cuda/bin/nvprof --output-profile parallel_offloading.prof -f ./parallel_offloading
==35547== NVPROF is profiling process 35547, command: ./parallel_offloading
Avg time: 5825 ms
==35547== Generated result file: /home/shiltian/Documents/project/multiple_streams/tests/multistreams/parallel_offloading.prof

We can see that 1081 vs 5825 ms, approximately 5.4x speedup.

[tianshilei1992]

In D74145#1865837, @jdoerfert wrote:

We will probably need "version 2" functions soon which take additional information, e.g., the stream to be used. I would suggest to test this as is and merge it before we go there. It should already allow overlap between threads that offload. The "version 2" will only shrink the overhead per thread. That said, we are working on the nowait support so there will be other changes soon anyway.

Yes, later we will take stream it previous used for data transfer into consideration when selecting stream for kernel, and other potential optimization.

[tianshilei1992]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

246

Sure

[ye-luo]

In D74145#1865837, @jdoerfert wrote:

We will probably need "version 2" functions soon which take additional information, e.g., the stream to be used. I would suggest to test this as is and merge it before we go there. It should already allow overlap between threads that offload. The "version 2" will only shrink the overhead per thread. That said, we are working on the nowait support so there will be other changes soon anyway.

@ye-luo Do you have a way to test this or do we need to fix the linker issue first?

My standalone code can be used to verify multi-stream concurrent execution and whether transfer and execution use the same stream by profiling with nvprof.

[jdoerfert]

In D74145#1865927, @ye-luo wrote:

In D74145#1865837, @jdoerfert wrote:

We will probably need "version 2" functions soon which take additional information, e.g., the stream to be used. I would suggest to test this as is and merge it before we go there. It should already allow overlap between threads that offload. The "version 2" will only shrink the overhead per thread. That said, we are working on the nowait support so there will be other changes soon anyway.

@ye-luo Do you have a way to test this or do we need to fix the linker issue first?

My standalone code can be used to verify multi-stream concurrent execution and whether transfer and execution use the same stream by profiling with nvprof.

What standalone code? Can you run it with this patch? The transfer will use a different stream for now but it should be OK for now. "version 2" will do the same stream.

[tianshilei1992]

Add a new test case to check map is working correct.

[jdoerfert]

I'm fine with this, anyone else?

[JonChesterfield]

We need to lose std::make_unique before landing as the C++11 = > C++14 move has proven contentious. Otherwise LGTM.

[tianshilei1992]

In D74145#1866192, @JonChesterfield wrote:

We need to lose std::make_unique before landing as the C++11 = > C++14 move has proven contentious. Otherwise LGTM.

I can use a macro here like #if __cplusplus < 201402L.

[tianshilei1992]

Add a backup statement in case that the library is not compiled with C++14

[jdoerfert]

In D74145#1866224, @tianshilei1992 wrote:

Add a backup statement in case that the library is not compiled with C++14

I would prefer not to do this. Let's wait till Monday and replies on the RFC.

[JonChesterfield]

I was thinking of going back to the explicit new, not an ifdef on c++ version. Then we can land this now and optionally revisit once the codebase moves to 14.

[ye-luo]

I tested the patch. The stream of H2D, D2H and compute behaves asynchronously as expected.

openmp/libomptarget/plugins/cuda/src/rtl.cpp

246

I don't like this choice. The hardware limit is 32 which is preferred. Users can play with environment variable if they need more.
On the nvprof, it is impossible to digest 256 streams from OpenMP plus other application streams.

[jdoerfert]

In D74145#1866382, @ye-luo wrote:

I tested the patch. The stream of H2D, D2H and compute behaves asynchronously as expected.

I do accept this pending D74258 and the C++14 RFC. If they go through the version of this patch that uses C++14 is fine.

We can discuss and modify the stream number afterwards as necessary (assuming we don't find a consensus now).
This patch is strictly positive so we should work from here.

openmp/libomptarget/plugins/cuda/src/rtl.cpp

246

@ye-luo Do you experience a downside to 256 streams?

There should not be a performance problem but it should help us to be future and backwards compatible.

[ye-luo]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

246

I don't have strong evidence about performance impact. I though more streams should cost the driver a bit more to monitor and schedule workload to the hardware.

[jdoerfert]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

246

I would expect, or maybe hope, that the driver just does the modulo internally. There is no point in tracking more than the number of hardware streams so why would they. To that end they can just do hw_stream = user_stream % num_hw_streams, which would make sense because it is portable (=backwards/future compatible).

[JonChesterfield]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

182

It looks like DeviceID should be unsigned here

[tianshilei1992]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

182

Well, yes, it should be. But if you take a look at what they're used, for example at line 725, you can see the declaration is int32_t device_id.

[jdoerfert]

I'll commit this one and D74258 later.

openmp/libomptarget/plugins/cuda/src/rtl.cpp

182

we make it an unsigned here. I can do that before I commit as well.

[grokos]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

182

Well, strictly speaking, device IDs in libomptarget are signed. E.g. the default device has an ID of -1 and the host device has ID -10. On the other hand, such negative values should never reach the plugin, if that ever happens then something is buggy in the base library. So it's really up to you to either keep the signed flavor or switch to unsigned.

[jdoerfert]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

182

signed + assertion(id >= 0) ?

[tianshilei1992]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

182

It would be better we put this check in each API call.

[jdoerfert]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

182

true, we add them in (a) different commit(s) though. Can you add the check to the assert you have below?
(Nit: you can also use int(NextStreamId.size()) to save some characters)

[jdoerfert]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

263

No ifdef needed anymore, C++14 is here.

[grokos]

The patch looks good to me. One final comment: is LIBOMPTARGET_NUM_STREAMS the most appropriate name for the new env var? Because it targets the CUDA plugin specifically, should we change the name to something like LIBOMPTARGET_CUDA_NUM_STREAMS?

[tianshilei1992]

In D74145#1871257, @grokos wrote:

The patch looks good to me. One final comment: is LIBOMPTARGET_NUM_STREAMS the most appropriate name for the new env var? Because it targets the CUDA plugin specifically, should we change the name to something like LIBOMPTARGET_CUDA_NUM_STREAMS?

My idea is the concept of stream is widely used in different platforms. They might use different terminology.

[jdoerfert]

In D74145#1871269, @tianshilei1992 wrote:

In D74145#1871257, @grokos wrote:

The patch looks good to me. One final comment: is LIBOMPTARGET_NUM_STREAMS the most appropriate name for the new env var? Because it targets the CUDA plugin specifically, should we change the name to something like LIBOMPTARGET_CUDA_NUM_STREAMS?

My idea is the concept of stream is widely used in different platforms. They might use different terminology.

Let's not complicate the name and just interpret "STREAMS" as whatever the equivalent of the platform is. That way we can have a single environment variable.

[protze.joachim]

The test started to fail recently

openmp/libomptarget/test/offloading/parallel_offloading_map.c

11

With D89523, this line does not compile and the test fails. As I understand D89523, this assignment is non-comforming. Using #define instead should fix the issue

[tianshilei1992]

In D74145#2369400, @protze.joachim wrote:

The test started to fail recently

I'll fix this issue. Thanks for your notification!

openmp/libomptarget/plugins/cuda/src/rtl.cpp

 //===----RTLs/cuda/src/rtl.cpp - Target RTLs Implementation ------- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // RTL for CUDA machine
 //
 //===----------------------------------------------------------------------===//
 
+#include <atomic>
 #include <cassert>
 #include <cstddef>
 #include <cuda.h>
 #include <list>
+#include <memory>
 #include <string>
 #include <vector>
 
 #include "omptargetplugin.h"
 
 #ifndef TARGET_NAME
 #define TARGET_NAME CUDA
 #endif

 
 /// List that contains all the kernels.
 /// FIXME: we may need this to be per device and per library.
 std::list<KernelTy> KernelsList;
 
 /// Class containing all the device information.
 class RTLDeviceInfoTy {
   std::vector<std::list<FuncOrGblEntryTy>> FuncGblEntries;
+  std::vector<std::unique_ptr<std::atomic_uint>> NextStreamId;
 
 public:
   int NumberOfDevices;
   std::vector<CUmodule> Modules;
   std::vector<CUcontext> Contexts;
+  std::vector<std::vector<CUstream>> Streams;
 
   // Device properties
   std::vector<int> ThreadsPerBlock;
   std::vector<int> BlocksPerGrid;
   std::vector<int> WarpSize;
 
   // OpenMP properties
   std::vector<int> NumTeams;
   std::vector<int> NumThreads;
 
   // OpenMP Environment properties
   int EnvNumTeams;
   int EnvTeamLimit;
+  int EnvNumStreams;
 
   // OpenMP Requires Flags
   int64_t RequiresFlags;
 
   //static int EnvNumThreads;
   static const int HardTeamLimit = 1<<16; // 64k
   static const int HardThreadLimit = 1024;
   static const int DefaultNumTeams = 128;
void clearOffloadEntriesTable(int32_t device_id) {
     assert(device_id < (int32_t)FuncGblEntries.size() &&
            "Unexpected device id!");
     FuncGblEntries[device_id].emplace_back();
     FuncOrGblEntryTy &E = FuncGblEntries[device_id].back();
     E.Entries.clear();
     E.Table.EntriesBegin = E.Table.EntriesEnd = 0;
   }
 
+  // Get the next stream on a given device in a round robin manner
+  CUstream &getNextStream(const int DeviceId) {
+    const int Id = NextStreamId[DeviceId]->fetch_add(1);
+    return Streams[DeviceId][Id % EnvNumStreams];
+  }
+
   RTLDeviceInfoTy() {
 #ifdef OMPTARGET_DEBUG
     if (char *envStr = getenv("LIBOMPTARGET_DEBUG")) {
       DebugLevel = std::stoi(envStr);
     }
 #endif // OMPTARGET_DEBUG
 
     DP("Start initializing CUDA\n");
#endif // OMPTARGET_DEBUG
 
     if (NumberOfDevices == 0) {
       DP("There are no devices supporting CUDA.\n");
       return;
     }
 
     FuncGblEntries.resize(NumberOfDevices);
     Contexts.resize(NumberOfDevices);
+    Streams.resize(NumberOfDevices);
+    NextStreamId.resize(NumberOfDevices);
     ThreadsPerBlock.resize(NumberOfDevices);
     BlocksPerGrid.resize(NumberOfDevices);
     WarpSize.resize(NumberOfDevices);
     NumTeams.resize(NumberOfDevices);
     NumThreads.resize(NumberOfDevices);
 
     // Get environment variables regarding teams
     char *envStr = getenv("OMP_TEAM_LIMIT");
     if (envStr) {
       // OMP_TEAM_LIMIT has been set
       EnvTeamLimit = std::stoi(envStr);
       DP("Parsed OMP_TEAM_LIMIT=%d\n", EnvTeamLimit);
     } else {
       EnvTeamLimit = -1;
     }
     envStr = getenv("OMP_NUM_TEAMS");
     if (envStr) {
       // OMP_NUM_TEAMS has been set
       EnvNumTeams = std::stoi(envStr);
       DP("Parsed OMP_NUM_TEAMS=%d\n", EnvNumTeams);
     } else {
       EnvNumTeams = -1;
     }
 
+    // By default let's create 256 streams per device
+    EnvNumStreams = 256;
+    envStr = getenv("LIBOMPTARGET_NUM_STREAMS");
+    if (envStr) {
+      EnvNumStreams = std::stoi(envStr);
+    }
+
+    // Initialize streams for each device
+    for (std::vector<CUstream> &S : Streams) {
+      S.resize(EnvNumStreams);
+    }
+
+    // Initialize the next stream id
+    for (std::unique_ptr<std::atomic_uint> &Ptr : NextStreamId) {
+#if __cplusplus < 201402L
+      Ptr = std::unique_ptr<std::atomic_uint>(new std::atomic_uint(0));
+#else
+      Ptr = std::make_unique<std::atomic_uint>(0);
+#endif
+    }
+
     // Default state.
     RequiresFlags = OMP_REQ_UNDEFINED;
   }
 
   ~RTLDeviceInfoTy() {
     // Close modules
     for (auto &module : Modules)
       if (module) {
         CUresult err = cuModuleUnload(module);
         if (err != CUDA_SUCCESS) {
           DP("Error when unloading CUDA module\n");
           CUDA_ERR_STRING(err);
         }
       }
 
+    // Destroy streams before contexts
+    for (int I = 0; I < NumberOfDevices; ++I) {
+      CUresult err = cuCtxSetCurrent(Contexts[I]);
+      if (err != CUDA_SUCCESS) {
+        DP("Error when setting current CUDA context\n");
+        CUDA_ERR_STRING(err);
+      }
+
+      for (auto &S : Streams[I])
+        if (S) {
+          err = cuStreamDestroy(S);
+          if (err != CUDA_SUCCESS) {
+            DP("Error when destroying CUDA stream\n");
+            CUDA_ERR_STRING(err);
+          }
+        }
+    }
+
     // Destroy contexts
     for (auto &ctx : Contexts)
       if (ctx) {
         CUresult err = cuCtxDestroy(ctx);
         if (err != CUDA_SUCCESS) {
           DP("Error when destroying CUDA context\n");
           CUDA_ERR_STRING(err);
         }
int32_t __tgt_rtl_init_device(int32_t device_id) {
   err = cuCtxCreate(&DeviceInfo.Contexts[device_id], CU_CTX_SCHED_BLOCKING_SYNC,
                     cuDevice);
   if (err != CUDA_SUCCESS) {
     DP("Error when creating a CUDA context\n");
     CUDA_ERR_STRING(err);
     return OFFLOAD_FAIL;
   }
 
+  err = cuCtxSetCurrent(DeviceInfo.Contexts[device_id]);
+  if (err != CUDA_SUCCESS) {
+    DP("Error when setting current CUDA context\n");
+    CUDA_ERR_STRING(err);
+  }
+
+  for (CUstream &Stream : DeviceInfo.Streams[device_id]) {
+    err = cuStreamCreate(&Stream, CU_STREAM_NON_BLOCKING);
+    if (err != CUDA_SUCCESS) {
+      DP("Error when creating CUDA stream\n");
+      CUDA_ERR_STRING(err);
+    }
+  }
+
   // Query attributes to determine number of threads/block and blocks/grid.
   int maxGridDimX;
   err = cuDeviceGetAttribute(&maxGridDimX, CU_DEVICE_ATTRIBUTE_MAX_GRID_DIM_X,
                              cuDevice);
   if (err != CUDA_SUCCESS) {
     DP("Error getting max grid dimension, use default\n");
     DeviceInfo.BlocksPerGrid[device_id] = RTLDeviceInfoTy::DefaultNumTeams;
   } else if (maxGridDimX <= RTLDeviceInfoTy::HardTeamLimit) {
if (e->size) {
       // check for to and link variables types:
       //  (DeviceInfo.RequiresFlags & OMP_REQ_UNIFIED_SHARED_MEMORY &&
       //   (e->flags & OMP_DECLARE_TARGET_LINK ||
       //    e->flags == OMP_DECLARE_TARGET_TO))
       if (DeviceInfo.RequiresFlags & OMP_REQ_UNIFIED_SHARED_MEMORY) {
         // If unified memory is present any target link or to variables
         // can access host addresses directly. There is no longer a
         // need for device copies.
         cuMemcpyHtoD(cuptr, e->addr, sizeof(void *));
         DP("Copy linked variable host address (" DPxMOD ")"
            "to device address (" DPxMOD ")\n",
           DPxPTR(*((void**)e->addr)), DPxPTR(cuptr));
       }
 
       DeviceInfo.addOffloadEntry(device_id, entry);
 
       continue;
int32_t __tgt_rtl_data_submit(int32_t device_id, void *tgt_ptr, void *hst_ptr,
   // Set the context we are using.
   CUresult err = cuCtxSetCurrent(DeviceInfo.Contexts[device_id]);
   if (err != CUDA_SUCCESS) {
     DP("Error when setting CUDA context\n");
     CUDA_ERR_STRING(err);
     return OFFLOAD_FAIL;
   }
 
-  err = cuMemcpyHtoD((CUdeviceptr)tgt_ptr, hst_ptr, size);
+  CUstream &Stream = DeviceInfo.getNextStream(device_id);
+
+  err = cuMemcpyHtoDAsync((CUdeviceptr)tgt_ptr, hst_ptr, size, Stream);
   if (err != CUDA_SUCCESS) {
     DP("Error when copying data from host to device. Pointers: host = " DPxMOD
-       ", device = " DPxMOD ", size = %" PRId64 "\n", DPxPTR(hst_ptr),
-       DPxPTR(tgt_ptr), size);
+       ", device = " DPxMOD ", size = %" PRId64 "\n",
+       DPxPTR(hst_ptr), DPxPTR(tgt_ptr), size);
+    CUDA_ERR_STRING(err);
+    return OFFLOAD_FAIL;
+  }
+
+  err = cuStreamSynchronize(Stream);
+  if (err != CUDA_SUCCESS) {
+    DP("Error when synchronizing async data transfer from host to device. "
+       "Pointers: host = " DPxMOD ", device = " DPxMOD ", size = %" PRId64 "\n",
+       DPxPTR(hst_ptr), DPxPTR(tgt_ptr), size);
     CUDA_ERR_STRING(err);
     return OFFLOAD_FAIL;
   }
+
   return OFFLOAD_SUCCESS;
 }
 
 int32_t __tgt_rtl_data_retrieve(int32_t device_id, void *hst_ptr, void *tgt_ptr,
     int64_t size) {
   // Set the context we are using.
   CUresult err = cuCtxSetCurrent(DeviceInfo.Contexts[device_id]);
   if (err != CUDA_SUCCESS) {
     DP("Error when setting CUDA context\n");
     CUDA_ERR_STRING(err);
     return OFFLOAD_FAIL;
   }
 
-  err = cuMemcpyDtoH(hst_ptr, (CUdeviceptr)tgt_ptr, size);
+  CUstream &Stream = DeviceInfo.getNextStream(device_id);
+
+  err = cuMemcpyDtoHAsync(hst_ptr, (CUdeviceptr)tgt_ptr, size, Stream);
   if (err != CUDA_SUCCESS) {
     DP("Error when copying data from device to host. Pointers: host = " DPxMOD
-        ", device = " DPxMOD ", size = %" PRId64 "\n", DPxPTR(hst_ptr),
-        DPxPTR(tgt_ptr), size);
+       ", device = " DPxMOD ", size = %" PRId64 "\n",
+       DPxPTR(hst_ptr), DPxPTR(tgt_ptr), size);
     CUDA_ERR_STRING(err);
     return OFFLOAD_FAIL;
   }
+
+  err = cuStreamSynchronize(Stream);
+  if (err != CUDA_SUCCESS) {
+    DP("Error when synchronizing async data transfer from device to host. "
+       "Pointers: host = " DPxMOD ", device = " DPxMOD ", size = %" PRId64 "\n",
+       DPxPTR(hst_ptr), DPxPTR(tgt_ptr), size);
+    CUDA_ERR_STRING(err);
+    return OFFLOAD_FAIL;
+  }
+
   return OFFLOAD_SUCCESS;
 }
 
 int32_t __tgt_rtl_data_delete(int32_t device_id, void *tgt_ptr) {
   // Set the context we are using.
   CUresult err = cuCtxSetCurrent(DeviceInfo.Contexts[device_id]);
   if (err != CUDA_SUCCESS) {
     DP("Error when setting CUDA context\n");
if (team_num <= 0) {
     cudaBlocksPerGrid = team_num;
     DP("Using requested number of teams %d\n", team_num);
   }
 
   // Run on the device.
   DP("Launch kernel with %d blocks and %d threads\n", cudaBlocksPerGrid,
      cudaThreadsPerBlock);
 
+  CUstream &Stream = DeviceInfo.getNextStream(device_id);
+
   err = cuLaunchKernel(KernelInfo->Func, cudaBlocksPerGrid, 1, 1,
-      cudaThreadsPerBlock, 1, 1, 0 /*bytes of shared memory*/, 0, &args[0], 0);
+                       cudaThreadsPerBlock, 1, 1, 0 /*bytes of shared memory*/,
+                       Stream, &args[0], 0);
   if (err != CUDA_SUCCESS) {
     DP("Device kernel launch failed!\n");
     CUDA_ERR_STRING(err);
     return OFFLOAD_FAIL;
   }
 
   DP("Launch of entry point at " DPxMOD " successful!\n",
       DPxPTR(tgt_entry_ptr));
 
-  CUresult sync_err = cuCtxSynchronize();
+  CUresult sync_err = cuStreamSynchronize(Stream);
   if (sync_err != CUDA_SUCCESS) {
     DP("Kernel execution error at " DPxMOD "!\n", DPxPTR(tgt_entry_ptr));
     CUDA_ERR_STRING(sync_err);
     return OFFLOAD_FAIL;
   } else {
     DP("Kernel execution at " DPxMOD " successful!\n", DPxPTR(tgt_entry_ptr));
   }
 

openmp/libomptarget/test/offloading/parallel_offloading_map.c

+// RUN: %libomptarget-compilexx-run-and-check-aarch64-unknown-linux-gnu
+// RUN: %libomptarget-compilexx-run-and-check-powerpc64-ibm-linux-gnu
+// RUN: %libomptarget-compilexx-run-and-check-powerpc64le-ibm-linux-gnu
+// RUN: %libomptarget-compilexx-run-and-check-x86_64-pc-linux-gnu
+#include <assert.h>
+#include <stdio.h>
+
+int main(int argc, char *argv[]) {
+  const int num_threads = 64, N = 128;
+  int array[num_threads] = {0};
+
+#pragma omp parallel for
+  for (int i = 0; i < num_threads; ++i) {
+    int tmp[N];
+
+    for (int j = 0; j < N; ++j) {
+      tmp[j] = i;
+    }
+
+#pragma omp target teams distribute parallel for map(tofrom : tmp)
+    for (int j = 0; j < N; ++j) {
+      tmp[j] += j;
+    }
+
+    for (int j = 0; j < N; ++j) {
+      array[i] += tmp[j];
+    }
+  }
+
+  // Verify
+  for (int i = 0; i < num_threads; ++i) {
+    const int ref = (0 + N - 1) * N / 2 + i * N;
+    assert(array[i] == ref);
+  }
+
+  printf("PASS\n");
+
+  return 0;
+}
+
+// CHECK: PASS