Skip to content

Commit

Permalink
Added "-fopenmp" cflags (#2783) (#3006)
Browse files Browse the repository at this point in the history
* Fixes #2783
- Added "-fopenmp" cflags
- Refactored cxx flags creation

* Fixed for macosx
  • Loading branch information
vfdev-5 authored Nov 16, 2020
1 parent dbdc726 commit 74de51d
Showing 1 changed file with 6 additions and 7 deletions.
13 changes: 6 additions & 7 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -175,7 +175,9 @@ def get_extensions():

define_macros = []

extra_compile_args = {}
extra_compile_args = {
'cxx': []
}
if (torch.cuda.is_available() and ((CUDA_HOME is not None) or is_rocm_pytorch)) \
or os.getenv('FORCE_CUDA', '0') == '1':
extension = CUDAExtension
Expand All @@ -190,16 +192,13 @@ def get_extensions():
else:
define_macros += [('WITH_HIP', None)]
nvcc_flags = []
extra_compile_args = {
'cxx': [],
'nvcc': nvcc_flags,
}
extra_compile_args['nvcc'] = nvcc_flags

if sys.platform == 'win32':
define_macros += [('torchvision_EXPORTS', None)]

extra_compile_args.setdefault('cxx', [])
extra_compile_args['cxx'].append('/MP')
elif sys.platform == 'linux':
extra_compile_args['cxx'].append('-fopenmp')

debug_mode = os.getenv('DEBUG', '0') == '1'
if debug_mode:
Expand Down

0 comments on commit 74de51d

Please sign in to comment.