2023-11-08 09:01:59 -07:00
if ( NOT INTERN_BUILD_ATEN_OPS )
return ( )
endif ( )
# Find modules
if ( NOT INTERN_BUILD_MOBILE )
list ( APPEND CMAKE_MODULE_PATH /usr/lib/x86_64-linux-gnu/ )
list ( APPEND CMAKE_LIBRARY_PATH /usr/lib/x86_64-linux-gnu/ /usr/lib/aarch64-linux-gnu/ )
endif ( )
list ( APPEND CMAKE_MODULE_PATH
$ { C M A K E _ C U R R E N T _ S O U R C E _ D I R } / . . / c m a k e / M o d u l e s
$ { C M A K E _ C U R R E N T _ S O U R C E _ D I R } / . . / c m a k e / p u b l i c
$ { C M A K E _ C U R R E N T _ S O U R C E _ D I R } / . . / c m a k e / M o d u l e s _ C U D A _ f i x )
cmake_policy ( SET CMP0012 NEW )
#############################################
set ( ATen_CPU_SRCS )
set ( ATen_CPU_TEST_SRCS )
set ( ATen_CPU_INCLUDE )
set ( ATen_THIRD_PARTY_INCLUDE )
set ( ATen_CUDA_CPP_SRCS )
set ( ATen_CUDA_CU_SRCS )
set ( ATen_CUDA_LINALG_SRCS )
set ( ATen_CUDA_SRCS_W_SORT_BY_KEY )
set ( ATen_CUDA_TEST_SRCS )
set ( ATen_CUDA_INCLUDE )
set ( ATen_NVRTC_STUB_SRCS )
set ( ATen_HIP_SRCS )
set ( ATen_HIP_SRCS_W_SORT_BY_KEY )
set ( ATen_HIP_TEST_SRCS )
set ( ATen_HIP_INCLUDE )
set ( ATen_MPS_SRCS )
set ( ATen_MPS_TEST_SRCS )
set ( ATen_VULKAN_TEST_SRCS )
set ( ATen_CPU_DEPENDENCY_LIBS )
set ( ATen_CUDA_DEPENDENCY_LIBS )
set ( ATen_HIP_DEPENDENCY_LIBS )
set ( ATen_PUBLIC_CUDA_DEPENDENCY_LIBS )
set ( ATen_PUBLIC_HIP_DEPENDENCY_LIBS )
set ( ATEN_INSTALL_BIN_SUBDIR "bin" CACHE PATH "ATen install binary subdirectory" )
set ( ATEN_INSTALL_LIB_SUBDIR "lib" CACHE PATH "ATen install library subdirectory" )
set ( ATEN_INSTALL_INCLUDE_SUBDIR "include" CACHE PATH "ATen install include subdirectory" )
set ( MEM_EFF_ATTENTION_CUDA_SOURCES )
if ( USE_CUDA )
list ( APPEND ATen_CUDA_INCLUDE ${ CUDA_INCLUDE_DIRS } )
endif ( )
set ( TH_LINK_STYLE STATIC )
set ( TH_CPU_INCLUDE
$ { C M A K E _ C U R R E N T _ S O U R C E _ D I R } / s r c
$ { C M A K E _ C U R R E N T _ B I N A R Y _ D I R } / s r c
$ { C M A K E _ B I N A R Y _ D I R } / a t e n / s r c )
list ( APPEND ATen_CPU_INCLUDE ${ TH_CPU_INCLUDE } )
if ( USE_VULKAN )
list ( APPEND ATen_CPU_INCLUDE ${ CMAKE_BINARY_DIR } /vulkan ${ CMAKE_CURRENT_SOURCE_DIR } /../third_party/VulkanMemoryAllocator )
endif ( )
# Find the HIP package, set the HIP paths, load the HIP CMake.
if ( USE_ROCM )
include ( LoadHIP )
if ( NOT PYTORCH_FOUND_HIP )
set ( USE_ROCM OFF )
endif ( )
endif ( )
# Both CUDA and ROCM are enabled and found. Report an error.
if ( USE_CUDA AND USE_ROCM )
message ( FATAL_ERROR "Both CUDA and ROCm are enabled and found. PyTorch can only be built with either of them. Please turn one off by using either USE_CUDA=OFF or USE_ROCM=OFF." )
endif ( )
if ( USE_ROCM )
# TODO: AT_HIP_ENABLED (change this once we represent HIP as HIP in
# ATen proper)
set ( AT_CUDA_ENABLED 1 )
add_subdirectory ( src/THH )
message ( "ROCm is enabled." )
elseif ( USE_CUDA )
set ( AT_CUDA_ENABLED 1 )
add_subdirectory ( src/THC )
else ( )
message ( "disabling CUDA because USE_CUDA is set false" )
set ( AT_CUDA_ENABLED 0 )
endif ( )
if ( NOT USE_NNPACK )
set ( AT_NNPACK_ENABLED 0 )
else ( )
set ( AT_NNPACK_ENABLED 1 )
endif ( )
2023-11-08 09:13:36 -07:00
if ( NOT USE_CUSPARSELT )
set ( AT_CUSPARSELT_ENABLED 0 )
else ( )
set ( AT_CUSPARSELT_ENABLED 1 )
endif ( )
2023-11-08 09:01:59 -07:00
list ( APPEND ATen_CPU_INCLUDE
2023-11-08 09:13:36 -07:00
$ { C M A K E _ C U R R E N T _ S O U R C E _ D I R } / s r c )
2023-11-08 09:01:59 -07:00
add_subdirectory ( src/ATen )
# Pass source, includes, and libs to parent
set ( ATen_CPU_SRCS ${ ATen_CPU_SRCS } PARENT_SCOPE )
set ( ATen_CORE_SRCS ${ ATen_CORE_SRCS } PARENT_SCOPE )
set ( ATen_CUDA_CU_SRCS ${ ATen_CUDA_CU_SRCS } PARENT_SCOPE )
set ( ATen_CUDA_CPP_SRCS ${ ATen_CUDA_CPP_SRCS } PARENT_SCOPE )
set ( ATen_CUDA_LINALG_SRCS ${ ATen_CUDA_LINALG_SRCS } PARENT_SCOPE )
set ( ATen_CUDA_SRCS_W_SORT_BY_KEY ${ ATen_CUDA_SRCS_W_SORT_BY_KEY } PARENT_SCOPE )
set ( ATen_CUDA_CU_SRCS_W_SORT_BY_KEY ${ ATen_CUDA_CU_SRCS_W_SORT_BY_KEY } PARENT_SCOPE )
set ( ATen_HIP_SRCS ${ ATen_HIP_SRCS } PARENT_SCOPE )
set ( ATen_MPS_SRCS ${ ATen_MPS_SRCS } PARENT_SCOPE )
set ( ATen_MPS_TEST_SRCS ${ ATen_MPS_TEST_SRCS } PARENT_SCOPE )
set ( ATen_HIP_SRCS_W_SORT_BY_KEY ${ ATen_HIP_SRCS_W_SORT_BY_KEY } PARENT_SCOPE )
set ( ATen_NVRTC_STUB_SRCS ${ ATen_NVRTC_STUB_SRCS } PARENT_SCOPE )
set ( ATen_CPU_TEST_SRCS ${ ATen_CPU_TEST_SRCS } PARENT_SCOPE )
set ( ATen_CUDA_TEST_SRCS ${ ATen_CUDA_TEST_SRCS } PARENT_SCOPE )
set ( ATen_HIP_TEST_SRCS ${ ATen_HIP_TEST_SRCS } PARENT_SCOPE )
set ( ATen_VULKAN_TEST_SRCS ${ ATen_VULKAN_TEST_SRCS } PARENT_SCOPE )
set ( ATen_MOBILE_BENCHMARK_SRCS ${ ATen_MOBILE_BENCHMARK_SRCS } PARENT_SCOPE )
set ( ATen_MOBILE_TEST_SRCS ${ ATen_MOBILE_TEST_SRCS } PARENT_SCOPE )
set ( ATen_VEC_TEST_SRCS ${ ATen_VEC_TEST_SRCS } PARENT_SCOPE )
set ( ATen_CPU_INCLUDE ${ ATen_CPU_INCLUDE } PARENT_SCOPE )
set ( ATen_CUDA_INCLUDE ${ ATen_CUDA_INCLUDE } PARENT_SCOPE )
set ( ATen_HIP_INCLUDE ${ ATen_HIP_INCLUDE } PARENT_SCOPE )
set ( ATen_THIRD_PARTY_INCLUDE ${ ATen_THIRD_PARTY_INCLUDE } PARENT_SCOPE )
set ( ATen_CPU_DEPENDENCY_LIBS ${ ATen_CPU_DEPENDENCY_LIBS } PARENT_SCOPE )
set ( ATen_CUDA_DEPENDENCY_LIBS ${ ATen_CUDA_DEPENDENCY_LIBS } PARENT_SCOPE )
set ( ATen_HIP_DEPENDENCY_LIBS ${ ATen_HIP_DEPENDENCY_LIBS } PARENT_SCOPE )
set ( ATen_CORE_TEST_SRCS ${ ATen_CORE_TEST_SRCS } PARENT_SCOPE )
2023-11-08 09:13:36 -07:00
set ( MEM_EFF_ATTENTION_CUDA_SOURCES ${ MEM_EFF_ATTENTION_CUDA_SOURCES } PARENT_SCOPE )