import logging
import os
from datetime import datetime
import shutil
import subprocess
import sys
from sysconfig import get_paths
from setuptools import setup, find_packages, Extension, Command
from setuptools.command.bdist_wheel import bdist_wheel
from setuptools.command.build_ext import build_ext
from setuptools.command.build_py import build_py
from setuptools.command.develop import develop
from setuptools.command.install import install
def req_file(filename):
with open(filename, encoding='utf-8') as f:
content = f.readlines()
return [x.strip() for x in content]
install_requires = []
ROOT_DIR = os.path.dirname(__file__)
def get_value_from_lines(lines: list[str], key: str) -> str:
for line in lines:
line = " ".join(line.split())
if key in line:
return line.split(":")[-1].strip()
return ""
def get_chip_type() -> str:
npu_smi = shutil.which("npu-smi")
if npu_smi is None:
logging.warning(
"npu-smi command not found, if this is an npu environment, "
"please check if npu driver is installed correctly."
)
return ""
try:
npu_info_lines = subprocess.check_output([npu_smi, "info", "-l"]).decode().strip().split("\n")
npu_id = int(get_value_from_lines(npu_info_lines, "NPU ID"))
board_info_lines = (
subprocess.check_output([npu_smi, "info", "-t", "board", "-i", str(npu_id)]).decode().strip().split("\n")
)
chip_name = get_value_from_lines(board_info_lines, "Chip Name")
if not chip_name:
chip_info_lines = (
subprocess.check_output([npu_smi, "info", "-t", "board", "-i", str(npu_id), "-c", "0"])
.decode()
.strip()
.split("\n")
)
else:
chip_info_lines = board_info_lines
chip_name = get_value_from_lines(chip_info_lines, "Chip Name")
chip_type = get_value_from_lines(chip_info_lines, "Chip Type")
npu_name = get_value_from_lines(chip_info_lines, "NPU Name")
if "910" in chip_name:
if chip_type:
assert not npu_name
return (chip_type + chip_name).lower()
else:
assert npu_name
return (chip_name + "_" + npu_name).lower()
elif "950" in chip_name:
assert npu_name
return (chip_name + "_" + npu_name).lower()
else:
raise ValueError(f"Unable to recognize chip name: {chip_name}, please manually set env SOC_VERSION")
except subprocess.CalledProcessError as e:
logging.warning("npu-smi call failed: %s. ACLNN custom operator compilation will be skipped.", e)
return ""
except FileNotFoundError:
logging.warning(
"npu-smi command not found, if this is an npu environment, "
"please check if npu driver is installed correctly."
)
return ""
_env_soc_version = os.environ.get("SOC_VERSION", "")
if not _env_soc_version:
_env_soc_version = get_chip_type()
if _env_soc_version:
logging.info("Auto-detected SOC_VERSION: %s", _env_soc_version)
else:
logging.warning(
"Could not determine chip type automatically. "
"ACLNN custom operator compilation will be skipped.\n"
"To enable aclnn compilation, set the 'SOC_VERSION' environment "
"variable to specify the target chip, for example:\n"
' - arch32: export SOC_VERSION="ascend910b1"\n'
' - arch32: export SOC_VERSION="ascend910_9391"\n'
' - arch35: export SOC_VERSION="<value starting with ascend950>"\n'
)
def _find_cmake_binary():
cmake_from_path = shutil.which("cmake")
if cmake_from_path and os.path.isfile(cmake_from_path):
try:
with open(cmake_from_path, "rb") as f:
head = f.read(128)
if not head.startswith(b"#!/") or b"python" not in head.lower():
return cmake_from_path
except OSError:
pass
for candidate in ["/usr/bin/cmake", "/usr/local/bin/cmake"]:
if os.path.isfile(candidate) and os.access(candidate, os.X_OK):
return candidate
if cmake_from_path:
return cmake_from_path
return "cmake"
def normalize_soc_version(value: str) -> str:
value = value.lower()
if value.startswith("ascend910b"):
return "ascend910b"
if value.startswith("ascend910_93"):
return "ascend910_93"
if value.startswith("ascend950"):
return "ascend950"
return ""
class CMakeExtension(Extension):
def __init__(self, name: str, cmake_lists_dir: str = ".", **kwargs) -> None:
super().__init__(name, sources=[], py_limited_api=False, **kwargs)
self.cmake_lists_dir = os.path.abspath(cmake_lists_dir)
class cmake_build_ext(build_ext):
did_config: dict[str, bool] = {}
_cmake_package_install_path = ""
_catlass_library_names = {
"libmindspeed_catlass_mindspeed_basic_matmul.so",
"libmindspeed_catlass_mindspeed_chunk_loss.so",
}
def _sync_catlass_libraries(self) -> None:
"""Make generated CATLASS DSOs visible to editable installs.
CMake installs the extension and CATLASS libraries into setuptools'
build tree. An editable install points Python at the source package,
and setuptools copies only the extension module there, so the CATLASS
DSOs would otherwise be missing from ``mindspeed_ops/lib``.
"""
source_lib = os.path.join(ROOT_DIR, "mindspeed_ops", "lib")
candidates = [
getattr(self, "_cmake_package_install_path", ""),
os.path.join(self.build_lib, "mindspeed_ops"),
]
copied = []
os.makedirs(source_lib, exist_ok=True)
for filename in self._catlass_library_names:
destination = os.path.join(source_lib, filename)
if os.path.isfile(destination):
os.remove(destination)
for package_dir in candidates:
if not package_dir or not os.path.isdir(os.path.join(package_dir, "lib")):
continue
for filename in os.listdir(os.path.join(package_dir, "lib")):
if filename not in self._catlass_library_names:
continue
source = os.path.join(package_dir, "lib", filename)
destination = os.path.join(source_lib, filename)
if os.path.abspath(source) == os.path.abspath(destination):
continue
shutil.copy2(source, destination)
copied.append(destination)
if copied:
for path in copied:
print(f"Copy CATLASS library: {path}")
def compute_num_jobs(self):
num_jobs = os.environ.get("MAX_JOBS")
if num_jobs is not None:
num_jobs = int(num_jobs)
else:
try:
num_jobs = len(os.sched_getaffinity(0))
except AttributeError:
num_jobs = os.cpu_count()
num_jobs = max(1, num_jobs)
return num_jobs
def configure(self, ext: CMakeExtension) -> None:
build_temp = self.build_temp
os.makedirs(build_temp, exist_ok=True)
source_dir = os.path.abspath(ROOT_DIR)
python_executable = sys.executable
cmake_args = [_find_cmake_binary()]
cmake_args += ["-DCMAKE_BUILD_TYPE=Release"]
cmake_args += ["-DCMAKE_EXPORT_COMPILE_COMMANDS=1"]
asc_home = os.environ.get("ASCEND_HOME_PATH")
if asc_home is None:
logging.warning(
"No ASCEND_HOME_PATH found in environment. "
"Using default path. Set ASCEND_HOME_PATH if you have a "
"custom CANN installation."
)
asc_home = "/usr/local/Ascend/ascend-toolkit/latest"
logging.info("ASCEND_HOME_PATH: %s", asc_home)
cmake_args += [f"-DASCEND_HOME_PATH={asc_home}"]
cmake_args += [f"-DPYTHON_EXECUTABLE={python_executable}"]
cmake_args += [f"-DPYTHON_INCLUDE_PATH={get_paths()['include']}"]
try:
pybind11_cmake_path = (
subprocess.check_output(
[python_executable, "-m", "pybind11", "--cmakedir"],
stderr=subprocess.PIPE,
)
.decode()
.strip()
)
except subprocess.CalledProcessError as e:
raise RuntimeError(
f"Failed to locate pybind11 CMake directory via "
f"'{python_executable} -m pybind11 --cmakedir'. "
f"Install pybind11 in this Python environment:\n"
f" {python_executable} -m pip install pybind11"
) from e
cmake_args += [f"-DCMAKE_PREFIX_PATH={pybind11_cmake_path}"]
package_install_path = os.path.dirname(os.path.abspath(self.get_ext_fullpath(ext.name)))
self._cmake_package_install_path = package_install_path
install_path = os.path.dirname(package_install_path)
cmake_args += [f"-DCMAKE_INSTALL_PREFIX={install_path}"]
cmake_args += [f"-DMINDSPEED_PACKAGE_INSTALL_PATH={package_install_path}"]
normalized_soc = normalize_soc_version(_env_soc_version)
cmake_args += [f"-DMINDSPEED_CATLASS_SOC={normalized_soc}"]
catlass_source = os.environ.get("CATLASS_SOURCE_DIR", "")
if catlass_source:
cmake_args += [f"-DCATLASS_SOURCE_DIR={catlass_source}"]
catlass_enabled = os.environ.get("MINDSPEED_BUILD_CATLASS", "0") == "1" and normalized_soc in {
"ascend910b",
"ascend910_93",
}
cmake_args += [f"-DMINDSPEED_CATLASS_ENABLE={'ON' if catlass_enabled else 'OFF'}"]
if os.environ.get("MINDSPEED_ALLOW_UNPINNED_CATLASS", "0") == "1":
cmake_args += ["-DMINDSPEED_ALLOW_UNPINNED_CATLASS=ON"]
try:
torch_npu_path = (
subprocess.check_output(
[python_executable, "-c", "import torch_npu; print(torch_npu.__file__)"],
stderr=subprocess.PIPE,
)
.decode()
.strip()
)
torch_npu_path = os.path.dirname(torch_npu_path)
except subprocess.CalledProcessError:
torch_npu_path = ""
logging.warning(
"torch_npu not found via import. "
"If torch_npu is installed at a non-standard location, "
"set TORCH_NPU_PATH manually."
)
if torch_npu_path:
cmake_args += [f"-DTORCH_NPU_PATH={torch_npu_path}"]
cmake_args += [source_dir]
logging.info("cmake config command: %s", cmake_args)
try:
subprocess.check_call(cmake_args, cwd=self.build_temp)
except subprocess.CalledProcessError as e:
raise RuntimeError(f"CMake configuration failed: {e}")
def build_extensions(self) -> None:
try:
subprocess.check_output([_find_cmake_binary(), "--version"])
except OSError as e:
raise RuntimeError(f"Cannot find CMake executable: {e}")
if not os.path.exists(self.build_temp):
os.makedirs(self.build_temp)
targets = []
def target_name(s: str) -> str:
return s.removeprefix("mindspeed_ops.")
for ext in self.extensions:
self.configure(ext)
targets.append(target_name(ext.name))
num_jobs = self.compute_num_jobs()
build_args = [
"--build",
".",
"-j",
str(num_jobs),
*[f"--target={name}" for name in targets],
]
try:
subprocess.check_call([_find_cmake_binary(), *build_args], cwd=self.build_temp)
except OSError as e:
raise RuntimeError(f"Build library failed: {e}")
install_args = [
_find_cmake_binary(),
"--install",
".",
]
try:
subprocess.check_call(install_args, cwd=self.build_temp)
except OSError as e:
raise RuntimeError(f"Install library failed: {e}")
src_cann = os.path.join(ROOT_DIR, "mindspeed_ops", "_cann_ops_custom")
dst_cann = os.path.join(self.build_lib, "mindspeed_ops", "_cann_ops_custom")
if os.path.isdir(src_cann):
if os.path.exists(dst_cann):
shutil.rmtree(dst_cann)
shutil.copytree(src_cann, dst_cann)
print(f"Copy: {src_cann} -> {dst_cann}")
self._sync_catlass_libraries()
def run(self):
if _env_soc_version:
self.run_command("build_aclnn")
super().run()
class build_and_install_aclnn(Command):
description = "Build and install AclNN by running build_aclnn.sh"
user_options = []
def initialize_options(self):
pass
def finalize_options(self):
pass
def run(self):
try:
print("Running bash build_aclnn.sh ...")
subprocess.check_call(
[shutil.which("bash") or "bash", "mindspeed_ops/csrc/build_aclnn.sh", ROOT_DIR, _env_soc_version]
)
print("build_aclnn.sh executed successfully!")
except subprocess.CalledProcessError as e:
print(f"Error running build_aclnn.sh: {e}")
raise SystemExit(e.returncode)
class custom_develop(develop):
def run(self):
self.run_command("build_ext")
super().run()
class custom_install(install):
def run(self):
self.run_command("build_ext")
install.run(self)
class dated_bdist_wheel(bdist_wheel):
"""bdist_wheel with a --date-tag option.
When set, the build date (YYYYMMDD) is stamped into the wheel filename
as the build tag, e.g. mindspeed_ops-0.1.0-20260827-cp311-...whl.
The version metadata itself is unchanged.
"""
date_tag = 0
build_number = None
user_options = bdist_wheel.user_options + [
("date-tag", None, "append the build date (YYYYMMDD) to the wheel filename"),
]
boolean_options = bdist_wheel.boolean_options + ["date-tag"]
def initialize_options(self):
super().initialize_options()
self.date_tag = 0
def finalize_options(self):
if self.date_tag and not self.build_number:
self.build_number = datetime.now().strftime("%Y%m%d")
super().finalize_options()
ext_modules = []
if _env_soc_version:
ext_modules = [CMakeExtension(name="mindspeed_ops.mindspeed_ops_C")]
cmdclass = {
"develop": custom_develop,
"build_py": build_py,
"build_aclnn": build_and_install_aclnn,
"build_ext": cmake_build_ext,
"install": custom_install,
"bdist_wheel": dated_bdist_wheel,
}
setup(
name="mindspeed_ops",
version=os.environ.get("MINDSPEED_OPS_VERSION", "26.2.0"),
description="Ascend-optimized custom operators for training tasks",
packages=find_packages(exclude=["tests", "tests.*"]),
license="Apache-2.0",
classifiers=[
"Programming Language :: Python :: 3",
"License :: OSI Approved :: Apache Software License",
"Operating System :: OS Independent",
],
python_requires='>=3.10',
install_requires=install_requires,
package_data={
"mindspeed_ops": [
"lib/*.so",
"lib/*.pyd",
"csrc/atb/*.cpp",
"csrc/atb/inc/*.h",
"csrc/atb/utils/*.cpp",
"csrc/flop_counter/*.cpp",
"csrc/flop_counter/*.h",
]
},
ext_modules=ext_modules,
cmdclass=cmdclass,
)
