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
diff --git a/cmake/Dependencies.cmake b/cmake/Dependencies.cmake
index 7ed9c2c..7372a2f 100644
--- a/cmake/Dependencies.cmake
+++ b/cmake/Dependencies.cmake
@@ -674,8 +674,11 @@
# 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)
diff --git a/setup.py b/setup.py
index ff71248..6af53ce 100644
--- a/setup.py
+++ b/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.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 @@
# 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/'
diff --git a/tools/setup_helpers/cmake.py b/tools/setup_helpers/cmake.py
index 3c6bd8d..7c62b7c 100644
--- a/tools/setup_helpers/cmake.py
+++ b/tools/setup_helpers/cmake.py
@@ -16,6 +16,7 @@
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 @@
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', "")
diff --git a/tools/setup_helpers/numpy_.py b/tools/setup_helpers/numpy_.py
new file mode 100644
index 0000000..b5f32cd
--- /dev/null
+++ b/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()
