xref: /aosp_15_r20/external/pytorch/torch/csrc/api/include/torch/xpu.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #pragma once
2 
3 #include <torch/csrc/Export.h>
4 
5 #include <cstddef>
6 #include <cstdint>
7 
8 namespace torch::xpu {
9 
10 /// Returns the number of XPU devices available.
11 size_t TORCH_API device_count();
12 
13 /// Returns true if at least one XPU device is available.
14 bool TORCH_API is_available();
15 
16 /// Sets the seed for the current GPU.
17 void TORCH_API manual_seed(uint64_t seed);
18 
19 /// Sets the seed for all available GPUs.
20 void TORCH_API manual_seed_all(uint64_t seed);
21 
22 /// Waits for all kernels in all streams on a XPU device to complete.
23 void TORCH_API synchronize(int64_t device_index);
24 
25 } // namespace torch::xpu
26