PeerToPeerAccess.h 1.4 KB

12345678910111213141516171819202122232425262728293031323334353637383940
  1. #if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
  2. #pragma once
  3. #include <c10/core/Device.h>
  4. #include <c10/cuda/CUDAMacros.h>
  5. #include <c10/macros/Macros.h>
  6. #include <cstdint>
  7. namespace c10::cuda {
  8. namespace detail {
  9. /// Initialize the peer-to-peer and fabric access caches.
  10. /// Must be called before any calls to get_p2p_access or get_fabric_access.
  11. /// @param num_devices The number of CUDA devices in the system.
  12. C10_CUDA_API void init_p2p_access_cache(int64_t num_devices);
  13. } // namespace detail
  14. /// Query if peer-to-peer access is available between two devices.
  15. /// @param source_dev The source device index.
  16. /// @param dest_dev The destination device index.
  17. /// @return true if P2P access is available, false otherwise.
  18. C10_CUDA_API bool get_p2p_access(
  19. c10::DeviceIndex source_dev,
  20. c10::DeviceIndex dest_dev);
  21. /// Query if GPU fabric (high-speed interconnect like NVLink/NVSwitch) is
  22. /// available for a device. This checks both hardware support and the ability
  23. /// to allocate/export/import memory with fabric handles.
  24. /// @param device The device index to check.
  25. /// @return true if fabric access is available, false otherwise.
  26. C10_CUDA_API bool get_fabric_access(c10::DeviceIndex device);
  27. } // namespace c10::cuda
  28. #else
  29. #error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
  30. #endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)