diff --git a/cmake/find_ms_internal_kernels_lib.cmake b/cmake/find_ms_internal_kernels_lib.cmake index 9eb7ff22235d0a722d429aba1b34c37532e12a89..cdc9b96178453601dfeef1619ad9484253c19e35 100644 --- a/cmake/find_ms_internal_kernels_lib.cmake +++ b/cmake/find_ms_internal_kernels_lib.cmake @@ -67,7 +67,13 @@ endif() # MindSpore Internal Kernels Path Detection # ============================================================================= -set(INTERNAL_KERNEL_ROOT_PATH "${CMAKE_SOURCE_DIR}/prebuild/ms_kernels_internal/ms_kernels_internal/internal_kernel") +if(DEFINED ENV{MS_CUSTOM_INTERNAL_KERNEL_HOME}) + set(INTERNAL_KERNEL_ROOT_PATH "$ENV{MS_CUSTOM_INTERNAL_KERNEL_HOME}/internal_kernel") + message(STATUS "Using environment variable MS_CUSTOM_INTERNAL_KERNEL_HOME=$ENV{MS_CUSTOM_INTERNAL_KERNEL_HOME} for internal kernel path") +else() + set(INTERNAL_KERNEL_ROOT_PATH "${CMAKE_SOURCE_DIR}/prebuild/ms_kernels_internal/ms_kernels_internal/internal_kernel") +endif() + set(INTERNAL_KERNEL_INC_PATH "${INTERNAL_KERNEL_ROOT_PATH}" "${INTERNAL_KERNEL_ROOT_PATH}/include") # Check if paths exist @@ -83,8 +89,12 @@ message(STATUS "INTERNAL_KERNEL_INC_PATH: ${INTERNAL_KERNEL_INC_PATH}") # ============================================================================= # Library Detection # ============================================================================= - -set(INTERNAL_KERNEL_LIB_PATH "${CMAKE_SOURCE_DIR}/prebuild/ms_kernels_internal/ms_kernels_internal/internal_kernel/lib") +if(DEFINED ENV{MS_CUSTOM_INTERNAL_KERNEL_HOME}) + set(INTERNAL_KERNEL_LIB_PATH "$ENV{MS_CUSTOM_INTERNAL_KERNEL_HOME}/internal_kernel/lib") + message(STATUS "Using environment variable MS_CUSTOM_INTERNAL_KERNEL_HOME=$ENV{MS_CUSTOM_INTERNAL_KERNEL_HOME} for internal kernel path") +else() + set(INTERNAL_KERNEL_LIB_PATH "${CMAKE_SOURCE_DIR}/prebuild/ms_kernels_internal/ms_kernels_internal/internal_kernel/lib") +endif() message(STATUS "INTERNAL_KERNEL_LIB_PATH: ${INTERNAL_KERNEL_LIB_PATH}") # Check for ms_kernels_internal library diff --git a/python/ms_custom_ops/__init__.py b/python/ms_custom_ops/__init__.py index d7af8202cdb895884e871a3a944b0604edeeffdc..20814da8491e4d6184279faf3dc67915a0cae452 100644 --- a/python/ms_custom_ops/__init__.py +++ b/python/ms_custom_ops/__init__.py @@ -2,6 +2,7 @@ import os import ctypes import mindspore + def _init_env(): """init env.""" current_path = os.path.dirname(os.path.abspath(__file__)) @@ -11,18 +12,28 @@ def _init_env(): os.environ["ASCEND_CUSTOM_OPP_PATH"] = env_path + ":" + origin_env_path else: os.environ["ASCEND_CUSTOM_OPP_PATH"] = env_path - + if os.getenv("ASDOPS_LOG_LEVEL") is None: os.environ["ASDOPS_LOG_LEVEL"] = "ERROR" if os.getenv("ASDOPS_LOG_TO_STDOUT") is None: os.environ["ASDOPS_LOG_TO_STDOUT"] = "1" - internal_lib_path = os.path.join(current_path, "prebuild/ms_kernels_internal/internal_kernel/lib/libms_kernels_internal_v2.so") + internal_lib_path = os.path.join( + current_path, + "prebuild/ms_kernels_internal/internal_kernel/lib/libms_kernels_internal_v2.so", + ) + + if not os.path.exists(internal_lib_path): + raise FileNotFoundError( + f"MS internal kernel library not found at {internal_lib_path}" + ) ctypes.CDLL(internal_lib_path) + _init_env() from .ms_custom_ops import * + # Import generated ops interfaces try: from .gen_ops_def import * @@ -40,21 +51,23 @@ __all__ = [] # Add ops from gen_ops_def if available try: from . import gen_ops_def - if hasattr(gen_ops_def, '__all__'): + + if hasattr(gen_ops_def, "__all__"): __all__.extend(gen_ops_def.__all__) else: # If no __all__ defined, add all public functions - __all__.extend([name for name in dir(gen_ops_def) if not name.startswith('_')]) + __all__.extend([name for name in dir(gen_ops_def) if not name.startswith("_")]) except ImportError: pass -# Add ops from gen_ops_prim if available +# Add ops from gen_ops_prim if available try: from . import gen_ops_prim - if hasattr(gen_ops_prim, '__all__'): + + if hasattr(gen_ops_prim, "__all__"): __all__.extend(gen_ops_prim.__all__) else: # If no __all__ defined, add all public functions - __all__.extend([name for name in dir(gen_ops_prim) if not name.startswith('_')]) + __all__.extend([name for name in dir(gen_ops_prim) if not name.startswith("_")]) except ImportError: pass diff --git a/scripts/prepare_prebuild.sh b/scripts/prepare_prebuild.sh index 4aa82741e72f2205e07891076e114038ce403ca4..36c1f7f8e9eb3f46ca69615bfb8280591e23a5cd 100644 --- a/scripts/prepare_prebuild.sh +++ b/scripts/prepare_prebuild.sh @@ -7,12 +7,16 @@ function prepare_ms_kernels_internal() if [ -d "prebuild/ms_kernels_internal/" ]; then rm -rf prebuild/ms_kernels_internal/ms_kernels_internal fi - type=aarch64 - uname -a | grep aarch64 &> /dev/null || type=x86_64 - tar -zxf prebuild/ms_kernels_internal/${type}/ms_kernels_internal.tar.gz -C prebuild/ms_kernels_internal - tar -zxf prebuild/ms_kernels_internal/${type}/ms_kernels_dependency.tar.gz -C prebuild/ms_kernels_internal - cp prebuild/ms_kernels_internal/ms_kernels_dependency/asdops prebuild/ms_kernels_internal/ms_kernels_internal/ -rf - rm -rf prebuild/ms_kernels_internal/ms_kernels_dependency + + if [ -z "${MS_CUSTOM_INTERNAL_KERNEL_HOME}" ]; then + echo "[INFO]:use default internal tar." + type=aarch64 + uname -a | grep aarch64 &> /dev/null || type=x86_64 + tar -zxf prebuild/ms_kernels_internal/${type}/ms_kernels_internal.tar.gz -C prebuild/ms_kernels_internal + tar -zxf prebuild/ms_kernels_internal/${type}/ms_kernels_dependency.tar.gz -C prebuild/ms_kernels_internal + cp prebuild/ms_kernels_internal/ms_kernels_dependency/asdops prebuild/ms_kernels_internal/ms_kernels_internal/ -rf + rm -rf prebuild/ms_kernels_internal/ms_kernels_dependency + fi } prepare_ms_kernels_internal diff --git a/setup.py b/setup.py index a8215d9b687b873fe9f6eb63644081baf63234d2..e3c778ce7db73339a25123cabc39e1f60b06132a 100644 --- a/setup.py +++ b/setup.py @@ -65,7 +65,7 @@ def get_requirements() -> List[str]: requirements_path = get_path(filename) if not os.path.exists(requirements_path): return [] - + with open(requirements_path) as f: requirements = f.read().strip().split("\n") resolved_requirements = [] @@ -90,12 +90,15 @@ def write_commit_id(): commit_info = "" try: commit_info += subprocess.check_output( - ["git", "rev-parse", "--abbrev-ref", "HEAD"]).decode("utf-8") + ["git", "rev-parse", "--abbrev-ref", "HEAD"] + ).decode("utf-8") commit_info += subprocess.check_output( - ["git", "log", "--abbrev-commit", "-1"]).decode("utf-8") + ["git", "log", "--abbrev-commit", "-1"] + ).decode("utf-8") except subprocess.CalledProcessError: - logger.warning("Can't get commit id information. " - "Please make sure git is available.") + logger.warning( + "Can't get commit id information. " "Please make sure git is available." + ) commit_info = "git is not available while building." with open("./python/ms_custom_ops/.commit_id", "w") as f: @@ -110,16 +113,23 @@ def get_version(): else: return "0.1.0" + version = get_version() + def _get_ascend_home_path(): return os.environ.get("ASCEND_HOME_PATH", "/usr/local/Ascend/ascend-toolkit/latest") + def _get_ascend_env_path(): - env_script_path = os.path.realpath(os.path.join(_get_ascend_home_path(), "..", "set_env.sh")) + env_script_path = os.path.realpath( + os.path.join(_get_ascend_home_path(), "..", "set_env.sh") + ) if not os.path.exists(env_script_path): - raise ValueError(f"The file '{env_script_path}' is not found, " - "please make sure environment variable 'ASCEND_HOME_PATH' is set correctly.") + raise ValueError( + f"The file '{env_script_path}' is not found, " + "please make sure environment variable 'ASCEND_HOME_PATH' is set correctly." + ) return env_script_path @@ -127,31 +137,36 @@ def generate_docs(): """Generate YAML documentation from Markdown sources.""" logger.info("Generating documentation...") doc_generator_script = os.path.join(ROOT_DIR, "scripts", "doc_generator.py") - + if not os.path.exists(doc_generator_script): - logger.warning(f"Documentation generator script not found: {doc_generator_script}") + logger.warning( + f"Documentation generator script not found: {doc_generator_script}" + ) return - + try: # Run the documentation generator result = subprocess.run( [sys.executable, doc_generator_script], cwd=ROOT_DIR, capture_output=True, - text=True + text=True, ) - + if result.returncode == 0: logger.info("Documentation generated successfully") if result.stdout: logger.info(f"Generator output: {result.stdout}") else: - logger.warning(f"Documentation generation failed with exit code {result.returncode}") + logger.warning( + f"Documentation generation failed with exit code {result.returncode}" + ) if result.stderr: logger.warning(f"Generator error: {result.stderr}") except Exception as e: logger.warning(f"Failed to run documentation generator: {e}") + class CustomBuildExt(build_ext): ROOT_DIR = os.path.abspath(os.path.dirname(__file__)) @@ -159,7 +174,7 @@ class CustomBuildExt(build_ext): """Override run method to include documentation generation.""" # Generate documentation before building extensions generate_docs() - + # Continue with normal build process super().run() @@ -173,12 +188,20 @@ class CustomBuildExt(build_ext): def _prepare_prebuild(): bash_cmd = f"bash scripts/prepare_prebuild.sh" try: - result = subprocess.run(bash_cmd, cwd=self.ROOT_DIR, text=True, shell=True, capture_output=False) + result = subprocess.run( + bash_cmd, + cwd=self.ROOT_DIR, + text=True, + shell=True, + capture_output=False, + ) if result.returncode != 0: logger.error("call prepare_prebuild.sh failed:") logger.error(result.stdout) # Print standard output logger.error(result.stderr) # Print error output - raise RuntimeError(f"call prepare_prebuild.sh failed with exit code {result.returncode}") + raise RuntimeError( + f"call prepare_prebuild.sh failed with exit code {result.returncode}" + ) except subprocess.CalledProcessError as e: raise RuntimeError(f"Failed to call prepare_prebuild.sh: {e}") @@ -198,7 +221,7 @@ class CustomBuildExt(build_ext): dst_dir = os.path.dirname(dst_so_path) package_path = os.path.join(dst_dir, package_name) os.makedirs(package_path, exist_ok=True) - + # Also prepare the Python package directory for generated files python_package_path = os.path.join(ROOT_DIR, "python", package_name) os.makedirs(python_package_path, exist_ok=True) @@ -207,7 +230,9 @@ class CustomBuildExt(build_ext): compile_cores = int(os.getenv("CMAKE_THREAD_NUM")) else: compile_cores = max(1, available_cores // 2) - logger.info(f"Available CPU cores: {available_cores}, using {compile_cores} cores for compilation") + logger.info( + f"Available CPU cores: {available_cores}, using {compile_cores} cores for compilation" + ) # Combine all cmake commands into one string cmake_cmd = ( f"source {env_script_path} && " @@ -225,12 +250,20 @@ class CustomBuildExt(build_ext): try: # Run the combined cmake command logger.info(f"Running combined CMake commands:\n{cmake_cmd}") - result = subprocess.run(cmake_cmd, cwd=self.ROOT_DIR, text=True, shell=True, capture_output=False) + result = subprocess.run( + cmake_cmd, + cwd=self.ROOT_DIR, + text=True, + shell=True, + capture_output=False, + ) if result.returncode != 0: logger.info("CMake commands failed:") logger.info(result.stdout) # Print standard output logger.info(result.stderr) # Print error output - raise RuntimeError(f"Combined CMake commands failed with exit code {result.returncode}") + raise RuntimeError( + f"Combined CMake commands failed with exit code {result.returncode}" + ) except subprocess.CalledProcessError as e: raise RuntimeError(f"Failed to build {so_name}: {e}") @@ -243,7 +276,9 @@ class CustomBuildExt(build_ext): logger.info(f"Copied {so_name} to {dst_so_path}") # Copy generated Python files to Python package directory - auto_generate_dir = os.path.join(build_extension_dir, ext_name + "_auto_generate") + auto_generate_dir = os.path.join( + build_extension_dir, ext_name + "_auto_generate" + ) if os.path.exists(auto_generate_dir): generated_files = ["gen_ops_def.py", "gen_ops_prim.py"] for gen_file in generated_files: @@ -251,9 +286,16 @@ class CustomBuildExt(build_ext): if os.path.exists(src_gen_path): dst_gen_path = os.path.join(package_path, gen_file) shutil.copy(src_gen_path, dst_gen_path) - replace_cmd = ["sed", "-i", "s/import ms_custom_ops/from . import ms_custom_ops/g", dst_gen_path] + replace_cmd = [ + "sed", + "-i", + "s/import ms_custom_ops/from . import ms_custom_ops/g", + dst_gen_path, + ] try: - result = subprocess.run(replace_cmd, cwd=self.ROOT_DIR, text=True, shell=False) + result = subprocess.run( + replace_cmd, cwd=self.ROOT_DIR, text=True, shell=False + ) except subprocess.CalledProcessError as e: raise RuntimeError(f"Failed to exec command {replace_cmd}: {e}") logger.info(f"Copied {gen_file} to {dst_gen_path}") @@ -263,9 +305,17 @@ class CustomBuildExt(build_ext): logger.warning(f"Auto-generate directory not found: {auto_generate_dir}") # Copy prebuild/ms_kernels_internal files to the package - prebuild_path = os.path.join(ROOT_DIR, "prebuild", "ms_kernels_internal/ms_kernels_internal") + if "MS_CUSTOM_INTERNAL_KERNEL_HOME" in os.environ: + prebuild_path = os.environ["MS_CUSTOM_INTERNAL_KERNEL_HOME"] + else: + prebuild_path = os.path.join( + ROOT_DIR, "prebuild", "ms_kernels_internal/ms_kernels_internal" + ) + if os.path.exists(prebuild_path): - package_prebuild_path = os.path.join(package_path, "prebuild", "ms_kernels_internal") + package_prebuild_path = os.path.join( + package_path, "prebuild", "ms_kernels_internal" + ) os.makedirs(package_prebuild_path, exist_ok=True) # Walk through the prebuild directory and copy all files @@ -285,22 +335,14 @@ class CustomBuildExt(build_ext): logger.info(f"Copied {src_file_path} to {dst_file_path}") - - write_commit_id() package_data = { - "": [ - "*.so", - "lib/*.so", - ".commit_id" - ], - "ms_custom_ops": [ - "gen_ops_def.py", - "gen_ops_prim.py" - ] + "": ["*.so", "lib/*.so", ".commit_id"], + "ms_custom_ops": ["gen_ops_def.py", "gen_ops_prim.py"], } + def _get_ext_modules(): ext_modules = [] if os.path.exists(_get_ascend_home_path()): @@ -308,14 +350,13 @@ def _get_ext_modules(): ext_modules.append(Extension("ms_custom_ops", sources=[])) return ext_modules + setup( name=package_name, version=version, author="MindSpore Team", license="Apache 2.0", - description=( - "MindSpore Custom Operations for Ascend NPU" - ), + description=("MindSpore Custom Operations for Ascend NPU"), long_description=read_readme(), long_description_content_type="text/markdown", url="https://gitee.com/mindspore/ms-custom-ops",