mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
Summary: This commit removes the dependency on `build_pytorch_libs.sh` by moving the remaining functionality that is not expressible in cmake into python. Removing the indirection through bash also removes over 300 lines of environment munging code that is incredibly hard to understand because it passes a lot of secret parameters through `os.env`. Pull Request resolved: https://github.com/pytorch/pytorch/pull/16289 Reviewed By: ezyang Differential Revision: D13821662 Pulled By: zdevito fbshipit-source-id: d658d26925e3b1169ac1e3d44a159cf8a1f0d9b1
284 lines
11 KiB
Python
284 lines
11 KiB
Python
from .setup_helpers.env import (IS_ARM, IS_DARWIN, IS_LINUX, IS_PPC, IS_WINDOWS,
|
|
DEBUG, REL_WITH_DEB_INFO, USE_MKLDNN,
|
|
check_env_flag, check_negative_env_flag, hotpatch_build_env_vars)
|
|
|
|
import os
|
|
import sys
|
|
import distutils
|
|
import distutils.sysconfig
|
|
from distutils.file_util import copy_file
|
|
from distutils.dir_util import copy_tree
|
|
from subprocess import check_call, call, check_output
|
|
from distutils.version import LooseVersion
|
|
from .setup_helpers.cuda import USE_CUDA, CUDA_HOME
|
|
from .setup_helpers.dist_check import USE_DISTRIBUTED, USE_GLOO_IBVERBS
|
|
from .setup_helpers.nccl import USE_SYSTEM_NCCL, NCCL_INCLUDE_DIR, NCCL_ROOT_DIR, NCCL_SYSTEM_LIB
|
|
from .setup_helpers.rocm import ROCM_HOME, ROCM_VERSION, USE_ROCM
|
|
from .setup_helpers.nnpack import USE_NNPACK
|
|
from .setup_helpers.qnnpack import USE_QNNPACK
|
|
from .setup_helpers.cudnn import CUDNN_INCLUDE_DIR, CUDNN_LIB_DIR, CUDNN_LIBRARY, USE_CUDNN
|
|
|
|
|
|
from pprint import pprint
|
|
from glob import glob
|
|
import multiprocessing
|
|
import shutil
|
|
|
|
|
|
def which(thefile):
|
|
path = os.environ.get("PATH", os.defpath).split(os.pathsep)
|
|
for dir in path:
|
|
name = os.path.join(dir, thefile)
|
|
if (os.path.exists(name) and os.access(name, os.F_OK | os.X_OK)
|
|
and not os.path.isdir(name)):
|
|
return name
|
|
return None
|
|
|
|
|
|
def cmake_version(cmd):
|
|
for line in check_output([cmd, '--version']).split('\n'):
|
|
if 'version' in line:
|
|
return LooseVersion(line.strip().split(' ')[2])
|
|
raise Exception('no version found')
|
|
|
|
|
|
def get_cmake_command():
|
|
cmake_command = 'cmake'
|
|
if IS_WINDOWS:
|
|
return cmake_command
|
|
cmake3 = which('cmake3')
|
|
if cmake3 is not None:
|
|
cmake = which('cmake')
|
|
if cmake is not None:
|
|
bare_version = cmake_version(cmake)
|
|
if bare_version < LooseVersion("3.5.0") and cmake_version(cmake3) > bare_version:
|
|
cmake_command = 'cmake3'
|
|
return cmake_command
|
|
|
|
|
|
def cmake_defines(lst, **kwargs):
|
|
for key in sorted(kwargs.keys()):
|
|
value = kwargs[key]
|
|
if value is not None:
|
|
lst.append('-D{}={}'.format(key, value))
|
|
|
|
|
|
# Ninja
|
|
try:
|
|
import ninja
|
|
USE_NINJA = True
|
|
except ImportError:
|
|
USE_NINJA = False
|
|
|
|
base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
torch_lib_dir = base_dir + "/torch/lib"
|
|
install_dir = base_dir + "/torch/lib/tmp_install"
|
|
build_type = "Release"
|
|
if DEBUG:
|
|
build_type = "Debug"
|
|
elif REL_WITH_DEB_INFO:
|
|
build_type = "RelWithDebInfo"
|
|
|
|
|
|
def mkdir_p(dir):
|
|
try:
|
|
os.makedirs(dir)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
def run_cmake(version,
|
|
cmake_python_library,
|
|
build_python,
|
|
build_test,
|
|
build_dir):
|
|
cmake_args = [
|
|
get_cmake_command(),
|
|
base_dir
|
|
]
|
|
if USE_NINJA:
|
|
cmake_args.append('-GNinja')
|
|
try:
|
|
import numpy as np
|
|
NUMPY_INCLUDE_DIR = np.get_include()
|
|
USE_NUMPY = True
|
|
except ImportError:
|
|
USE_NUMPY = False
|
|
NUMPY_INCLUDE_DIR = None
|
|
|
|
cflags = os.getenv('CFLAGS') or ""
|
|
ldflags = os.getenv('LDFLAGS') or ""
|
|
if IS_DARWIN:
|
|
ldflags += " -Wl,-rpath,@loader_path"
|
|
elif USE_ROCM:
|
|
ldflags += " -Wl,-rpath,\\\\\\$ORIGIN"
|
|
elif IS_WINDOWS:
|
|
cflags += " /EHa"
|
|
else:
|
|
ldflags += " -Wl,-rpath,$ORIGIN"
|
|
|
|
# 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()
|
|
my_env['PYTORCH_PYTHON'] = sys.executable
|
|
if USE_CUDNN:
|
|
my_env['CUDNN_LIBRARY'] = CUDNN_LIBRARY
|
|
my_env['CUDNN_INCLUDE_DIR'] = CUDNN_INCLUDE_DIR
|
|
if USE_CUDA:
|
|
my_env['CUDA_BIN_PATH'] = CUDA_HOME
|
|
|
|
mkdir_p(install_dir)
|
|
mkdir_p(build_dir)
|
|
|
|
cmake_defines(
|
|
cmake_args,
|
|
PYTHON_EXECUTABLE=sys.executable,
|
|
PYTHON_LIBRARY=cmake_python_library,
|
|
PYTHON_INCLUDE_DIR=distutils.sysconfig.get_python_inc(),
|
|
BUILDING_WITH_TORCH_LIBS="ON",
|
|
TORCH_BUILD_VERSION=version,
|
|
CMAKE_BUILD_TYPE=build_type,
|
|
BUILD_TORCH="ON",
|
|
BUILD_PYTHON=build_python,
|
|
BUILD_SHARED_LIBS=os.getenv("BUILD_SHARED_LIBS") or "ON",
|
|
BUILD_BINARY=check_env_flag('BUILD_BINARY'),
|
|
BUILD_TEST=build_test,
|
|
INSTALL_TEST=build_test,
|
|
BUILD_CAFFE2_OPS=not check_negative_env_flag('BUILD_CAFFE2_OPS'),
|
|
ONNX_NAMESPACE=os.getenv("ONNX_NAMESPACE") or "onnx_torch",
|
|
USE_CUDA=USE_CUDA,
|
|
USE_DISTRIBUTED=USE_DISTRIBUTED,
|
|
USE_FBGEMM=not (check_env_flag('NO_FBGEMM') or check_negative_env_flag('USE_FBGEMM')),
|
|
USE_NUMPY=USE_NUMPY,
|
|
NUMPY_INCLUDE_DIR=NUMPY_INCLUDE_DIR,
|
|
USE_SYSTEM_NCCL=USE_SYSTEM_NCCL,
|
|
NCCL_INCLUDE_DIR=NCCL_INCLUDE_DIR,
|
|
NCCL_ROOT_DIR=NCCL_ROOT_DIR,
|
|
NCCL_SYSTEM_LIB=NCCL_SYSTEM_LIB,
|
|
CAFFE2_STATIC_LINK_CUDA=check_env_flag('USE_CUDA_STATIC_LINK'),
|
|
USE_ROCM=USE_ROCM,
|
|
USE_NNPACK=USE_NNPACK,
|
|
USE_LEVELDB=check_env_flag('USE_LEVELDB'),
|
|
USE_LMDB=check_env_flag('USE_LMDB'),
|
|
USE_OPENCV=check_env_flag('USE_OPENCV'),
|
|
USE_QNNPACK=USE_QNNPACK,
|
|
USE_TENSORRT=check_env_flag('USE_TENSORRT'),
|
|
USE_FFMPEG=check_env_flag('USE_FFMPEG'),
|
|
USE_SYSTEM_EIGEN_INSTALL="OFF",
|
|
USE_MKLDNN=USE_MKLDNN,
|
|
NCCL_EXTERNAL=USE_CUDA,
|
|
CMAKE_INSTALL_PREFIX=install_dir,
|
|
CMAKE_C_FLAGS=cflags,
|
|
CMAKE_CXX_FLAGS=cflags,
|
|
CMAKE_EXE_LINKER_FLAGS=ldflags,
|
|
CMAKE_SHARED_LINKER_FLAGS=ldflags,
|
|
THD_SO_VERSION="1",
|
|
CMAKE_PREFIX_PATH=os.getenv('CMAKE_PREFIX_PATH') or distutils.sysconfig.get_python_lib(),
|
|
BLAS=os.getenv('BLAS'),
|
|
CUDA_NVCC_EXECUTABLE=os.getenv('CUDA_NVCC_EXECUTABLE'),
|
|
USE_REDIS=os.getenv('USE_REDIS'),
|
|
USE_GLOG=os.getenv('USE_GLOG'),
|
|
USE_GFLAGS=os.getenv('USE_GFLAGS'))
|
|
|
|
if USE_GLOO_IBVERBS:
|
|
cmake_defines(cmake_args, USE_IBVERBS="1", USE_GLOO_IBVERBS="1")
|
|
|
|
expected_wrapper = '/usr/local/opt/ccache/libexec'
|
|
if IS_DARWIN and os.path.exists(expected_wrapper):
|
|
cmake_defines(cmake_args,
|
|
CMAKE_C_COMPILER="{}/gcc".format(expected_wrapper),
|
|
CMAKE_CXX_COMPILER="{}/g++".format(expected_wrapper))
|
|
pprint(cmake_args)
|
|
check_call(['printenv'])
|
|
check_call(cmake_args, cwd=build_dir, env=my_env)
|
|
|
|
|
|
def copy_files(build_test):
|
|
def copy_all(pattern, dst):
|
|
for file in glob(pattern):
|
|
if os.path.isdir(file):
|
|
copy_tree(file, dst, update=True)
|
|
else:
|
|
copy_file(file, dst, update=True)
|
|
|
|
shutil.rmtree(install_dir + '/lib/cmake', ignore_errors=True)
|
|
shutil.rmtree(install_dir + '/lib/python', ignore_errors=True)
|
|
copy_all(install_dir + '/lib/*', torch_lib_dir)
|
|
if os.path.exists(install_dir + '/lib64'):
|
|
copy_all(install_dir + '/lib64/*', torch_lib_dir)
|
|
copy_file(base_dir + '/aten/src/THNN/generic/THNN.h', torch_lib_dir, update=True)
|
|
copy_file(base_dir + '/aten/src/THCUNN/generic/THCUNN.h', torch_lib_dir, update=True)
|
|
|
|
copy_tree(install_dir + '/include', torch_lib_dir + '/include', update=True)
|
|
if os.path.exists(install_dir + '/bin/'):
|
|
copy_all(install_dir + '/bin/*', torch_lib_dir)
|
|
|
|
if build_test:
|
|
# Copy the test files to pytorch/caffe2 manually
|
|
# They were built in pytorch/torch/lib/tmp_install/test
|
|
# Why do we do this? So, setup.py has this section called 'package_data' which
|
|
# you need to specify to include non-default files (usually .py files).
|
|
# package_data takes a map from 'python package' to 'globs of files to
|
|
# include'. By 'python package', it means a folder with an __init__.py file
|
|
# that's not excluded in the find_packages call earlier in setup.py. So to
|
|
# include our cpp_test into the site-packages folder in
|
|
# site-packages/caffe2/cpp_test, we have to copy the cpp_test folder into the
|
|
# root caffe2 folder and then tell setup.py to include them. Having another
|
|
# folder like site-packages/caffe2_cpp_test would also be possible by adding a
|
|
# caffe2_cpp_test folder to pytorch with an __init__.py in it.
|
|
mkdir_p(base_dir + '/caffe2/cpp_test/')
|
|
copy_tree(install_dir + '/test', base_dir + '/caffe2/cpp_test', update=True)
|
|
|
|
|
|
def build_caffe2(version,
|
|
cmake_python_library,
|
|
build_python,
|
|
rerun_cmake,
|
|
build_dir):
|
|
build_test = not check_negative_env_flag('BUILD_TEST')
|
|
if rerun_cmake or not os.path.exists('build/CMakeCache.txt'):
|
|
run_cmake(version,
|
|
cmake_python_library,
|
|
build_python,
|
|
build_test,
|
|
build_dir)
|
|
|
|
if IS_WINDOWS:
|
|
if USE_NINJA:
|
|
# sccache will fail if all cores are used for compiling
|
|
j = max(1, multiprocessing.cpu_count() - 1)
|
|
check_call(['cmake', '--build', '.', '--target', 'install', '--config', build_type, '--', '-j', str(j)],
|
|
cwd=build_dir)
|
|
else:
|
|
check_call(['msbuild', 'INSTALL.vcxproj', '/p:Configuration={}'.format(build_type)],
|
|
cwd=build_dir)
|
|
else:
|
|
if USE_NINJA:
|
|
check_call(['ninja', 'install'], cwd=build_dir)
|
|
else:
|
|
check_call(['make', '-j', str(multiprocessing.cpu_count()), 'install'], cwd=build_dir)
|
|
|
|
# in cmake, .cu compilation involves generating certain intermediates
|
|
# such as .cu.o and .cu.depend, and these intermediates finally get compiled
|
|
# into the final .so.
|
|
# Ninja updates build.ninja's timestamp after all dependent files have been built,
|
|
# and re-kicks cmake on incremental builds if any of the dependent files
|
|
# have a timestamp newer than build.ninja's timestamp.
|
|
# There is a cmake bug with the Ninja backend, where the .cu.depend files
|
|
# are still compiling by the time the build.ninja timestamp is updated,
|
|
# so the .cu.depend file's newer timestamp is screwing with ninja's incremental
|
|
# build detector.
|
|
# This line works around that bug by manually updating the build.ninja timestamp
|
|
# after the entire build is finished.
|
|
if os.path.exists('build/build.ninja'):
|
|
os.utime('build/build.ninja', None)
|
|
|
|
if build_python:
|
|
for proto_file in glob('build/caffe2/proto/*.py'):
|
|
if proto_file != 'build/caffe2/proto/__init__.py':
|
|
shutil.copyfile(proto_file, "caffe2/proto/" + os.path.basename(proto_file))
|
|
|
|
copy_files(build_test)
|