xref: /aosp_15_r20/external/pytorch/test/cpp/c10d/CUDATest.hpp (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
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