mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
c326533999
Fixes #ISSUE_NUMBER Pull Request resolved: https://github.com/pytorch/pytorch/pull/131521 Approved by: https://github.com/jataylo, https://github.com/pruthvistony, https://github.com/malfet
80 lines
2.6 KiB
CMake
80 lines
2.6 KiB
CMake
|
|
set(AOT_INDUCTOR_TEST_ROOT ${TORCH_ROOT}/test/cpp/aoti_inference)
|
|
|
|
# Build custom TorchScript op for AOTInductor
|
|
add_library(aoti_custom_class SHARED aoti_custom_class.cpp)
|
|
set_target_properties(aoti_custom_class PROPERTIES
|
|
LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR})
|
|
if(USE_CUDA)
|
|
target_compile_definitions(aoti_custom_class PRIVATE USE_CUDA)
|
|
elseif(USE_ROCM)
|
|
target_compile_definitions(aoti_custom_class PRIVATE USE_ROCM)
|
|
endif()
|
|
# Link against LibTorch
|
|
target_link_libraries(aoti_custom_class torch)
|
|
|
|
# the custom command that generates the TorchScript module
|
|
add_custom_command(
|
|
OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/script_data.pt
|
|
${CMAKE_CURRENT_BINARY_DIR}/script_model_cpu.pt
|
|
${CMAKE_CURRENT_BINARY_DIR}/script_model_cuda.pt
|
|
COMMAND python ${AOT_INDUCTOR_TEST_ROOT}/compile_model.py
|
|
DEPENDS compile_model.py
|
|
)
|
|
add_custom_target(aoti_script_model ALL
|
|
DEPENDS ${CMAKE_CURRENT_BINARY_DIR}/script_data.pt
|
|
DEPENDS ${CMAKE_CURRENT_BINARY_DIR}/script_model_cpu.pt
|
|
DEPENDS ${CMAKE_CURRENT_BINARY_DIR}/script_model_cuda.pt
|
|
)
|
|
add_dependencies(aoti_script_model aoti_custom_class)
|
|
|
|
# Build the cpp gtest binary containing the cpp-only tests.
|
|
set(INDUCTOR_TEST_SRCS
|
|
${AOT_INDUCTOR_TEST_ROOT}/test.cpp
|
|
)
|
|
|
|
add_executable(test_aoti_inference
|
|
${TORCH_ROOT}/test/cpp/common/main.cpp
|
|
${INDUCTOR_TEST_SRCS}
|
|
data.pt
|
|
script_data.pt
|
|
script_model_cpu.pt
|
|
script_model_cuda.pt
|
|
)
|
|
add_dependencies(test_aoti_inference aoti_custom_class aoti_script_model)
|
|
|
|
# TODO temporary until we can delete the old gtest polyfills.
|
|
target_compile_definitions(test_aoti_inference PRIVATE USE_GTEST)
|
|
|
|
# Define a custom command to generate the library
|
|
add_custom_command(
|
|
OUTPUT data.pt
|
|
COMMAND python ${AOT_INDUCTOR_TEST_ROOT}/test.py
|
|
DEPENDS ${AOT_INDUCTOR_TEST_ROOT}/test.py
|
|
)
|
|
|
|
target_link_libraries(test_aoti_inference PRIVATE
|
|
torch
|
|
gtest
|
|
-Wl,--no-as-needed aoti_custom_class
|
|
)
|
|
|
|
if(USE_CUDA)
|
|
target_include_directories(test_aoti_inference PRIVATE ${ATen_CUDA_INCLUDE})
|
|
target_compile_definitions(test_aoti_inference PRIVATE USE_CUDA)
|
|
elseif(USE_ROCM)
|
|
target_include_directories(test_aoti_inference PRIVATE ${ATen_HIP_INCLUDE})
|
|
target_compile_definitions(test_aoti_inference PRIVATE USE_ROCM)
|
|
endif()
|
|
target_compile_definitions(test_aoti_inference PRIVATE
|
|
CMAKE_CURRENT_BINARY_DIR=${CMAKE_CURRENT_BINARY_DIR}
|
|
)
|
|
|
|
if(INSTALL_TEST)
|
|
install(TARGETS test_aoti_inference DESTINATION bin)
|
|
# Install PDB files for MSVC builds
|
|
if(MSVC AND BUILD_SHARED_LIBS)
|
|
install(FILES $<TARGET_PDB_FILE:test_aoti_inference> DESTINATION bin OPTIONAL)
|
|
endif()
|
|
endif()
|