ROCm: enable cpp_extensions.load/load_inline (pytorch#35897)

Summary:
This enables cpp_extensions.load/load_inline. This works by hipify-ing cuda sources.
Also enable tests.
CuDNN/MIOpen extensions aren't yet supported, I propose to not do this in this PR.
Pull Request resolved: pytorch#35897

Differential Revision: D20983279

Pulled By: ezyang

fbshipit-source-id: a5d0f5ac592d04488a6a46522c58e2ee0a6fd57c

Author: t-vi

test/run_test.py

@@ -121,7 +121,6 @@
     'distributed/rpc/test_dist_optimizer_spawn',
     'distributed/rpc/test_rpc_spawn',
     'test_cpp_extensions_aot_ninja',
-    'test_cpp_extensions_jit',
     'test_determination',
     'test_multiprocessing',
     'test_jit_simple',

test/test_cpp_extensions_jit.py

@@ -12,12 +12,13 @@
 import torch
 import torch.backends.cudnn
 import torch.utils.cpp_extension
-from torch.utils.cpp_extension import CUDA_HOME
+from torch.utils.cpp_extension import CUDA_HOME, ROCM_HOME
 
 
 TEST_CUDA = torch.cuda.is_available() and CUDA_HOME is not None
 TEST_CUDNN = False
-if TEST_CUDA:
+TEST_ROCM = torch.cuda.is_available() and torch.version.hip is not None and ROCM_HOME is not None
+if TEST_CUDA and torch.version.cuda is not None:  # the skip CUDNN test for ROCm
     CUDNN_HEADER_EXISTS = os.path.isfile(os.path.join(CUDA_HOME, "include/cudnn.h"))
     TEST_CUDNN = (
         TEST_CUDA and CUDNN_HEADER_EXISTS and torch.backends.cudnn.is_available()
@@ -109,6 +110,7 @@ def test_jit_cuda_extension(self):
             ],
             extra_cuda_cflags=["-O2"],
             verbose=True,
+            keep_intermediates=False,
         )
 
         x = torch.zeros(100, device="cuda", dtype=torch.float32)
@@ -184,6 +186,7 @@ def _check_cuobjdump_output(expected_values, is_ptx=False):
                 os.environ['TORCH_CUDA_ARCH_LIST'] = old_envvar
 
     @unittest.skipIf(not TEST_CUDA, "CUDA not found")
+    @unittest.skipIf(TEST_ROCM, "disabled on rocm")
     def test_jit_cuda_archflags(self):
         # Test a number of combinations:
         #   - the default for the machine we're testing on

torch/utils/cpp_extension.py

@@ -15,6 +15,8 @@
 import torch
 from .file_baton import FileBaton
 from ._cpp_extension_versioner import ExtensionVersioner
+from .hipify import hipify_python
+from .hipify.hipify_python import get_hip_file_path, GeneratedFileCleaner
 
 from setuptools.command.build_ext import build_ext
 
@@ -814,7 +816,8 @@ def load(name,
          build_directory=None,
          verbose=False,
          with_cuda=None,
-         is_python_module=True):
+         is_python_module=True,
+         keep_intermediates=True):
     '''
     Loads a PyTorch C++ extension just-in-time (JIT).
 
@@ -895,7 +898,8 @@ def load(name,
         build_directory or _get_build_directory(name, verbose),
         verbose,
         with_cuda,
-        is_python_module)
+        is_python_module,
+        keep_intermediates=keep_intermediates)
 
 
 def load_inline(name,
@@ -910,7 +914,8 @@ def load_inline(name,
                 verbose=False,
                 with_cuda=None,
                 is_python_module=True,
-                with_pytorch_error_handling=True):
+                with_pytorch_error_handling=True,
+                keep_intermediates=True):
     '''
     Loads a PyTorch C++ extension just-in-time (JIT) from string sources.
 
@@ -1043,7 +1048,8 @@ def load_inline(name,
         build_directory,
         verbose,
         with_cuda,
-        is_python_module)
+        is_python_module,
+        keep_intermediates=keep_intermediates)
 
 
 def _jit_compile(name,
@@ -1055,7 +1061,11 @@ def _jit_compile(name,
                  build_directory,
                  verbose,
                  with_cuda,
-                 is_python_module):
+                 is_python_module,
+                 keep_intermediates=True):
+    if with_cuda is None:
+        with_cuda = any(map(_is_cuda_file, sources))
+    with_cudnn = any(['cudnn' in f for f in extra_ldflags or []])
     old_version = JIT_EXTENSION_VERSIONER.get_version(name)
     version = JIT_EXTENSION_VERSIONER.bump_version_if_changed(
         name,
@@ -1074,16 +1084,27 @@ def _jit_compile(name,
         baton = FileBaton(os.path.join(build_directory, 'lock'))
         if baton.try_acquire():
             try:
-                _write_ninja_file_and_build_library(
-                    name=name,
-                    sources=sources,
-                    extra_cflags=extra_cflags or [],
-                    extra_cuda_cflags=extra_cuda_cflags or [],
-                    extra_ldflags=extra_ldflags or [],
-                    extra_include_paths=extra_include_paths or [],
-                    build_directory=build_directory,
-                    verbose=verbose,
-                    with_cuda=with_cuda)
+                with GeneratedFileCleaner(keep_intermediates=keep_intermediates) as clean_ctx:
+                    if IS_HIP_EXTENSION and (with_cuda or with_cudnn):
+                        hipify_python.hipify(
+                            project_directory=build_directory,
+                            output_directory=build_directory,
+                            includes=os.path.join(build_directory, '*'),
+                            extra_files=[os.path.abspath(s) for s in sources],
+                            show_detailed=verbose,
+                            is_pytorch_extension=True,
+                            clean_ctx=clean_ctx
+                        )
+                    _write_ninja_file_and_build_library(
+                        name=name,
+                        sources=sources,
+                        extra_cflags=extra_cflags or [],
+                        extra_cuda_cflags=extra_cuda_cflags or [],
+                        extra_ldflags=extra_ldflags or [],
+                        extra_include_paths=extra_include_paths or [],
+                        build_directory=build_directory,
+                        verbose=verbose,
+                        with_cuda=with_cuda)
             finally:
                 baton.release()
         else:
@@ -1231,10 +1252,10 @@ def _prepare_ldflags(extra_ldflags, with_cuda, verbose):
         extra_ldflags.append('-L{}'.format(lib_path))
         extra_ldflags.append('-lc10')
         if with_cuda:
-            extra_ldflags.append('-lc10_cuda')
+            extra_ldflags.append('-lc10_hip' if IS_HIP_EXTENSION else '-lc10_cuda')
         extra_ldflags.append('-ltorch_cpu')
         if with_cuda:
-            extra_ldflags.append('-ltorch_cuda')
+            extra_ldflags.append('-ltorch_hip' if IS_HIP_EXTENSION else '-ltorch_cuda')
         extra_ldflags.append('-ltorch')
         extra_ldflags.append('-ltorch_python')
 
@@ -1465,7 +1486,15 @@ def _write_ninja_file_to_build_library(path,
     else:
         cflags = common_cflags + ['-fPIC', '-std=c++14'] + extra_cflags
 
-    if with_cuda:
+    if with_cuda and IS_HIP_EXTENSION:
+        cuda_flags = ['-DWITH_HIP'] + cflags + COMMON_HIPCC_FLAGS
+        cuda_flags += extra_cuda_cflags
+        cuda_flags += _get_rocm_arch_flags(cuda_flags)
+        sources = [s if not _is_cuda_file(s) else
+                   os.path.abspath(os.path.join(
+                       path, get_hip_file_path(os.path.relpath(s, path))))
+                   for s in sources]
+    elif with_cuda:
         cuda_flags = common_cflags + COMMON_NVCC_FLAGS + _get_cuda_arch_flags()
         if IS_WINDOWS:
             for flag in COMMON_MSVC_FLAGS:
@@ -1568,7 +1597,11 @@ def sanitize_flags(flags):
     config = ['ninja_required_version = 1.3']
     config.append('cxx = {}'.format(compiler))
     if with_cuda:
-        config.append('nvcc = {}'.format(_join_cuda_home('bin', 'nvcc')))
+        if IS_HIP_EXTENSION:
+            nvcc = _join_rocm_home('bin', 'hipcc')
+        else:
+            nvcc = _join_cuda_home('bin', 'nvcc')
+        config.append('nvcc = {}'.format(nvcc))
 
     flags = ['cflags = {}'.format(' '.join(cflags))]
     flags.append('post_cflags = {}'.format(' '.join(post_cflags)))

torch/utils/hipify/hipify_python.py

@@ -72,6 +72,45 @@ class bcolors:
     BOLD = '\033[1m'
     UNDERLINE = '\033[4m'
 
+# To the programmer, the output of hipify most likely are intermediates.
+# This class allows users of hipify to ask for a cleanup by running the
+# hipify and compilation in a with instantiating this context manager class
+# with keep_intermediates=False.
+# The main usecase is the cpp_extensions, specifically the load method.
+# It is a good idea to keep intermediates (in case of errors or to
+# not recompile unchanged files), but in cases where you don't want to
+# keep them (e.g. in the CI), this can be used to remove files.
+class GeneratedFileCleaner:
+    """Context Manager to clean up generated files"""
+    def __init__(self, keep_intermediates=False):
+        self.keep_intermediates = keep_intermediates
+        self.files_to_clean = set()
+        self.dirs_to_clean = []
+
+    def __enter__(self):
+        return self
+
+    def open(self, fn, *args):
+        if not os.path.exists(fn):
+            self.files_to_clean.add(os.path.abspath(fn))
+        return open(fn, *args)
+
+    def makedirs(self, dn, exist_ok=False):
+        parent, n = os.path.split(dn)
+        if not n:
+            parent, n = os.path.split(parent)
+        if parent and n and not os.path.exists(parent):
+            self.makedirs(parent, exist_ok=True)
+        if not os.path.isdir(dn) or not exist_ok:
+            os.mkdir(dn)
+            self.dirs_to_clean.append(os.path.abspath(dn))
+
+    def __exit__(self, type, value, traceback):
+        if not self.keep_intermediates:
+            for f in self.files_to_clean:
+                os.unlink(f)
+            for d in self.dirs_to_clean[::-1]:
+                os.rmdir(d)
 
 def matched_files_iter(root_path, includes=('*',), ignores=(), extensions=(), out_of_place_only=False):
     def _fnmatch(filepath, patterns):
@@ -120,19 +159,24 @@ def preprocess(
         show_detailed=False,
         show_progress=True,
         hip_clang_launch=False,
-        is_pytorch_extension=False):
+        is_pytorch_extension=False,
+        clean_ctx=None):
     """
     Call preprocessor on selected files.
 
     Arguments)
         show_detailed - Show a detailed summary of the transpilation process.
     """
 
+    if clean_ctx is None:
+        clean_ctx = GeneratedFileCleaner(keep_intermediates=True)
+
     # Preprocessing statistics.
     stats = {"unsupported_calls": [], "kernel_launches": []}
 
     for filepath in all_files:
-        result = preprocessor(output_directory, filepath, stats, hip_clang_launch, is_pytorch_extension)
+        result = preprocessor(output_directory, filepath, stats, hip_clang_launch, is_pytorch_extension, clean_ctx)
+
         # Show what happened
         if show_progress:
             print(
@@ -606,15 +650,15 @@ 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, is_pytorch_extension):
+def preprocessor(output_directory, filepath, stats, hip_clang_launch, is_pytorch_extension, clean_ctx):
     """ Executes the CUDA -> HIP conversion on the specified file. """
     fin_path = os.path.join(output_directory, filepath)
     with open(fin_path, 'r') as fin:
         output_source = fin.read()
 
     fout_path = os.path.join(output_directory, get_hip_file_path(filepath))
     if not os.path.exists(os.path.dirname(fout_path)):
-        os.makedirs(os.path.dirname(fout_path))
+        clean_ctx.makedirs(os.path.dirname(fout_path))
 
     # unsupported_calls statistics reporting is broken atm
     def pt_repl(m):
@@ -675,7 +719,7 @@ def repl(m):
         with open(fout_path, 'r') as fout_old:
             do_write = fout_old.read() != output_source
     if do_write:
-        with open(fout_path, 'w') as fout:
+        with clean_ctx.open(fout_path, 'w') as fout:
             fout.write(output_source)
         return "ok"
     else:
@@ -776,11 +820,13 @@ def hipify(
     extensions=(".cu", ".cuh", ".c", ".cc", ".cpp", ".h", ".in", ".hpp"),
     output_directory="",
     includes=(),
+    extra_files=(),
     out_of_place_only=False,
     ignores=(),
     show_progress=True,
     hip_clang_launch=False,
     is_pytorch_extension=False,
+    clean_ctx=None
 ):
     if project_directory == "":
         project_directory = os.getcwd()
@@ -802,6 +848,8 @@ def hipify(
     all_files = list(matched_files_iter(output_directory, includes=includes,
                                         ignores=ignores, extensions=extensions,
                                         out_of_place_only=out_of_place_only))
+    all_files_set = set(all_files)
+    all_files += [f for f in extra_files if f not in all_files_set]
 
     # Start Preprocessor
     preprocess(
@@ -810,4 +858,5 @@ def hipify(
         show_detailed=show_detailed,
         show_progress=show_progress,
         hip_clang_launch=hip_clang_launch,
-        is_pytorch_extension=is_pytorch_extension)
+        is_pytorch_extension=is_pytorch_extension,
+        clean_ctx=clean_ctx)