frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-11-11 22:34:53 +08:00
Code Issues Packages Projects Releases Wiki Activity

Refactor hotpatch_vars and apply it to libtorch (#14976)

Browse Source 
Summary:
Fixes #14801.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/14976

Differential Revision: D13485381

Pulled By: soumith

fbshipit-source-id: 0af3c2e1b90988d56f6f85632328d1e4b788ffd2
This commit is contained in:
peter
2018-12-16 21:50:43 -08:00
committed by Facebook Github Bot
parent 656b565a0f
commit d71fac20eb
3 changed files with 32 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
tools/setup_helpers/env.py
View File

@ -27,3 +27,27 @@ def gather_paths(env_vars):
def lib_paths_from_base(base_path):
return [os.path.join(base_path, s) for s in ['lib/x64', 'lib', 'lib64']]
def hotpatch_var(var, prefix='USE_'):
if check_env_flag('NO_' + var):
os.environ[prefix + var] = '0'
elif check_negative_env_flag('NO_' + var):
os.environ[prefix + var] = '1'
elif check_env_flag('WITH_' + var):
os.environ[prefix + var] = '1'
elif check_negative_env_flag('WITH_' + var):
os.environ[prefix + var] = '0'
def hotpatch_build_env_vars():
# Before we run the setup_helpers, let's look for NO_* and WITH_*
# variables and hotpatch environment with the USE_* equivalent
use_env_vars = ['CUDA', 'CUDNN', 'FBGEMM', 'MIOPEN', 'MKLDNN', 'NNPACK', 'DISTRIBUTED',
'OPENCV', 'TENSORRT', 'QNNPACK', 'FFMPEG', 'SYSTEM_NCCL',
'GLOO_IBVERBS']
list(map(hotpatch_var, use_env_vars))
# Also hotpatch a few with BUILD_* equivalent
build_env_vars = ['BINARY', 'TEST', 'CAFFE2_OPS']
[hotpatch_var(v, 'BUILD_') for v in build_env_vars]