frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-21 05:34:18 +08:00
Code Issues Packages Projects Releases Wiki Activity

setup.py and cmake improvements (#5269)

Browse Source 
* Document env vars and properly propagate MAX_JOBS down.

Signed-off-by: Edward Z. Yang <ezyang@fb.com>

* Apply CFLAGS and LDFLAGS environment variables to cmake builds.

Signed-off-by: Edward Z. Yang <ezyang@fb.com>

* Test that running built program works; fixes #5151.

Signed-off-by: Edward Z. Yang <ezyang@fb.com>

* CMake CR.

Signed-off-by: Edward Z. Yang <ezyang@fb.com>
This commit is contained in:
Edward Z. Yang
2018-02-20 16:55:57 -05:00
committed by GitHub
parent 5ce46be17c
commit 031412a14b
3 changed files with 155 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
aten/CMakeLists.txt
View File

@ -218,6 +218,26 @@ ENDIF(UNIX)
INCLUDE (CheckIncludeFile)
INCLUDE (CheckCSourceCompiles)
INCLUDE (CheckCSourceRuns)
# Check that our programs run. This is different from the native CMake compiler
# check, which just tests if the program compiles and links. This is important
# because with ASAN you might need to help the compiled library find some
# dynamic libraries.
CHECK_C_SOURCE_RUNS("
int main() { return 0; }
" COMPILER_WORKS)
IF(NOT COMPILER_WORKS)
# Force cmake to retest next time around
unset(COMPILER_WORKS CACHE)
MESSAGE(FATAL_ERROR
"Could not run a simple program built with your compiler. "
"If you are trying to use -fsanitize=address, make sure "
"libasan is properly installed on your system (you can confirm "
"if the problem is this by attempting to build and run a "
"small program.)")
ENDIF()
CHECK_INCLUDE_FILE(cpuid.h HAVE_CPUID_H)
# Check for a cpuid intrinsic
IF(HAVE_CPUID_H)

93
setup.py
View File

@ -1,3 +1,83 @@
# Welcome to the PyTorch setup.py.
#
# Environment variables you are probably interestd in:
#
# DEBUG
# build with -O0 and -g (debug symbols)
#
# MAX_JOBS
# maximum number of compile jobs we should use to compile your code
#
# NO_CUDA
# disables CUDA build
#
# CFLAGS
# flags to apply to both C and C++ files to be compiled (a quirk of setup.py
# which we have faithfully adhered to in our build system is that CFLAGS
# also applies to C++ files, in contrast to the default behavior of autogoo
# and cmake build systems.)
#
# CC
# the C/C++ compiler to use (NB: the CXX flag has no effect for distutils
# compiles, because distutils always uses CC to compile, even for C++
# files.
#
# Environment variables for feature toggles:
#
# NO_CUDNN
# disables the cuDNN build
#
# NO_NNPACK
# disables NNPACK build
#
# NO_DISTRIBUTED
# disables THD (distributed) build
#
# NO_SYSTEM_NCCL
# disables use of system-wide nccl (we will use our submoduled
# copy in torch/lib/nccl)
#
# WITH_GLOO_IBVERBS
# toggle features related to distributed support
#
# WITH_SCALARS
# build with native support for scalars (zero-dim tensors)
#
# PYTORCH_BINARY_BUILD
# toggle static linking against libstdc++, used when we're building
# binaries for distribution
#
# PYTORCH_BUILD_VERSION
# PYTORCH_BUILD_NUMBER
# specify the version of PyTorch, rather than the hard-coded version
# in this file; used when we're building binaries for distribution
#
# Environment variables we respect (these environment variables are
# conventional and are often understood/set by other software.)
#
# CUDA_HOME (Linux/OS X)
# CUDA_PATH (Windows)
# specify where CUDA is installed; usually /usr/local/cuda or
# /usr/local/cuda-x.y
#
# CUDNN_LIB_DIR
# CUDNN_INCLUDE_DIR
# CUDNN_LIBRARY
# specify where cuDNN is installed
#
# NCCL_ROOT_DIR
# NCCL_LIB_DIR
# NCCL_INCLUDE_DIR
# specify where nccl is installed
#
# NVTOOLSEXT_PATH (Windows only)
# specify where nvtoolsext is installed
#
# LIBRARY_PATH
# LD_LIBRARY_PATH
# we will search for libraries in these paths
from setuptools import setup, Extension, distutils, Command, find_packages
import setuptools.command.build_ext
import setuptools.command.install
@ -9,6 +89,7 @@ import distutils.command.clean
import platform
import subprocess
import shutil
import multiprocessing
import sys
import os
import json
@ -38,6 +119,11 @@ if 'WITH_SCALARS' not in os.environ:
os.environ['WITH_SCALARS'] = '1'
WITH_SCALARS = check_env_flag('WITH_SCALARS')
NUM_JOBS = multiprocessing.cpu_count()
max_jobs = os.getenv("MAX_JOBS")
if max_jobs is not None:
NUM_JOBS = min(NUM_JOBS, int(max_jobs))
try:
import ninja
WITH_NINJA = True
@ -63,11 +149,7 @@ if not WITH_NINJA:
def _single_compile(obj):
src, ext = build[obj]
self._compile(obj, src, ext, cc_args, extra_postargs, pp_opts)
num_jobs = multiprocessing.cpu_count()
max_jobs = os.getenv("MAX_JOBS")
if max_jobs is not None:
num_jobs = min(num_jobs, int(max_jobs))
multiprocessing.pool.ThreadPool(num_jobs).map(_single_compile, objects)
multiprocessing.pool.ThreadPool(NUM_JOBS).map(_single_compile, objects)
return objects
distutils.ccompiler.CCompiler.compile = parallelCCompile
@ -120,6 +202,7 @@ def build_libs(libs):
build_libs_cmd = ['bash', 'torch/lib/build_libs.sh']
my_env = os.environ.copy()
my_env["PYTORCH_PYTHON"] = sys.executable
my_env["NUM_JOBS"] = str(NUM_JOBS)
if not IS_WINDOWS:
if WITH_NINJA:
my_env["CMAKE_GENERATOR"] = '-GNinja'

66
torch/lib/build_libs.sh
View File

@ -6,7 +6,7 @@
# called standalone to compile the libraries outside of the overall PyTorch
# build process.
set -e
set -ex
# Options for building only a subset of the libraries
WITH_CUDA=0
@ -27,6 +27,24 @@ if [[ "$1" == "--with-gloo-ibverbs" ]]; then
shift
fi
# Save user specified env vars, we will manually propagate them
# to cmake. We copy distutils semantics, referring to
# cpython/Lib/distutils/sysconfig.py as the source of truth
USER_CFLAGS=""
USER_LDFLAGS=""
if [[ -n "$LDFLAGS" ]]; then
USER_LDFLAGS="$USER_LDFLAGS $LDFLAGS"
fi
if [[ -n "$CFLAGS" ]]; then
USER_CFLAGS="$USER_CFLAGS $CFLAGS"
USER_LDFLAGS="$USER_LDFLAGS $CFLAGS"
fi
if [[ -n "$CPPFLAGS" ]]; then
# Unlike distutils, NOT modifying CXX
USER_C_CFLAGS="$USER_CFLAGS $CPPFLAGS"
USER_LDFLAGS="$USER_LDFLAGS $CPPFLAGS"
fi
cd "$(dirname "$0")/../.."
PWD=`printf "%q\n" "$(pwd)"`
BASE_DIR="$PWD"
@ -72,6 +90,9 @@ CUDA_NVCC_FLAGS=$C_FLAGS
if [[ $CUDA_DEBUG -eq 1 ]]; then
CUDA_NVCC_FLAGS="$CUDA_NVCC_FLAGS -g -G"
fi
if [ -z "$NUM_JOBS" ]; then
NUM_JOBS="$(getconf _NPROCESSORS_ONLN)"
fi
# Used to build an individual library, e.g. build TH
function build() {
@ -89,10 +110,10 @@ function build() {
${CMAKE_GENERATOR} \
-DTorch_FOUND="1" \
-DCMAKE_INSTALL_PREFIX="$INSTALL_DIR" \
-DCMAKE_C_FLAGS="$BUILD_C_FLAGS" \
-DCMAKE_CXX_FLAGS="$BUILD_C_FLAGS $CPP_FLAGS" \
-DCMAKE_EXE_LINKER_FLAGS="$LDFLAGS" \
-DCMAKE_SHARED_LINKER_FLAGS="$LDFLAGS" \
-DCMAKE_C_FLAGS="$BUILD_C_FLAGS $USER_CFLAGS" \
-DCMAKE_CXX_FLAGS="$BUILD_C_FLAGS $CPP_FLAGS $USER_CFLAGS" \
-DCMAKE_EXE_LINKER_FLAGS="$LDFLAGS $USER_LDFLAGS" \
-DCMAKE_SHARED_LINKER_FLAGS="$LDFLAGS $USER_LDFLAGS" \
-DCMAKE_INSTALL_LIBDIR="$INSTALL_DIR/lib" \
-DCUDA_NVCC_FLAGS="$CUDA_NVCC_FLAGS" \
-DCMAKE_PREFIX_PATH="$INSTALL_DIR" \
@ -119,7 +140,7 @@ function build() {
-DCMAKE_BUILD_TYPE=$([ $DEBUG ] && echo Debug || echo Release) \
${@:2} \
-DCMAKE_EXPORT_COMPILE_COMMANDS=1
${CMAKE_INSTALL} -j$(getconf _NPROCESSORS_ONLN)
${CMAKE_INSTALL} -j"$NUM_JOBS"
cd ../..
local lib_prefix=$INSTALL_DIR/lib/lib$1
@ -144,8 +165,9 @@ function build_nccl() {
${CMAKE_GENERATOR} \
-DCMAKE_BUILD_TYPE=Release \
-DCMAKE_INSTALL_PREFIX="$INSTALL_DIR" \
-DCMAKE_C_FLAGS="$C_FLAGS" \
-DCMAKE_CXX_FLAGS="$C_FLAGS $CPP_FLAGS"
-DCMAKE_C_FLAGS="$C_FLAGS $USER_CFLAGS" \
-DCMAKE_CXX_FLAGS="$C_FLAGS $CPP_FLAGS $USER_CFLAGS" \
-DCMAKE_SHARED_LINKER_FLAGS="$USER_LDFLAGS"
${CMAKE_INSTALL}
mkdir -p ${INSTALL_DIR}/lib
cp "lib/libnccl.so.1" "${INSTALL_DIR}/lib/libnccl.so.1"
@ -168,17 +190,23 @@ function build_aten() {
cd build/aten
${CMAKE_VERSION} ../../../../aten \
${CMAKE_GENERATOR} \
-DCMAKE_BUILD_TYPE=$([ $DEBUG ] && echo Debug || echo Release) \
-DNO_CUDA=$((1-$WITH_CUDA)) \
-DNO_NNPACK=$((1-$WITH_NNPACK)) \
-DCUDNN_INCLUDE_DIR=$CUDNN_INCLUDE_DIR \
-DCUDNN_LIB_DIR=$CUDNN_LIB_DIR \
-DCUDNN_LIBRARY=$CUDNN_LIBRARY \
-DATEN_NO_CONTRIB=1 \
-DCMAKE_INSTALL_PREFIX="$INSTALL_DIR" \
-DCMAKE_EXPORT_COMPILE_COMMANDS=1
# purpusefully not passing C_FLAGS for the same reason as above
${CMAKE_INSTALL} -j$(getconf _NPROCESSORS_ONLN)
-DCMAKE_BUILD_TYPE=$([ $DEBUG ] && echo Debug || echo Release) \
-DNO_CUDA=$((1-$WITH_CUDA)) \
-DNO_NNPACK=$((1-$WITH_NNPACK)) \
-DCUDNN_INCLUDE_DIR=$CUDNN_INCLUDE_DIR \
-DCUDNN_LIB_DIR=$CUDNN_LIB_DIR \
-DCUDNN_LIBRARY=$CUDNN_LIBRARY \
-DATEN_NO_CONTRIB=1 \
-DCMAKE_INSTALL_PREFIX="$INSTALL_DIR" \
-DCMAKE_EXPORT_COMPILE_COMMANDS=1 \
-DCMAKE_C_FLAGS="$USER_CFLAGS" \
-DCMAKE_CXX_FLAGS="$USER_CFLAGS" \
-DCMAKE_EXE_LINKER_FLAGS="$USER_LDFLAGS" \
-DCMAKE_SHARED_LINKER_FLAGS="$USER_LDFLAGS"
# STOP!!! Are you trying to add a C or CXX flag? Add it
# to aten/CMakeLists.txt, not here. We need the vanilla
# cmake build to work.
${CMAKE_INSTALL} -j"$NUM_JOBS"
cd ../..
}