Skip to content

Commit 74de51d

Browse files
authored
Added "-fopenmp" cflags (#2783) (#3006)
* Fixes #2783 - Added "-fopenmp" cflags - Refactored cxx flags creation * Fixed for macosx
1 parent dbdc726 commit 74de51d

File tree

1 file changed

+6
-7
lines changed

1 file changed

+6
-7
lines changed

setup.py

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -175,7 +175,9 @@ def get_extensions():
175175

176176
define_macros = []
177177

178-
extra_compile_args = {}
178+
extra_compile_args = {
179+
'cxx': []
180+
}
179181
if (torch.cuda.is_available() and ((CUDA_HOME is not None) or is_rocm_pytorch)) \
180182
or os.getenv('FORCE_CUDA', '0') == '1':
181183
extension = CUDAExtension
@@ -190,16 +192,13 @@ def get_extensions():
190192
else:
191193
define_macros += [('WITH_HIP', None)]
192194
nvcc_flags = []
193-
extra_compile_args = {
194-
'cxx': [],
195-
'nvcc': nvcc_flags,
196-
}
195+
extra_compile_args['nvcc'] = nvcc_flags
197196

198197
if sys.platform == 'win32':
199198
define_macros += [('torchvision_EXPORTS', None)]
200-
201-
extra_compile_args.setdefault('cxx', [])
202199
extra_compile_args['cxx'].append('/MP')
200+
elif sys.platform == 'linux':
201+
extra_compile_args['cxx'].append('-fopenmp')
203202

204203
debug_mode = os.getenv('DEBUG', '0') == '1'
205204
if debug_mode:

0 commit comments

Comments
 (0)