[ROCm] Added support for pytorch extensions to use HIP (pytorch#32669)

Summary:
This pull request has changes for:
1. Enabling a torch module with HIP code to be compiled by cpp_extensions.py
2. Fixes for hipify module to be able to be used by a torch extension

cc: ezyang iotamudelta jeffdaily
Pull Request resolved: pytorch#32669

Differential Revision: D20033893

Pulled By: zou3519

fbshipit-source-id: fd6ddc8cdcd3930f41008636bb2bc9dd26cdb008

Author: ashish

test/cpp_extensions/setup.py

@@ -3,7 +3,7 @@
 import os
 from setuptools import setup
 from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDAExtension
-from torch.utils.cpp_extension import CUDA_HOME
+from torch.utils.cpp_extension import CUDA_HOME, ROCM_HOME
 
 if sys.platform == 'win32':
     vc_version = os.getenv('VCToolsVersion', '')
@@ -35,6 +35,22 @@
         extra_compile_args={'cxx': CXX_FLAGS,
                             'nvcc': ['-O2']})
     ext_modules.append(extension)
+elif torch.cuda.is_available() and ROCM_HOME is not None:
+    from torch.utils.hipify import hipify_python
+    this_dir = os.path.dirname(os.path.abspath(__file__))
+    hipify_python.hipify(
+        project_directory=this_dir,
+        output_directory=this_dir,
+        includes="./*",
+        show_detailed=True,
+        is_pytorch_extension=True,)
+    extension = CUDAExtension(
+        'torch_test_cpp_extension.cuda', [
+            'cuda_extension.cpp',
+            'hip/hip_extension_kernel.hip',
+            'hip/hip_extension_kernel2.hip',
+        ])
+    ext_modules.append(extension)
 
 setup(
     name='torch_test_cpp_extension',

test/run_test.py

@@ -22,8 +22,8 @@
 TESTS = [
     'test_autograd',
     'test_complex',
-    'test_cpp_extensions_aot',
     'test_cpp_extensions_aot_no_ninja',
+    'test_cpp_extensions_aot_ninja',
     'test_cpp_extensions_jit',
     'distributed/test_c10d',
     'distributed/test_c10d_spawn',
@@ -93,8 +93,7 @@
 ]
 
 ROCM_BLACKLIST = [
-    'test_cpp_extensions_aot',
-    'test_cpp_extensions_aot_no_ninja',
+    'test_cpp_extensions_aot_ninja',
     'test_cpp_extensions_jit',
     'test_multiprocessing',
     'distributed/rpc/test_rpc_spawn',
@@ -129,7 +128,7 @@
 Ninja (https://ninja-build.org) is required for some of the C++ extensions
 tests, but it could not be found. Install ninja with `pip install ninja`
 or `conda install ninja`. Alternatively, disable said tests with
-`run_test.py --exclude test_cpp_extensions_aot test_cpp_extensions_jit`.
+`run_test.py --exclude test_cpp_extensions_aot_ninja test_cpp_extensions_jit`.
 """
 
 
@@ -199,8 +198,8 @@ def _test_cpp_extensions_aot(executable, test_module, test_directory, options, u
         os.environ['PYTHONPATH'] = python_path
 
 
-def test_cpp_extensions_aot(executable, test_module, test_directory, options):
-    return _test_cpp_extensions_aot(executable, test_module, test_directory,
+def test_cpp_extensions_aot_ninja(executable, test_module, test_directory, options):
+    return _test_cpp_extensions_aot(executable, 'test_cpp_extensions_aot', test_directory,
                                     options, use_ninja=True)
 
 
@@ -261,8 +260,8 @@ def test_distributed(executable, test_module, test_directory, options):
 
 CUSTOM_HANDLERS = {
     'test_cuda_primary_ctx': test_cuda_primary_ctx,
-    'test_cpp_extensions_aot': test_cpp_extensions_aot,
     'test_cpp_extensions_aot_no_ninja': test_cpp_extensions_aot_no_ninja,
+    'test_cpp_extensions_aot_ninja': test_cpp_extensions_aot_ninja,
     'distributed/test_distributed': test_distributed,
 }
 
@@ -430,8 +429,8 @@ def get_selected_tests(options):
     if sys.platform == 'win32' and not options.ignore_win_blacklist:
         target_arch = os.environ.get('VSCMD_ARG_TGT_ARCH')
         if target_arch != 'x64':
-            WINDOWS_BLACKLIST.append('cpp_extensions_aot')
             WINDOWS_BLACKLIST.append('cpp_extensions_aot_no_ninja')
+            WINDOWS_BLACKLIST.append('cpp_extensions_aot_ninja')
             WINDOWS_BLACKLIST.append('cpp_extensions_jit')
             WINDOWS_BLACKLIST.append('jit')
             WINDOWS_BLACKLIST.append('jit_fuser')

test/test_cpp_extensions_aot.py

@@ -14,17 +14,17 @@
 except ImportError:
     raise RuntimeError(
         "test_cpp_extensions_aot.py cannot be invoked directly. Run "
-        "`python run_test.py -i test_cpp_extensions_aot` instead."
+        "`python run_test.py -i test_cpp_extensions_aot_ninja` instead."
     )
 
 
 class TestCppExtensionAOT(common.TestCase):
     """Tests ahead-of-time cpp extensions
 
-    NOTE: run_test.py's test_cpp_extensions_aot_no_ninja target
-    also runs this test case, but with ninja disabled. If you are debugging
+    NOTE: run_test.py's test_cpp_extensions_aot_ninja target
+    also runs this test case, but with ninja enabled. If you are debugging
     a test failure here from the CI, check the logs for which target
-    (test_cpp_extensions_aot vs test_cpp_extensions_aot_no_ninja)
+    (test_cpp_extensions_aot_no_ninja vs test_cpp_extensions_aot_ninja)
     failed.
     """
 

torch/utils/cpp_extension.py

@@ -50,6 +50,44 @@ def _find_cuda_home():
         print("No CUDA runtime is found, using CUDA_HOME='{}'".format(cuda_home))
     return cuda_home
 
+def _find_rocm_home():
+    '''Finds the ROCm install path.'''
+    # Guess #1
+    rocm_home = os.environ.get('ROCM_HOME') or os.environ.get('ROCM_PATH')
+    if rocm_home is None:
+        # Guess #2
+        try:
+            hipcc = subprocess.check_output(
+                ['which', 'hipcc']).decode().rstrip('\r\n')
+            # this will be either <ROCM_HOME>/hip/bin/hipcc or <ROCM_HOME>/bin/hipcc
+            rocm_home = os.path.dirname(os.path.dirname(hipcc))
+            if os.path.basename(rocm_home) == 'hip':
+                rocm_home = os.path.dirname(rocm_home)
+        except Exception:
+            # Guess #3
+            rocm_home = '/opt/rocm'
+            if not os.path.exists(rocm_home):
+                rocm_home = None
+    if rocm_home and torch.version.hip is None:
+        print("No ROCm runtime is found, using ROCM_HOME='{}'".format(rocm_home))
+    return rocm_home
+
+
+def _join_rocm_home(*paths):
+    '''
+    Joins paths with ROCM_HOME, or raises an error if it ROCM_HOME is not set.
+
+    This is basically a lazy way of raising an error for missing $ROCM_HOME
+    only once we need to get any ROCm-specific path.
+    '''
+    if ROCM_HOME is None:
+        raise EnvironmentError('ROCM_HOME environment variable is not set. '
+                               'Please set it to your ROCm install root.')
+    elif IS_WINDOWS:
+        raise EnvironmentError('Building PyTorch extensions using '
+                               'ROCm and Windows is not supported.')
+    return os.path.join(ROCM_HOME, *paths)
+
 
 MINIMUM_GCC_VERSION = (4, 9, 0)
 MINIMUM_MSVC_VERSION = (19, 0, 24215)
@@ -85,6 +123,9 @@ def _find_cuda_home():
 
                               !! WARNING !!
 '''
+ROCM_HOME = _find_rocm_home()
+MIOPEN_HOME = _join_rocm_home('miopen') if ROCM_HOME else None
+IS_HIP_EXTENSION = True if ((ROCM_HOME is not None) and (torch.version.hip is not None)) else False
 CUDA_HOME = _find_cuda_home()
 CUDNN_HOME = os.environ.get('CUDNN_HOME') or os.environ.get('CUDNN_PATH')
 # PyTorch releases have the version pattern major.minor.patch, whereas when
@@ -101,6 +142,14 @@ def _find_cuda_home():
     '--expt-relaxed-constexpr'
 ]
 
+COMMON_HIPCC_FLAGS = [
+    '-fPIC',
+    '-D__HIP_PLATFORM_HCC__=1',
+    '-DCUDA_HAS_FP16=1',
+    '-D__HIP_NO_HALF_OPERATORS__=1',
+    '-D__HIP_NO_HALF_CONVERSIONS__=1',
+]
+
 JIT_EXTENSION_VERSIONER = ExtensionVersioner()
 
 
@@ -243,12 +292,15 @@ def __init__(self, *args, **kwargs):
         super(BuildExtension, self).__init__(*args, **kwargs)
         self.no_python_abi_suffix = kwargs.get("no_python_abi_suffix", False)
 
-        self.use_ninja = kwargs.get('use_ninja', True)
+        self.use_ninja = kwargs.get('use_ninja', False if IS_HIP_EXTENSION else True)
         if self.use_ninja:
             # 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 not _is_ninja_available():
+            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():
                 warnings.warn(msg.format('we could not find ninja.'))
                 self.use_ninja = False
 
@@ -259,8 +311,8 @@ def build_extensions(self):
             self._define_torch_extension_name(extension)
             self._add_gnu_cpp_abi_flag(extension)
 
-        # Register .cu and .cuh as valid source extensions.
-        self.compiler.src_extensions += ['.cu', '.cuh']
+        # Register .cu, .cuh and .hip as valid source extensions.
+        self.compiler.src_extensions += ['.cu', '.cuh', '.hip']
         # Save the original _compile method for later.
         if self.compiler.compiler_type == 'msvc':
             self.compiler._cpp_extensions += ['.cu', '.cuh']
@@ -289,15 +341,20 @@ def unix_wrap_single_compile(obj, src, ext, cc_args, extra_postargs, pp_opts):
             try:
                 original_compiler = self.compiler.compiler_so
                 if _is_cuda_file(src):
-                    nvcc = _join_cuda_home('bin', 'nvcc')
+                    nvcc = (_join_rocm_home('bin', 'hipcc') if IS_HIP_EXTENSION else _join_cuda_home('bin', 'nvcc'))
                     if not isinstance(nvcc, list):
                         nvcc = [nvcc]
                     self.compiler.set_executable('compiler_so', nvcc)
                     if isinstance(cflags, dict):
                         cflags = cflags['nvcc']
-                    cflags = unix_cuda_flags(cflags)
+                    if IS_HIP_EXTENSION:
+                        cflags = cflags + _get_rocm_arch_flags(cflags)
+                    else:
+                        cflags = unix_cuda_flags(cflags)
                 elif isinstance(cflags, dict):
                     cflags = cflags['cxx']
+                if IS_HIP_EXTENSION:
+                    cflags = cflags + COMMON_HIPCC_FLAGS
                 append_std14_if_no_std_present(cflags)
 
                 original_compile(obj, src, ext, cc_args, cflags, pp_opts)
@@ -649,13 +706,17 @@ def CUDAExtension(name, sources, *args, **kwargs):
     kwargs['library_dirs'] = library_dirs
 
     libraries = kwargs.get('libraries', [])
-    libraries.append('cudart')
     libraries.append('c10')
-    libraries.append('c10_cuda')
     libraries.append('torch')
     libraries.append('torch_cpu')
-    libraries.append('torch_cuda')
     libraries.append('torch_python')
+    if IS_HIP_EXTENSION:
+        libraries.append('c10_hip')
+        libraries.append('torch_hip')
+    else:
+        libraries.append('cudart')
+        libraries.append('c10_cuda')
+        libraries.append('torch_cuda')
     kwargs['libraries'] = libraries
 
     include_dirs = kwargs.get('include_dirs', [])
@@ -689,7 +750,12 @@ def include_paths(cuda=False):
         os.path.join(lib_include, 'TH'),
         os.path.join(lib_include, 'THC')
     ]
-    if cuda:
+    if cuda and IS_HIP_EXTENSION:
+        paths.append(os.path.join(lib_include, 'THH'))
+        paths.append(_join_rocm_home('include'))
+        if MIOPEN_HOME is not None:
+            paths.append(os.path.join(MIOPEN_HOME, 'include'))
+    elif cuda:
         cuda_home_include = _join_cuda_home('include')
         # if we have the Debian/Ubuntu packages for cuda, we get /usr as cuda home.
         # but gcc doesn't like having /usr/include passed explicitly
@@ -718,7 +784,10 @@ def library_paths(cuda=False):
     lib_path = os.path.join(torch_path, 'lib')
     paths.append(lib_path)
 
-    if cuda:
+    if cuda and IS_HIP_EXTENSION:
+        lib_dir = 'lib'
+        paths.append(_join_rocm_home(lib_dir))
+    elif cuda:
         if IS_WINDOWS:
             lib_dir = 'lib/x64'
         else:
@@ -1251,6 +1320,22 @@ def _get_cuda_arch_flags(cflags=None):
     return list(set(flags))
 
 
+def _get_rocm_arch_flags(cflags=None):
+    # If cflags is given, there may already be user-provided arch flags in it
+    # (from `extra_compile_args`)
+    if cflags is not None:
+        for flag in cflags:
+            if 'amdgpu-target' in flag:
+                return ['-fno-gpu-rdc']
+    return [
+        '--amdgpu-target=gfx803',
+        '--amdgpu-target=gfx900',
+        '--amdgpu-target=gfx906',
+        '--amdgpu-target=gfx908',
+        '-fno-gpu-rdc'
+    ]
+
+
 def _get_build_directory(name, verbose):
     root_extensions_directory = os.environ.get('TORCH_EXTENSIONS_DIR')
     if root_extensions_directory is None:
@@ -1567,4 +1652,7 @@ def _join_cuda_home(*paths):
 
 
 def _is_cuda_file(path):
-    return os.path.splitext(path)[1] in ['.cu', '.cuh']
+    valid_ext = ['.cu', '.cuh']
+    if IS_HIP_EXTENSION:
+        valid_ext.append('.hip')
+    return os.path.splitext(path)[1] in valid_ext

torch/utils/hipify/hipify_python.py

@@ -119,7 +119,8 @@ def preprocess(
         all_files,
         show_detailed=False,
         show_progress=True,
-        hip_clang_launch=False):
+        hip_clang_launch=False,
+        is_pytorch_extension=False):
     """
     Call preprocessor on selected files.
 
@@ -131,7 +132,7 @@ def preprocess(
     stats = {"unsupported_calls": [], "kernel_launches": []}
 
     for filepath in all_files:
-        result = preprocessor(output_directory, filepath, stats, hip_clang_launch)
+        result = preprocessor(output_directory, filepath, stats, hip_clang_launch, is_pytorch_extension)
         # Show what happened
         if show_progress:
             print(
@@ -605,7 +606,7 @@ def pattern(self):
 RE_THC_GENERIC_FILE = re.compile(r'#define THC_GENERIC_FILE "([^"]+)"')
 RE_CU_SUFFIX = re.compile(r'\.cu\b')  # be careful not to pick up .cuh
 
-def preprocessor(output_directory, filepath, stats, hip_clang_launch):
+def preprocessor(output_directory, filepath, stats, hip_clang_launch, is_pytorch_extension):
     """ Executes the CUDA -> HIP conversion on the specified file. """
     fin_path = os.path.join(output_directory, filepath)
     with open(fin_path, 'r') as fin:
@@ -616,14 +617,18 @@ def preprocessor(output_directory, filepath, stats, hip_clang_launch):
         os.makedirs(os.path.dirname(fout_path))
 
     # unsupported_calls statistics reporting is broken atm
-    if is_pytorch_file(filepath):
-        def pt_repl(m):
-            return PYTORCH_MAP[m.group(0)]
+    def pt_repl(m):
+        return PYTORCH_MAP[m.group(0)]
+
+    if is_pytorch_extension:
         output_source = RE_PYTORCH_PREPROCESSOR.sub(pt_repl, output_source)
     else:
-        def c2_repl(m):
-            return CAFFE2_MAP[m.group(0)]
-        output_source = RE_CAFFE2_PREPROCESSOR.sub(c2_repl, output_source)
+        if is_pytorch_file(filepath):
+            output_source = RE_PYTORCH_PREPROCESSOR.sub(pt_repl, output_source)
+        else:
+            def c2_repl(m):
+                return CAFFE2_MAP[m.group(0)]
+            output_source = RE_CAFFE2_PREPROCESSOR.sub(c2_repl, output_source)
 
     # Header rewrites
     def mk_repl(templ):
@@ -775,6 +780,7 @@ def hipify(
     ignores=(),
     show_progress=True,
     hip_clang_launch=False,
+    is_pytorch_extension=False,
 ):
     if project_directory == "":
         project_directory = os.getcwd()
@@ -803,4 +809,5 @@ def hipify(
         all_files,
         show_detailed=show_detailed,
         show_progress=show_progress,
-        hip_clang_launch=hip_clang_launch)
+        hip_clang_launch=hip_clang_launch,
+        is_pytorch_extension=is_pytorch_extension)