Star0
0
代码
代码
Issues4
Pull Requests
项目讨论
Wiki
分析
Star0
0
已开启
【初始化】torch_npu初始化流程分析 #1
hbhu_bin创建于  3月12日
H
hbhu_bin成员
3月12日 创建

torch_npu 初始化流程详解

目录

  1. 概述
  2. 自动加载机制
  3. 完整初始化流程
  4. torch._C vs torch_npu._C
  5. PyInit__C 调用时机
  6. Python C 扩展机制
  7. _initExtension 调用流程
  8. aclInit 调用时机
  9. _lazy_init 触发时机
  10. 常见问题解答

概述

torch_npu 是华为昇腾 NPU 的 PyTorch 扩展库，通过 PyTorch 的插件机制实现无缝集成。本文档详细分析 torch_npu 的初始化流程，包括：

  • 自动加载机制
  • Python C 扩展原理
  • 设备初始化流程
  • ACL (Ascend Computing Language) 初始化时机
  • 延迟初始化策略

自动加载机制

Entry Points 机制

# setup.py:753-760
entry_points={
    'console_scripts': [
        'torch_npu_run = torch_npu.distributed.run:_main',
    ],
    'torch.backends': [
        'torch_npu = torch_npu:_autoload',
    ],
}

这是 PyTorch 2.0+ 提供的设备后端自动加载机制。当用户执行 import torch 时，PyTorch 会自动扫描已安装包中注册的 torch.backends 入口点，并调用对应的函数。

自动加载流程

import torch
    │
    ▼
PyTorch 扫描 torch.backends entry points
    │
    │  通过 importlib.metadata 查找已安装的包
    │  发现: torch_npu = torch_npu:_autoload
    │
    ▼
调用 torch_npu._autoload()
    │
    │  torch_npu/__init__.py:333-335
    │
    ▼
执行 torch_npu/__init__.py 完整初始化流程
    │
    ├── 注册 NPU 设备
    ├── 应用 Monkey Patches
    └── 初始化 C++ 扩展

_autoload 函数实现

# torch_npu/__init__.py:333-335
def _autoload():
    # 恢复环境变量，子进程需要继承其值
    os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = ORG_AUTOLOAD

这个函数非常简单，只是恢复了环境变量。真正的初始化在模块加载时完成。

完整初始化流程

步骤概览

步骤 代码位置 作用
1 __init__.py:14-15 防止循环依赖
2 __init__.py:22-38 检查加速器冲突
3 __init__.py:40-66 加载核心模块
4 __init__.py:209-215 注册 NPU 设备
5 __init__.py:217-221 应用 Monkey Patches
6 __init__.py:238 初始化 C++ 扩展

步骤 1: 防止循环依赖

# torch_npu/__init__.py:14-15
ORG_AUTOLOAD = os.getenv("TORCH_DEVICE_BACKEND_AUTOLOAD", "1")
os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0"

在 import torch 之前临时禁用自动加载，避免循环依赖。

步骤 2: 检查加速器冲突

# torch_npu/__init__.py:22-38
acc = torch._C._get_accelerator()
if acc.type != "cpu":
    raise RuntimeError(
        f"Two accelerators cannot be used at the same time "
        f"in PyTorch: npu and {acc.type}. You can install "
        f"the cpu version of PyTorch to use your npu device, "
        f"or use the {acc.type} device with "
        f"'export TORCH_DEVICE_BACKEND_AUTOLOAD=0'."
    )

确保系统中没有其他加速器（如 CUDA）同时启用。

步骤 3: 加载核心模块

# torch_npu/__init__.py:40-66
import torch_npu.npu       # 核心 NPU 功能
import torch_npu._C        # C++ 扩展模块
import torch_npu.npu.amp   # 自动混合精度
import torch_npu.optim     # 优化器
import torch_npu.dynamo    # Dynamo 支持

步骤 4: 注册 NPU 设备（最关键！）

# torch_npu/__init__.py:209-215
# 将 PyTorch 预留的 privateuse1 设备重命名为 'npu'
torch.utils.rename_privateuse1_backend("npu")

# 注册设备模块，使得可以访问 torch.npu
torch._register_device_module('npu', torch_npu.npu)

# 自动生成设备相关方法
unsupported_dtype = [torch.quint8, torch.quint4x2, torch.quint2x4, torch.qint32, torch.qint8]
torch.utils.generate_methods_for_privateuse1_backend(
    for_tensor=True,      # 生成 torch.Tensor.npu()
    for_module=True,      # 生成 torch.npu.* 方法
    for_storage=True,     # 生成存储相关方法
    unsupported_dtype=unsupported_dtype
)

这行代码的效果：

  • torch.Tensor.npu() - 将张量移动到 NPU
  • torch.npu.is_available() - 检查 NPU 是否可用
  • torch.npu.device_count() - 获取 NPU 数量
  • torch.npu.current_device() - 获取当前设备

步骤 5: 应用 Monkey Patches

# torch_npu/__init__.py:217-221
_apply_patches(all_monkey_patches)  # 应用各种补丁
_apply_class_patches()              # 类级别补丁
_asd_patch()                        # ASD 补丁
_except_handler.patch_excepthook()  # 异常处理补丁

步骤 6: 初始化 C++ 扩展

# torch_npu/__init__.py:238
torch_npu._C._initExtension()

torch._C vs torch_npu._C

这两个是完全不同的模块，分别属于不同的包：

对比表格

特性 torch._C torch_npu._C
所属包 PyTorch 核心 torch_npu 扩展
定义位置 PyTorch 源码 torch_npu 源码
入口函数 PyInit__C() PyInit__C()
模块名 torch._C torch_npu._C
功能 PyTorch 核心功能 NPU 相关功能

torch._C（PyTorch 核心）

torch._C
├── 核心张量操作
├── 自动求导引擎
├── JIT 编译器
├── CUDA 支持
├── 分布式通信
├── 内存管理
└── 其他 PyTorch 核心功能

典型使用：

import torch

# torch._C 是 PyTorch 的内部模块
torch._C._is_device_backend_autoload_enabled()
torch._C._get_accelerator()
torch._C._cuda_getDeviceCount()

torch_npu._C（NPU 扩展）

torch_npu._C
├── NPU 设备管理
│   ├── _npu_init()
│   ├── _npu_setDevice()
│   ├── _npu_getDevice()
│   └── _npu_synchronize()
├── NPU 内存管理
│   ├── _npu_emptyCache()
│   ├── _npu_memoryStats()
│   └── _npu_getDeviceMemories()
├── NPU Stream/Event
│   ├── _npu_getCurrentStream()
│   └── _npu_setStream()
├── NPU 配置
│   ├── _npu_setOption()
│   └── _npu_getOption()
├── Python 绑定初始化
│   └── _initExtension()
└── 其他 NPU 特有功能

典型使用：

import torch_npu

# torch_npu._C 是 torch_npu 的内部模块
torch_npu._C._npu_init()
torch_npu._C._npu_setDevice(0)
torch_npu._C._npu_synchronize()
torch_npu._C._initExtension()

加载时机

import torch
    │
    ├── 加载 torch._C (PyTorch 核心)
    │
    └── 扫描 torch.backends entry points
            │
            └── 调用 torch_npu._autoload()
                    │
                    └── import torch_npu
                            │
                            └── import torch_npu._C (NPU 扩展)

总结

模块 说明
torch._C PyTorch 核心 C++ 扩展，包含所有基础功能
torch_npu._C torch_npu 扩展 C++ 模块，包含 NPU 特有功能

两者虽然入口函数名相同（PyInit__C），但模块名不同，是独立的两个模块，互不冲突。

PyInit__C 调用时机

触发位置

# torch_npu/__init__.py:66
import torch_npu._C  # 触发 PyInit__C()

调用链

import torch_npu._C
    │
    ▼
Python 解释器查找模块
    │
    │  在 torch_npu 包中找到 _C 扩展模块
    │  (_C.cpython-310-x86_64-linux-gnu.so)
    │
    ▼
dlopen() 加载共享库
    │
    │  Python 使用 dlopen 加载 .so 文件
    │
    ▼
dlsym() 查找 PyInit__C 符号
    │
    │  Python 约定：模块名 xxx 的初始化函数为 PyInit_xxx
    │
    ▼
调用 PyInit__C()
    │
    ▼
initModule() 创建模块对象

PyInit__C 函数实现

// torch_npu/csrc/InitNpuBindings.cpp:216-219
PyMODINIT_FUNC PyInit__C(void)
{
    return initModule();
}

initModule() 函数

// torch_npu/csrc/InitNpuBindings.cpp:165-214
PyObject* initModule()
{
    // 1. 初始化线程数
    at::internal::lazy_init_num_threads();

    // 2. 注册 Python 方法
    AddPyMethodDefs(methods, TorchNpuMethods);
    AddPyMethodDefs(methods, TorchSanitizerMethods);
    AddPyMethodDefs(methods, THNPModule_get_methods());
    AddPyMethodDefs(methods, torch_npu::utils::npu_extension_functions());

    // 3. 创建 Python 模块 torch_npu._C
    static struct PyModuleDef torchnpu_module = {
        PyModuleDef_HEAD_INIT,
        "torch_npu._C",  // 模块名
        nullptr,
        -1,
        methods.data()
    };
    module = PyModule_Create(&torchnpu_module);

    // 4. 初始化核心对象
    THNPStream_init(module);   // Stream 对象
    THNPEvent_init(module);    // Event 对象
    THNPGraph_init(module);    // Graph 对象
    THNPMemPool_init(module);  // MemPool 对象

    // 5. 注册设备属性和内存
    RegisterNPUDeviceProperties(module);
    BindGetDeviceProperties(module);
    RegisterNPUDeviceMemories(module);

    // 6. 初始化通信和 profiler
    initCommMethods();
    torch_npu::profiler::initMstx(module);

    return module;
}

Python C 扩展机制

基本原理

Python 解释器是用 C 语言实现的（CPython），它提供了 C API 让开发者可以用 C/C++ 编写扩展模块。

Python 代码
     │
     ▼
┌─────────────────────────────────────┐
│  CPython 解释器                      │
│  - 查找模块                          │
│  - 调用 PyInit_xxx() 初始化          │
│  - 通过 PyMethodDef 查找方法         │
│  - 调用对应的 C 函数                 │
└─────────────────────────────────────┘
     │
     ▼
C/C++ 代码执行

核心组件

1. 模块定义 (PyModuleDef)

static struct PyModuleDef module_def = {
    PyModuleDef_HEAD_INIT,    // 必须的头部
    "模块名",                  // 模块名称
    nullptr,                   // 模块文档
    -1,                        // 模块状态大小
    methods                    // 方法数组
};

2. 方法定义 (PyMethodDef)

static PyMethodDef methods[] = {
    // Python 方法名, C 函数指针, 调用方式, 文档
    {"_initExtension", (PyCFunction)THPModule_initExtension, METH_NOARGS, nullptr},
    {"_npu_init", (PyCFunction)THNPModule_initExtension, METH_NOARGS, nullptr},
    {"_npu_setDevice", (PyCFunction)THNPModule_setDevice_wrap, METH_O, nullptr},
    {nullptr, nullptr, 0, nullptr}  // 结束标记
};

3. 模块初始化函数 (PyInit_xxx)

// 命名规则：PyInit_模块名
// 模块名 mymodule → PyInit_mymodule
// 模块名 torch._C → PyInit_torch__C（点号替换为下划线）
PyMODINIT_FUNC PyInit_mymodule(void)
{
    return PyModule_Create(&module_def);
}

4. C 函数实现

// C 函数签名必须符合 Python 要求
// 参数：self (模块或实例), args (位置参数)
static PyObject* THPModule_initExtension(PyObject* self, PyObject* noargs)
{
    // 执行 C++ 代码
    _initialize_python_bindings();

    // 返回 Python None
    Py_RETURN_NONE;
}

参数传递方式 (METH_xxx)

标志 说明 使用场景
METH_NOARGS 无参数 func()
METH_O 单个参数 func(arg)
METH_VARARGS 位置参数 func(*args)
METH_KEYWORDS 关键字参数 func(**kwargs)
METH_STATIC 静态方法 Class.func()
METH_CLASS 类方法 Class.func(cls)

调用流程

Python: module.func(args)
    │
    ▼
CPython 解释器
    │
    ├── 1. 在模块中查找 "func" 方法
    │
    ├── 2. 在 PyMethodDef 数组中找到匹配项
    │
    ├── 3. 根据 METH_xxx 解析参数
    │
    └── 4. 调用对应的 C 函数
    │
    ▼
C++ 函数执行
    │
    ▼
返回 PyObject* (Python 对象)

常用 CPython API

API 作用
PyModule_Create() 创建模块对象
PyModule_AddObject() 向模块添加对象
PyArg_ParseTuple() 解析位置参数
PyArg_ParseTupleAndKeywords() 解析位置和关键字参数
Py_BuildValue() 构建返回值
Py_RETURN_NONE 返回 None
PyLong_FromLong() 创建整数对象
PyUnicode_FromString() 创建字符串对象
PyErr_SetString() 设置异常
PyImport_ImportModule() 导入模块

完整示例

// ============== C++ 代码 ==============

// 1. 定义 C 函数
static PyObject* add_numbers(PyObject* self, PyObject* args)
{
    long a, b;

    // 解析参数：两个整数
    if (!PyArg_ParseTuple(args, "ll", &a, &b)) {
        return nullptr;  // 解析失败，返回 NULL 表示异常
    }

    // 计算结果
    long result = a + b;

    // 返回 Python 整数对象
    return PyLong_FromLong(result);
}

// 2. 方法定义
static PyMethodDef MyMethods[] = {
    {"add", add_numbers, METH_VARARGS, "Add two integers"},
    {nullptr, nullptr, 0, nullptr}
};

// 3. 模块定义
static struct PyModuleDef mymodule = {
    PyModuleDef_HEAD_INIT,
    "mymodule",
    "My example module",
    -1,
    MyMethods
};

// 4. 初始化函数
PyMODINIT_FUNC PyInit_mymodule(void)
{
    return PyModule_Create(&mymodule);
}

# ============== Python 代码 ==============
import mymodule

result = mymodule.add(3, 5)  # 调用 C 函数
print(result)  # 8

_initExtension 调用流程

完整调用链

torch_npu._C._initExtension()
    │
    ▼
Python 查找 torch_npu._C 模块中的 _initExtension 方法
    │
    ▼
调用 C++ 函数 THPModule_initExtension()
    │
    │  torch_npu/csrc/utils/TensorType.cpp:328-334
    │
    ▼
_initialize_python_bindings()
    │
    │  torch_npu/csrc/utils/TensorType.cpp:266-291
    │
    ▼
初始化 NPU 张量类型绑定
    │
    ├── torch.npu.FloatTensor
    ├── torch.npu.DoubleTensor
    ├── torch.npu.HalfTensor
    └── ...

代码详解

1. Python 调用入口

# torch_npu/__init__.py:238
torch_npu._C._initExtension()

2. C++ 方法注册

// torch_npu/csrc/utils/TensorType.cpp:337-340
static PyMethodDef TorchNpuExtensionMethods[] = {
    {"_initExtension", (PyCFunction)THPModule_initExtension, METH_NOARGS, nullptr},
    {nullptr, nullptr, 0, nullptr}
};

3. 方法添加到模块

// torch_npu/csrc/InitNpuBindings.cpp:176
AddPyMethodDefs(methods, torch_npu::utils::npu_extension_functions());

其中 npu_extension_functions() 返回 TorchNpuExtensionMethods 数组：

// torch_npu/csrc/utils/TensorType.cpp:342-345
PyMethodDef *npu_extension_functions()
{
    return TorchNpuExtensionMethods;
}

4. C++ 实现函数

// torch_npu/csrc/utils/TensorType.cpp:328-334
static PyObject *THPModule_initExtension(PyObject *_unused, PyObject *noargs)
{
    HANDLE_TH_ERRORS
    _initialize_python_bindings();  // 核心：初始化 Python 绑定
    Py_RETURN_NONE;
    END_HANDLE_TH_ERRORS
}

5. 初始化 Python 绑定

// torch_npu/csrc/utils/TensorType.cpp:266-291
void _initialize_python_bindings()
{
    // 1. 初始化 NPU ATen 类型
    initialize_npu_aten_types(tensor_types);

    // 2. 初始化 Python 元类
    py_initialize_metaclass(metaclass);

    // 3. 获取张量字典
    auto tensor_dict = get_tensor_dict();

    // 4. 初始化每个 Python 类型对象
    // torch.npu.FloatTensor, torch.npu.DoubleTensor 等
    for (auto &tensor_type : tensor_types) {
        py_initialize_tensor_type(tensor_type.py_type, tensor_type.name, tensor_dict.get());
    }

    // 5. 将类型对象添加到模块
    py_bind_tensor_types(tensor_types);
}

时序图

import torch_npu._C
    │
    ▼
PyInit__C() 执行
    │
    ├── AddPyMethodDefs(methods, npu_extension_functions())
    │       │
    │       └── 注册 _initExtension 方法
    │
    ▼
torch_npu._C 模块创建完成
    │
    ▼
执行 torch_npu/__init__.py
    │
    ▼
torch_npu._C._initExtension()  ← Python 调用
    │
    ▼
THPModule_initExtension()  ← C++ 函数
    │
    ▼
_initialize_python_bindings()
    │
    ├── 初始化 torch.npu.FloatTensor
    ├── 初始化 torch.npu.DoubleTensor
    ├── 初始化 torch.npu.HalfTensor
    └── ...
    │
    ▼
返回 None

与 _npu_init 的区别

函数 作用 是否初始化 ACL
_initExtension() 初始化 Python 张量类型绑定 ❌ 否
_npu_init() 初始化 NPU 设备（调用 aclInit） ✅ 是 
# torch_npu/__init__.py

# 第238行：初始化 Python 绑定（不初始化 ACL）
torch_npu._C._initExtension()

# _npu_init 在 _lazy_init() 中调用（首次使用 NPU 时）
# torch_npu/npu/__init__.py:271
torch_npu._C._npu_init()

aclInit 调用时机

完整调用链

torch.ones(5, device="npu")
    │
    ▼
torch_npu.npu._lazy_init()
    │
    │  torch_npu/npu/__init__.py:242-282
    │
    ▼
torch_npu._C._npu_init()
    │
    │  Python -> C++ 绑定
    │
    ▼
THNPModule_initExtension()
    │
    │  torch_npu/csrc/npu/Module.cpp:662-697
    │
    │  核心调用:
    │  c10_npu::NpuSysCtrl::GetInstance().Initialize()
    │
    ▼
NpuSysCtrl::Initialize()
    │
    │  torch_npu/csrc/core/npu/sys_ctrl/npu_sys_ctrl.cpp
    │
    ▼
★ aclInit(json_path)  ← ACL 初始化！

NpuSysCtrl::Initialize() 详细流程

// torch_npu/csrc/core/npu/sys_ctrl/npu_sys_ctrl.cpp:145-248
NpuSysCtrl::SysStatus NpuSysCtrl::Initialize(int device_id)
{
    // 1. 双重检查锁定，防止重复初始化
    if (init_flag_) {
        return INIT_SUCC;
    }
    std::lock_guard<std::mutex> lock(init_mutex_);
    if (init_flag_) {
        return INIT_SUCC;
    }

    // 2. 获取 ACL 配置文件路径
    std::string json_path = GetAclConfigJsonPath();
    const char *json_path_ptr = json_path == "" ? nullptr : json_path.c_str();

    // ★★★ 3. 调用 aclInit - 这是 ACL 初始化的关键接口！★★★
    auto init_ret = aclInit(json_path_ptr);  // 第157行

    // 处理重复初始化的情况
    if (init_ret == ACL_ERROR_REPEAT_INITIALIZE) {
        repeat_init_acl_flag_ = false;
        ASCEND_LOGI("acl has allready init by other component.");
    } else if (init_ret != ACL_ERROR_NONE) {
        NPU_CHECK_ERROR(init_ret, "aclInit");
    }

    // 4. 初始化 Dump 功能
    if (c10_npu::option::OptionsManager::CheckAclDumpDateEnable()) {
        NPU_CHECK_ERROR(aclmdlInitDump());
    }

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

    // 6. 设置设备
    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_));
        }
    }

    // 7. 获取 SoC 版本
    auto soc_name = c10_npu::acl::AclGetSocName();
    c10_npu::SetSocVersion(soc_name);

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

    // 9. 初始化算子编译模式
    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();
    }

    // 10. 设置确定性计算
    if (!c10_npu::is_lazy_set_device()) {
        SetDeterministicFromLevel();
        NPU_CHECK_ERROR(c10_npu::acl::AclrtSetOpExecuteTimeOut(kMaxOpExecuteTimeOut));
    }

    init_flag_ = true;
    return INIT_SUCC;
}

ACL 配置文件

// torch_npu/csrc/core/npu/sys_ctrl/npu_sys_ctrl.cpp:108-130
std::string GetAclConfigJsonPath()
{
    std::string npu_path = GetTorchNpuFile();
    std::string json_path = "";

    // 根据是否延迟设置设备选择不同的配置文件
    if (c10_npu::is_lazy_set_device()) {
        json_path = npu_path.append("torch_npu/acl_default.json");
    } else {
        json_path = npu_path.append("torch_npu/acl.json");
    }

    return json_path_str;
}

配置文件位置：

  • torch_npu/acl.json - 标准配置
  • torch_npu/acl_default.json - 延迟初始化配置

初始化时序图

import torch
    │
    ▼
PyTorch 扫描 torch.backends entry points
    │
    ▼
调用 torch_npu._autoload()
    │
    ▼
执行 torch_npu/__init__.py
    │
    ├── import torch_npu._C
    │       │
    │       ▼
    │   PyInit__C() 执行
    │       │
    │       ▼
    │   torch_npu._C 模块创建完成
    │
    ├── torch_npu._C._initExtension()
    │       │
    │       ▼
    │   初始化 Python 张量类型绑定
    │
    │   ← 此时 aclInit() 尚未调用！
    │
    ▼
torch_npu 模块加载完成
    │
    ▼
用户代码: x = torch.ones(5, device="npu")
    │
    ▼
torch_npu.npu._lazy_init()
    │
    ▼
torch_npu._C._npu_init()
    │
    ▼
THNPModule_initExtension()
    │
    ▼
NpuSysCtrl::Initialize()
    │
    ▼
★ aclInit(json_path)  ← ACL 真正初始化！
    │
    ▼
aclrtSetDevice() / 设置设备
    │
    ▼
NPU 设备可用

_lazy_init 触发时机

延迟初始化函数

# torch_npu/npu/__init__.py:242-282
def _lazy_init():
    global _initialized, _original_pid, _queued_calls
    if _initialized or hasattr(_tls, 'is_initializing'):
        return
    with _initialization_lock:
        if _initialized:
            return

        # 核心：调用 C++ 初始化函数
        torch_npu._C._npu_init()

        _original_pid = os.getpid()
        _tls.is_initializing = True
        try:
            _queue_call(_queued_calls)
        finally:
            delattr(_tls, 'is_initializing')
        _initialized = True

Python 层触发点

1. 设备相关操作

函数 文件位置 触发场景
synchronize() npu/utils.py:70 同步 NPU 设备
current_device() npu/utils.py:95 获取当前设备
device.__enter__() npu/utils.py:117 进入设备上下文

2. 内存相关操作

函数 文件位置 触发场景
mem_get_info() npu/__init__.py:495 获取内存信息
memory_stats() npu/memory.py:139 获取内存统计
empty_cache() npu/memory.py:869 清空缓存

3. 随机数相关操作

函数 文件位置 触发场景
get_rng_state_offset() npu/__init__.py:349 获取 RNG 状态偏移
manual_seed() npu/random.py:23 设置随机种子
seed() npu/random.py:164 生成随机种子

4. 设备属性查询

函数 文件位置 触发场景
get_device_properties() npu/__init__.py:485 获取设备属性
utilization() npu/__init__.py:566 获取设备利用率

5. 其他操作

函数 文件位置 触发场景
init() npu/__init__.py:239 显式初始化
ipc_collect() npu/utils.py:84 IPC 内存回收

C++ 层触发点

// torch_npu/csrc/utils/LazyInit.cpp:13-31
void npu_lazy_init()
{
    pybind11::gil_scoped_acquire g;
    if (!npu_run_yet) {
        // 导入 Python 模块
        auto module = THPObjectPtr(PyImport_ImportModule("torch_npu.npu"));
        if (!module) {
            throw python_error();
        }
        // 调用 Python 的 _lazy_init 方法
        auto res = THPObjectPtr(PyObject_CallMethod(module.get(), "_lazy_init", ""));
        if (!res) {
            throw python_error();
        }
        npu_run_yet = true;
    }
}

C++ 层调用位置：

// torch_npu/csrc/core/npu/NPUHooksInterface.cpp:28-33
void NPUHooksInterface::init() const
{
    torch_npu::utils::npu_lazy_init();
}

// torch_npu/csrc/npu/Module.cpp
PyObject* THNPModule_setDevice_wrap(...) {
    torch_npu::utils::npu_lazy_init();
    // ...
}

PyObject* THNPModule_getDevice_wrap(...) {
    torch_npu::utils::npu_lazy_init();
    // ...
}

典型触发场景

import torch
# 此时 _lazy_init() 尚未调用

# 场景 1: 创建 NPU 张量
x = torch.ones(5, device="npu")
#       ↓
#  Tensor.__new__ → NPUHooksInterface::init() → _lazy_init()

# 场景 2: 查询设备
torch.npu.current_device()
#       ↓
#  current_device() → _lazy_init()

# 场景 3: 同步设备
torch.npu.synchronize()
#       ↓
#  synchronize() → _lazy_init()

# 场景 4: 设置随机种子
torch.npu.manual_seed(42)
#       ↓
#  manual_seed() → _lazy_init()

# 场景 5: 显式初始化
torch.npu.init()
#       ↓
#  init() → _lazy_init()

常见问题解答

Q1: torch.tensor([1]) 会触发 NPU 初始化吗？

答案：不会。

torch.tensor([1]) 默认在 CPU 上创建张量，不会触发 NPU 初始化。

import torch

# 情况1: 不触发 NPU 初始化
x = torch.tensor([1])
print(x.device)  # cpu

# 情况2: 触发 NPU 初始化
y = torch.tensor([1], device="npu")
print(y.device)  # npu:0

# 情况3: 先创建 CPU 张量，再移动到 NPU（也会触发）
z = torch.tensor([1]).npu()
print(z.device)  # npu:0

Q2: 为什么采用延迟初始化？

  1. 快速启动：import torch 时不需要初始化 ACL，加快启动速度
  2. 资源节约：只在真正需要时才占用 NPU 资源
  3. 兼容性：允许在没有 NPU 的环境中导入 torch_npu
  4. 灵活性：用户可以在初始化前设置环境变量和配置

Q3: 如何判断 NPU 是否已初始化？

import torch

# 方法1: 检查初始化状态
print(torch.npu.is_initialized())

# 方法2: 检查是否可用（会触发初始化）
print(torch.npu.is_available())

Q4: torch._C 和 torch_npu._C 有什么区别？

模块 说明
torch._C PyTorch 核心 C++ 扩展，包含所有基础功能
torch_npu._C torch_npu 扩展 C++ 模块，包含 NPU 特有功能

两者是独立的模块，互不冲突。

Q5: _initExtension 和 _npu_init 有什么区别？

函数 作用 是否初始化 ACL
_initExtension() 初始化 Python 张量类型绑定 ❌ 否
_npu_init() 初始化 NPU 设备（调用 aclInit） ✅ 是

Q6: 初始化流程总结

阶段 函数 作用
1. 模块加载 torch_npu._autoload() 触发 torch_npu 模块加载
2. C++ 模块初始化 PyInit__C() 创建 torch_npu._C 模块
3. Python 绑定 _initExtension() 初始化张量类型绑定
4. 设备注册 rename_privateuse1_backend("npu") 注册 NPU 设备类型
5. 延迟初始化 torch_npu.npu._lazy_init() 首次使用 NPU 时触发
6. C++ 初始化 THNPModule_initExtension() 调用底层初始化
7. ACL 初始化 aclInit() 初始化 Ascend Computing Language
8. 设备设置 aclrtSetDevice() 设置当前 NPU 设备
9. 内存分配器 NPUCachingAllocator::init() 初始化内存管理

附录：关键文件路径

文件 路径 说明
入口点配置 setup.py:753-760 定义 torch.backends entry point
模块初始化 torch_npu/__init__.py Python 层初始化
设备注册 torch_npu/__init__.py:209-215 注册 NPU 设备
C++ 模块入口 torch_npu/csrc/InitNpuBindings.cpp PyInit__C 函数
方法注册 torch_npu/csrc/utils/TensorType.cpp _initExtension 方法
延迟初始化 torch_npu/npu/__init__.py:242-282 _lazy_init 函数
NPU 初始化 torch_npu/csrc/npu/Module.cpp:662-697 THNPModule_initExtension
ACL 初始化 torch_npu/csrc/core/npu/sys_ctrl/npu_sys_ctrl.cpp NpuSysCtrl::Initialize
C++ 桥接 torch_npu/csrc/utils/LazyInit.cpp npu_lazy_init 函数
Guard 实现 torch_npu/csrc/core/npu/impl/NPUGuardImpl.cpp 设备操作实现

参考资料

likedislike
Hhbhu_bin成员
3月12日 修改了issue 的描述
Hhbhu_bin成员
3月12日 修改了issue 的描述