setup.py and cmake improvements (#5269)

* 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>

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)

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()
+        multiprocessing.pool.ThreadPool(NUM_JOBS).map(_single_compile, objects)
-        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)
 
         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'

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_C_FLAGS="$BUILD_C_FLAGS $USER_CFLAGS" \
-              -DCMAKE_CXX_FLAGS="$BUILD_C_FLAGS $CPP_FLAGS" \
+              -DCMAKE_CXX_FLAGS="$BUILD_C_FLAGS $CPP_FLAGS $USER_CFLAGS" \
-              -DCMAKE_EXE_LINKER_FLAGS="$LDFLAGS" \
+              -DCMAKE_EXE_LINKER_FLAGS="$LDFLAGS $USER_LDFLAGS" \
-              -DCMAKE_SHARED_LINKER_FLAGS="$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_C_FLAGS="$C_FLAGS $USER_CFLAGS" \
-               -DCMAKE_CXX_FLAGS="$C_FLAGS $CPP_FLAGS"
+               -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) \
+      -DCMAKE_BUILD_TYPE=$([ $DEBUG ] && echo Debug || echo Release) \
-  -DNO_CUDA=$((1-$WITH_CUDA)) \
+      -DNO_CUDA=$((1-$WITH_CUDA)) \
-  -DNO_NNPACK=$((1-$WITH_NNPACK)) \
+      -DNO_NNPACK=$((1-$WITH_NNPACK)) \
-  -DCUDNN_INCLUDE_DIR=$CUDNN_INCLUDE_DIR \
+      -DCUDNN_INCLUDE_DIR=$CUDNN_INCLUDE_DIR \
-  -DCUDNN_LIB_DIR=$CUDNN_LIB_DIR \
+      -DCUDNN_LIB_DIR=$CUDNN_LIB_DIR \
-  -DCUDNN_LIBRARY=$CUDNN_LIBRARY \
+      -DCUDNN_LIBRARY=$CUDNN_LIBRARY \
-  -DATEN_NO_CONTRIB=1 \
+      -DATEN_NO_CONTRIB=1 \
-  -DCMAKE_INSTALL_PREFIX="$INSTALL_DIR" \
+      -DCMAKE_INSTALL_PREFIX="$INSTALL_DIR" \
-  -DCMAKE_EXPORT_COMPILE_COMMANDS=1
+      -DCMAKE_EXPORT_COMPILE_COMMANDS=1 \
-  # purpusefully not passing C_FLAGS for the same reason as above
+      -DCMAKE_C_FLAGS="$USER_CFLAGS" \
-  ${CMAKE_INSTALL} -j$(getconf _NPROCESSORS_ONLN)
+      -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 ../..
 }