提交提案之前，请先检索仓库内是否已有相同的提案，如已有请在同一提案中进行讨论。

性能优化具体描述

PYTORCH_HAL_BASED_NPU_CHECK 环境变量与 is_available() 分支优化方案

1. 需求背景与目标

当前 torch.npu.is_available() 在 NPU 尚未初始化时调用耗时较长。该接口通过 device_count() 检查设备，而 device_count() 在初始化前不会缓存结果，连续调用时会重复加载 libascend_hal.so、查询设备数量并解析 ASCEND_RT_VISIBLE_DEVICES。当框架或业务代码频繁查询 NPU 可用性时，重复 HAL 探测会累积为明显的额外耗时。

本方案的目标是降低 torch.npu.is_available() 默认路径的调用耗时，同时保证 NPU 可用性判断结果正确，并为需要原有 HAL 探测语义的场景保留显式开关。

2. 方案设计

2.1 新增环境变量

项目	说明
名称	PYTORCH_HAL_BASED_NPU_CHECK
生效条件	变量值严格等于字符串 "1"
未设置或其他值	is_available() 使用 _npu_getDeviceCount() 执行默认 Runtime 检查
HAL 失败	沿用 device_count() 既有逻辑，HAL 返回负值时回退 Runtime API
HAL 返回 0	0 是合法设备数，不触发回退，is_available() 返回 False

环境变量命名和严格匹配语义对齐 CUDA 的 PYTORCH_NVML_BASED_CUDA_CHECK

环境变量取值说明

取值	行为
"1"	开启 HAL 路径，调用 device_count()
其他（包括未设置）	关闭 HAL 路径，调用 _npu_getDeviceCount()

判断条件为 os.getenv("PYTORCH_HAL_BASED_NPU_CHECK") == "1"，因此该开关默认关闭，且只有字符串 "1" 会使其生效。

2.2 调用路径

torch.npu.is_available()
  ├─ torch_npu 未包含 NPU 能力 → False
  ├─ PYTORCH_HAL_BASED_NPU_CHECK == "1"
  │    └─ device_count() > 0
  │         ├─ ascend_hal 探测成功 → 使用 HAL 结果
  │         └─ ascend_hal 探测失败（返回负值）→ Runtime API
  └─ 其他取值或未设置
       └─ torch_npu._C._npu_getDeviceCount() > 0

2.3 代码分支改进

改进前：
def is_available():
    if (not hasattr(torch_npu._C, '_npu_setDevice')):
        return False
    return device_count() > 0
    
改进后：
def _hal_based_avail() -> bool:
    return os.getenv("PYTORCH_HAL_BASED_NPU_CHECK") == "1"

def is_available():
    if not hasattr(torch_npu._C, "_npu_setDevice"):
        return False
    if _hal_based_avail():
        return device_count() > 0
    return torch_npu._C._npu_getDeviceCount() > 0

3. 与 CUDA 的对齐情况

本方案与 CUDA 侧 is_available() 的环境变量分支设计一致，环境变量值严格等于 "1" 时启用替代检测路径，否则默认使用 Runtime 设备计数。

torch.cuda.is_available()
  ├─ PyTorch 未编译 CUDA 支持 → False
  ├─ PYTORCH_NVML_BASED_CUDA_CHECK == "1"
  │    └─ device_count() > 0
  │         ├─ NVML 探测成功 → 使用 NVML 结果
  │         └─ NVML 探测失败（返回负值）→ CUDA Runtime API
  └─ 其他取值或未设置
       └─ torch._C._cuda_getDeviceCount() > 0

cuda侧源码

def _nvml_based_avail() -> bool:
    return os.getenv("PYTORCH_NVML_BASED_CUDA_CHECK") == "1"

def is_available() -> bool:
    if not _is_compiled():
        return False
    if _nvml_based_avail():
        return device_count() > 0
    else:
        return torch._C._cuda_getDeviceCount() > 0

4. CI 单元测试设计

测试文件：test/test_npu_hal_based_avail.py。

• HAL 路径测试：设置 PYTORCH_HAL_BASED_NPU_CHECK=1，验证 is_available() 根据 device_count() 的结果返回 True 或 False，且不调用 Runtime 设备计数接口。
• 默认 Runtime 路径测试：环境变量未设置或取其他值时，验证 is_available() 根据 _npu_getDeviceCount() 的结果返回 True 或 False，且不调用 HAL 设备计数接口。

5. 后续项

5.1 fork poisoning

CUDA 默认 Runtime 查询会初始化驱动，因此其测试会检查进程是否进入 bad fork 状态。NPU 的 _npu_getDeviceCount() 调用后 torch_npu.npu.is_initialized() 仍为 False，现有 bad fork 标志无法反映 CANN 底层是否存在 fork 副作用，因此不直接照搬 CUDA 的毒进程断言。

本方案改用实际行为验证：分别在父进程未查询可用性和调用默认 Runtime is_available() 后 fork 子进程，两组各执行 3 次相同的 NPU 初始化和基础张量操作，均成功退出。该结果说明在本次测试环境和操作范围内，未观察到 _npu_getDeviceCount() 引入 fork poisoning；不作为所有 CANN、驱动和工作负载的通用保证。

5.2 可见设备环境变量

默认 Runtime 路径会缓存首次查询得到的非零设备数。因此，ASCEND_RT_VISIBLE_DEVICES 应在进程启动前设置；同一进程内动态修改该变量，不保证对后续 is_available() 调用生效。

5.3 C++ noexcept 错误路径

当前 c10_npu::device_count() noexcept 内部错误处理可能抛出异常，从而触发 std::terminate。该问题原本存在于 HAL 失败后的 Runtime 回退路径；默认改走 Runtime 后，其暴露范围扩大。建议单独修复为捕获异常、记录告警并返回 0，与 CUDA 的容错语义对齐。

性能劣化说明

其他相关讨论

环境信息

-操作系统：openEuler 24.03
-昇腾硬件：910B
-CANN版本：9.2.0
-软件版本：Python 3.10、torch_npu 2.7.1

欢迎加入社区，感谢您对社区的贡献 🎉!