From 7b28fb8df13f7a43c75f9f987ebf3aa693ee90aa Mon Sep 17 00:00:00 2001 From: Li Baoming <41820386+baominghelly@users.noreply.github.com> Date: Mon, 14 Sep 2026 17:23:32 +0800 Subject: [PATCH] fix(ascend): detect exact SoC version --- .../ascend/custom/cmake/detect_soc.cmake | 62 +++++++++++-- tests/test_ascend_soc_detection.py | 90 +++++++++++++++++++ 2 files changed, 143 insertions(+), 9 deletions(-) create mode 100644 tests/test_ascend_soc_detection.py diff --git a/src/native/ascend/custom/cmake/detect_soc.cmake b/src/native/ascend/custom/cmake/detect_soc.cmake index 548ead0aa..bf2d4765c 100644 --- a/src/native/ascend/custom/cmake/detect_soc.cmake +++ b/src/native/ascend/custom/cmake/detect_soc.cmake @@ -1,9 +1,9 @@ -# Auto-detect the Ascend SOC version from `npu-smi info`. +# Auto-detect the Ascend SOC version from torch_npu or `npu-smi info`. # -# `infiniops_detect_soc()` parses the first `910*` / `310*` entry -# in `npu-smi info` and writes `Ascend` into the named variable in -# the caller's scope. Falls back to `Ascend910B4` when detection fails -# (no NPU on the host, `npu-smi` missing, output format mismatch). +# `infiniops_detect_soc()` first asks torch_npu for the exact device +# name, then falls back to a model-qualified `910*` / `310*` token from +# `npu-smi info`. Falls back to `Ascend910B4` when detection fails (no NPU +# on the host, missing tools, or an output format without a model suffix). # # Called from both `src/CMakeLists.txt` (outer `pip install` build, to # forward `SOC_VERSION` to the standalone `build.sh` invocation) and @@ -11,10 +11,54 @@ # by that `build.sh`). function(infiniops_detect_soc out_var) - execute_process( - COMMAND bash -c "npu-smi info 2>/dev/null | awk '/910B|910A|310/ {for (i=1;i<=NF;i++) if ($i ~ /^(910|310)/) {print \"Ascend\" $i; exit}}'" - OUTPUT_VARIABLE _detected - OUTPUT_STRIP_TRAILING_WHITESPACE) + set(_detected "") + set(_python "") + + if(DEFINED _TORCH_PYTHON AND NOT "${_TORCH_PYTHON}" STREQUAL "") + set(_python "${_TORCH_PYTHON}") + elseif(DEFINED Python_EXECUTABLE AND NOT "${Python_EXECUTABLE}" STREQUAL "") + set(_python "${Python_EXECUTABLE}") + else() + find_program(_python NAMES python3 python) + endif() + + if(_python) + execute_process( + COMMAND "${_python}" -c "import torch, torch_npu; print(torch.npu.get_device_name(0))" + RESULT_VARIABLE _torch_npu_result + OUTPUT_VARIABLE _torch_npu_output + ERROR_QUIET + OUTPUT_STRIP_TRAILING_WHITESPACE) + + if(_torch_npu_result EQUAL 0 AND + _torch_npu_output MATCHES "^Ascend(910|310)[A-Za-z0-9_]*$") + set(_detected "${_torch_npu_output}") + endif() + endif() + + if(NOT _detected) + execute_process( + COMMAND npu-smi info + RESULT_VARIABLE _npu_smi_result + OUTPUT_VARIABLE _npu_smi_output + ERROR_QUIET + OUTPUT_STRIP_TRAILING_WHITESPACE) + + if(_npu_smi_result EQUAL 0) + string(REGEX MATCH + "(Ascend)?(910|310)[A-Za-z_][A-Za-z0-9_]*" + _npu_smi_soc + "${_npu_smi_output}") + + if(_npu_smi_soc) + if(_npu_smi_soc MATCHES "^Ascend") + set(_detected "${_npu_smi_soc}") + else() + set(_detected "Ascend${_npu_smi_soc}") + endif() + endif() + endif() + endif() if(_detected) set(${out_var} "${_detected}" PARENT_SCOPE) diff --git a/tests/test_ascend_soc_detection.py b/tests/test_ascend_soc_detection.py new file mode 100644 index 000000000..855442ba1 --- /dev/null +++ b/tests/test_ascend_soc_detection.py @@ -0,0 +1,90 @@ +import os +import stat +import subprocess +from pathlib import Path + +import pytest + + +DETECT_SOC = ( + Path(__file__).parents[1] + / "src" + / "native" + / "ascend" + / "custom" + / "cmake" + / "detect_soc.cmake" +) + + +def _write_executable(path, contents): + path.write_text(contents) + path.chmod(path.stat().st_mode | stat.S_IXUSR) + + +def _detect_soc(tmp_path, python_output=None, npu_smi_output=""): + if os.name == "nt": + pytest.skip("the detector invokes POSIX Ascend command-line tools") + + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + python = bin_dir / "python3" + npu_smi = bin_dir / "npu-smi" + + if python_output is None: + _write_executable(python, "#!/bin/sh\nexit 1\n") + else: + _write_executable(python, f"#!/bin/sh\nprintf '%s\\n' '{python_output}'\n") + + _write_executable( + npu_smi, + "#!/bin/sh\ncat <<'EOF'\n" + npu_smi_output + "\nEOF\n", + ) + + driver = tmp_path / "detect.cmake" + driver.write_text( + f'set(Python_EXECUTABLE "{python.as_posix()}")\n' + f'include("{DETECT_SOC.as_posix()}")\n' + "infiniops_detect_soc(DETECTED_SOC)\n" + 'message(STATUS "DETECTED_SOC=${DETECTED_SOC}")\n' + ) + env = os.environ.copy() + env["PATH"] = f"{bin_dir}{os.pathsep}{env['PATH']}" + result = subprocess.run( + ["cmake", "-P", str(driver)], + capture_output=True, + text=True, + env=env, + check=True, + ) + + return result.stdout + + +def test_detect_soc_prefers_exact_torch_npu_name(tmp_path): + output = _detect_soc( + tmp_path, + python_output="Ascend910_9362", + npu_smi_output="| 7 Ascend910 | 0 0 / 0 3109 / 65536 |", + ) + + assert "DETECTED_SOC=Ascend910_9362" in output + + +def test_detect_soc_ignores_hbm_numbers_in_npu_smi(tmp_path): + output = _detect_soc( + tmp_path, + npu_smi_output="| 7 Ascend910 | 0 0 / 0 3109 / 65536 |", + ) + + assert "DETECTED_SOC=Ascend910B4" in output + assert "Ascend3109" not in output + + +def test_detect_soc_accepts_model_qualified_npu_smi_name(tmp_path): + output = _detect_soc( + tmp_path, + npu_smi_output="| 0 910B4 | 0 0 / 0 2789 / 32768 |", + ) + + assert "DETECTED_SOC=Ascend910B4" in output