PeerToPeerAccess.h 1.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142
  1. #if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
  2. #pragma once
  3. #include <c10/core/Device.h>
  4. #include <c10/cuda/PeerToPeerAccess.h>
  5. #include <c10/macros/Macros.h>
  6. #include <cstdint>
  7. namespace at::cuda {
  8. namespace detail {
  9. /// Initialize the peer-to-peer and fabric access caches.
  10. /// Forwards to c10::cuda::detail::init_p2p_access_cache.
  11. /// @param num_devices The number of CUDA devices in the system.
  12. inline void init_p2p_access_cache(int64_t num_devices) {
  13. c10::cuda::detail::init_p2p_access_cache(num_devices);
  14. }
  15. } // namespace detail
  16. /// Query if peer-to-peer access is available between two devices.
  17. /// This wrapper ensures CUDA lazy initialization before forwarding to c10.
  18. /// @param source_dev The source device index.
  19. /// @param dest_dev The destination device index.
  20. /// @return true if P2P access is available, false otherwise.
  21. TORCH_CUDA_CPP_API bool get_p2p_access(
  22. c10::DeviceIndex source_dev,
  23. c10::DeviceIndex dest_dev);
  24. /// Query if GPU fabric (high-speed interconnect) is available for a device.
  25. /// This wrapper ensures CUDA lazy initialization before forwarding to c10.
  26. /// @param device The device index to check.
  27. /// @return true if fabric access is available, false otherwise.
  28. TORCH_CUDA_CPP_API bool get_fabric_access(c10::DeviceIndex device);
  29. } // namespace at::cuda
  30. #else
  31. #error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
  32. #endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)