From 031412a14b7c20e9030b00c26f30c5524a0eb028 Mon Sep 17 00:00:00 2001 From: "Edward Z. Yang" Date: Tue, 20 Feb 2018 16:55:57 -0500 Subject: [PATCH] setup.py and cmake improvements (#5269) * Document env vars and properly propagate MAX_JOBS down. Signed-off-by: Edward Z. Yang * Apply CFLAGS and LDFLAGS environment variables to cmake builds. Signed-off-by: Edward Z. Yang * Test that running built program works; fixes #5151. Signed-off-by: Edward Z. Yang * CMake CR. Signed-off-by: Edward Z. Yang --- aten/CMakeLists.txt | 20 +++++++++ setup.py | 93 ++++++++++++++++++++++++++++++++++++++--- torch/lib/build_libs.sh | 66 ++++++++++++++++++++--------- 3 files changed, 155 insertions(+), 24 deletions(-) diff --git a/aten/CMakeLists.txt b/aten/CMakeLists.txt index 9dd0a8010d9..64690837a54 100644 --- a/aten/CMakeLists.txt +++ b/aten/CMakeLists.txt @@ -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) diff --git a/setup.py b/setup.py index 20d74a879cc..9d0c5427603 100644 --- a/setup.py +++ b/setup.py @@ -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' diff --git a/torch/lib/build_libs.sh b/torch/lib/build_libs.sh index 63fb011f32b..4ebeaaefdbe 100755 --- a/torch/lib/build_libs.sh +++ b/torch/lib/build_libs.sh @@ -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 ../.. }