pytorch

1
"Manages CMake."
2

3
from __future__ import annotations
4

5
import multiprocessing
6
import os
7
import platform
8
import sys
9
import sysconfig
10
from distutils.version import LooseVersion
11
from subprocess import CalledProcessError, check_call, check_output
12
from typing import Any, cast
13

14
from . import which
15
from .cmake_utils import CMakeValue, get_cmake_cache_variables_from_file
16
from .env import BUILD_DIR, check_negative_env_flag, IS_64BIT, IS_DARWIN, IS_WINDOWS
17

18

19
def _mkdir_p(d: str) -> None:
20
    try:
21
        os.makedirs(d, exist_ok=True)
22
    except OSError as e:
23
        raise RuntimeError(
24
            f"Failed to create folder {os.path.abspath(d)}: {e.strerror}"
25
        ) from e
26

27

28
# Ninja
29
# Use ninja if it is on the PATH. Previous version of PyTorch required the
30
# ninja python package, but we no longer use it, so we do not have to import it
31
USE_NINJA = not check_negative_env_flag("USE_NINJA") and which("ninja") is not None
32
if "CMAKE_GENERATOR" in os.environ:
33
    USE_NINJA = os.environ["CMAKE_GENERATOR"].lower() == "ninja"
34

35

36
class CMake:
37
    "Manages cmake."
38

39
    def __init__(self, build_dir: str = BUILD_DIR) -> None:
40
        self._cmake_command = CMake._get_cmake_command()
41
        self.build_dir = build_dir
42

43
    @property
44
    def _cmake_cache_file(self) -> str:
45
        r"""Returns the path to CMakeCache.txt.
46

47
        Returns:
48
          string: The path to CMakeCache.txt.
49
        """
50
        return os.path.join(self.build_dir, "CMakeCache.txt")
51

52
    @staticmethod
53
    def _get_cmake_command() -> str:
54
        "Returns cmake command."
55

56
        cmake_command = "cmake"
57
        if IS_WINDOWS:
58
            return cmake_command
59
        cmake3_version = CMake._get_version(which("cmake3"))
60
        cmake_version = CMake._get_version(which("cmake"))
61

62
        _cmake_min_version = LooseVersion("3.18.0")
63
        if all(
64
            ver is None or ver < _cmake_min_version
65
            for ver in [cmake_version, cmake3_version]
66
        ):
67
            raise RuntimeError("no cmake or cmake3 with version >= 3.18.0 found")
68

69
        if cmake3_version is None:
70
            cmake_command = "cmake"
71
        elif cmake_version is None:
72
            cmake_command = "cmake3"
73
        else:
74
            if cmake3_version >= cmake_version:
75
                cmake_command = "cmake3"
76
            else:
77
                cmake_command = "cmake"
78
        return cmake_command
79

80
    @staticmethod
81
    def _get_version(cmd: str | None) -> Any:
82
        "Returns cmake version."
83

84
        if cmd is None:
85
            return None
86
        for line in check_output([cmd, "--version"]).decode("utf-8").split("\n"):
87
            if "version" in line:
88
                return LooseVersion(line.strip().split(" ")[2])
89
        raise RuntimeError("no version found")
90

91
    def run(self, args: list[str], env: dict[str, str]) -> None:
92
        "Executes cmake with arguments and an environment."
93

94
        command = [self._cmake_command] + args
95
        print(" ".join(command))
96
        try:
97
            check_call(command, cwd=self.build_dir, env=env)
98
        except (CalledProcessError, KeyboardInterrupt) as e:
99
            # This error indicates that there was a problem with cmake, the
100
            # Python backtrace adds no signal here so skip over it by catching
101
            # the error and exiting manually
102
            sys.exit(1)
103

104
    @staticmethod
105
    def defines(args: list[str], **kwargs: CMakeValue) -> None:
106
        "Adds definitions to a cmake argument list."
107
        for key, value in sorted(kwargs.items()):
108
            if value is not None:
109
                args.append(f"-D{key}={value}")
110

111
    def get_cmake_cache_variables(self) -> dict[str, CMakeValue]:
112
        r"""Gets values in CMakeCache.txt into a dictionary.
113
        Returns:
114
          dict: A ``dict`` containing the value of cached CMake variables.
115
        """
116
        with open(self._cmake_cache_file) as f:
117
            return get_cmake_cache_variables_from_file(f)
118

119
    def generate(
120
        self,
121
        version: str | None,
122
        cmake_python_library: str | None,
123
        build_python: bool,
124
        build_test: bool,
125
        my_env: dict[str, str],
126
        rerun: bool,
127
    ) -> None:
128
        "Runs cmake to generate native build files."
129

130
        if rerun and os.path.isfile(self._cmake_cache_file):
131
            os.remove(self._cmake_cache_file)
132

133
        ninja_build_file = os.path.join(self.build_dir, "build.ninja")
134
        if os.path.exists(self._cmake_cache_file) and not (
135
            USE_NINJA and not os.path.exists(ninja_build_file)
136
        ):
137
            # Everything's in place. Do not rerun.
138
            return
139

140
        args = []
141
        if USE_NINJA:
142
            # Avoid conflicts in '-G' and the `CMAKE_GENERATOR`
143
            os.environ["CMAKE_GENERATOR"] = "Ninja"
144
            args.append("-GNinja")
145
        elif IS_WINDOWS:
146
            generator = os.getenv("CMAKE_GENERATOR", "Visual Studio 16 2019")
147
            supported = ["Visual Studio 16 2019", "Visual Studio 17 2022"]
148
            if generator not in supported:
149
                print("Unsupported `CMAKE_GENERATOR`: " + generator)
150
                print("Please set it to one of the following values: ")
151
                print("\n".join(supported))
152
                sys.exit(1)
153
            args.append("-G" + generator)
154
            toolset_dict = {}
155
            toolset_version = os.getenv("CMAKE_GENERATOR_TOOLSET_VERSION")
156
            if toolset_version is not None:
157
                toolset_dict["version"] = toolset_version
158
                curr_toolset = os.getenv("VCToolsVersion")
159
                if curr_toolset is None:
160
                    print(
161
                        "When you specify `CMAKE_GENERATOR_TOOLSET_VERSION`, you must also "
162
                        "activate the vs environment of this version. Please read the notes "
163
                        "in the build steps carefully."
164
                    )
165
                    sys.exit(1)
166
            if IS_64BIT:
167
                if platform.machine() == "ARM64":
168
                    args.append("-A ARM64")
169
                else:
170
                    args.append("-Ax64")
171
                    toolset_dict["host"] = "x64"
172
            if toolset_dict:
173
                toolset_expr = ",".join([f"{k}={v}" for k, v in toolset_dict.items()])
174
                args.append("-T" + toolset_expr)
175

176
        base_dir = os.path.dirname(
177
            os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
178
        )
179
        install_dir = os.path.join(base_dir, "torch")
180

181
        _mkdir_p(install_dir)
182
        _mkdir_p(self.build_dir)
183

184
        # Store build options that are directly stored in environment variables
185
        build_options: dict[str, CMakeValue] = {}
186

187
        # Build options that do not start with "BUILD_", "USE_", or "CMAKE_" and are directly controlled by env vars.
188
        # This is a dict that maps environment variables to the corresponding variable name in CMake.
189
        additional_options = {
190
            # Key: environment variable name. Value: Corresponding variable name to be passed to CMake. If you are
191
            # adding a new build option to this block: Consider making these two names identical and adding this option
192
            # in the block below.
193
            "_GLIBCXX_USE_CXX11_ABI": "GLIBCXX_USE_CXX11_ABI",
194
            "CUDNN_LIB_DIR": "CUDNN_LIBRARY",
195
            "USE_CUDA_STATIC_LINK": "CAFFE2_STATIC_LINK_CUDA",
196
        }
197
        additional_options.update(
198
            {
199
                # Build options that have the same environment variable name and CMake variable name and that do not start
200
                # with "BUILD_", "USE_", or "CMAKE_". If you are adding a new build option, also make sure you add it to
201
                # CMakeLists.txt.
202
                var: var
203
                for var in (
204
                    "UBSAN_FLAGS",
205
                    "BLAS",
206
                    "WITH_BLAS",
207
                    "CUDA_HOST_COMPILER",
208
                    "CUDA_NVCC_EXECUTABLE",
209
                    "CUDA_SEPARABLE_COMPILATION",
210
                    "CUDNN_LIBRARY",
211
                    "CUDNN_INCLUDE_DIR",
212
                    "CUDNN_ROOT",
213
                    "EXPERIMENTAL_SINGLE_THREAD_POOL",
214
                    "INSTALL_TEST",
215
                    "JAVA_HOME",
216
                    "INTEL_MKL_DIR",
217
                    "INTEL_OMP_DIR",
218
                    "MKL_THREADING",
219
                    "MKLDNN_CPU_RUNTIME",
220
                    "MSVC_Z7_OVERRIDE",
221
                    "CAFFE2_USE_MSVC_STATIC_RUNTIME",
222
                    "Numa_INCLUDE_DIR",
223
                    "Numa_LIBRARIES",
224
                    "ONNX_ML",
225
                    "ONNX_NAMESPACE",
226
                    "ATEN_THREADING",
227
                    "WERROR",
228
                    "OPENSSL_ROOT_DIR",
229
                    "STATIC_DISPATCH_BACKEND",
230
                    "SELECTED_OP_LIST",
231
                    "TORCH_CUDA_ARCH_LIST",
232
                    "TRACING_BASED",
233
                    "PYTHON_LIB_REL_PATH",
234
                )
235
            }
236
        )
237

238
        # Aliases which are lower priority than their canonical option
239
        low_priority_aliases = {
240
            "CUDA_HOST_COMPILER": "CMAKE_CUDA_HOST_COMPILER",
241
            "CUDAHOSTCXX": "CUDA_HOST_COMPILER",
242
            "CMAKE_CUDA_HOST_COMPILER": "CUDA_HOST_COMPILER",
243
            "CMAKE_CUDA_COMPILER": "CUDA_NVCC_EXECUTABLE",
244
            "CUDACXX": "CUDA_NVCC_EXECUTABLE",
245
        }
246
        for var, val in my_env.items():
247
            # We currently pass over all environment variables that start with "BUILD_", "USE_", and "CMAKE_". This is
248
            # because we currently have no reliable way to get the list of all build options we have specified in
249
            # CMakeLists.txt. (`cmake -L` won't print dependent options when the dependency condition is not met.) We
250
            # will possibly change this in the future by parsing CMakeLists.txt ourselves (then additional_options would
251
            # also not be needed to be specified here).
252
            true_var = additional_options.get(var)
253
            if true_var is not None:
254
                build_options[true_var] = val
255
            elif var.startswith(("BUILD_", "USE_", "CMAKE_")) or var.endswith(
256
                ("EXITCODE", "EXITCODE__TRYRUN_OUTPUT")
257
            ):
258
                build_options[var] = val
259

260
            if var in low_priority_aliases:
261
                key = low_priority_aliases[var]
262
                if key not in build_options:
263
                    build_options[key] = val
264

265
        # The default value cannot be easily obtained in CMakeLists.txt. We set it here.
266
        py_lib_path = sysconfig.get_path("purelib")
267
        cmake_prefix_path = build_options.get("CMAKE_PREFIX_PATH", None)
268
        if cmake_prefix_path:
269
            build_options["CMAKE_PREFIX_PATH"] = (
270
                py_lib_path + ";" + cast(str, cmake_prefix_path)
271
            )
272
        else:
273
            build_options["CMAKE_PREFIX_PATH"] = py_lib_path
274

275
        # Some options must be post-processed. Ideally, this list will be shrunk to only one or two options in the
276
        # future, as CMake can detect many of these libraries pretty comfortably. We have them here for now before CMake
277
        # integration is completed. They appear here not in the CMake.defines call below because they start with either
278
        # "BUILD_" or "USE_" and must be overwritten here.
279
        build_options.update(
280
            {
281
                # Note: Do not add new build options to this dict if it is directly read from environment variable -- you
282
                # only need to add one in `CMakeLists.txt`. All build options that start with "BUILD_", "USE_", or "CMAKE_"
283
                # are automatically passed to CMake; For other options you can add to additional_options above.
284
                "BUILD_PYTHON": build_python,
285
                "BUILD_TEST": build_test,
286
                # Most library detection should go to CMake script, except this one, which Python can do a much better job
287
                # due to NumPy's inherent Pythonic nature.
288
                "USE_NUMPY": not check_negative_env_flag("USE_NUMPY"),
289
            }
290
        )
291

292
        # Options starting with CMAKE_
293
        cmake__options = {
294
            "CMAKE_INSTALL_PREFIX": install_dir,
295
        }
296

297
        # We set some CMAKE_* options in our Python build code instead of relying on the user's direct settings. Emit an
298
        # error if the user also attempts to set these CMAKE options directly.
299
        specified_cmake__options = set(build_options).intersection(cmake__options)
300
        if len(specified_cmake__options) > 0:
301
            print(
302
                ", ".join(specified_cmake__options)
303
                + " should not be specified in the environment variable. They are directly set by PyTorch build script."
304
            )
305
            sys.exit(1)
306
        build_options.update(cmake__options)
307

308
        CMake.defines(
309
            args,
310
            Python_EXECUTABLE=sys.executable,
311
            TORCH_BUILD_VERSION=version,
312
            **build_options,
313
        )
314

315
        expected_wrapper = "/usr/local/opt/ccache/libexec"
316
        if IS_DARWIN and os.path.exists(expected_wrapper):
317
            if "CMAKE_C_COMPILER" not in build_options and "CC" not in os.environ:
318
                CMake.defines(args, CMAKE_C_COMPILER=f"{expected_wrapper}/gcc")
319
            if "CMAKE_CXX_COMPILER" not in build_options and "CXX" not in os.environ:
320
                CMake.defines(args, CMAKE_CXX_COMPILER=f"{expected_wrapper}/g++")
321

322
        for env_var_name in my_env:
323
            if env_var_name.startswith("gh"):
324
                # github env vars use utf-8, on windows, non-ascii code may
325
                # cause problem, so encode first
326
                try:
327
                    my_env[env_var_name] = str(my_env[env_var_name].encode("utf-8"))
328
                except UnicodeDecodeError as e:
329
                    shex = ":".join(f"{ord(c):02x}" for c in my_env[env_var_name])
330
                    print(
331
                        f"Invalid ENV[{env_var_name}] = {shex}",
332
                        file=sys.stderr,
333
                    )
334
                    print(e, file=sys.stderr)
335
        # According to the CMake manual, we should pass the arguments first,
336
        # and put the directory as the last element. Otherwise, these flags
337
        # may not be passed correctly.
338
        # Reference:
339
        # 1. https://cmake.org/cmake/help/latest/manual/cmake.1.html#synopsis
340
        # 2. https://stackoverflow.com/a/27169347
341
        args.append(base_dir)
342
        self.run(args, env=my_env)
343

344
    def build(self, my_env: dict[str, str]) -> None:
345
        "Runs cmake to build binaries."
346

347
        from .env import build_type
348

349
        build_args = [
350
            "--build",
351
            ".",
352
            "--target",
353
            "install",
354
            "--config",
355
            build_type.build_type_string,
356
        ]
357

358
        # Determine the parallelism according to the following
359
        # priorities:
360
        # 1) MAX_JOBS environment variable
361
        # 2) If using the Ninja build system, delegate decision to it.
362
        # 3) Otherwise, fall back to the number of processors.
363

364
        # Allow the user to set parallelism explicitly. If unset,
365
        # we'll try to figure it out.
366
        max_jobs = os.getenv("MAX_JOBS")
367

368
        if max_jobs is not None or not USE_NINJA:
369
            # Ninja is capable of figuring out the parallelism on its
370
            # own: only specify it explicitly if we are not using
371
            # Ninja.
372

373
            # This lists the number of processors available on the
374
            # machine. This may be an overestimate of the usable
375
            # processors if CPU scheduling affinity limits it
376
            # further. In the future, we should check for that with
377
            # os.sched_getaffinity(0) on platforms that support it.
378
            max_jobs = max_jobs or str(multiprocessing.cpu_count())
379

380
            # This ``if-else'' clause would be unnecessary when cmake
381
            # 3.12 becomes minimum, which provides a '-j' option:
382
            # build_args += ['-j', max_jobs] would be sufficient by
383
            # then. Until then, we use "--" to pass parameters to the
384
            # underlying build system.
385
            build_args += ["--"]
386
            if IS_WINDOWS and not USE_NINJA:
387
                # We are likely using msbuild here
388
                build_args += [f"/p:CL_MPCount={max_jobs}"]
389
            else:
390
                build_args += ["-j", max_jobs]
391
        self.run(build_args, my_env)
392