wehub-resource-sync
87 lines
3.1 KiB
Python
87 lines
3.1 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import re
|
|
import os
|
|
try:
|
|
import torch_npu
|
|
except ImportError as e:
|
|
pass
|
|
|
|
try:
|
|
# is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed
|
|
# if successful this also means we're doing a local install and not JIT compile path
|
|
from op_builder import __deepspeed__ # noqa: F401 # type: ignore
|
|
from op_builder.builder import OpBuilder
|
|
except ImportError:
|
|
from deepspeed.ops.op_builder.builder import OpBuilder
|
|
|
|
|
|
class NPUOpBuilder(OpBuilder):
|
|
_ascend_path = None
|
|
_torch_npu_path = None
|
|
_cann_version = None
|
|
|
|
def __init__(self, name):
|
|
super().__init__(name)
|
|
self._ascend_path = self.installed_cann_path()
|
|
self._torch_npu_path = os.path.join(os.path.dirname(os.path.abspath(torch_npu.__file__)))
|
|
try:
|
|
self._cann_version = self.installed_cann_version(self.name)
|
|
except Exception:
|
|
print(f"{self.name} ascend_cann is missing, npu ops cannot be compiled!")
|
|
|
|
def cann_defs(self):
|
|
if self._cann_version:
|
|
return '-D__ENABLE_CANN__'
|
|
return '-D__DISABLE_CANN__'
|
|
|
|
def installed_cann_path(self):
|
|
if "ASCEND_HOME_PATH" in os.environ or os.path.exists(os.environ["ASCEND_HOME_PATH"]):
|
|
return os.environ["ASCEND_HOME_PATH"]
|
|
return None
|
|
|
|
def installed_cann_version(self, name=""):
|
|
ascend_path = self.installed_cann_path()
|
|
assert ascend_path is not None, "CANN_HOME does not exist, unable to compile NPU op(s)"
|
|
cann_version = ""
|
|
for dirpath, _, filenames in os.walk(os.path.realpath(ascend_path)):
|
|
if cann_version:
|
|
break
|
|
install_files = [file for file in filenames if re.match(r"ascend_.*_install\.info", file)]
|
|
if install_files:
|
|
filepath = os.path.join(dirpath, install_files[0])
|
|
with open(filepath, "r") as f:
|
|
for line in f:
|
|
if line.find("version") != -1:
|
|
cann_version = line.strip().split("=")[-1]
|
|
break
|
|
return cann_version
|
|
|
|
def include_paths(self):
|
|
paths = super().include_paths()
|
|
paths += [os.path.join(self._ascend_path, 'include'), os.path.join(self._torch_npu_path, 'include')]
|
|
return paths
|
|
|
|
def cxx_args(self):
|
|
args = super().cxx_args()
|
|
args += ['-O3', '-std=c++17', '-g', '-Wno-reorder', '-fopenmp']
|
|
args += ['-fstack-protector-all', '-Wl,-z,relro,-z,now,-z,noexecstack', '-Wl,--disable-new-dtags,--rpath']
|
|
args += [
|
|
self.cann_defs(),
|
|
self.cpu_arch(),
|
|
self.simd_width(), '-L' + os.path.join(self._ascend_path, 'lib64'),
|
|
'-L' + os.path.join(self._torch_npu_path, 'lib')
|
|
]
|
|
return args
|
|
|
|
def extra_ldflags(self):
|
|
flags = super().extra_ldflags()
|
|
flags += [
|
|
'-L' + os.path.join(self._ascend_path, 'lib64'), '-lascendcl',
|
|
'-L' + os.path.join(self._torch_npu_path, 'lib'), '-ltorch_npu'
|
|
]
|
|
return flags
|