1 | #include "CUDATest.hpp" |
2 | #include <ATen/cuda/Exceptions.h> |
3 | |
4 | namespace c10d { |
5 | namespace test { |
6 | |
7 | namespace { |
8 | __global__ void waitClocks(const uint64_t count) { |
9 | clock_t start = clock64(); |
10 | clock_t offset = 0; |
11 | while (offset < count) { |
12 | offset = clock() - start; |
13 | } |
14 | } |
15 | |
16 | } // namespace |
17 | |
18 | void cudaSleep(at::cuda::CUDAStream& stream, uint64_t clocks) { |
19 | waitClocks<<<1, 1, 0, stream.stream()>>>(clocks); |
20 | C10_CUDA_KERNEL_LAUNCH_CHECK(); |
21 | } |
22 | |
23 | int cudaNumDevices() { |
24 | int n = 0; |
25 | C10_CUDA_CHECK_WARN(cudaGetDeviceCount(&n)); |
26 | return n; |
27 | } |
28 | |
29 | } // namespace test |
30 | } // namespace c10d |
31 | |