ITKIMPACT/itk-module-init.cmake at main · InsightSoftwareConsortium/ITKIMPACT · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
#==========================================================================
#
#   Copyright NumFOCUS
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#          https://www.apache.org/licenses/LICENSE-2.0.txt
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.
#
#==========================================================================

# LibTorch is provided by the installed `torch` Python package (the libtorch it bundles),
# not a separately downloaded C++ distribution. One wheel then inherits CPU or CUDA from
# whatever torch the user installed, and we never ship a second multi-GB copy of LibTorch.
# Auto-locate it via `torch.utils.cmake_prefix_path` unless the caller already pointed
# CMAKE_PREFIX_PATH / Torch_DIR at a LibTorch (e.g. a manual C++ distribution).
#
# torch is a hard dependency, so if the `torch` package is missing we pip-install the CPU
# build (best effort) rather than fail outright — this lets the standard ITK CI (which does
# not provision LibTorch) build the module out of the box. Turn it off with
# -DIMPACT_AUTO_INSTALL_TORCH=OFF, or point -DTorch_DIR / -DCMAKE_PREFIX_PATH at a manual
# (e.g. CUDA) LibTorch, which always takes precedence.
option(IMPACT_AUTO_INSTALL_TORCH
  "If the torch Python package is missing, pip-install the CPU build so find_package(Torch) succeeds." ON)

if(NOT Torch_DIR AND NOT DEFINED Torch_ROOT)
  # Prefer the interpreter CMake was configured with (e.g. the cp3xx wheel-build python in
  # the manylinux container) — a bare `python3` on PATH can be an old system Python for which
  # a current torch has no wheel. Fall back to a python3 on PATH only if none was provided.
  if(DEFINED Python3_EXECUTABLE AND EXISTS "${Python3_EXECUTABLE}")
    set(IMPACT_Python_EXECUTABLE "${Python3_EXECUTABLE}")
  elseif(DEFINED PYTHON_EXECUTABLE AND EXISTS "${PYTHON_EXECUTABLE}")
    set(IMPACT_Python_EXECUTABLE "${PYTHON_EXECUTABLE}")
  else()
    find_program(IMPACT_Python_EXECUTABLE NAMES python3 python)
  endif()
  if(IMPACT_Python_EXECUTABLE)
    execute_process(
      COMMAND "${IMPACT_Python_EXECUTABLE}" -c "import torch; print(torch.utils.cmake_prefix_path)"
      OUTPUT_VARIABLE IMPACT_Torch_cmake_prefix
      OUTPUT_STRIP_TRAILING_WHITESPACE
      RESULT_VARIABLE IMPACT_Torch_query
      ERROR_QUIET)

    if(NOT IMPACT_Torch_query EQUAL 0 AND IMPACT_AUTO_INSTALL_TORCH)
      message(STATUS "IMPACT: the torch Python package was not found; installing the CPU build via pip ...")
      # An environment created without pip (e.g. `uv venv`) still ships ensurepip in the
      # standard library, so bootstrap pip rather than failing on `-m pip`.
      execute_process(
        COMMAND "${IMPACT_Python_EXECUTABLE}" -m pip --version
        RESULT_VARIABLE IMPACT_pip_present OUTPUT_QUIET ERROR_QUIET)
      if(NOT IMPACT_pip_present EQUAL 0)
        execute_process(
          COMMAND "${IMPACT_Python_EXECUTABLE}" -m ensurepip --upgrade
          RESULT_VARIABLE IMPACT_ensurepip OUTPUT_QUIET ERROR_QUIET)
      endif()
      execute_process(
        COMMAND "${IMPACT_Python_EXECUTABLE}" -m pip install "torch==2.12.*"
                --index-url https://download.pytorch.org/whl/cpu
        RESULT_VARIABLE IMPACT_Torch_pip)
      if(NOT IMPACT_Torch_pip EQUAL 0)
        # Some distributions ship an externally-managed Python (PEP 668); retry allowing it.
        execute_process(
          COMMAND "${IMPACT_Python_EXECUTABLE}" -m pip install --break-system-packages "torch==2.12.*"
                  --index-url https://download.pytorch.org/whl/cpu
          RESULT_VARIABLE IMPACT_Torch_pip)
      endif()
      if(IMPACT_Torch_pip EQUAL 0)
        execute_process(
          COMMAND "${IMPACT_Python_EXECUTABLE}" -c "import torch; print(torch.utils.cmake_prefix_path)"
          OUTPUT_VARIABLE IMPACT_Torch_cmake_prefix
          OUTPUT_STRIP_TRAILING_WHITESPACE
          RESULT_VARIABLE IMPACT_Torch_query
          ERROR_QUIET)
      endif()
    endif()

    if(IMPACT_Torch_query EQUAL 0 AND IMPACT_Torch_cmake_prefix)
      message(STATUS "IMPACT: using LibTorch from the installed torch package: ${IMPACT_Torch_cmake_prefix}")
      list(APPEND CMAKE_PREFIX_PATH "${IMPACT_Torch_cmake_prefix}")
    else()
      # Say what to do here: the bare find_package(Torch) failure below names a
      # TorchConfig.cmake the user has no reason to know how to get.
      message(WARNING
        "IMPACT: no usable torch for '${IMPACT_Python_EXECUTABLE}', and it could not be "
        "installed automatically. LibTorch ships inside the torch Python package, so either\n"
        "  * install it first:  ${IMPACT_Python_EXECUTABLE} -m pip install 'torch==2.12.*'\n"
        "    (add --index-url https://download.pytorch.org/whl/cpu for the smaller CPU build)\n"
        "  * or point CMake at an existing one:  -DTorch_DIR=<prefix>/share/cmake/Torch\n"
        "Under `pip install`, the build runs in an isolated environment with no pip of its "
        "own, so the automatic install cannot run; torch is declared as a build requirement "
        "for that case -- use --no-build-isolation only if you provide torch yourself.")
    endif()
  endif()
endif()

find_package(Torch REQUIRED)

# The module is mostly header-only (templates), plus a small compiled Impact library
# (src/) for the torch-free ImpactModelConfiguration facade. Its public templates still
# instantiate LibTorch in client code, so LibTorch is propagated as an extra library
# clients must link, in addition to the Impact library itself.
set(Impact_LIBRARIES Impact ${TORCH_LIBRARIES})

set(Impact_WRAP_INCLUDE_DIRS
  ${TORCH_INCLUDE_DIRS}
)

set(Impact_SYSTEM_INCLUDE_DIRS ${TORCH_INCLUDE_DIRS})
set(Impact_SYSTEM_LIBRARY_DIRS ${Torch_LIBRARY_DIRS})

# When this module is consumed by an app, make it re-find Torch so the imported
# torch targets exist downstream.
set(Impact_EXPORT_CODE_INSTALL "
set(Torch_DIR \"${Torch_DIR}\")
find_package(Torch REQUIRED)
")
set(Impact_EXPORT_CODE_BUILD "
if(NOT ITK_BINARY_DIR)
  set(Torch_DIR \"${Torch_DIR}\")
  find_package(Torch REQUIRED)
endif()
")