Fixing reading of FBGEMM from env variables

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

Reviewed By: orionr

Differential Revision: D13406778

Pulled By: pjh5

fbshipit-source-id: 2265f01170fb7969cbdf4e44ca6ef183f5d8017d

diff --git a/setup.py b/setup.py
index 6507d71..aa501af 100644 (file)
--- a/setup.py
+++ b/setup.py
@@ -172,11 +172,11 @@ build_env_vars = ['BINARY', 'TEST', 'CAFFE2_OPS']
 from tools.setup_helpers.cuda import USE_CUDA, CUDA_HOME, CUDA_VERSION
 from tools.setup_helpers.build import (BUILD_BINARY, BUILD_TEST,
                                        BUILD_CAFFE2_OPS, USE_LEVELDB,
-                                       USE_LMDB, USE_OPENCV, USE_TENSORRT, USE_FFMPEG)
+                                       USE_LMDB, USE_OPENCV, USE_TENSORRT,
+                                       USE_FFMPEG, USE_FBGEMM)
 from tools.setup_helpers.rocm import USE_ROCM, ROCM_HOME, ROCM_VERSION
 from tools.setup_helpers.cudnn import (USE_CUDNN, CUDNN_LIBRARY,
                                        CUDNN_LIB_DIR, CUDNN_INCLUDE_DIR)
-from tools.setup_helpers.fbgemm import USE_FBGEMM
 from tools.setup_helpers.miopen import (USE_MIOPEN, MIOPEN_LIBRARY,
                                         MIOPEN_LIB_DIR, MIOPEN_INCLUDE_DIR)
 from tools.setup_helpers.nccl import USE_NCCL, USE_SYSTEM_NCCL, NCCL_LIB_DIR, \

diff --git a/tools/build_pytorch_libs.sh b/tools/build_pytorch_libs.sh
index f1b98b0..8ab575c 100755 (executable)
--- a/tools/build_pytorch_libs.sh
+++ b/tools/build_pytorch_libs.sh
@@ -52,7 +52,7 @@ fi
 
 # Options for building only a subset of the libraries
 USE_CUDA=0
-USE_FBGEMM=1
+USE_FBGEMM=0
 USE_ROCM=0
 USE_NNPACK=0
 USE_MKLDNN=0

diff --git a/tools/setup_helpers/build.py b/tools/setup_helpers/build.py
index 81e8255..252da64 100644 (file)
--- a/tools/setup_helpers/build.py
+++ b/tools/setup_helpers/build.py
@@ -8,3 +8,4 @@ USE_LMDB = check_env_flag('USE_LMDB')
 USE_OPENCV = check_env_flag('USE_OPENCV')
 USE_TENSORRT = check_env_flag('USE_TENSORRT')
 USE_FFMPEG = check_env_flag('USE_FFMPEG')
+USE_FBGEMM = not (check_env_flag('NO_FBGEMM') or check_negative_env_flag('USE_FBGEMM'))

diff --git a/tools/setup_helpers/fbgemm.py b/tools/setup_helpers/fbgemm.py
deleted file mode 100644 (file)
index bd30bcc..0000000
--- a/tools/setup_helpers/fbgemm.py
+++ /dev/null
@@ -1,12 +0,0 @@
-from .env import check_env_flag
-
-USE_FBGEMM = True
-
-if check_env_flag('NO_FBGEMM'):
-    USE_FBGEMM = False
-else:
-    # Enable FBGEMM if explicitly enabled
-    if check_env_flag('USE_FBGEMM'):
-        USE_FBGEMM = True
-    else:
-        USE_FBGEMM = False