set(NATIVERT_TEST_ROOT ${TORCH_ROOT}/test/cpp/nativert)

file(GLOB_RECURSE NATIVERT_ALL_TEST_FILES "${NATIVERT_TEST_ROOT}/test_*.cpp")

# Build the cpp gtest binary containing the cpp-only tests.
set(NATIVERT_TEST_SRCS
  ${NATIVERT_ALL_TEST_FILES}
  ${TORCH_ROOT}/torch/nativert/ModelRunner.cpp
  ${TORCH_ROOT}/torch/nativert/graph/TensorMeta.cpp
  ${TORCH_ROOT}/torch/nativert/graph/Graph.cpp
  ${TORCH_ROOT}/torch/nativert/graph/GraphPasses.cpp
  ${TORCH_ROOT}/torch/nativert/graph/GraphSignature.cpp
  ${TORCH_ROOT}/torch/nativert/graph/GraphUtils.cpp
  ${TORCH_ROOT}/torch/nativert/graph/Serialization.cpp
  ${TORCH_ROOT}/torch/nativert/executor/OpKernel.cpp
  ${TORCH_ROOT}/torch/nativert/executor/PlacementUtils.cpp
  ${TORCH_ROOT}/torch/nativert/executor/Weights.cpp
  ${TORCH_ROOT}/torch/nativert/common/FileUtil.cpp
  ${TORCH_ROOT}/torch/nativert/executor/memory/FunctionSchema.cpp
  ${TORCH_ROOT}/torch/nativert/executor/ExecutionPlanner.cpp
  ${TORCH_ROOT}/torch/nativert/detail/ITree.cpp
  ${TORCH_ROOT}/torch/nativert/executor/ExecutionFrame.cpp
  ${TORCH_ROOT}/torch/nativert/kernels/C10Kernel.cpp
  ${TORCH_ROOT}/torch/nativert/executor/memory/GreedyBySize.cpp
  ${TORCH_ROOT}/torch/nativert/executor/memory/Bump.cpp
  ${TORCH_ROOT}/torch/nativert/executor/memory/DisjointStorageGroups.cpp
  ${TORCH_ROOT}/torch/nativert/executor/memory/LayoutPlanner.cpp
  ${TORCH_ROOT}/torch/nativert/executor/memory/LayoutManager.cpp
  ${TORCH_ROOT}/torch/nativert/executor/memory/AliasAnalyzer.cpp
  ${TORCH_ROOT}/torch/nativert/executor/Executor.cpp
  ${TORCH_ROOT}/torch/nativert/kernels/KernelFactory.cpp
  ${TORCH_ROOT}/torch/nativert/executor/ConstantFolder.cpp
  ${TORCH_ROOT}/torch/nativert/executor/GraphExecutorBase.cpp
  ${TORCH_ROOT}/torch/nativert/executor/SerialGraphExecutor.cpp
  ${TORCH_ROOT}/torch/nativert/executor/ParallelGraphExecutor.cpp
  ${TORCH_ROOT}/torch/nativert/kernels/AutoFunctionalizeKernel.cpp
  ${TORCH_ROOT}/torch/nativert/kernels/CallTorchBindKernel.cpp
  ${TORCH_ROOT}/torch/nativert/kernels/HigherOrderKernel.cpp
  ${TORCH_ROOT}/torch/nativert/graph/passes/SubgraphRewriter.cpp
  ${TORCH_ROOT}/torch/nativert/graph/passes/pass_manager/GraphPasses.cpp
  ${TORCH_ROOT}/torch/nativert/graph/passes/pass_manager/PassManager.cpp
  ${TORCH_ROOT}/torch/nativert/kernels/KernelHandlerRegistry.cpp
  ${TORCH_ROOT}/torch/nativert/executor/triton/CpuTritonKernelManager.cpp
  ${TORCH_ROOT}/torch/nativert/kernels/TritonKernel.cpp
  ${TORCH_ROOT}/torch/nativert/executor/DelegateExecutor.cpp
  ${TORCH_ROOT}/torch/nativert/executor/AOTInductorDelegateExecutor.cpp
  ${TORCH_ROOT}/torch/nativert/kernels/ETCallDelegateKernel.cpp
  ${TORCH_ROOT}/torch/csrc/inductor/aoti_torch/oss_proxy_executor.cpp
)

if(USE_CUDA OR USE_ROCM)
  list(APPEND NATIVERT_TEST_SRCS ${TORCH_ROOT}/torch/nativert/executor/triton/CudaTritonKernelManager.cpp)
  list(APPEND NATIVERT_TEST_SRCS ${TORCH_ROOT}/torch/nativert/executor/AOTInductorModelContainerCudaShim.cpp)
endif()

add_executable(test_nativert
  ${TORCH_ROOT}/test/cpp/common/main.cpp
  ${NATIVERT_TEST_SRCS}
)

if(MSVC)
  target_compile_definitions(test_nativert PRIVATE NATIVERT_MSVC_TEST)
endif()

# TODO temporary until we can delete the old gtest polyfills.
target_compile_definitions(test_nativert PRIVATE USE_GTEST)

set(NATIVERT_TEST_DEPENDENCIES torch gtest_main)

target_link_libraries(test_nativert PRIVATE ${CMAKE_DL_LIBS})
target_link_libraries(test_nativert PRIVATE ${NATIVERT_TEST_DEPENDENCIES})
target_link_libraries(test_nativert PRIVATE fmt::fmt-header-only)
target_include_directories(test_nativert PRIVATE ${ATen_CPU_INCLUDE})

if(USE_CUDA)
  target_compile_definitions(test_nativert PRIVATE USE_CUDA)
elseif(USE_ROCM)
  target_link_libraries(test_nativert PRIVATE
    hiprtc::hiprtc
    hip::amdhip64
    ${TORCH_CUDA_LIBRARIES})

  target_compile_definitions(test_nativert PRIVATE USE_ROCM)
endif()

if(INSTALL_TEST)
  set_target_properties(test_nativert PROPERTIES INSTALL_RPATH "${CMAKE_INSTALL_RPATH}:${_rpath_portable_origin}/../lib")
  install(TARGETS test_nativert DESTINATION bin)
  # Install PDB files for MSVC builds
  if(MSVC AND BUILD_SHARED_LIBS)
    install(FILES $<TARGET_PDB_FILE:test_nativert> DESTINATION bin OPTIONAL)
  endif()
endif()