XPUFunctions.h 1.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445
  1. #pragma once
  2. #include <c10/core/Device.h>
  3. #include <c10/xpu/XPUDeviceProp.h>
  4. #include <c10/xpu/XPUMacros.h>
  5. // The naming convention used here matches the naming convention of torch.xpu
  6. namespace c10::xpu {
  7. // Log a warning only once if no devices are detected.
  8. C10_XPU_API DeviceIndex device_count();
  9. // Throws an error if no devices are detected.
  10. C10_XPU_API DeviceIndex device_count_ensure_non_zero();
  11. C10_XPU_API DeviceIndex current_device();
  12. C10_XPU_API void set_device(DeviceIndex device);
  13. C10_XPU_API DeviceIndex exchange_device(DeviceIndex device);
  14. C10_XPU_API DeviceIndex maybe_exchange_device(DeviceIndex to_device);
  15. C10_XPU_API sycl::device& get_raw_device(DeviceIndex device);
  16. C10_XPU_API sycl::context& get_device_context();
  17. C10_XPU_API void get_device_properties(
  18. DeviceProp* device_prop,
  19. DeviceIndex device);
  20. C10_XPU_API DeviceIndex get_device_idx_from_pointer(void* ptr);
  21. static inline void check_device_index(DeviceIndex device_index) {
  22. TORCH_CHECK(
  23. device_index >= 0 && device_index < c10::xpu::device_count(),
  24. "The device index is out of range. It must be in [0, ",
  25. static_cast<int>(c10::xpu::device_count()),
  26. "), but got ",
  27. static_cast<int>(device_index),
  28. ".");
  29. }
  30. } // namespace c10::xpu