Move redundant code that checks NumPy during build to a helper module and add an option to disable building with NumPy

Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/21417

Reviewed By: ezyang

Differential Revision: D15694357

Pulled By: fmassa

fbshipit-source-id: bc1bda23349ba4531f19619fa4adecb846225c20

cmake/Dependencies.cmake

@@ -674,8 +674,11 @@ if(BUILD_PYTHON)
   # don't want to overwrite it because we trust python more than cmake
   if (NUMPY_INCLUDE_DIR)
     set(NUMPY_FOUND ON)
-  else()
+  elseif(USE_NUMPY)
     find_package(NumPy)
+    if(NOT NUMPY_FOUND)
+      message(WARNING "NumPy could not be found. Not building with NumPy. Suppress this warning with -DUSE_NUMPY=OFF")
+    endif()
   endif()
 
   if(PYTHONINTERP_FOUND AND PYTHONLIBS_FOUND)

setup.py

@@ -33,6 +33,9 @@
 #   USE_FBGEMM=0
 #     disables the FBGEMM build
 #
+#   USE_NUMPY=0
+#     disables the NumPy build
+#
 #   BUILD_TEST=0
 #     disables the test build
 #
@@ -187,6 +190,7 @@ from tools.setup_helpers.cudnn import USE_CUDNN, CUDNN_LIBRARY, CUDNN_INCLUDE_DI
 from tools.setup_helpers.rocm import USE_ROCM
 from tools.setup_helpers.miopen import USE_MIOPEN, MIOPEN_LIBRARY, MIOPEN_INCLUDE_DIR
 from tools.setup_helpers.nccl import USE_NCCL, USE_SYSTEM_NCCL, NCCL_SYSTEM_LIB, NCCL_INCLUDE_DIR
+from tools.setup_helpers.numpy_ import USE_NUMPY
 from tools.setup_helpers.dist_check import USE_DISTRIBUTED
 ################################################################################
 # Parameters parsed from environment
@@ -619,14 +623,6 @@ main_sources = ["torch/csrc/stub.cpp"]
 # before libcaffe2.so in the linker command.
 main_link_args.extend(CAFFE2_LIBS)
 
-try:
-    import numpy as np
-except ImportError:
-    USE_NUMPY = False
-else:
-    NUMPY_INCLUDE_DIR = np.get_include()
-    USE_NUMPY = True
-
 if USE_CUDA:
     if IS_WINDOWS:
         cuda_lib_path = CUDA_HOME + '/lib/x64/'

tools/setup_helpers/cmake.py

@@ -16,6 +16,7 @@ from .cuda import USE_CUDA
 from .dist_check import USE_DISTRIBUTED, USE_GLOO_IBVERBS
 from .nccl import (USE_SYSTEM_NCCL, NCCL_INCLUDE_DIR, NCCL_ROOT_DIR,
                    NCCL_SYSTEM_LIB, USE_NCCL)
+from .numpy_ import USE_NUMPY, NUMPY_INCLUDE_DIR
 from .rocm import USE_ROCM
 from .nnpack import USE_NNPACK
 from .qnnpack import USE_QNNPACK
@@ -110,14 +111,6 @@ def run(version,
         if IS_64BIT:
             cmake_args.append('-Ax64')
             cmake_args.append('-Thost=x64')
-    try:
-        import numpy as np
-    except ImportError:
-        USE_NUMPY = False
-        NUMPY_INCLUDE_DIR = None
-    else:
-        NUMPY_INCLUDE_DIR = np.get_include()
-        USE_NUMPY = True
 
     cflags = os.getenv('CFLAGS', "") + " " + os.getenv('CPPFLAGS', "")
     ldflags = os.getenv('LDFLAGS', "")

tools/setup_helpers/numpy_.py

@@ -0,0 +1,19 @@
+"NumPy helper."
+
+from .env import check_negative_env_flag
+
+
+# Set USE_NUMPY to what the user wants, because even if we fail here, cmake
+# will check for the presence of NumPy again (`cmake/Dependencies.cmake`).
+USE_NUMPY = not check_negative_env_flag('USE_NUMPY')
+NUMPY_INCLUDE_DIR = None
+
+if USE_NUMPY:
+    try:
+        import numpy as np
+    except ImportError:
+        pass
+    else:
+        # To reach here, the user must has not disabled NumPy build and the
+        # NumPy library is present in the system.
+        NUMPY_INCLUDE_DIR = np.get_include()