背景概述

在深度学习框架的生态系统中，设备管理模块是连接上层框架与底层硬件运行时的关键桥梁。随着昇腾 AI 处理器（Ascend NPU）在企业级 AI 训练和推理场景中的广泛应用，如何高效、稳定地管理 NPU 设备资源成为开发者面临的核心挑战之一。

torch_npu 作为 PyTorch 框架的昇腾 NPU 适配层，其设备管理模块承担着设备生命周期管理、计算流调度、事件同步、内存分配等核心职责，并通过简洁的 Python API 向开发者暴露易用的操作接口。本文从架构设计、核心流程、关键实现到使用示例，对设备管理模块进行全方位解析，帮助开发者理解其工作原理，并为相关系统优化与日常开发提供参考。

1. 模块概述

设备管理模块是 torch_npu 的核心基础设施，负责昇腾 NPU 设备的生命周期管理，是连接 PyTorch 框架与 CANN（Compute Architecture for Neural Networks）底层运行时的桥梁。

1.1 核心职责

• 设备生命周期管理：设备的初始化、切换、重置与释放
• 流管理（Stream Management）：创建和管理计算流，支持异步执行
• 事件管理（Event Management）：实现流间同步与性能计时
• 上下文管理（Context Management）：管理设备上下文，支持多设备协同场景
• 选项配置（Option Configuration）：提供灵活的运行时配置机制
• 设备属性查询：获取设备名称、内存信息、利用率等属性

1.2 模块组成

设备管理模块由多个职责清晰的子模块协作组成，主要组件包括：

• 系统控制（NpuSysCtrl）：单例式系统控制器，负责运行时初始化、清理，以及释放回调的注册与统一调度
• 设备函数（NPUFunctions）：提供设备数量查询、设备切换、设备同步等核心操作
• 流管理（NPUStream）：负责计算流的创建、复用与调度，支持多优先级与同步模式
• 事件管理（NPUEvent）：实现流间同步、事件查询与性能计时
• 设备守卫（NPUGuard / NPUGuardImpl）：以 RAII 方式自动管理设备上下文
• 选项管理（OptionsManager / OptionRegister）：提供运行时配置项的注册与查询
• ACL 接口封装（AclInterface）：封装 CANN 底层 ACL 接口调用

各组件相互协作，共同支撑 PyTorch 上层接口对 NPU 设备的统一访问。

2. 架构设计

设备管理模块采用清晰的分层架构，自顶向下依次为 Python API 层（torch_npu.npu）、C++ 核心层（c10_npu）、ACL 接口封装层（AclInterface）与 CANN 运行时（aclrt* API），整体结构如下。

2.1 整体架构图

┌─────────────────────────────────────────────────────────────────────────────┐
│                           Python API Layer                                   │
│  ┌─────────────┐ ┌─────────────┐ ┌─────────────┐ ┌─────────────────────────┐│
│  │   device()  │ │synchronize()│ │ set_device()│ │ current_device() / ...  ││
│  └──────┬──────┘ └──────┬──────┘ └──────┬──────┘ └────────────┬────────────┘│
│         │               │               │                     │             │
│  ┌──────┴───────────────┴───────────────┴─────────────────────┴────────────┐│
│  │                    torch_npu._C (Python Bindings)                       ││
│  └─────────────────────────────────────┬───────────────────────────────────┘│
└────────────────────────────────────────┼────────────────────────────────────┘
                                         │
┌────────────────────────────────────────┼────────────────────────────────────┐
│                           C++ Core Layer                                    │
│  ┌─────────────────────────────────────┴───────────────────────────────────┐│
│  │                         NpuSysCtrl (Singleton)                          ││
│  │   - Initialize() / Finalize()                                           ││
│  │   - LazyInitialize()                                                    ││
│  │   - RegisterReleaseFn()                                                 ││
│  └──────────────┬─────────────────────────────────────┬────────────────────┘│
│                 │                                     │                     │
│  ┌──────────────▼──────────────┐   ┌──────────────────▼─────────────────┐   │
│  │      NPUFunctions           │   │        NPUGuardImpl                │   │
│  │  - GetDevice() / SetDevice()│   │  - exchangeDevice()                │   │
│  │  - device_count()           │   │  - getDevice() / setDevice()       │   │
│  │  - current_device()         │   │  - getStream() / exchangeStream()  │   │
│  │  - device_synchronize()     │   │  - Event operations                │   │
│  └──────────────┬──────────────┘   └──────────────────┬─────────────────┘   │
│                 │                                     │                     │
│  ┌──────────────▼──────────────┐   ┌──────────────────▼─────────────────┐   │
│  │       NPUStream             │   │         NPUEvent                   │   │
│  │  - stream pool management   │   │  - record() / synchronize()        │   │
│  │  - getCurrentNPUStream()    │   │  - elapsed_time() / query()        │   │
│  │  - setCurrentNPUStream()    │   │  - IPC support                     │   │
│  └──────────────┬──────────────┘   └────────────────────────────────────┘   │
│                 │                                                           │
│  ┌──────────────▼──────────────────────────────────────────────────────────┐│
│  │                    OptionsManager / OptionRegister                      ││
│  │   - Environment variables configuration                                  ││
│  │   - Runtime options management                                           ││
│  └─────────────────────────────────────┬───────────────────────────────────┘│
└────────────────────────────────────────┼────────────────────────────────────┘
                                         │
┌────────────────────────────────────────┼────────────────────────────────────┐
│                    ACL Interface Layer                                      │
│  ┌─────────────────────────────────────┴───────────────────────────────────┐│
│  │                        AclInterface                                     ││
│  │   - aclrtSetDevice() / aclrtGetDevice()                                 ││
│  │   - aclrtCreateStream() / aclrtSynchronizeStream()                      ││
│  │   - aclrtCreateEvent() / aclrtRecordEvent()                             ││
│  │   - aclrtMalloc() / aclrtFree()                                         ││
│  └─────────────────────────────────────┬───────────────────────────────────┘│
└────────────────────────────────────────┼────────────────────────────────────┘
                                         │
┌────────────────────────────────────────▼────────────────────────────────────┐
│                    CANN Runtime (Ascend Computing Platform)                 │
│  ┌────────────────────────────────────────────────────────────────────────┐ │
│  │                     Ascend NPU Hardware                                │ │
│  └────────────────────────────────────────────────────────────────────────┘ │
└─────────────────────────────────────────────────────────────────────────────┘

2.2 核心模块交互关系

┌────────────────┐     uses      ┌──────────────────┐
│  NPUGuard      │──────────────▶│  NPUGuardImpl    │
└────────────────┘               └────────┬─────────┘
                                          │
                                          │ calls
                                          ▼
┌────────────────┐    manages    ┌──────────────────┐
│  NpuSysCtrl    │──────────────▶│  NPUFunctions    │
│  (Singleton)   │               └────────┬─────────┘
└───────┬────────┘                        │
        │                                 │ uses
        │ initializes                     ▼
        │                        ┌──────────────────┐
        │                        │   NPUStream      │
        │                        │   (Stream Pool)  │
        │                        └────────┬─────────┘
        │                                 │
        ▼                                 │ synchronizes
┌────────────────┐                        ▼
│  ACL Runtime   │◀──────────────┌──────────────────┐
│  aclrt*()      │               │    NPUEvent      │
└────────────────┘               └──────────────────┘

2.3 核心类关系（UML）

2.4 设备函数模块类图

2.5 流管理模块类图

3. 核心流程

设备管理模块的关键流程包括设备初始化、设备切换与资源释放。

3.1 设备初始化流程

                    ┌─────────────────────────┐
                    │  torch_npu.npu.init()   │
                    │  (Python Layer)         │
                    └───────────┬─────────────┘
                                │
                                ▼
                    ┌─────────────────────────┐
                    │    _lazy_init()         │
                    │  检查 _initialized 标志  │
                    └───────────┬─────────────┘
                                │
                    ┌───────────┴───────────┐
                    │                       │
               已初始化                   未初始化
                    │                       │
                    ▼                       ▼
              ┌──────────┐    ┌───────────────────────────┐
              │  返回    │    │ torch_npu._C._npu_init()  │
              └──────────┘    └───────────┬───────────────┘
                                          │
                                          ▼
                              ┌───────────────────────────┐
                              │ NpuSysCtrl::Initialize()  │
                              │ (C++ Singleton)           │
                              └───────────┬───────────────┘
                                          │
                    ┌─────────────────────┼─────────────────────┐
                    │                     │                     │
                    ▼                     ▼                     ▼
          ┌─────────────────┐  ┌─────────────────┐  ┌─────────────────┐
          │   aclInit()     │  │ SetDevice(0)    │  │ InitAllocator() │
          │ (ACL Runtime)   │  │ (默认设备)       │  │ (内存分配器)    │
          └─────────────────┘  └─────────────────┘  └─────────────────┘
                    │                     │                     │
                    └─────────────────────┼─────────────────────┘
                                          │
                                          ▼
                              ┌───────────────────────────┐
                              │  SetSocVersion()          │
                              │  设置芯片版本信息           │
                              └───────────┬───────────────┘
                                          │
                                          ▼
                              ┌───────────────────────────┐
                              │  InitAclops()             │
                              │  初始化算子库              │
                              └───────────┬───────────────┘
                                          │
                                          ▼
                              ┌───────────────────────────┐
                              │  SetThreadAffinity()      │
                              │  设置线程亲和性            │
                              └───────────┬───────────────┘
                                          │
                                          ▼
                              ┌───────────────────────────┐
                              │  init_flag_ = true        │
                              │  标记初始化完成            │
                              └───────────────────────────┘

详细初始化步骤包括：

1. ACL 初始化：调用 aclInit() 初始化 CANN 运行时
2. Dump 配置：如果启用了 ACL dump，初始化 dump 功能
3. 内存分配器初始化：初始化 NPUCachingAllocator 和 NPUWorkspaceAllocator
4. 设备设置：
   • 非延迟模式：直接设置设备（默认设备 0）
   • 延迟模式（CANN >= 8.3.RC1）：延迟到首次使用时设置
5. SoC 版本设置：获取并设置当前芯片版本
6. 浮点溢出模式：根据芯片特性设置 INF/NAN 模式
7. 算子初始化：初始化 ACLOP 算子库
8. 线程亲和性：设置主线程的 CPU 亲和性

3.2 设备切换流程

                    ┌─────────────────────────┐
                    │ torch.npu.set_device(n)│
                    │ 或 with device(n):      │
                    └───────────┬─────────────┘
                                │
                                ▼
                    ┌─────────────────────────┐
                    │  _get_device_index()    │
                    │  解析设备索引            │
                    └───────────┬─────────────┘
                                │
                                ▼
                    ┌─────────────────────────┐
                    │ torch_npu._C._npu_      │
                    │ setDevice(device_id)    │
                    └───────────┬─────────────┘
                                │
                                ▼
                    ┌─────────────────────────┐
                    │  c10_npu::SetDevice()   │
                    │  (设备函数模块)          │
                    └───────────┬─────────────┘
                                │
                    ┌───────────┴───────────┐
                    │                       │
            local_device == device   local_device != device
                    │                       │
                    ▼                       ▼
              ┌──────────┐    ┌───────────────────────────┐
              │  返回    │    │  SetThreadAffinity()      │
              │ ACL_OK   │    │  (如需要绑定CPU核心)        │
              └──────────┘    └───────────┬───────────────┘
                                          │
                                          ▼
                              ┌───────────────────────────┐
                              │   aclrtSetDevice(device)  │
                              │   (ACL Runtime Call)      │
                              └───────────┬───────────────┘
                                          │
                              ┌───────────┴───────────┐
                              │                       │
                           成功                     失败
                              │                       │
                              ▼                       ▼
                    ┌─────────────────┐    ┌──────────────────┐
                    │ 更新 local_device│    │  抛出异常         │
                    │ 记录到 used_devices│   │  NPU_CHECK_ERROR │
                    └─────────────────┘    └──────────────────┘

3.3 资源释放流程

                    ┌─────────────────────────┐
                    │  Python 进程退出        │
                    │  或显式调用 finalize    │
                    └───────────┬─────────────┘
                                │
                                ▼
                    ┌─────────────────────────┐
                    │ NpuSysCtrl::Finalize()  │
                    └───────────┬─────────────┘
                                │
                                ▼
                    ┌─────────────────────────┐
                    │  检查 init_flag_        │
                    │  如果未初始化则返回      │
                    └───────────┬─────────────┘
                                │
                                ▼
                    ┌─────────────────────────┐
                    │ 注册最后释放函数:        │
                    │ - ClearEvent()          │
                    │ - DestroyUsedStreams()  │
                    │ - ResetUsedDevices()    │
                    │ - aclFinalize()         │
                    └───────────┬─────────────┘
                                │
                                ▼
                    ┌─────────────────────────┐
                    │  init_flag_ = false     │
                    └───────────┬─────────────┘
                                │
                                ▼
                    ┌─────────────────────────┐
                    │  Finalize Dump (如启用)  │
                    └───────────┬─────────────┘
                                │
                                ▼
              ┌─────────────────────────────────────┐
              │  按优先级执行 release_fn_ 回调函数    │
              │  Priority: First -> Middle -> Last  │
              └─────────────────────────────────────┘
                                │
    ┌───────────────────────────┼───────────────────────────┐
    │                           │                           │
    ▼                           ▼                           ▼
┌─────────────┐        ┌─────────────────┐        ┌─────────────────┐
│ PriorityFirst│        │ PriorityMiddle  │        │ PriorityLast    │
│ 高优先级释放  │        │ 中优先级释放     │        │ 低优先级释放     │
└─────────────┘        └─────────────────┘        │ - ClearEvent    │
                                                   │ - DestroyStreams │
                                                   │ - ResetDevices   │
                                                   │ - aclFinalize    │
                                                   └─────────────────┘

4. 核心实现解析

4.1 NpuSysCtrl 实现

NpuSysCtrl 是一个单例类，负责整个 NPU 系统的初始化和清理工作。

4.1.1 单例获取

NpuSysCtrl &NpuSysCtrl::GetInstance()
{
    static NpuSysCtrl instance;
    return instance;
}

4.1.2 初始化流程

NpuSysCtrl::SysStatus NpuSysCtrl::Initialize(int device_id)
{
    if (init_flag_) {
        return INIT_SUCC;
    }
    std::lock_guard lock(init_mutex_);
    if (init_flag_) {
        return INIT_SUCC;
    }

    // 1. 获取 ACL 配置文件路径并初始化 ACL
    std::string json_path = GetAclConfigJsonPath();
    const char *json_path_ptr = json_path == "" ? nullptr : json_path.c_str();
    auto init_ret = aclInit(json_path_ptr);

    // 处理重复初始化的情况
    if (init_ret == ACL_ERROR_REPEAT_INITIALIZE) {
        repeat_init_acl_flag_ = false;  // 标记 ACL 非本模块初始化
    } else if (init_ret != ACL_ERROR_NONE) {
        NPU_CHECK_ERROR(init_ret, "aclInit");
    }

    // 2. 初始化内存分配器
    c10_npu::NPUCachingAllocator::init();
    c10_npu::NPUWorkspaceAllocator::init();

    // 3. 设置设备 (根据是否延迟初始化模式)
    if (!c10_npu::is_lazy_set_device()) {
        auto ret = aclrtGetDevice(&device_id_);
        if (ret != ACL_ERROR_NONE) {
            device_id_ = (device_id == -1) ? 0 : device_id;
            NPU_CHECK_ERROR(c10_npu::SetDevice(device_id_));
        }
    }

    // 4. 设置 SoC 版本
    auto soc_name = c10_npu::acl::AclGetSocName();
    c10_npu::SetSocVersion(soc_name);

    // 5. 设置浮点溢出模式
    if (!c10_npu::is_lazy_set_device()) {
        if (c10_npu::IsSupportInfNan()) {
            c10_npu::acl::AclrtSetDeviceSatMode(ACL_RT_OVERFLOW_MODE_INFNAN);
        } else {
            c10_npu::acl::AclrtSetDeviceSatMode(ACL_RT_OVERFLOW_MODE_SATURATION);
        }
    }

    // 6. 初始化算子
    auto acl_op_init_mode = c10_npu::option::OptionsManager::GetAclOpInitMode();
    if (acl_op_init_mode == 0) {
        at_npu::aclops::InitAclops();
    } else {
        at_npu::aclops::InitializeJitCompilationMode();
    }

    // 7. 执行延迟注册的回调
    for (const auto &iter: lazy_fn_) {
        const auto &call_ = iter.first;
        const auto &in = iter.second;
        call_(in);
    }
    lazy_fn_.clear();

    // 8. 设置线程亲和性
    SetMainThread();
    if (SetThreadAffinityInInitialize()) {
        SetThreadAffinity(device_id_);
    }

    init_flag_ = true;
    return INIT_SUCC;
}

4.1.3 延迟初始化（LazyInitialize）

// 仅在 CANN >= 8.3.RC1 版本生效
NpuSysCtrl::SysStatus NpuSysCtrl::LazyInitialize(int device_id)
{
    if (!c10_npu::is_lazy_set_device()) {
        return INIT_SUCC;  // 非延迟模式直接返回
    }

    if (lazy_init_flag_) {
        return INIT_SUCC;
    }
    std::lock_guard lock(lazy_init_mutex_);
    if (lazy_init_flag_) {
        return INIT_SUCC;
    }

    auto ret = aclrtGetDevice(&device_id_);

    // 设置浮点溢出模式
    if (c10_npu::IsSupportInfNan()) {
        c10_npu::acl::AclrtSetDeviceSatMode(ACL_RT_OVERFLOW_MODE_INFNAN);
    } else {
        c10_npu::acl::AclrtSetDeviceSatMode(ACL_RT_OVERFLOW_MODE_SATURATION);
    }

    // 设置系统参数
    NPU_CHECK_ERROR(at_npu::native::AclrtCtxSetSysParamOpt(
        aclSysParamOpt::ACL_OPT_DETERMINISTIC, 0));
    NPU_CHECK_ERROR(c10_npu::acl::AclrtSetOpExecuteTimeOut(kMaxOpExecuteTimeOut));

    lazy_init_flag_ = true;
    return INIT_SUCC;
}

4.1.4 资源释放

NpuSysCtrl::SysStatus NpuSysCtrl::Finalize()
{
    if (!init_flag_) {
        return FINALIZE_SUCC;
    }

    // 注册最后执行的清理函数
    this->RegisterReleaseFn(
        [=]() -> void {
            c10_npu::NPUEventManager::GetInstance().ClearEvent();
            NPU_CHECK_WARN(c10_npu::DestroyUsedStreams());
            NPU_CHECK_WARN(c10_npu::ResetUsedDevices());
            // 只有当 ACL 是由本模块初始化时才调用 aclFinalize
            if (repeat_init_acl_flag_) {
                NPU_CHECK_WARN(aclFinalize());
            }
        },
        ReleasePriority::PriorityLast);

    init_flag_ = false;

    // 按优先级执行所有注册的释放函数
    for (const auto &iter: release_fn_) {
        const auto &fn_vec = iter.second;
        for (const auto &fn: fn_vec) {
            fn();
        }
    }
    release_fn_.clear();

    return FINALIZE_SUCC;
}

4.2 NPUFunctions 实现

4.2.1 全局状态变量

namespace c10_npu {

static uint32_t dev_count = 0;                              // 设备总数
static thread_local int local_device = -1;                  // 线程本地设备缓存
static std::unordered_map used_devices; // 已使用设备的上下文
std::recursive_mutex mtx;                                   // 保护 used_devices 的互斥锁
thread_local int targetDeviceIndex = -1;                    // 目标设备索引 (延迟设置)

4.2.2 获取设备

aclError GetDevice(int32_t *device)
{
    // 优先使用目标设备索引 (延迟设置场景)
    if (targetDeviceIndex >= 0) {
        *device = targetDeviceIndex;
        return ACL_ERROR_NONE;
    }

    // 使用线程本地缓存
    if (local_device >= 0) {
        *device = local_device;
        return ACL_ERROR_NONE;
    }

    // 从 ACL 获取当前设备
    aclError err = aclrtGetDevice(device);
    if (err != ACL_ERROR_NONE) {
        CHECK_AND_THROW_ERROR_WITH_SPECIFIC_MESSAGE(err);
    }

    // 非延迟模式下更新缓存
    if (!is_lazy_set_device()) {
        if (err == ACL_ERROR_NONE) {
            local_device = *device;
        }
    }

    // 上下文为空时默认返回设备 0
    if (err == ACL_ERROR_RT_CONTEXT_NULL) {
        *device = 0;
        return ACL_ERROR_NONE;
    }
    return err;
}

4.2.3 设置设备

aclError SetDevice(c10::DeviceIndex device)
{
    TORCH_CHECK(device >= 0, "device id must be positive!", PTA_ERROR(ErrCode::VALUE));
    targetDeviceIndex = -1;  // 清除延迟设置标记

    // 已在当前设备上，直接返回
    if (local_device == device) {
        return ACL_ERROR_NONE;
    }

    // 设置线程 CPU 亲和性
    if (c10_npu::NeedMainThreadBind()) {
        c10_npu::SetThreadAffinity(device);
    }

    // 调用 ACL 设置设备
    aclError err = aclrtSetDevice(device);
    if (err == ACL_ERROR_NONE) {
        local_device = device;
        std::lock_guard lock(mtx);
        // 记录设备上下文
        if (used_devices.find(local_device) == used_devices.end()) {
            NPU_CHECK_ERROR(aclrtGetCurrentContext(&used_devices[local_device]));
        }
    }
    return err;
}

4.2.4 设备同步

void device_synchronize()
{
    NPU_CHECK_ERROR(c10_npu::acl::AclrtSynchronizeDeviceWithTimeout());
#ifndef BUILD_LIBTORCH
    // 触发 Python 层的同步回调 (用于性能分析等)
    const c10_npu::impl::PyCallbackTrigger* trigger =
        c10_npu::impl::NPUTrace::getTrace();
    if (C10_UNLIKELY(trigger)) {
        trigger->traceNpuDeviceSynchronization();
    }
#endif
}

4.3 设备属性获取

4.3.1 Python 层属性获取

def device_count() -> int:
    """返回可用 NPU 设备数量"""
    global _cached_device_count
    if _cached_device_count is not None:
        return _cached_device_count
    ascend_hal_count = _device_count_ascend_hal()
    r = torch_npu._C._npu_getDeviceCount() if ascend_hal_count < 0 else ascend_hal_count
    if _initialized:
        _cached_device_count = r
    return r

def get_device_name(device_name=None):
    """获取设备名称"""
    device_id = _get_device_index(device_name, optional=True)
    if device_id < 0 or device_id >= device_count():
        raise AssertionError("Invalid device id")
    return torch_npu._C._npu_getDeviceName()

def get_device_properties(device_name=None):
    """获取设备属性"""
    device_id = _get_device_index(device_name, optional=True)
    torch_npu.npu._lazy_init()
    return torch_npu._C._npu_getDeviceProperties(device_id)

def mem_get_info(device=None):
    """获取设备内存信息"""
    if device is None:
        device = torch_npu.npu.current_device()
    device_id = _get_device_index(device)
    torch_npu.npu._lazy_init()
    device_prop = torch_npu._C._npu_getDeviceMemories(device_id)
    return device_prop.free_memory, device_prop.total_memory

def utilization(device=None):
    """获取设备利用率"""
    device_id = _get_device_index(device, optional=True)
    torch_npu.npu._lazy_init()
    return torch_npu._C._npu_getDeviceUtilizationRate(device_id)

4.3.2 SoC 版本检测

设备驱动通过 CANN 运行时接口获取当前芯片型号，框架内部按公开型号（Ascend 910 系列、Ascend 310 系列、Ascend 950 等）维护版本信息，用于控制浮点溢出模式、BF16 支持等与芯片能力相关的运行时特性，并支持通过环境变量强制开启或关闭部分特性。具体的型号枚举与映射逻辑属于实现细节，此处不再展开。

4.4 流管理实现

4.4.1 流池初始化

static void initGlobalStreamState()
{
    num_npus = c10_npu::device_count();
    AT_ASSERTM(num_npus <= C10_COMPILE_TIME_MAX_NPUS,
               "Number of NPU devices exceeds compiled maximum");

    int device_id = 0;
    auto ret = c10_npu::GetDevice(&device_id);

    // 初始化默认流
    default_streams[device_id].device_index = device_id;
    for (const auto p : c10::irange(kMaxStreamPriorities)) {
        npu_counters[p][device_id] = 0;
    }

    auto& default_streamsi = default_streams[device_id];
    NPU_CHECK_ERROR(acl::AclrtCreateStreamWithConfig(
        &default_streamsi.stream, 0,
        (ACL_STREAM_FAST_LAUNCH | ACL_STREAM_FAST_SYNC)));

    // 初始化任务队列
    if (c10_npu::option::OptionsManager::GetTaskQueueEnable()) {
        default_streamsi.repo->InitRepo(device_id);
    }

    // 初始化辅助流
    secondary_streams[device_id].device_index = device_id;
    NPU_CHECK_ERROR(acl::AclrtCreateStreamWithConfig(
        &secondary_streams[device_id].stream, 0,
        (ACL_STREAM_FAST_LAUNCH | ACL_STREAM_FAST_SYNC)));
}

4.4.2 获取当前流

NPUStream getCurrentNPUStream(c10::DeviceIndex device_index)
{
    initNPUStreamsOnce();  // 确保流已初始化
    if (device_index == -1) {
        device_index = current_device();
    }
    check_npu(device_index);
    return NPUStream_fromInternals(current_streams[device_index]);
}

void setCurrentNPUStream(NPUStream stream)
{
    initNPUStreamsOnce();
    auto ptr = NPUStream_internals(stream);
    AT_ASSERT(ptr, PTA_ERROR(ErrCode::PTR));
    if (current_streams[ptr->device_index]->stream != ptr->stream) {
        ASCEND_LOGI("Exchange NPU current stream from %p to %p",
            current_streams[ptr->device_index]->stream, ptr->stream);
    }
    current_streams[ptr->device_index] = ptr;
}

5. 设计特点与使用建议

5.1 设计优势

优势	说明
线程安全设计	使用 thread_local 变量缓存设备 ID，减少跨线程竞争；使用 recursive_mutex 保护共享状态
延迟初始化	支持延迟设备设置，避免不必要的设备初始化开销
RAII 资源管理	NPUGuard 和 OptionalNPUGuard 使用 RAII 模式自动管理设备上下文切换
单例模式	NpuSysCtrl 使用单例模式确保全局唯一的系统控制器
流池复用	预创建流池，避免频繁创建销毁流的开销
优先级释放	资源释放支持优先级，确保依赖关系正确
完善的错误处理	统一的错误码和异常处理机制
PyTorch 兼容	实现 DeviceGuardImplInterface，与 PyTorch 设备抽象无缝集成
可配置性强	通过 OptionsManager 支持丰富的运行时配置

5.2 设计特点与演进方向

设计特点	说明	演进方向
最大设备数编译期确定	通过编译期常量控制最大设备数，流池等结构在编译期即可确定大小	可结合部署形态规划运行期动态扩展能力
全局互斥锁保护共享状态	已用设备上下文等共享状态由互斥锁保护，保证多线程访问一致	可进一步细化锁粒度，降低多设备高并发场景的竞争
流池预创建与复用	预创建固定数量的计算流并复用，减少频繁创建销毁流的开销	可支持按负载动态调整流池大小
依赖 CANN ACL 运行时	通过标准 ACL 接口与 CANN 运行时对接，接口层保持稳定	随 CANN 版本演进持续适配新增能力
初始化顺序有约束	设备相关操作需在初始化完成后调用，框架提供显式初始化入口	可通过更完善的状态管理进一步简化使用
设备数量启动时确定	设备资源在进程启动阶段确定，调度逻辑简洁	可结合产品形态规划设备热插拔支持

面向后续版本，可从性能、可维护性与功能三个方面持续演进：

• 性能优化：细粒度锁（按设备维度拆分互斥锁）、流池大小可配置、设备切换缓存策略优化，进一步降低高并发场景的开销
• 可维护性优化：以显式初始化状态机替代多标志位组合、统一运行时配置入口、关键路径增加结构化日志，提升可观测性
• 功能增强：设备健康检查、资源配额管理、设备间数据迁移等能力可结合生态需求逐步规划

5.3 使用注意事项

• 多进程场景：多进程场景推荐使用 spawn 启动方式；框架已提供 after-fork 处理，可确保子进程中的设备状态正确
• 流与事件资源：建议在长生命周期进程中复用流与事件对象，避免不必要的资源增长；进程退出时框架会统一完成资源清理
• 初始化方式：框架支持常规与延迟两种初始化模式，建议同一进程内保持一致的使用方式，避免因初始化时机差异带来状态不一致

6. Python API 使用示例

设备管理模块不仅提供底层 C++ 核心能力，还通过 Python API 向开发者暴露简洁、直观的操作接口。以下示例覆盖设备操作、上下文管理、流与事件同步、内存控制等核心功能，并给出多设备训练场景下的实践建议。

6.1 基本设备操作

以下示例展示了如何通过 torch_npu 的 Python API 完成设备可用性检查、初始化、切换以及属性查询等基础操作：

import torch
import torch_npu

# 1. 检查 NPU 可用性
if torch_npu.npu.is_available():
    print(f"NPU 可用，设备数量: {torch_npu.npu.device_count()}")
else:
    print("NPU 不可用")

# 2. 初始化 NPU (通常自动完成)
torch_npu.npu.init()

# 3. 获取/设置当前设备
current = torch_npu.npu.current_device()
print(f"当前设备: {current}")

torch_npu.npu.set_device(0)
print(f"切换后设备: {torch_npu.npu.current_device()}")

# 4. 获取设备信息
device_name = torch_npu.npu.get_device_name(0)
print(f"设备名称: {device_name}")

properties = torch_npu.npu.get_device_properties(0)
print(f"设备属性: {properties}")

free_mem, total_mem = torch_npu.npu.mem_get_info(0)
print(f"内存信息: {free_mem / 1e9:.2f}GB / {total_mem / 1e9:.2f}GB")

# 5. 获取设备利用率
util = torch_npu.npu.utilization(0)
print(f"设备利用率: {util}%")

# 6. 获取 CANN 版本
cann_version = torch_npu.npu.utils.get_cann_version()
print(f"CANN 版本: {cann_version}")

6.2 设备上下文管理

设备上下文管理是编写多设备代码的基础。torch_npu 提供了与 PyTorch 一致的 device 和 device_of 上下文管理器，支持嵌套使用：

import torch
import torch_npu

# 1. 使用 with 语句管理设备上下文
with torch_npu.npu.device(0):
    tensor_on_0 = torch.randn(1000, 1000, device='npu')
    print(f"张量在设备: {tensor_on_0.device}")

with torch_npu.npu.device(1):
    tensor_on_1 = torch.randn(1000, 1000, device='npu')
    print(f"张量在设备: {tensor_on_1.device}")

# 2. 使用 device_of 跟随张量设备
with torch_npu.npu.device_of(tensor_on_0):
    result = tensor_on_0 + 1
    print(f"结果在设备: {result.device}")

# 3. 嵌套设备上下文
with torch_npu.npu.device(0):
    print(f"外层设备: {torch_npu.npu.current_device()}")
    with torch_npu.npu.device(1):
        print(f"内层设备: {torch_npu.npu.current_device()}")
    print(f"返回外层: {torch_npu.npu.current_device()}")

6.3 流管理

流（Stream）机制是实现异步计算和并行执行的关键。以下示例展示了流的创建、切换和同步方法：

import torch
import torch_npu

# 1. 获取默认流和当前流
default_stream = torch_npu.npu.default_stream()
current_stream = torch_npu.npu.current_stream()
print(f"默认流: {default_stream}")
print(f"当前流: {current_stream}")

# 2. 创建新流
new_stream = torch_npu.npu.Stream()
print(f"新建流: {new_stream}")

# 3. 使用流上下文
with torch_npu.npu.stream(new_stream):
    # 这里的操作在 new_stream 上执行
    tensor = torch.randn(1000, 1000, device='npu')
    result = tensor @ tensor.T

# 4. 设置当前流
torch_npu.npu.set_stream(new_stream)

# 5. 流同步
new_stream.synchronize()

6.4 事件同步

事件（Event）用于流间同步和性能计时。以下示例展示了如何利用事件测量算子执行时间：

import torch
import torch_npu
import time

# 1. 创建事件
start_event = torch_npu.npu.Event(enable_timing=True)
end_event = torch_npu.npu.Event(enable_timing=True)

# 2. 使用事件计时
tensor = torch.randn(10000, 10000, device='npu')

start_event.record()
result = tensor @ tensor.T
end_event.record()

# 3. 等待事件完成
end_event.synchronize()

# 4. 计算耗时
elapsed_time = start_event.elapsed_time(end_event)
print(f"矩阵乘法耗时: {elapsed_time:.2f} ms")

# 5. 查询事件状态
print(f"事件是否完成: {end_event.query()}")

6.5 设备同步

设备同步用于确保所有异步操作完成，在调试和性能分析场景中尤为重要：

import torch
import torch_npu

# 1. 同步当前设备
torch_npu.npu.synchronize()

# 2. 同步指定设备
torch_npu.npu.synchronize(device=0)

# 3. 使用同步调试模式
# "default" - 不警告
# "warn" - 同步时发出警告
# "error" - 同步时抛出错误
torch_npu.npu.set_sync_debug_mode("warn")

# 执行同步操作时会有警告
torch_npu.npu.synchronize()

# 获取当前同步调试模式
mode = torch_npu.npu.get_sync_debug_mode()
print(f"同步调试模式: {mode}")

# 恢复默认
torch_npu.npu.set_sync_debug_mode("default")

6.6 内存管理

内存管理接口帮助开发者监控和优化显存使用，避免内存溢出或碎片化问题：

import torch
import torch_npu

# 1. 查看内存统计
stats = torch_npu.npu.memory_stats()
print(f"内存统计: {stats}")

# 2. 获取已分配内存
allocated = torch_npu.npu.memory_allocated()
print(f"已分配内存: {allocated / 1e9:.2f} GB")

# 3. 获取缓存内存
cached = torch_npu.npu.memory_cached()
print(f"缓存内存: {cached / 1e9:.2f} GB")

# 4. 获取最大内存
max_allocated = torch_npu.npu.max_memory_allocated()
print(f"最大分配内存: {max_allocated / 1e9:.2f} GB")

# 5. 清空缓存
torch_npu.npu.empty_cache()

# 6. 重置统计
torch_npu.npu.reset_peak_memory_stats()
torch_npu.npu.reset_accumulated_memory_stats()

# 7. 设置内存分数限制
torch_npu.npu.set_per_process_memory_fraction(0.8, device=0)

6.7 设备恢复

以下示例展示了设备停止与重启接口的用法，帮助开发者在需要时回收或重建设备资源：

import torch
import torch_npu

# 1. 停止设备
# 注意: 这会停止设备上的所有计算任务
# torch_npu.npu.stop_device(device_id=0)

# 2. 重启设备
# rebuild_all_resources=True 会重建所有资源
# torch_npu.npu.restart_device(device_id=0, rebuild_all_resources=True)

6.8 高级配置

torch_npu 提供了一系列高级配置接口，用于满足特定场景下的性能调优和功能定制需求：

import torch
import torch_npu

# 1. 检查 BF16 支持
bf16_supported = torch_npu.npu.is_bf16_supported()
print(f"BF16 支持: {bf16_supported}")

# 2. 检查 INF/NAN 模式支持
from torch_npu.npu.utils import is_support_inf_nan
inf_nan_supported = is_support_inf_nan()
print(f"INF/NAN 模式支持: {inf_nan_supported}")

# 3. 设置算子超时
torch_npu.npu.set_op_timeout_ms(timeout=30000)  # 30 秒

# 4. Dump 配置
torch_npu.npu.init_dump()
torch_npu.npu.set_dump("/path/to/dump_config.json")
# ... 执行需要 dump 的操作 ...
torch_npu.npu.finalize_dump()

# 5. 使用一致性算法模式
torch_npu.npu.use_consistent_algorithms(is_enable=True)
is_enabled = torch_npu.npu.are_consistent_algorithms_enable()
print(f"一致性算法模式: {is_enabled}")

# 6. P2P 访问检查
can_access = torch_npu.npu.can_device_access_peer(device_id=0, peer_device_id=1)
print(f"设备 0 可以访问设备 1: {can_access}")

6.9 多设备训练的组织方式

以下示例展示了以设备上下文为单位组织训练流程的方式：将单卡训练逻辑封装为可按设备索引执行的函数，多设备并行可通过多进程方式在各设备上分别启动训练任务：

import torch
import torch.nn as nn
import torch_npu

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(1024, 512)
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, 10)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        return self.fc3(x)

def train_on_device(device_id):
    """在指定设备上训练"""
    with torch_npu.npu.device(device_id):
        # 初始化模型
        model = SimpleModel().npu()
        optimizer = torch.optim.Adam(model.parameters())
        criterion = nn.CrossEntropyLoss()

        # 创建数据
        data = torch.randn(32, 1024, device='npu')
        labels = torch.randint(0, 10, (32,), device='npu')

        # 训练步骤
        for step in range(100):
            optimizer.zero_grad()
            outputs = model(data)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            if step % 10 == 0:
                print(f"Device {device_id}, Step {step}, Loss: {loss.item():.4f}")

        # 同步
        torch_npu.npu.synchronize()

    return model

# 在多个设备上并行训练 (使用多进程)
if __name__ == "__main__":
    import multiprocessing as mp

    num_devices = torch_npu.npu.device_count()
    print(f"使用 {num_devices} 个设备进行训练")

    # 简单示例: 在设备 0 上训练
    model = train_on_device(0)
    print("训练完成!")

7. 总结

torch_npu 的设备管理模块是连接 PyTorch 与昇腾 NPU 硬件的关键基础设施，通过以下机制确保了 NPU 设备的高效管理：

• 分层架构：Python API → C++ 核心 → ACL 接口 → CANN 运行时
• 线程安全：使用线程本地缓存和互斥锁保护共享状态
• 资源复用：流池、事件池等机制减少资源创建开销
• RAII 模式：Guard 类自动管理设备上下文生命周期
• 可配置性：OptionsManager 支持丰富的运行时配置

该模块与 PyTorch 的设备抽象接口紧密集成，使得用户可以使用熟悉的 PyTorch API 操作 NPU 设备，降低了迁移成本。结合文中的 Python API 示例，开发者可以快速上手 torch_npu 的设备管理功能，并将其应用于实际的模型训练和推理任务中。通过持续的优化与演进，该模块将更好地支撑大规模 AI 训练和推理场景。