xref: /aosp_15_r20/external/pytorch/test/cpp/c10d/CMakeLists.txt (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1if(USE_CUDA)
2  add_library(c10d_cuda_test CUDATest.cu)
3  target_include_directories(c10d_cuda_test PRIVATE $<BUILD_INTERFACE:${TORCH_SRC_DIR}/csrc/distributed>)
4  target_link_libraries(c10d_cuda_test torch_cuda)
5  add_dependencies(c10d_cuda_test torch_cuda)
6endif()
7
8function(c10d_add_test test_src)
9  get_filename_component(test_name ${test_src} NAME_WE)
10  add_executable(${test_name} "${test_src}")
11  target_include_directories(${test_name} PRIVATE $<BUILD_INTERFACE:${TORCH_SRC_DIR}/csrc/distributed>)
12  target_link_libraries(${test_name} ${ARGN})
13  if(NOT WIN32)
14    target_link_libraries(${test_name} pthread)
15  endif()
16  add_test(NAME ${test_name} COMMAND $<TARGET_FILE:${test_name}>)
17endfunction()
18
19c10d_add_test(BackoffTest.cpp torch_cpu gtest_main)
20c10d_add_test(FileStoreTest.cpp torch_cpu gtest_main)
21c10d_add_test(TCPStoreTest.cpp torch_cpu gtest_main)
22if(INSTALL_TEST)
23  install(TARGETS FileStoreTest DESTINATION bin)
24  install(TARGETS TCPStoreTest DESTINATION bin)
25endif()
26if(NOT WIN32)
27  c10d_add_test(HashStoreTest.cpp torch_cpu gtest_main)
28  if(INSTALL_TEST)
29    install(TARGETS HashStoreTest DESTINATION bin)
30  endif()
31endif()
32
33if(USE_CUDA)
34  if(USE_GLOO AND USE_C10D_GLOO)
35    c10d_add_test(ProcessGroupGlooTest.cpp torch_cpu c10d_cuda_test gtest_main)
36    if(INSTALL_TEST)
37      install(TARGETS ProcessGroupGlooTest DESTINATION bin)
38    endif()
39    c10d_add_test(ProcessGroupGlooAsyncTest.cpp torch_cpu c10d_cuda_test gtest_main)
40  endif()
41  if(USE_NCCL AND USE_C10D_NCCL)
42    # NCCL is a private dependency of libtorch, but the tests include some
43    # private headers of libtorch, which in turn include NCCL. As a hacky
44    # alternative to making NCCL a public dependency of libtorch, we make it
45    # a private dependency of the tests as well.
46    c10d_add_test(
47      ProcessGroupNCCLTest.cpp
48      torch_cpu c10d_cuda_test gtest_main __caffe2_nccl)
49    c10d_add_test(
50      ProcessGroupNCCLErrorsTest.cpp
51      torch_cpu c10d_cuda_test gtest_main __caffe2_nccl)
52    if(INSTALL_TEST)
53      install(TARGETS ProcessGroupNCCLTest DESTINATION bin)
54      install(TARGETS ProcessGroupNCCLErrorsTest DESTINATION bin)
55      install(TARGETS c10d_cuda_test DESTINATION lib)
56    endif()
57  endif()
58  if(USE_UCC AND USE_C10D_UCC)
59    # UCC is a private dependency of libtorch, but the tests include some
60    # private headers of libtorch, which in turn include UCC. As a hacky
61    # alternative to making UCC a public dependency of libtorch, we make it
62    # a private dependency of the tests as well.
63    c10d_add_test(
64      ProcessGroupUCCTest.cpp
65      torch_cpu c10d_cuda_test gtest_main __caffe2_ucc)
66    if(INSTALL_TEST)
67      install(TARGETS ProcessGroupUCCTest DESTINATION bin)
68      install(TARGETS c10d_cuda_test DESTINATION lib)
69    endif()
70  endif()
71else()
72  if(USE_GLOO AND USE_C10D_GLOO)
73    c10d_add_test(ProcessGroupGlooTest.cpp torch_cpu gtest_main)
74  endif()
75endif()
76
77if(USE_MPI AND USE_C10D_MPI)
78  add_definitions(-DMPIEXEC=${MPIEXEC})
79  # MPI is a private dependency of libtorch, but the tests include some
80  # private headers of libtorch, which in turn include MPI. As a hacky
81  # alternative to making MPI a public dependency of libtorch, we make it
82  # a private dependency of the tests as well.
83  c10d_add_test(ProcessGroupMPITest.cpp torch_cpu MPI::MPI_CXX)
84  if(INSTALL_TEST)
85    install(TARGETS ProcessGroupMPITest DESTINATION bin)
86  endif()
87endif()
88
89if(LINUX AND USE_GLOO AND USE_C10D_GLOO)
90  add_executable(example_allreduce example/allreduce.cpp)
91  target_include_directories(example_allreduce PRIVATE $<BUILD_INTERFACE:${TORCH_SRC_DIR}/csrc/distributed>)
92  target_link_libraries(example_allreduce pthread torch_cpu)
93  if(USE_CUDA)
94    target_link_libraries(example_allreduce torch_cuda)
95  endif()
96endif()
97