[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

1113

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

1098–1099

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

1098–1099

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

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

class DeviceRTLTy {
 
   std::vector<DeviceDataTy> DeviceData;
   std::vector<std::vector<CUmodule>> Modules;
 
   /// Vector of flags indicating the initalization status of all associated
   /// devices.
   std::vector<bool> InitializedFlags;
 
+  enum class PeerAccessState : uint8_t { Unkown, Enabled, Yes, No };
+  std::vector<std::vector<PeerAccessState>> PeerAccessMatrix;
+  std::mutex PeerAccessMatrixLock;
+
   /// A class responsible for interacting with device native runtime library to
   /// allocate and free memory.
   class CUDADeviceAllocatorTy : public DeviceAllocatorTy {
     std::unordered_map<void *, TargetAllocTy> HostPinnedAllocs;
 
   public:
     void *allocate(size_t Size, void *, TargetAllocTy Kind) override {
       if (Size == 0)
if (NumberOfDevices == 0) {
       DP("There are no devices supporting CUDA.\n");
       return;
     }
 
     DeviceData.resize(NumberOfDevices);
     Modules.resize(NumberOfDevices);
     StreamPool.resize(NumberOfDevices);
     EventPool.resize(NumberOfDevices);
+    PeerAccessMatrix.resize(NumberOfDevices);
+    for (auto &V : PeerAccessMatrix)
+      V.resize(NumberOfDevices, PeerAccessState::Unkown);
 
     // Get environment variables regarding teams
     if (const char *EnvStr = getenv("OMP_TEAM_LIMIT")) {
       // OMP_TEAM_LIMIT has been set
       EnvTeamLimit = std::stoi(EnvStr);
       DP("Parsed OMP_TEAM_LIMIT=%d\n", EnvTeamLimit);
     }
     if (const char *EnvStr = getenv("OMP_TEAMS_THREAD_LIMIT")) {
if (Err != CUDA_SUCCESS) {
       CUDA_ERR_STRING(Err);
       return OFFLOAD_FAIL;
     }
 
     return OFFLOAD_SUCCESS;
   }
 
   int dataExchange(int SrcDevId, const void *SrcPtr, int DstDevId, void *DstPtr,
-                   int64_t Size, __tgt_async_info *AsyncInfo) const {
+                   int64_t Size, __tgt_async_info *AsyncInfo) {
     assert(AsyncInfo && "AsyncInfo is nullptr");
 
     CUresult Err;
     CUstream Stream = getStream(SrcDevId, AsyncInfo);
 
     // If they are two devices, we try peer to peer copy first
     if (SrcDevId != DstDevId) {
+      std::lock_guard<std::mutex> LG(PeerAccessMatrixLock);
+
+      switch (PeerAccessMatrix[SrcDevId][DstDevId]) {
+      case PeerAccessState::No: {
+        REPORT("Peer access from %" PRId32 " to %" PRId32
+               " is not supported. Fall back to D2D memcpy.\n",
+               SrcDevId, DstDevId);
+        return memcpyDtoD(SrcPtr, DstPtr, Size, Stream);
+      }
+      case PeerAccessState::Unkown: {
         int CanAccessPeer = 0;
         Err = cuDeviceCanAccessPeer(&CanAccessPeer, SrcDevId, DstDevId);
         if (Err != CUDA_SUCCESS) {
           REPORT("Error returned from cuDeviceCanAccessPeer. src = %" PRId32
-               ", dst = %" PRId32 "\n",
+                 ", dst = %" PRId32 ". Fall back to D2D memcpy.\n",
                  SrcDevId, DstDevId);
           CUDA_ERR_STRING(Err);
+          PeerAccessMatrix[SrcDevId][DstDevId] = PeerAccessState::No;
           return memcpyDtoD(SrcPtr, DstPtr, Size, Stream);
         }
 
         if (!CanAccessPeer) {
-        DP("P2P memcpy not supported so fall back to D2D memcpy");
+          REPORT("P2P access from %d to %d is not supported. Fall back to D2D "
+                 "memcpy.\n",
+                 SrcDevId, DstDevId);
+          PeerAccessMatrix[SrcDevId][DstDevId] = PeerAccessState::No;
           return memcpyDtoD(SrcPtr, DstPtr, Size, Stream);
         }
 
+        LLVM_FALLTHROUGH;
+      }
+      case PeerAccessState::Yes: {
         Err = cuCtxEnablePeerAccess(DeviceData[DstDevId].Context, 0);
         if (Err != CUDA_SUCCESS) {
           REPORT("Error returned from cuCtxEnablePeerAccess. src = %" PRId32
-               ", dst = %" PRId32 "\n",
+                 ", dst = %" PRId32 ". Fall back to D2D memcpy.\n",
                  SrcDevId, DstDevId);
           CUDA_ERR_STRING(Err);
+          PeerAccessMatrix[SrcDevId][DstDevId] = PeerAccessState::No;
           return memcpyDtoD(SrcPtr, DstPtr, Size, Stream);
         }
 
-      Err = cuMemcpyPeerAsync((CUdeviceptr)DstPtr, DeviceData[DstDevId].Context,
-                              (CUdeviceptr)SrcPtr, DeviceData[SrcDevId].Context,
-                              Size, Stream);
+        PeerAccessMatrix[SrcDevId][DstDevId] = PeerAccessState::Enabled;
+
+        LLVM_FALLTHROUGH;
+      }
+      case PeerAccessState::Enabled: {
+        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",
+           ", src_id =%" PRId32 ", dst_ptr = " DPxMOD ", dst_id =%" PRId32
+           ". Fall back to D2D memcpy.\n",
            DPxPTR(SrcPtr), SrcDevId, DPxPTR(DstPtr), DstDevId);
         CUDA_ERR_STRING(Err);
+
+        return memcpyDtoD(SrcPtr, DstPtr, Size, Stream);
+      }
+      default:
+        REPORT("Unknown PeerAccessState %d.\n",
+               int(PeerAccessMatrix[SrcDevId][DstDevId]));
+        return OFFLOAD_FAIL;
+      }
     }
 
     return memcpyDtoD(SrcPtr, DstPtr, Size, Stream);
   }
 
   int dataDelete(const int DeviceId, void *TgtPtr) {
     if (UseMemoryManager)
       return MemoryManagers[DeviceId]->free(TgtPtr);
 
     return DeviceAllocators[DeviceId].free(TgtPtr);
   }
 
   int runTargetTeamRegion(const int DeviceId, void *TgtEntryPtr, void **TgtArgs,
                           ptrdiff_t *TgtOffsets, const int ArgNum,
                           const int TeamNum, const int ThreadLimit,
                           const unsigned int LoopTripCount,
                           __tgt_async_info *AsyncInfo) const {
     // All args are references.
     std::vector<void *> Args(ArgNum);
     std::vector<void *> Ptrs(ArgNum);
 
     for (int I = 0; I < ArgNum; ++I) {
       Ptrs[I] = (void *)((intptr_t)TgtArgs[I] + TgtOffsets[I]);

 
 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");