"vscode:/vscode.git/clone" did not exist on "5ed704ec8c4e68f1bc846ab4e3c9e355585d62da"
Commit 40816696 authored by wenjh's avatar wenjh
Browse files

Make hygon backend installable


Signed-off-by: wenjh's avatarwenjh <wenjh@sugon.com>
parent 73d959a4
......@@ -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:
target_dir = install_dir / "transformer_engine" / lib_dir
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"):
......
......@@ -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__,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment