Enable ahead of time compilation for HIPExtensions using ninja (#37800)

Summary:
This pull request enables ahead of time compilation of HIPExtensions with ninja by setting appropriate compilation flags for ROCm environment. Also, this enables the unit test for testing cuda_extensions on ROCm as well as removing test for ahead of time compilation of extensions with ninja from ROCM_BLACKLIST

ezyang jeffdaily
Pull Request resolved: https://github.com/pytorch/pytorch/pull/37800

Differential Revision: D21408148

Pulled By: soumith

fbshipit-source-id: 146f4ffb3418f3534e6ce86805d3fe9c3eae84e1
This commit is contained in:
ashishfarmer
2020-05-05 20:51:43 -07:00
committed by Facebook GitHub Bot
parent 70f375becf
commit 402f635bbe
3 changed files with 6 additions and 7 deletions

View File

@ -324,10 +324,7 @@ class BuildExtension(build_ext, object):
# Test if we can use ninja. Fallback otherwise.
msg = ('Attempted to use ninja as the BuildExtension backend but '
'{}. Falling back to using the slow distutils backend.')
if IS_HIP_EXTENSION:
warnings.warn(msg.format('HIP extensions is not supported yet for ninja.'))
self.use_ninja = False
elif not _is_ninja_available():
if not _is_ninja_available():
warnings.warn(msg.format('we could not find ninja.'))
self.use_ninja = False
@ -435,7 +432,11 @@ class BuildExtension(build_ext, object):
cuda_post_cflags = extra_postargs['nvcc']
else:
cuda_post_cflags = list(extra_postargs)
cuda_post_cflags = unix_cuda_flags(cuda_post_cflags)
if IS_HIP_EXTENSION:
cuda_post_cflags = cuda_post_cflags + _get_rocm_arch_flags(cuda_post_cflags)
cuda_post_cflags = cuda_post_cflags + COMMON_HIPCC_FLAGS
else:
cuda_post_cflags = unix_cuda_flags(cuda_post_cflags)
append_std14_if_no_std_present(cuda_post_cflags)
_write_ninja_file_and_compile_objects(