Update

guangyey · guangyey · commit e927ca92c7fa · 2026-03-09T13:18:49.000Z
[ghstack-poisoned]
diff --git a/aten/src/ATen/xpu/PeerToPeerAccess.h b/aten/src/ATen/xpu/PeerToPeerAccess.h
@@ -1,15 +1,24 @@
 #pragma once
 
-#include <c10/core/Device.h>
-#include <c10/macros/Macros.h>
+#include <c10/xpu/PeerToPeerAccess.h>
+
+#include <ATen/Context.h>
 
 namespace at::xpu {
 namespace detail {
-void init_p2p_access_cache(c10::DeviceIndex num_devices);
+// Initialize the peer-to-peer access cache for XPU devices.
+inline void init_p2p_access_cache(c10::DeviceIndex num_devices) {
+  c10::xpu::detail::init_p2p_access_cache(num_devices);
+}
 } // namespace detail
 
-TORCH_XPU_API bool get_p2p_access(
+// Query if peer-to-peer access is available between two devices.
+// This wrapper ensures XPU lazy initialization before forwarding to c10.
+inline bool get_p2p_access(
     c10::DeviceIndex dev,
-    c10::DeviceIndex dev_to_access);
+    c10::DeviceIndex dev_to_access) {
+  at::globalContext().lazyInitDevice(c10::DeviceType::XPU);
+  return c10::xpu::get_p2p_access(dev, dev_to_access);
+}
 
 } // namespace at::xpu
diff --git a/c10/xpu/CMakeLists.txt b/c10/xpu/CMakeLists.txt
@@ -15,12 +15,14 @@ configure_file(
     ${CMAKE_BINARY_DIR}/c10/xpu/impl/xpu_cmake_macros.h)
 
 set(C10_XPU_SRCS
+    PeerToPeerAccess.cpp
     XPUCachingAllocator.cpp
     XPUFunctions.cpp
     XPUStream.cpp
     impl/XPUGuardImpl.cpp
 )
 set(C10_XPU_HEADERS
+    PeerToPeerAccess.h
     XPUCachingAllocator.h
     XPUDeviceProp.h
     XPUException.h
diff --git a/c10/xpu/PeerToPeerAccess.cpp b/c10/xpu/PeerToPeerAccess.cpp
@@ -1,11 +1,8 @@
-#include <ATen/xpu/PeerToPeerAccess.h>
-#include <ATen/xpu/XPUContext.h>
-
-#include <c10/util/Exception.h>
 #include <c10/util/irange.h>
+#include <c10/xpu/PeerToPeerAccess.h>
 #include <c10/xpu/XPUCachingAllocator.h>
 
-namespace at::xpu {
+namespace c10::xpu {
 
 // p2pAccessEnabled_ is a flattened 2D matrix of size [num_devices x
 // num_devices].
@@ -35,8 +32,6 @@ void init_p2p_access_cache(c10::DeviceIndex num_devices) {
 } // namespace detail
 
 bool get_p2p_access(c10::DeviceIndex dev, c10::DeviceIndex dev_to_access) {
-  at::globalContext().lazyInitDevice(c10::DeviceType::XPU);
-
   check_device_index(dev);
   check_device_index(dev_to_access);
 
@@ -60,4 +55,4 @@ bool get_p2p_access(c10::DeviceIndex dev, c10::DeviceIndex dev_to_access) {
   return static_cast<bool>(cache);
 }
 
-} // namespace at::xpu
+} // namespace c10::xpu
diff --git a/c10/xpu/PeerToPeerAccess.h b/c10/xpu/PeerToPeerAccess.h
@@ -0,0 +1,18 @@
+#pragma once
+
+#include <c10/core/Device.h>
+#include <c10/macros/Macros.h>
+#include <c10/xpu/XPUMacros.h>
+
+namespace c10::xpu {
+namespace detail {
+// Initialize the peer-to-peer access cache for XPU devices.
+C10_XPU_API void init_p2p_access_cache(c10::DeviceIndex num_devices);
+} // namespace detail
+
+// Query if peer-to-peer access is available between two devices.
+C10_XPU_API bool get_p2p_access(
+    c10::DeviceIndex dev,
+    c10::DeviceIndex dev_to_access);
+
+} // namespace c10::xpu
