Make hygon backend installable

Signed-off-by: wenjh <wenjh@sugon.com>

build_tools/build_ext.py

@@ -145,7 +145,10 @@ def get_build_ext(
             # For editable/inplace builds this is not a concern as
             # the SOs will be in a local directory anyway.
             if not self.inplace:
+                if not bool(int(os.getenv("NVTE_HYGON_BACKEND", "0"))):
                     target_dir = install_dir / "transformer_engine" / lib_dir
+                else:
+                    target_dir = install_dir / "transformer_engine_hygon" / lib_dir
                 target_dir.mkdir(exist_ok=True, parents=True)
 
                 for ext in Path(self.build_lib).glob("*.so"):

setup.py

@@ -50,6 +50,17 @@ if rocm_build():
 else:
     archs = cuda_archs()
 
+if bool(int(os.getenv("NVTE_HYGON_BACKEND", "0"))):
+    common_dir = current_file_path / "transformer_engine" / "common"
+    common_copy = current_file_path / "transformer_engine_hygon" / "common"
+    if common_copy.exists():
+        shutil.rmtree(common_copy)
+    shutil.copytree(common_dir, common_copy)
+    csrc_dir = current_file_path / "transformer_engine" / "pytorch" / "csrc"
+    csrc_copy = current_file_path / "transformer_engine_hygon" / "pytorch" / "csrc"
+    if csrc_copy.exists():
+        shutil.rmtree(csrc_copy)
+    shutil.copytree(csrc_dir, csrc_copy)
 
 class TimedBdist(bdist_wheel):
     """Helper class to measure build time"""
@@ -297,16 +308,6 @@ if __name__ == "__main__":
         )
     else:
         # Configure package of hygon backend for TransformerEngine-FL
-        common_dir = current_file_path / "transformer_engine" / "common"
-        common_copy = current_file_path / "transformer_engine_hygon" / "common"
-        if common_copy.exists():
-            shutil.rmtree(common_copy)
-        shutil.copytree(common_dir, common_copy)
-        csrc_dir = current_file_path / "transformer_engine" / "pytorch" / "csrc"
-        csrc_copy = current_file_path / "transformer_engine_hygon" / "pytorch" / "csrc"
-        if csrc_copy.exists():
-            shutil.rmtree(csrc_copy)
-        shutil.copytree(csrc_dir, csrc_copy)
         setuptools.setup(
             name="transformer_engine_hygon",
             version=__version__,