file(GLOB OP_SRC print_summary.cc comm.cc tabulate_multi_device.cc)

add_library(deepmd_op_pt MODULE ${OP_SRC})
# link: libdeepmd libtorch
target_link_libraries(deepmd_op_pt PRIVATE ${TORCH_LIBRARIES})
if(${OP_CXX_ABI_PT} EQUAL ${OP_CXX_ABI})
  target_link_libraries(deepmd_op_pt PRIVATE ${LIB_DEEPMD})
else()
  target_link_libraries(deepmd_op_pt PRIVATE ${LIB_DEEPMD}_compat_cxxabi)
endif()
remove_definitions(-D_GLIBCXX_USE_CXX11_ABI=${OP_CXX_ABI})
target_compile_definitions(
  deepmd_op_pt
  PUBLIC "$<$<COMPILE_LANGUAGE:CXX>:_GLIBCXX_USE_CXX11_ABI=${OP_CXX_ABI_PT}>")
if(APPLE)
  set_target_properties(deepmd_op_pt PROPERTIES INSTALL_RPATH "@loader_path")
else()
  set_target_properties(deepmd_op_pt PROPERTIES INSTALL_RPATH "$ORIGIN")
endif()

find_package(MPI)
if(MPI_FOUND)
  include(CheckCXXSymbolExists)
  set(CMAKE_REQUIRED_INCLUDES ${MPI_CXX_INCLUDE_DIRS})
  set(CMAKE_REQUIRED_LIBRARIES ${MPI_CXX_LIBRARIES})
  check_cxx_symbol_exists(MPIX_Query_cuda_support "mpi.h" CUDA_AWARE)
  if(NOT CUDA_AWARE)
    check_cxx_symbol_exists(MPIX_Query_cuda_support "mpi.h;mpi-ext.h" OMP_CUDA)
    if(NOT OMP_CUDA)
      target_compile_definitions(deepmd_op_pt PRIVATE NO_CUDA_AWARE)
    endif()
  endif()
  target_link_libraries(deepmd_op_pt PRIVATE MPI::MPI_CXX)
  target_compile_definitions(deepmd_op_pt PRIVATE USE_MPI)
endif()
if(CMAKE_TESTING_ENABLED)
  target_link_libraries(deepmd_op_pt PRIVATE coverage_config)
endif()

if(BUILD_PY_IF)
  install(TARGETS deepmd_op_pt DESTINATION deepmd/lib/)
else(BUILD_PY_IF)
  install(TARGETS deepmd_op_pt DESTINATION lib/)
endif(BUILD_PY_IF)
