Invoke hipify from op builder for JIT extension builds (#1807)

Co-authored-by: Jeff Rasley <jerasley@microsoft.com>

.github/workflows/main.yml

@@ -127,7 +127,9 @@ jobs:
       # Runs a set of commands using the runners shell
       - name: Install deepspeed
         run: |
-          pip install .[dev,1bit,autotuning]
+          pip uninstall -y deepspeed
+          sudo /opt/conda/bin/pip uninstall -y deepspeed
+          sudo /opt/conda/bin/pip install .[dev,1bit,autotuning]
           python -c "from deepspeed.env_report import cli_main; cli_main()"
           #ds_report
       # Runs a set of commands using the runners shell

MANIFEST.in

@@ -1,4 +1,4 @@
 include *.txt README.md
 recursive-include requirements *.txt
-recursive-include deepspeed *.cpp *.h *.cu *.tr *.cuh *.cc *.json
+recursive-include deepspeed *.cpp *.h *.cu *.hip *.tr *.cuh *.cc *.json
 recursive-include csrc *.cpp *.h *.cu *.tr *.cuh *.cc

csrc/transformer/inference/includes/cublas_wrappers.h

@@ -7,7 +7,6 @@
 #include <cuda_runtime.h>
 #include <mma.h>
 #include <stdio.h>
-#include "cublas_wrappers.h"
 
 int cublas_gemm_ex(cublasHandle_t handle,
                    cublasOperation_t transa,

op_builder/builder.py

@@ -126,6 +126,9 @@ def sources(self):
         '''
         pass
 
+    def hipify_extension(self):
+        pass
+
     @staticmethod
     def assert_torch_info(torch_info):
         install_torch_version = torch_info['version']
@@ -606,6 +609,21 @@ def builder(self):
             cuda_ext.sources = sources
         return cuda_ext
 
+    def hipify_extension(self):
+        if self.is_rocm_pytorch():
+            from torch.utils.hipify import hipify_python
+            hipify_python.hipify(
+                project_directory=os.getcwd(),
+                output_directory=os.getcwd(),
+                header_include_dirs=self.include_paths(),
+                includes=[os.path.join(os.getcwd(),
+                                       '*')],
+                extra_files=[os.path.abspath(s) for s in self.sources()],
+                show_detailed=True,
+                is_pytorch_extension=True,
+                hipify_extra_files_only=True,
+            )
+
     def cxx_args(self):
         if sys.platform == "win32":
             return ['-O2']

setup.py

@@ -156,6 +156,10 @@ def op_enabled(op_name):
         reqs = builder.python_requirements()
         install_requires += builder.python_requirements()
 
+    # if op is compatible but install is not enabled (JIT mode)
+    if is_rocm_pytorch and op_compatible and not op_enabled(op_name):
+        builder.hipify_extension()
+
     # If op install enabled, add builder to extensions
     if op_enabled(op_name) and op_compatible:
         assert torch_available, f"Unable to pre-compile {op_name}, please first install torch"