1#include "CUDATest.hpp"
2#include <ATen/cuda/Exceptions.h>
3
4namespace c10d {
5namespace test {
6
7namespace {
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
18void cudaSleep(at::cuda::CUDAStream& stream, uint64_t clocks) {
19 waitClocks<<<1, 1, 0, stream.stream()>>>(clocks);
20 C10_CUDA_KERNEL_LAUNCH_CHECK();
21}
22
23int cudaNumDevices() {
24 int n = 0;
25 C10_CUDA_CHECK_WARN(cudaGetDeviceCount(&n));
26 return n;
27}
28
29} // namespace test
30} // namespace c10d
31