1 2set(AOT_INDUCTOR_TEST_ROOT ${TORCH_ROOT}/test/cpp/aoti_inference) 3 4# Build custom TorchScript op for AOTInductor 5add_library(aoti_custom_class SHARED aoti_custom_class.cpp) 6set_target_properties(aoti_custom_class PROPERTIES 7 LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}) 8if(USE_CUDA) 9 target_compile_definitions(aoti_custom_class PRIVATE USE_CUDA) 10elseif(USE_ROCM) 11 target_compile_definitions(aoti_custom_class PRIVATE USE_ROCM) 12endif() 13# Link against LibTorch 14target_link_libraries(aoti_custom_class torch) 15 16# the custom command that generates the TorchScript module 17add_custom_command( 18 OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/script_data.pt 19 ${CMAKE_CURRENT_BINARY_DIR}/script_model_cpu.pt 20 ${CMAKE_CURRENT_BINARY_DIR}/script_model_cuda.pt 21 COMMAND python ${AOT_INDUCTOR_TEST_ROOT}/compile_model.py 22 DEPENDS compile_model.py 23) 24add_custom_target(aoti_script_model ALL 25 DEPENDS ${CMAKE_CURRENT_BINARY_DIR}/script_data.pt 26 DEPENDS ${CMAKE_CURRENT_BINARY_DIR}/script_model_cpu.pt 27 DEPENDS ${CMAKE_CURRENT_BINARY_DIR}/script_model_cuda.pt 28) 29add_dependencies(aoti_script_model aoti_custom_class) 30 31# Build the cpp gtest binary containing the cpp-only tests. 32set(INDUCTOR_TEST_SRCS 33 ${AOT_INDUCTOR_TEST_ROOT}/test.cpp 34) 35 36add_executable(test_aoti_inference 37 ${TORCH_ROOT}/test/cpp/common/main.cpp 38 ${INDUCTOR_TEST_SRCS} 39 data.pt 40 script_data.pt 41 script_model_cpu.pt 42 script_model_cuda.pt 43) 44add_dependencies(test_aoti_inference aoti_custom_class aoti_script_model) 45 46# TODO temporary until we can delete the old gtest polyfills. 47target_compile_definitions(test_aoti_inference PRIVATE USE_GTEST) 48 49# Define a custom command to generate the library 50add_custom_command( 51 OUTPUT data.pt 52 COMMAND python ${AOT_INDUCTOR_TEST_ROOT}/test.py 53 DEPENDS ${AOT_INDUCTOR_TEST_ROOT}/test.py 54) 55 56target_link_libraries(test_aoti_inference PRIVATE 57 torch 58 gtest 59 -Wl,--no-as-needed aoti_custom_class 60) 61 62if(USE_CUDA) 63 target_include_directories(test_aoti_inference PRIVATE ${ATen_CUDA_INCLUDE}) 64 target_compile_definitions(test_aoti_inference PRIVATE USE_CUDA) 65elseif(USE_ROCM) 66 target_include_directories(test_aoti_inference PRIVATE ${ATen_HIP_INCLUDE}) 67 target_compile_definitions(test_aoti_inference PRIVATE USE_ROCM) 68endif() 69target_compile_definitions(test_aoti_inference PRIVATE 70 CMAKE_CURRENT_BINARY_DIR=${CMAKE_CURRENT_BINARY_DIR} 71) 72 73if(INSTALL_TEST) 74 install(TARGETS test_aoti_inference DESTINATION bin) 75 # Install PDB files for MSVC builds 76 if(MSVC AND BUILD_SHARED_LIBS) 77 install(FILES $<TARGET_PDB_FILE:test_aoti_inference> DESTINATION bin OPTIONAL) 78 endif() 79endif() 80