1 | #pragma once |
---|---|
2 | |
3 | #include <c10/cuda/CUDAStream.h> |
4 | |
5 | namespace c10d { |
6 | namespace test { |
7 | |
8 | #ifdef _WIN32 |
9 | #define EXPORT_TEST_API __declspec(dllexport) |
10 | #else |
11 | #define EXPORT_TEST_API |
12 | #endif |
13 | |
14 | EXPORT_TEST_API void cudaSleep(at::cuda::CUDAStream& stream, uint64_t clocks); |
15 | |
16 | EXPORT_TEST_API int cudaNumDevices(); |
17 | |
18 | } // namespace test |
19 | } // namespace c10d |
20 |