Enable torchvision build with ROCm on Windows (#6)

* changing flags according to clang

* update changes

* fix hipcc run on windows

* changed include dirs conversion for rocm windows

* fixed clang compilation with space in include paths

* fix hip kernels build

* restruct CXX and MSVC flags

* change cxx flags

* change include headers in hipified source files

* change cuda to hip in include paths

* cleaning code

* fixed win_wrap_single_compile

* fix paths in hipify

* format code

* removing unnecessary flags

* cleaning code

* fix cleaning

* cleaning code

* cleaning code

* restructured ninja commands

* renaming hipcc function

* added set runtime lib function

* add set hipcc runtime library function

* format code

Author: tvukovic-amd

torch/utils/cpp_extension.py

@@ -173,9 +173,6 @@ def _join_rocm_home(*paths) -> str:
     if ROCM_HOME is None:
         raise OSError('ROCM_HOME environment variable is not set. '
                       'Please set it to your ROCm install root.')
-    elif IS_WINDOWS:
-        raise OSError('Building PyTorch extensions using '
-                      'ROCm and Windows is not supported.')
     return os.path.join(ROCM_HOME, *paths)
 
 def _join_sycl_home(*paths) -> str:
@@ -270,12 +267,14 @@ def _join_sycl_home(*paths) -> str:
 ]
 
 COMMON_HIP_FLAGS = [
-    '-fPIC',
     '-D__HIP_PLATFORM_AMD__=1',
     '-DUSE_ROCM=1',
     '-DHIPBLAS_V2',
 ]
 
+if not IS_WINDOWS:
+    COMMON_HIP_FLAGS.append('-fPIC')
+
 COMMON_HIPCC_FLAGS = [
     '-DCUDA_HAS_FP16=1',
     '-D__HIP_NO_HALF_OPERATORS__=1',
@@ -511,6 +510,19 @@ def _check_cuda_version(compiler_name: str, compiler_version: TorchVersion) -> N
                 f'Please make sure to use an adequate version of {compiler_name} ({version_bound_str}).'
             )
 
+# Specify Visual Studio C runtime library for hipcc
+def _set_hipcc_runtime_lib(is_standalone):
+    debug = os.getenv("DEBUG")
+    if is_standalone:
+        if debug:
+            COMMON_HIP_FLAGS.append('-fms-runtime-lib=static_dbg')
+        else:
+            COMMON_HIP_FLAGS.append('-fms-runtime-lib=static')
+    else:
+        if debug:
+            COMMON_HIP_FLAGS.append('-fms-runtime-lib=dll_dbg')
+        else:
+            COMMON_HIP_FLAGS.append('-fms-runtime-lib=dll')
 
 def _append_sycl_std_if_no_std_present(cflags):
     if not any(flag.startswith('-sycl-std=') for flag in cflags):
@@ -831,6 +843,9 @@ def unix_wrap_ninja_compile(sources,
         def win_cuda_flags(cflags):
             return (COMMON_NVCC_FLAGS +
                     cflags + _get_cuda_arch_flags(cflags))
+        
+        def win_hip_flags(cflags):
+            return (COMMON_HIPCC_FLAGS + COMMON_HIP_FLAGS + cflags + _get_rocm_arch_flags(cflags))
 
         def win_wrap_single_compile(sources,
                                     output_dir=None,
@@ -868,19 +883,25 @@ def spawn(cmd):
                     src = src_list[0]
                     obj = obj_list[0]
                     if _is_cuda_file(src):
-                        nvcc = _join_cuda_home('bin', 'nvcc')
+                        if IS_HIP_EXTENSION:
+                            nvcc = _get_hipcc_path()
+                        else:
+                            nvcc = _join_cuda_home('bin', 'nvcc')     
                         if isinstance(self.cflags, dict):
                             cflags = self.cflags['nvcc']
                         elif isinstance(self.cflags, list):
                             cflags = self.cflags
                         else:
                             cflags = []
 
-                        cflags = win_cuda_flags(cflags) + ['-std=c++17', '--use-local-env']
+                        if IS_HIP_EXTENSION:
+                            cflags = win_hip_flags(cflags)   
+                        else:
+                            cflags = win_cuda_flags(cflags) + ['-std=c++17', '--use-local-env']
+                            for ignore_warning in MSVC_IGNORE_CUDAFE_WARNINGS:
+                                cflags = ['-Xcudafe', '--diag_suppress=' + ignore_warning] + cflags
                         for flag in COMMON_MSVC_FLAGS:
                             cflags = ['-Xcompiler', flag] + cflags
-                        for ignore_warning in MSVC_IGNORE_CUDAFE_WARNINGS:
-                            cflags = ['-Xcudafe', '--diag_suppress=' + ignore_warning] + cflags
                         cmd = [nvcc, '-c', src, '-o', obj] + include_list + cflags
                     elif isinstance(self.cflags, dict):
                         cflags = COMMON_MSVC_FLAGS + self.cflags['cxx']
@@ -909,7 +930,6 @@ def win_wrap_ninja_compile(sources,
                                    extra_preargs=None,
                                    extra_postargs=None,
                                    depends=None):
-
             if not self.compiler.initialized:
                 self.compiler.initialize()
             output_dir = os.path.abspath(output_dir)
@@ -926,14 +946,20 @@ def win_wrap_ninja_compile(sources,
                 self.compiler._setup_compile(output_dir, macros,
                                              include_dirs, sources,
                                              depends, extra_postargs)
+            # Replace space with \ when using hipcc (hipcc passes includes to clang without ""s so clang sees space in include paths as new argument)
+            if IS_HIP_EXTENSION:
+                pp_opts = ["-I{}".format(s[2:].replace(" ", "\\")) if s.startswith('-I') else s for s in pp_opts]              
             common_cflags = extra_preargs or []
             cflags = []
             if debug:
                 cflags.extend(self.compiler.compile_options_debug)
             else:
                 cflags.extend(self.compiler.compile_options)
-            common_cflags.extend(COMMON_MSVC_FLAGS)
-            cflags = cflags + common_cflags + pp_opts
+            cflags = cflags + common_cflags + pp_opts + COMMON_MSVC_FLAGS
+            if IS_HIP_EXTENSION:
+                common_cflags.extend(COMMON_HIP_FLAGS)
+            else:
+                common_cflags.extend(COMMON_MSVC_FLAGS) 
             with_cuda = any(map(_is_cuda_file, sources))
 
             # extra_postargs can be either:
@@ -943,25 +969,31 @@ def win_wrap_ninja_compile(sources,
                 post_cflags = extra_postargs['cxx']
             else:
                 post_cflags = list(extra_postargs)
+            if IS_HIP_EXTENSION:
+                post_cflags = COMMON_HIP_FLAGS + post_cflags
             append_std17_if_no_std_present(post_cflags)
 
             cuda_post_cflags = None
             cuda_cflags = None
             if with_cuda:
-                cuda_cflags = ['-std=c++17', '--use-local-env']
+                cuda_cflags = ['-std=c++17']
                 for common_cflag in common_cflags:
                     cuda_cflags.append('-Xcompiler')
                     cuda_cflags.append(common_cflag)
-                for ignore_warning in MSVC_IGNORE_CUDAFE_WARNINGS:
-                    cuda_cflags.append('-Xcudafe')
-                    cuda_cflags.append('--diag_suppress=' + ignore_warning)
+                if not IS_HIP_EXTENSION:
+                    cuda_cflags.append('--use-local-env')
+                    for ignore_warning in MSVC_IGNORE_CUDAFE_WARNINGS:
+                        cuda_cflags.append('-Xcudafe')
+                        cuda_cflags.append('--diag_suppress=' + ignore_warning)
                 cuda_cflags.extend(pp_opts)
                 if isinstance(extra_postargs, dict):
                     cuda_post_cflags = extra_postargs['nvcc']
                 else:
                     cuda_post_cflags = list(extra_postargs)
-                cuda_post_cflags = win_cuda_flags(cuda_post_cflags)
-
+                if IS_HIP_EXTENSION:
+                    cuda_post_cflags = win_hip_flags(cuda_post_cflags)  
+                else:
+                    cuda_post_cflags = win_cuda_flags(cuda_post_cflags)
             cflags = _nt_quote_args(cflags)
             post_cflags = _nt_quote_args(post_cflags)
             if with_cuda:
@@ -990,7 +1022,6 @@ def win_wrap_ninja_compile(sources,
 
             # Return *all* object filenames, not just the ones we just built.
             return objects
-
         # Monkey-patch the _compile or compile method.
         # https://github.com/python/cpython/blob/dc0284ee8f7a270b6005467f26d8e5773d76e959/Lib/distutils/ccompiler.py#L511
         if self.compiler.compiler_type == 'msvc':
@@ -2076,7 +2107,12 @@ def _jit_compile(name,
 
     return _import_module_from_library(name, build_directory, is_python_module)
 
-
+def _get_hipcc_path():
+    if IS_WINDOWS:
+        return _join_rocm_home('bin', 'hipcc.bat')
+    else:
+        return _join_rocm_home('bin', 'hipcc')
+    
 def _write_ninja_file_and_compile_objects(
         sources: list[str],
         objects,
@@ -2479,6 +2515,7 @@ def _run_ninja_build(build_directory: str, verbose: bool, error_prefix: str) ->
         stdout_fileno = 1
         subprocess.run(
             command,
+            shell=IS_WINDOWS and IS_HIP_EXTENSION,
             stdout=stdout_fileno if verbose else subprocess.PIPE,
             stderr=subprocess.STDOUT,
             cwd=build_directory,
@@ -2573,7 +2610,8 @@ def _write_ninja_file_to_build_library(path,
     common_cflags += [f"{x}" for x in _get_glibcxx_abi_build_flags()]
 
     if IS_WINDOWS:
-        cflags = common_cflags + COMMON_MSVC_FLAGS + ['/std:c++17'] + extra_cflags
+        cflags = common_cflags + ['/std:c++17'] + extra_cflags
+        clfags += COMMON_HIP_FLAGS if IS_HIP_EXTENSION else COMMON_MSVC_FLAGS
         cflags = _nt_quote_args(cflags)
     else:
         cflags = common_cflags + ['-fPIC', '-std=c++17'] + extra_cflags
@@ -2632,6 +2670,7 @@ def object_file_path(source_file: str) -> str:
 
     objects = [object_file_path(src) for src in sources]
     ldflags = ([] if is_standalone else [SHARED_FLAG]) + extra_ldflags
+    _set_hipcc_runtime_lib(is_standalone)
 
     # The darwin linker needs explicit consent to ignore unresolved symbols.
     if IS_MACOS:
@@ -2724,7 +2763,7 @@ def sanitize_flags(flags):
             nvcc = os.getenv("PYTORCH_NVCC")    # user can set nvcc compiler with ccache using the environment variable here
         else:
             if IS_HIP_EXTENSION:
-                nvcc = _join_rocm_home('bin', 'hipcc')
+                nvcc = _get_hipcc_path()
             else:
                 nvcc = _join_cuda_home('bin', 'nvcc')
         config.append(f'nvcc = {nvcc}')
@@ -2753,9 +2792,11 @@ def sanitize_flags(flags):
     # See https://ninja-build.org/build.ninja.html for reference.
     compile_rule = ['rule compile']
     if IS_WINDOWS:
+        compiler_name = "$cxx" if IS_HIP_EXTENSION else "cl"
         compile_rule.append(
-            '  command = cl /showIncludes $cflags -c $in /Fo$out $post_cflags')
-        compile_rule.append('  deps = msvc')
+                f'  command = {compiler_name} /showIncludes $cflags -c $in /Fo$out $post_cflags')
+        if not IS_HIP_EXTENSION:
+            compile_rule.append('  deps = msvc')
     else:
         compile_rule.append(
             '  command = $cxx -MMD -MF $out.d $cflags -c $in -o $out $post_cflags')

torch/utils/hipify/hipify_python.py

@@ -139,7 +139,6 @@ def __exit__(self, type, value, traceback):
             for d in self.dirs_to_clean[::-1]:
                 os.rmdir(d)
 
-
 # Follow UNIX convention for paths to use '/' instead of '\\' on Windows
 def _to_unix_path(path: str) -> str:
     return path.replace(os.sep, '/')
@@ -830,6 +829,7 @@ def preprocessor(
         show_progress: bool) -> HipifyResult:
     """ Executes the CUDA -> HIP conversion on the specified file. """
     fin_path = os.path.abspath(os.path.join(output_directory, filepath))
+    filepath = _to_unix_path(filepath) 
     hipify_result = HIPIFY_FINAL_RESULT[fin_path]
     if filepath not in all_files:
         hipify_result.hipified_path = None
@@ -932,8 +932,8 @@ def repl(m):
                         return templ.format(os.path.relpath(header_fout_path if header_fout_path is not None
                                                             else header_filepath, header_dir))
                 hipified_header_filepath = HIPIFY_FINAL_RESULT[header_filepath].hipified_path
-                return templ.format(os.path.relpath(hipified_header_filepath if hipified_header_filepath is not None
-                                                    else header_filepath, header_dir))
+                return templ.format(_to_unix_path(os.path.relpath(hipified_header_filepath if hipified_header_filepath is not None
+                                                    else header_filepath, header_dir)))
 
             return m.group(0)
         return repl