mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
c4ff49f4c7
This patch enables Pytorch build from source with Ninja and 'Visual Studio 16 2019' CMake generator on Windows on Arm. Tests: - Build from source: 'python setup.py develop'. - Run simple Pytorch example: passed - python test\test_torch.py: -- same results as on x64 -- Ran 1344 tests, failures=2 Pull Request resolved: https://github.com/pytorch/pytorch/pull/72424
89 lines
3.3 KiB
Python
89 lines
3.3 KiB
Python
import os
|
|
import platform
|
|
from glob import glob
|
|
import shutil
|
|
from typing import Dict, Optional
|
|
|
|
from .setup_helpers.env import IS_64BIT, IS_WINDOWS, check_negative_env_flag
|
|
from .setup_helpers.cmake import USE_NINJA, CMake
|
|
|
|
from setuptools import distutils # type: ignore[import]
|
|
|
|
def _overlay_windows_vcvars(env: Dict[str, str]) -> Dict[str, str]:
|
|
vc_arch = 'x64' if IS_64BIT else 'x86'
|
|
|
|
if platform.machine() == 'ARM64':
|
|
vc_arch = 'x64_arm64'
|
|
|
|
# First Win11 Windows on Arm build version that supports x64 emulation
|
|
# is 10.0.22000.
|
|
win11_1st_version = (10, 0, 22000)
|
|
current_win_version = tuple(int(version_part) for version_part in
|
|
platform.version().split('.'))
|
|
if current_win_version < win11_1st_version:
|
|
vc_arch = 'x86_arm64'
|
|
print("Warning: 32-bit toolchain will be used, but 64-bit linker "
|
|
"is recommended to avoid out-of-memory linker error!")
|
|
print("Warning: Please consider upgrading to Win11, where x64 "
|
|
"emulation is enabled!")
|
|
|
|
vc_env: Dict[str, str] = distutils._msvccompiler._get_vc_env(vc_arch)
|
|
# Keys in `_get_vc_env` are always lowercase.
|
|
# We turn them into uppercase before overlaying vcvars
|
|
# because OS environ keys are always uppercase on Windows.
|
|
# https://stackoverflow.com/a/7797329
|
|
vc_env = {k.upper(): v for k, v in vc_env.items()}
|
|
for k, v in env.items():
|
|
uk = k.upper()
|
|
if uk not in vc_env:
|
|
vc_env[uk] = v
|
|
return vc_env
|
|
|
|
|
|
def _create_build_env() -> Dict[str, str]:
|
|
# XXX - our cmake file sometimes looks at the system environment
|
|
# and not cmake flags!
|
|
# you should NEVER add something to this list. It is bad practice to
|
|
# have cmake read the environment
|
|
my_env = os.environ.copy()
|
|
if 'CUDA_HOME' in my_env: # Keep CUDA_HOME. This env variable is still used in other part.
|
|
my_env['CUDA_BIN_PATH'] = my_env['CUDA_HOME']
|
|
elif IS_WINDOWS: # we should eventually make this as part of FindCUDA.
|
|
cuda_win = glob('C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v*.*')
|
|
if len(cuda_win) > 0:
|
|
my_env['CUDA_BIN_PATH'] = cuda_win[0]
|
|
|
|
if IS_WINDOWS and USE_NINJA:
|
|
# When using Ninja under Windows, the gcc toolchain will be chosen as
|
|
# default. But it should be set to MSVC as the user's first choice.
|
|
my_env = _overlay_windows_vcvars(my_env)
|
|
my_env.setdefault('CC', 'cl')
|
|
my_env.setdefault('CXX', 'cl')
|
|
return my_env
|
|
|
|
|
|
def build_caffe2(
|
|
version: Optional[str],
|
|
cmake_python_library: Optional[str],
|
|
build_python: bool,
|
|
rerun_cmake: bool,
|
|
cmake_only: bool,
|
|
cmake: CMake,
|
|
) -> None:
|
|
my_env = _create_build_env()
|
|
build_test = not check_negative_env_flag('BUILD_TEST')
|
|
cmake.generate(version,
|
|
cmake_python_library,
|
|
build_python,
|
|
build_test,
|
|
my_env,
|
|
rerun_cmake)
|
|
if cmake_only:
|
|
return
|
|
cmake.build(my_env)
|
|
if build_python:
|
|
caffe2_proto_dir = os.path.join(cmake.build_dir, 'caffe2', 'proto')
|
|
for proto_file in glob(os.path.join(caffe2_proto_dir, '*.py')):
|
|
if proto_file != os.path.join(caffe2_proto_dir, '__init__.py'):
|
|
shutil.copy(proto_file, os.path.join('caffe2', 'proto'))
|