1from __future__ import annotations 2 3import os 4import platform 5import shutil 6from glob import glob 7 8from setuptools import distutils # type: ignore[import] 9 10from .setup_helpers.cmake import CMake, USE_NINJA 11from .setup_helpers.env import check_negative_env_flag, IS_64BIT, IS_WINDOWS 12 13 14def _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 50def _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 74def 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