[OpenMP][CUDA] Fix the issue that P2P memcpy doesn't work

[tianshilei1992]

Summary

This patch fixes the issue that P2P memcpy doesn't work. The root cause is we didn't set current context when calling the API function. In addition, a matrix to track the states of each pair of devices is also added such that we only need to query and configure the device once.

[tianshilei1992]

add blank line

[tianshilei1992]

add set context before d2d in case of any failure when setting destination device

[jdoerfert]

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

1064

Do we want to do this every time? Is this costly the second time around? Should we keep a flag?

[tianshilei1992]

add a state matrix

[tianshilei1992]

add lock

[tianshilei1992]

lazy initialization

[tianshilei1992]

remove unnecessary changes

[tianshilei1992]

remove intermediate state

[tianshilei1992]

ping

[tianshilei1992]

ping

[tianshilei1992]

rebase and ping again

[tianshilei1992]

new week, new ping :-)

[jdoerfert]

LG

[HaohaiWen]

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

1079–1080

This default label will cause warning/error...
error: default label in switch which covers all enumeration values [-Werror,-Wcovered-switch-default

[tianshilei1992]

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

1079–1080

Thanks for pointing that out. I'll push a fix pretty soon.

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

if (SrcDevId != DstDevId) {
       if (Err != CUDA_SUCCESS) {
         REPORT("Error returned from cuCtxEnablePeerAccess. src = %" PRId32
                ", dst = %" PRId32 "\n",
                SrcDevId, DstDevId);
         CUDA_ERR_STRING(Err);
         return memcpyDtoD(SrcPtr, DstPtr, Size, Stream);
       }
 
+      // Switch to destination context to enable peer access.
+      if (setContext(DstDevId) != OFFLOAD_SUCCESS)
+        return OFFLOAD_FAIL;
+
+      Err = cuCtxEnablePeerAccess(DeviceData[SrcDevId].Context, 0);
+      if (Err != CUDA_SUCCESS) {
+        REPORT("Error returned from cuCtxEnablePeerAccess. src = %" PRId32
+               ", dst = %" PRId32 "\n",
+               SrcDevId, DstDevId);
+        CUDA_ERR_STRING(Err);
+        // Switch back to source context to issue D2D.
+        if (setContext(SrcDevId) != OFFLOAD_SUCCESS)
+          return OFFLOAD_FAIL;
+        return memcpyDtoD(SrcPtr, DstPtr, Size, Stream);
+      }
+
+      // Switch back to source context to issue memcpy.
+      if (setContext(SrcDevId) != OFFLOAD_SUCCESS)
+        return OFFLOAD_FAIL;
+
       Err = cuMemcpyPeerAsync((CUdeviceptr)DstPtr, DeviceData[DstDevId].Context,
                               (CUdeviceptr)SrcPtr, DeviceData[SrcDevId].Context,
                               Size, Stream);
       if (Err == CUDA_SUCCESS)
         return OFFLOAD_SUCCESS;
 
       DP("Error returned from cuMemcpyPeerAsync. src_ptr = " DPxMOD
          ", src_id =%" PRId32 ", dst_ptr = " DPxMOD ", dst_id =%" PRId32 "\n",
          DPxPTR(SrcPtr), SrcDevId, DPxPTR(DstPtr), DstDevId);
       CUDA_ERR_STRING(Err);
     }
 
     return memcpyDtoD(SrcPtr, DstPtr, Size, Stream);
   }
 
   int dataDelete(const int DeviceId, void *TgtPtr) {
     if (UseMemoryManager)
       return MemoryManagers[DeviceId]->free(TgtPtr);
 
if (!DeviceInfo->Device) {
       } else {
         cuGetErrorString(Err, ErrStr);
         return OFFLOAD_FAIL;
       }
     }
     return OFFLOAD_SUCCESS;
   }
 
-  int setContext(int DeviceId) {
+  int setContext(int DeviceId) const {
     assert(InitializedFlags[DeviceId] && "Device is not initialized");
 
     CUresult Err = cuCtxSetCurrent(DeviceData[DeviceId].Context);
     if (!checkResult(Err, "error returned from cuCtxSetCurrent"))
       return OFFLOAD_FAIL;
 
     return OFFLOAD_SUCCESS;
   }

 
 int32_t __tgt_rtl_data_exchange_async(int32_t src_dev_id, void *src_ptr,
                                       int dst_dev_id, void *dst_ptr,
                                       int64_t size,
                                       __tgt_async_info *AsyncInfo) {
   assert(DeviceRTL.isValidDeviceId(src_dev_id) && "src_dev_id is invalid");
   assert(DeviceRTL.isValidDeviceId(dst_dev_id) && "dst_dev_id is invalid");
   assert(AsyncInfo && "AsyncInfo is nullptr");
-  // NOTE: We don't need to set context for data exchange as the device contexts
-  // are passed to CUDA function directly.
+
+  if (DeviceRTL.setContext(src_dev_id) != OFFLOAD_SUCCESS)
+    return OFFLOAD_FAIL;
+
   return DeviceRTL.dataExchange(src_dev_id, src_ptr, dst_dev_id, dst_ptr, size,
                                 AsyncInfo);
 }
 
 int32_t __tgt_rtl_data_exchange(int32_t src_dev_id, void *src_ptr,
                                 int32_t dst_dev_id, void *dst_ptr,
                                 int64_t size) {
   assert(DeviceRTL.isValidDeviceId(src_dev_id) && "src_dev_id is invalid");