pytorch

1
from __future__ import annotations
2

3
import os
4
import platform
5
import shutil
6
from glob import glob
7

8
from setuptools import distutils  # type: ignore[import]
9

10
from .setup_helpers.cmake import CMake, USE_NINJA
11
from .setup_helpers.env import check_negative_env_flag, IS_64BIT, IS_WINDOWS
12

13

14
def _overlay_windows_vcvars(env: dict[str, str]) -> dict[str, str]:
15
    vc_arch = "x64" if IS_64BIT else "x86"
16

17
    if platform.machine() == "ARM64":
18
        vc_arch = "x64_arm64"
19

20
        # First Win11 Windows on Arm build version that supports x64 emulation
21
        # is 10.0.22000.
22
        win11_1st_version = (10, 0, 22000)
23
        current_win_version = tuple(
24
            int(version_part) for version_part in platform.version().split(".")
25
        )
26
        if current_win_version < win11_1st_version:
27
            vc_arch = "x86_arm64"
28
            print(
29
                "Warning: 32-bit toolchain will be used, but 64-bit linker "
30
                "is recommended to avoid out-of-memory linker error!"
31
            )
32
            print(
33
                "Warning: Please consider upgrading to Win11, where x64 "
34
                "emulation is enabled!"
35
            )
36

37
    vc_env: dict[str, str] = distutils._msvccompiler._get_vc_env(vc_arch)
38
    # Keys in `_get_vc_env` are always lowercase.
39
    # We turn them into uppercase before overlaying vcvars
40
    # because OS environ keys are always uppercase on Windows.
41
    # https://stackoverflow.com/a/7797329
42
    vc_env = {k.upper(): v for k, v in vc_env.items()}
43
    for k, v in env.items():
44
        uk = k.upper()
45
        if uk not in vc_env:
46
            vc_env[uk] = v
47
    return vc_env
48

49

50
def _create_build_env() -> dict[str, str]:
51
    # XXX - our cmake file sometimes looks at the system environment
52
    # and not cmake flags!
53
    # you should NEVER add something to this list. It is bad practice to
54
    # have cmake read the environment
55
    my_env = os.environ.copy()
56
    if (
57
        "CUDA_HOME" in my_env
58
    ):  # Keep CUDA_HOME. This env variable is still used in other part.
59
        my_env["CUDA_BIN_PATH"] = my_env["CUDA_HOME"]
60
    elif IS_WINDOWS:  # we should eventually make this as part of FindCUDA.
61
        cuda_win = glob("C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v*.*")
62
        if len(cuda_win) > 0:
63
            my_env["CUDA_BIN_PATH"] = cuda_win[0]
64

65
    if IS_WINDOWS and USE_NINJA:
66
        # When using Ninja under Windows, the gcc toolchain will be chosen as
67
        # default. But it should be set to MSVC as the user's first choice.
68
        my_env = _overlay_windows_vcvars(my_env)
69
        my_env.setdefault("CC", "cl")
70
        my_env.setdefault("CXX", "cl")
71
    return my_env
72

73

74
def build_caffe2(
75
    version: str | None,
76
    cmake_python_library: str | None,
77
    build_python: bool,
78
    rerun_cmake: bool,
79
    cmake_only: bool,
80
    cmake: CMake,
81
) -> None:
82
    my_env = _create_build_env()
83
    build_test = not check_negative_env_flag("BUILD_TEST")
84
    cmake.generate(
85
        version, cmake_python_library, build_python, build_test, my_env, rerun_cmake
86
    )
87
    if cmake_only:
88
        return
89
    cmake.build(my_env)
90
    if build_python:
91
        caffe2_proto_dir = os.path.join(cmake.build_dir, "caffe2", "proto")
92
        for proto_file in glob(os.path.join(caffe2_proto_dir, "*.py")):
93
            if proto_file != os.path.join(caffe2_proto_dir, "__init__.py"):
94
                shutil.copy(proto_file, os.path.join("caffe2", "proto"))
95