xref: /aosp_15_r20/external/pytorch/torch/CMakeLists.txt (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1# This file used to build libtorch.so.
2# Now it only builds the Torch python bindings.
3
4if(NOT CAFFE2_CMAKE_BUILDING_WITH_MAIN_REPO)
5  cmake_minimum_required(VERSION 3.18 FATAL_ERROR)
6  project(torch CXX C)
7  find_package(torch REQUIRED)
8  option(USE_CUDA "Use CUDA" ON)
9  set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
10endif()
11
12if(NOT BUILD_PYTHON)
13  return()
14endif()
15
16set(TORCH_SRC_DIR "${CMAKE_CURRENT_SOURCE_DIR}")
17set(TORCH_ROOT "${TORCH_SRC_DIR}/..")
18
19if(NOT TORCH_INSTALL_LIB_DIR)
20  set(TORCH_INSTALL_LIB_DIR lib)
21endif()
22
23if(MSVC)
24    set(LIBSHM_SUBDIR libshm_windows)
25else()
26    set(LIBSHM_SUBDIR libshm)
27endif()
28
29set(LIBSHM_SRCDIR ${TORCH_SRC_DIR}/lib/${LIBSHM_SUBDIR})
30add_subdirectory(${LIBSHM_SRCDIR})
31
32
33# Generate files
34set(TOOLS_PATH "${TORCH_ROOT}/tools")
35
36
37set(TORCH_PYTHON_SRCS
38    ${GENERATED_THNN_CXX}
39    ${GENERATED_CXX_PYTHON}
40    )
41append_filelist("libtorch_python_core_sources" TORCH_PYTHON_SRCS)
42
43# NB: This has to match the condition under which the JIT test directory
44#     is included (at the time of writing that's in caffe2/CMakeLists.txt).
45if(BUILD_TEST)
46    add_definitions(-DBUILDING_TESTS)
47    list(APPEND TORCH_PYTHON_SRCS
48      ${TORCH_ROOT}/test/cpp/jit/torch_python_test.cpp
49      )
50endif()
51
52set(TORCH_PYTHON_INCLUDE_DIRECTORIES
53    ${PYTHON_INCLUDE_DIR}
54
55    ${TORCH_ROOT}
56    ${TORCH_ROOT}/aten/src
57    ${TORCH_ROOT}/aten/src/TH
58
59    ${CMAKE_BINARY_DIR}
60    ${CMAKE_BINARY_DIR}/aten/src
61    ${CMAKE_BINARY_DIR}/caffe2/aten/src
62    ${CMAKE_BINARY_DIR}/third_party
63    ${CMAKE_BINARY_DIR}/third_party/onnx
64
65    ${TORCH_ROOT}/third_party/valgrind-headers
66
67    ${TORCH_ROOT}/third_party/gloo
68    ${TORCH_ROOT}/third_party/onnx
69    ${TORCH_ROOT}/third_party/flatbuffers/include
70    ${TORCH_ROOT}/third_party/kineto/libkineto/include
71    ${TORCH_ROOT}/third_party/cpp-httplib
72    ${TORCH_ROOT}/third_party/nlohmann/include
73
74    ${TORCH_SRC_DIR}/csrc
75    ${TORCH_SRC_DIR}/csrc/api/include
76    ${TORCH_SRC_DIR}/lib
77    )
78
79list(APPEND TORCH_PYTHON_INCLUDE_DIRECTORIES ${LIBSHM_SRCDIR})
80
81set(TORCH_PYTHON_LINK_LIBRARIES
82    Python::Module
83    pybind::pybind11
84    opentelemetry::api
85    httplib
86    nlohmann
87    shm
88    fmt::fmt-header-only
89    ATEN_CPU_FILES_GEN_LIB)
90
91if(USE_ASAN AND TARGET Sanitizer::address)
92  list(APPEND TORCH_PYTHON_LINK_LIBRARIES Sanitizer::address)
93endif()
94if(USE_ASAN AND TARGET Sanitizer::undefined)
95  list(APPEND TORCH_PYTHON_LINK_LIBRARIES Sanitizer::undefined)
96endif()
97if(USE_TSAN AND TARGET Sanitizer::thread)
98  list(APPEND TORCH_PYTHON_LINK_LIBRARIES Sanitizer::thread)
99endif()
100
101set(TORCH_PYTHON_COMPILE_DEFINITIONS)
102
103set(TORCH_PYTHON_COMPILE_OPTIONS)
104
105set(TORCH_PYTHON_LINK_FLAGS "")
106
107if(MSVC)
108    string(APPEND TORCH_PYTHON_LINK_FLAGS " /NODEFAULTLIB:LIBCMT.LIB")
109    list(APPEND TORCH_PYTHON_LINK_LIBRARIES ${PYTHON_LIBRARIES} onnx_library)
110    if(NOT CMAKE_BUILD_TYPE MATCHES "Release")
111      string(APPEND TORCH_PYTHON_LINK_FLAGS " /DEBUG:FULL")
112    endif()
113elseif(APPLE)
114    string(APPEND TORCH_PYTHON_LINK_FLAGS " -undefined dynamic_lookup")
115else()
116    list(APPEND TORCH_PYTHON_COMPILE_OPTIONS
117      -fno-strict-aliasing
118      -Wno-strict-aliasing)
119endif()
120
121if(USE_ITT)
122  list(APPEND TORCH_PYTHON_SRCS
123    ${TORCH_SRC_DIR}/csrc/itt.cpp
124  )
125  list(APPEND TORCH_PYTHON_COMPILE_DEFINITIONS USE_ITT)
126endif()
127
128if(USE_CUDA)
129    include(${TORCH_ROOT}/cmake/public/cuda.cmake)
130    append_filelist("libtorch_python_cuda_core_sources" TORCH_PYTHON_SRCS)
131    list(APPEND TORCH_PYTHON_SRCS ${GENERATED_THNN_CXX_CUDA})
132
133    list(APPEND TORCH_PYTHON_COMPILE_DEFINITIONS USE_CUDA)
134    if(USE_CUDNN)
135        list(APPEND TORCH_PYTHON_LINK_LIBRARIES torch::cudnn)
136        list(APPEND TORCH_PYTHON_COMPILE_DEFINITIONS USE_CUDNN)
137    endif()
138    if(USE_CUSPARSELT)
139        list(APPEND TORCH_PYTHON_LINK_LIBRARIES torch::cusparselt)
140        list(APPEND TORCH_PYTHON_COMPILE_DEFINITIONS USE_CUSPARSELT)
141    endif()
142
143    if(TARGET torch::nvtx3)
144      list(APPEND TORCH_PYTHON_LINK_LIBRARIES torch::nvtx3)
145    else()
146      if(TARGET torch::nvtoolsext)
147        list(APPEND TORCH_PYTHON_LINK_LIBRARIES torch::nvtoolsext)
148      endif()
149    endif()
150endif()
151
152if(USE_ROCM)
153    append_filelist("libtorch_python_cuda_core_sources" TORCH_PYTHON_SRCS)
154    list(APPEND TORCH_PYTHON_SRCS ${GENERATED_THNN_CXX_CUDA})
155
156    list(APPEND TORCH_PYTHON_COMPILE_DEFINITIONS
157      USE_ROCM
158      __HIP_PLATFORM_AMD__
159      )
160    list(APPEND TORCH_PYTHON_LINK_LIBRARIES ${ROCM_ROCTX_LIB})
161endif()
162
163if(USE_XPU)
164    include(${TORCH_ROOT}/cmake/public/xpu.cmake)
165    append_filelist("libtorch_python_xpu_sources" TORCH_PYTHON_SRCS)
166
167    list(APPEND TORCH_PYTHON_COMPILE_DEFINITIONS USE_XPU)
168endif()
169
170if(USE_CUDNN OR USE_ROCM)
171    list(APPEND TORCH_PYTHON_SRCS
172      ${TORCH_SRC_DIR}/csrc/cuda/shared/cudnn.cpp
173      )
174    if(USE_STATIC_CUDNN)
175        set_source_files_properties(
176          ${TORCH_SRC_DIR}/csrc/cuda/shared/cudnn.cpp
177          PROPERTIES COMPILE_DEFINITIONS "USE_STATIC_CUDNN"
178        )
179    endif()
180endif()
181
182if(USE_CUSPARSELT)
183    list(APPEND TORCH_PYTHON_SRCS ${TORCH_SRC_DIR}/csrc/cuda/shared/cusparselt.cpp)
184endif()
185
186if(USE_MPS)
187    list(APPEND TORCH_PYTHON_SRCS ${MPS_PYTHON_SRCS})
188endif()
189
190if(USE_VALGRIND AND NOT WIN32)
191    list(APPEND TORCH_PYTHON_COMPILE_DEFINITIONS USE_VALGRIND)
192endif()
193
194# In the most recent CMake versions, a new 'TRANSFORM' subcommand of 'list' allows much of the boilerplate of defining the lists
195# of type stub files to be omitted.
196# For compatibility with older CMake versions, we omit it for now, but leave it as a comment in case compatibility with the older
197# CMake versions is eventually dropped.
198# set(Modules
199#     __init__
200#     activation
201#     adaptive
202#     batchnorm
203#     container
204#     conv
205#     distance
206#     dropout
207#     fold
208#     instancenorm
209#     linear
210#     loss
211#     module
212#     normalization
213#     padding
214#     pixelshuffle
215#     pooling
216#     rnn
217#     sparse
218#     upsampling
219# )
220# list(TRANSFORM Modules PREPEND "${TORCH_SRC_DIR}/nn/modules/")
221add_custom_target(torch_python_stubs DEPENDS
222    "${TORCH_SRC_DIR}/_C/__init__.pyi"
223    "${TORCH_SRC_DIR}/_C/_VariableFunctions.pyi"
224    "${TORCH_SRC_DIR}/nn/functional.pyi"
225    "${TORCH_SRC_DIR}/utils/data/datapipes/datapipe.pyi"
226)
227
228file(GLOB_RECURSE torchgen_python "${PROJECT_SOURCE_DIR}/torchgen/*.py")
229file(GLOB_RECURSE autograd_python "${TOOLS_PATH}/autograd/*.py")
230file(GLOB_RECURSE pyi_python "${TOOLS_PATH}/pyi/*.py")
231add_custom_command(
232    OUTPUT
233    "${TORCH_SRC_DIR}/_C/__init__.pyi"
234    "${TORCH_SRC_DIR}/_C/_VariableFunctions.pyi"
235    "${TORCH_SRC_DIR}/nn/functional.pyi"
236    COMMAND
237    "${Python_EXECUTABLE}" -mtools.pyi.gen_pyi
238      --native-functions-path "aten/src/ATen/native/native_functions.yaml"
239      --tags-path "aten/src/ATen/native/tags.yaml"
240      --deprecated-functions-path "tools/autograd/deprecated.yaml"
241    DEPENDS
242      "${TORCH_SRC_DIR}/_C/__init__.pyi.in"
243      "${TORCH_SRC_DIR}/_C/_VariableFunctions.pyi.in"
244      "${TORCH_SRC_DIR}/nn/functional.pyi.in"
245      "${TORCH_ROOT}/aten/src/ATen/native/native_functions.yaml"
246      "${TORCH_ROOT}/aten/src/ATen/native/tags.yaml"
247      "${TORCH_ROOT}/tools/autograd/deprecated.yaml"
248      ${pyi_python}
249      ${autograd_python}
250      ${torchgen_python}
251    WORKING_DIRECTORY
252    "${TORCH_ROOT}"
253)
254file(GLOB_RECURSE datapipe_files "${TORCH_SRC_DIR}/utils/data/datapipes/*.py")
255add_custom_command(
256    OUTPUT
257    "${TORCH_SRC_DIR}/utils/data/datapipes/datapipe.pyi"
258    COMMAND
259    "${Python_EXECUTABLE}" ${TORCH_SRC_DIR}/utils/data/datapipes/gen_pyi.py
260    DEPENDS
261    "${TORCH_SRC_DIR}/utils/data/datapipes/datapipe.pyi.in"
262    ${datapipe_files}
263    WORKING_DIRECTORY
264    "${TORCH_ROOT}"
265)
266if(USE_DISTRIBUTED)
267    if(WIN32)
268      append_filelist("libtorch_python_distributed_core_sources" TORCH_PYTHON_SRCS)
269    else()
270      append_filelist("libtorch_python_distributed_sources" TORCH_PYTHON_SRCS)
271    endif()
272    # Disable certain warnings for GCC-9.X
273    if(CMAKE_COMPILER_IS_GNUCXX)
274      set_source_files_properties(${TORCH_SRC_DIR}/csrc/distributed/autograd/init.cpp PROPERTIES COMPILE_FLAGS "-Wno-cast-function-type")
275      set_source_files_properties(${TORCH_SRC_DIR}/csrc/distributed/rpc/testing/init.cpp PROPERTIES COMPILE_FLAGS "-Wno-cast-function-type")
276      set_source_files_properties(${TORCH_SRC_DIR}/csrc/distributed/c10d/init.cpp PROPERTIES COMPILE_FLAGS "-Wno-cast-function-type")
277    endif()
278    # NCCL is a private dependency of libtorch, but libtorch_python includes
279    # some private headers of libtorch, which in turn include NCCL. As a hacky
280    # alternative to making NCCL a public dependency of libtorch, we make it
281    # a private dependency of libtorch_python as well.
282    if(USE_NCCL)
283      list(APPEND TORCH_PYTHON_LINK_LIBRARIES __caffe2_nccl)
284    endif()
285    # Same for MPI.
286    if(USE_MPI)
287      list(APPEND TORCH_PYTHON_LINK_LIBRARIES MPI::MPI_CXX)
288    endif()
289    list(APPEND TORCH_PYTHON_COMPILE_DEFINITIONS USE_C10D)
290
291endif()
292
293if(USE_NCCL AND NOT WIN32)
294    list(APPEND TORCH_PYTHON_SRCS
295      ${TORCH_SRC_DIR}/csrc/cuda/python_nccl.cpp)
296    list(APPEND TORCH_PYTHON_COMPILE_DEFINITIONS USE_NCCL)
297endif()
298
299if(NOT MSVC)
300  # cudaProfilerInitialize must go away
301  set_source_files_properties(${TORCH_SRC_DIR}/csrc/cuda/shared/cudart.cpp PROPERTIES COMPILE_FLAGS "-Wno-deprecated-declarations")
302endif()
303
304# coreml
305if(USE_COREML_DELEGATE)
306  list(APPEND TORCH_PYTHON_SRCS ${TORCH_SRC_DIR}/csrc/jit/backends/coreml/cpp/backend.cpp)
307  list(APPEND TORCH_PYTHON_SRCS ${TORCH_SRC_DIR}/csrc/jit/backends/coreml/cpp/preprocess.cpp)
308endif()
309
310
311add_library(torch_python SHARED ${TORCH_PYTHON_SRCS})
312
313if(CAFFE2_USE_MKL AND BUILD_LIBTORCHLESS)
314
315  # Use the RPATH of the linked libraries
316  set(CMAKE_INSTALL_RPATH_USE_LINK_PATH TRUE)
317  # we need to explicitly link caffe2::mkl in order to have the
318  # correct RPATH in torch_python for the split build
319  target_link_libraries(torch_python PRIVATE caffe2::mkl)
320endif()
321
322add_dependencies(torch_python onnx_proto)
323# Avoid numpy for the DEPLOY build
324if(USE_NUMPY)
325  target_link_libraries(torch_python PRIVATE Python::NumPy)
326  target_compile_definitions(torch_python PRIVATE USE_NUMPY)
327endif()
328
329if(USE_CUFILE AND NOT USE_ROCM)
330  target_compile_definitions(torch_python PRIVATE USE_CUFILE)
331endif()
332
333if(HAVE_SOVERSION)
334  set_target_properties(torch_python PROPERTIES
335      VERSION ${TORCH_VERSION} SOVERSION ${TORCH_SOVERSION})
336endif()
337
338# in case of the split build we need to add compile definitions
339if(BUILD_LIBTORCHLESS)
340
341  if(USE_UCC)
342    target_link_libraries(torch_python PRIVATE __caffe2_ucc)
343    target_compile_definitions(torch_python PRIVATE USE_UCC)
344  endif()
345
346  if(USE_UCC AND USE_C10D_UCC)
347    target_compile_definitions(torch_python PRIVATE USE_C10D_UCC)
348  endif()
349
350  if(USE_NCCL AND USE_C10D_NCCL)
351    target_compile_definitions(torch_python PRIVATE USE_C10D_NCCL)
352  endif()
353
354  if(USE_DISTRIBUTED)
355    target_compile_definitions(torch_python PRIVATE USE_DISTRIBUTED)
356  endif()
357
358  if(USE_MPI AND USE_C10D_MPI)
359    target_compile_definitions(torch_python PRIVATE USE_C10D_MPI)
360  endif()
361
362  if(USE_GLOO AND USE_C10D_GLOO)
363    target_compile_definitions(torch_python PRIVATE USE_C10D_GLOO)
364  endif()
365
366  if(NOT WIN32)
367    target_compile_definitions(torch_python PRIVATE USE_RPC)
368  endif()
369
370  if(USE_TENSORPIPE)
371    target_compile_definitions(torch_python PRIVATE USE_TENSORPIPE)
372  endif()
373
374  set(EXPERIMENTAL_SINGLE_THREAD_POOL "0" CACHE STRING
375  "Experimental option to use a single thread pool for inter- and intra-op parallelism")
376  if("${EXPERIMENTAL_SINGLE_THREAD_POOL}")
377    target_compile_definitions(torch_python PRIVATE "-DAT_EXPERIMENTAL_SINGLE_THREAD_POOL=1")
378  endif()
379
380endif()
381
382add_dependencies(torch_python torch_python_stubs)
383add_dependencies(torch_python flatbuffers)
384
385
386if(USE_PRECOMPILED_HEADERS)
387  target_precompile_headers(torch_python PRIVATE
388      "$<$<COMPILE_LANGUAGE:CXX>:ATen/ATen.h>")
389endif()
390
391# Required workaround for generated sources
392# See https://samthursfield.wordpress.com/2015/11/21/cmake-dependencies-between-targets-and-files-and-custom-commands/#custom-commands-in-different-directories
393add_dependencies(torch_python generate-torch-sources)
394set_source_files_properties(
395    ${GENERATED_THNN_SOURCES}
396    ${GENERATED_CXX_PYTHON}
397    PROPERTIES GENERATED TRUE
398    )
399
400# Disable certain warnings for GCC-9.X
401if(CMAKE_COMPILER_IS_GNUCXX)
402  set_source_files_properties(${TORCH_SRC_DIR}/csrc/Module.cpp PROPERTIES COMPILE_FLAGS "-Wno-cast-function-type")
403  set_source_files_properties(${TORCH_SRC_DIR}/csrc/autograd/python_variable.cpp PROPERTIES COMPILE_FLAGS "-Wno-cast-function-type")
404endif()
405
406# Preserve CUDA_GENCODE flags
407if(USE_CUDA)
408  torch_cuda_get_nvcc_gencode_flag(_ARCH_FLAGS)
409  set_source_files_properties(${TORCH_SRC_DIR}/csrc/cuda/Module.cpp PROPERTIES COMPILE_FLAGS "-DCUDA_ARCH_FLAGS=\"${_ARCH_FLAGS_readable}\"")
410endif()
411
412# Preserve HIP arch flags
413if(USE_ROCM)
414  string(REPLACE ";" " " PYTORCH_ROCM_ARCH_readable "${PYTORCH_ROCM_ARCH}")
415  set_source_files_properties(${TORCH_SRC_DIR}/csrc/cuda/Module.cpp PROPERTIES COMPILE_FLAGS "-DCUDA_ARCH_FLAGS=\"${PYTORCH_ROCM_ARCH_readable}\"")
416endif()
417
418target_compile_definitions(torch_python PRIVATE "-DTHP_BUILD_MAIN_LIB")
419
420target_link_libraries(torch_python PRIVATE ${TORCH_LIB} ${TORCH_PYTHON_LINK_LIBRARIES})
421
422target_compile_definitions(torch_python PRIVATE ${TORCH_PYTHON_COMPILE_DEFINITIONS})
423
424target_compile_options(torch_python PRIVATE ${TORCH_PYTHON_COMPILE_OPTIONS})
425
426target_include_directories(torch_python PUBLIC ${TORCH_PYTHON_INCLUDE_DIRECTORIES})
427
428if(USE_UCC)
429  target_link_libraries(torch_python PRIVATE __caffe2_ucc)
430  target_compile_definitions(torch_python PRIVATE USE_UCC)
431endif()
432
433if(BUILD_ONEDNN_GRAPH)
434  target_compile_definitions(torch_python PRIVATE "-DBUILD_ONEDNN_GRAPH")
435  if(NOT BUILD_LIBTORCHLESS)
436    target_compile_definitions(torch_cpu PRIVATE "-DBUILD_ONEDNN_GRAPH")
437  endif()
438endif()
439
440if(${CMAKE_BUILD_TYPE} STREQUAL "RelWithAssert")
441   # Workaround numerous decret-without-a-gil warnings from JIT
442   # see https://github.com/pytorch/pytorch/issues/130073
443   target_compile_definitions(torch_python PRIVATE "-DPYBIND11_NO_ASSERT_GIL_HELD_INCREF_DECREF")
444endif()
445
446if(NOT TORCH_PYTHON_LINK_FLAGS STREQUAL "")
447    set_target_properties(torch_python PROPERTIES LINK_FLAGS ${TORCH_PYTHON_LINK_FLAGS})
448endif()
449
450install(TARGETS torch_python DESTINATION "${TORCH_INSTALL_LIB_DIR}")
451
452# Generate torch/version.py from the appropriate CMake cache variables.
453if(${CMAKE_BUILD_TYPE} STREQUAL "Debug")
454  set(TORCH_VERSION_DEBUG 1)
455else()
456  set(TORCH_VERSION_DEBUG 0)
457endif()
458
459add_custom_command(
460  OUTPUT ${TORCH_SRC_DIR}/version.py
461  COMMAND "${CMAKE_COMMAND}" -E touch "${TOOLS_PATH}/generate_torch_version.py"
462  COMMAND
463    "${Python_EXECUTABLE}" "${TOOLS_PATH}/generate_torch_version.py"
464      --is-debug=${TORCH_VERSION_DEBUG}
465      --cuda-version=${CUDA_VERSION}
466      --hip-version=${HIP_VERSION}
467  DEPENDS ${TOOLS_PATH}/generate_torch_version.py
468  WORKING_DIRECTORY ${TORCH_ROOT}
469)
470add_custom_target(
471  gen_torch_version ALL
472  DEPENDS ${TORCH_SRC_DIR}/version.py
473)
474add_dependencies(torch_python gen_torch_version)
475
476# Skip building this library under MacOS, since it is currently failing to build on Mac
477# Github issue #61930
478if(NOT ${CMAKE_SYSTEM_NAME} MATCHES "Darwin")
479  # Add Android Nnapi delegate library
480  add_library(nnapi_backend SHARED
481          ${TORCH_SRC_DIR}/csrc/jit/backends/nnapi/nnapi_backend_lib.cpp
482          ${TORCH_SRC_DIR}/csrc/jit/backends/nnapi/nnapi_backend_preprocess.cpp
483          )
484  # Pybind11 requires explicit linking of the torch_python library
485  if(BUILD_LIBTORCHLESS)
486    target_link_libraries(nnapi_backend PRIVATE ${TORCH_LIB} torch_python pybind::pybind11)
487  else()
488    target_link_libraries(nnapi_backend PRIVATE torch torch_python pybind::pybind11)
489  endif()
490endif()
491
492set(TORCH_PYTHON_COMPILE_OPTIONS ${TORCH_PYTHON_COMPILE_OPTIONS} PARENT_SCOPE)
493set(TORCH_PYTHON_LINK_FLAGS ${TORCH_PYTHON_LINK_FLAGS} PARENT_SCOPE)
494