Some follow-up work on p2p

NVIDIA · Oct 28, 2024 · a4be366 · a4be366
1 parent 8f1691d
commit a4be366
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 8 deletions.
diff --git a/cudax/include/cuda/experimental/__device/all_devices.cuh b/cudax/include/cuda/experimental/__device/all_devices.cuh
@@ -209,14 +209,10 @@ _CCCL_NODISCARD inline ::std::vector<device_ref> device_ref::get_peers() const
     // group of peers is needed (for cases other than peer access control)
     if (__other_dev != *this)
     {
-      int __can_access;
-      _CCCL_TRY_CUDA_API(
-        ::cudaDeviceCanAccessPeer,
-        "Could not query if device can be peer accessed",
-        &__can_access,
-        get(),
-        __other_dev.get());
-      if (__can_access)
+      // While in almost all practical applications peer access should be symmetrical,
+      // it is possible to build a system with one directional peer access, check
+      // both ways here just to be safe
+      if (can_peer_access_to(__other_dev) && __other_dev.can_peer_access_to(*this))
       {
         __result.push_back(__other_dev);
       }

diff --git a/cudax/include/cuda/experimental/__device/device_ref.cuh b/cudax/include/cuda/experimental/__device/device_ref.cuh
@@ -119,6 +119,27 @@ public:
     return __name;
   }
 
+  // TODO not sure if this is the best name :(
+  //! @brief Queries if its possible for this device to directly access specified device's memory.
+  //!
+  //! If this function returns true, device supplied to this call can be passed into enable_peer_access
+  //! on memory resource or pool that manages memory on this device. It will make allocations from that
+  //! pool accessible by this device.
+  //!
+  //! @param __other_dev Device to query the peer access
+  //! @return true if its possible for this device to access the specified device's memory
+  bool can_peer_access_to(device_ref __other_dev) const
+  {
+    int __can_access;
+    _CCCL_TRY_CUDA_API(
+      ::cudaDeviceCanAccessPeer,
+      "Could not query if device can be peer accessed",
+      &__can_access,
+      get(),
+      __other_dev.get());
+    return __can_access;
+  }
+
   const arch_traits_t& arch_traits() const;
 
   // TODO this might return some more complex type in the future

diff --git a/cudax/test/device/device_smoke.cu b/cudax/test/device/device_smoke.cu
@@ -315,6 +315,13 @@ TEST_CASE("global devices vector", "[device]")
     CUDAX_REQUIRE(cudax::devices.size() - 1 == (*std::prev(cudax::devices.end())).get());
     CUDAX_REQUIRE(cudax::devices.size() - 1 == std::prev(cudax::devices.end())->get());
     CUDAX_REQUIRE(cudax::devices.size() - 1 == cudax::devices.end()[-1].get());
+
+    auto peers = cudax::devices[0].get_peers();
+    for (auto peer : peers)
+    {
+      CUDAX_REQUIRE(cudax::devices[0].can_peer_access_to(peer))
+      CUDAX_REQUIRE(peer.can_peer_access_to(cudax::devices[0]));
+    }
   }
 
   try