AsNumpy 快速上手指南仅改 import 即可将 NumPy 代码迁移到昇腾 NPU【免费下载链接】asnumpy哈尔滨工业大学计算学部苏统华、王甜甜老师团队联合华为CANN团队开发的华为昇腾NPU原生Numpy仓库项目地址: https://gitcode.com/cann/asnumpy本篇技术指南围绕 CANN / asnumpy 仓库的 docs/quick_start.md 展开系统讲解如何以最小改动只改 import、增加数据搬运把既有 NumPy 科学计算代码迁移到昇腾 910B NPU 上运行。读完本文你将掌握 AsNumpy 的环境要求、from_numpy/to_numpy双向数据搬运、与 NumPy 完全对齐的算子 API、设备切换方法以及从源码层理解其自动初始化、自动释放的资源管理机制。核心理念只改 import其余不变AsNumpy 的设计哲学可以浓缩为一句话change the import, keep the rest只改导入其余保持不变。它对外暴露的顶层 API 与 NumPy 同名同签名例如ap.add、ap.multiply、ap.sum、ap.mean同时通过ap.ndarray这个镜像numpy.ndarray的数据结构承载 NPU 上的张量数据。从 src/asnumpy/init.py 的_LAZY_MAPPING可以看到asnumpy顶层导出的算子覆盖了创建数组zeros、ones、full、linspace等、数学运算add、multiply、divide、exp2、sqrt等、逻辑运算all、any、equal、logical_and等、线性代数dot、matmul、vdot、einsum等、排序统计sort、mean与随机数random等 NumPy 常用子集且全部采用懒加载lazy import方式按需初始化避免拖慢包导入速度。这套同名同签名设计意味着你现有的 NumPy 代码基本不需要改写函数调用只需将数据从 CPU 内存搬到 NPU 显存计算完成后再搬回即可享受昇腾 NPU 的硬件加速。环境准备Prerequisites在开始迁移之前需要确认以下三项前提条件与仓库 README.md 中 Installation 一节列出的构建要求一致条件要求AsNumpy 安装已完成安装安装步骤见 README.md 安装章节昇腾硬件Ascend 910B NPUCANN 版本CANN 8.2.RC1.alpha003 及以上Python 版本Python 3.9此外若需要从源码构建README.md 明确要求 GCC 11.2、CMake 3.26并在构建前设置 CANN 环境变量export ASCEND_TOOLKIT_HOME/usr/local/Ascend/ascend-toolkit/latest推荐使用uv或pip完成安装详见 README.md安装完成后可以用如下最小片段验证环境是否就绪import asnumpy as ap arr ap.ones((1000, 1000), dtypeap.float32) print(arr.shape) # (1000, 1000)值得注意ap.float32等 dtype 别名直接解析为 NumPy 自身的类型对象ap.float32 is np.float32AsNumpy 并没有自建一套 dtype 系统。这一设计在 src/asnumpy/init.py 的_NUMPY_DTYPE_NAMES中有明确注释说明好处是 dtype 语义与 NumPy 完全一致、零心智负担。核心概念NPUArray 与 CPU↔NPU 双向传输AsNumpy 迁移模型的核心是两种数组与三个动作NumPy 数组numpy.ndarray驻留在 CPU 内存负责数据生成、加载与结果回读AsNumpy 数组asnumpy.ndarray即 NPUArray驻留在 NPU 设备显存负责高性能计算ap.ndarray.from_numpy(x)将 NumPy 数组拷贝到 NPU返回 NPUArraynpu_arr.to_numpy()将 NPU 计算结果拷贝回 CPU 内存返回 NumPy 数组其余所有ap.*函数在 NPUArray 上执行API 与 NumPy 一致。从 src/asnumpy/utils.py 可以看到from_numpy与to_numpy是类方法/实例方法的对称实现底层调用 pybind11 绑定的_core.ndarray。该公共ndarray类还实现了 NumPy 的__array_ufunc__src/asnumpy/utils.py与__array_function__src/asnumpy/utils.py协议当你在np.sin(npu_array)或np.add(a, b)中混入了 AsNumpy 数组时NumPy 会自动把调用转发给对应的asnumpy实现这种协议级兼容是迁移成本极低的重要原因。对应回归测试见 tests/asnumpy_tests/creation_tests/test_move_wrapper.py它验证了zeros、sin、add、sum的结果类型与数值均与 NumPy 基准一致。需要留意内存语义asnumpy.ndarray的构造函数对公共ndarray采用深拷贝语义src/asnumpy/utils.py因此ndarray(existing_array)不会消耗源数组而_core内部算子的私有结果则通过移动语义move semantics一次性包装进公共ndarray避免重复的设备内存拷贝。这些细节由框架内部保证使用者无需干预。迁移对照同一段代码的 CPU 与 NPU 写法原文档给出了一个直观的并排对照这里完整保留并逐步拆解。以 20000×20000 的矩阵乘法后求和为例NumPyCPU写法import numpy as np rows, cols 20000, 20000 m1 np.random.normal(0, 1, (rows, cols)) m2 np.random.normal(0, 1, (rows, cols)) # Compute on CPU product np.multiply(m1, m2) result np.sum(product) print(result)AsNumpyNPU写法import numpy as np import asnumpy as ap rows, cols 20000, 20000 m1 np.random.normal(0, 1, (rows, cols)) m2 np.random.normal(0, 1, (rows, cols)) # Transfer to NPU m1_npu ap.ndarray.from_numpy(m1) m2_npu ap.ndarray.from_numpy(m2) # Compute on NPU product ap.multiply(m1_npu, m2_npu) result ap.sum(product) print(result.to_numpy())两段代码的唯一差异点有三个其余完全一致多了一行import asnumpy as ap数据生成后用ap.ndarray.from_numpy搬到 NPU计算函数名从np.*换成ap.*或混用见下文协议兼容说明最终结果用result.to_numpy()回读。其中np.multiply/np.sum替换为ap.multiply/ap.sum后product与result都驻留在 NPU 上整条计算链不经过 CPU 回传这正是性能优势的来源——数据搬运只发生在这条链的起点和终点。端到端示例完整跑通一个加法→乘法→归约流水原文档提供了一个可直接复制运行的端到端示例这里完整保留并补充逐行注释与验证说明import numpy as np import asnumpy as ap # AsNumpy auto-initializes the NPU device on import # and releases it on exit (no manual init/finalize needed) # 1. Create data on CPU (NumPy) np_a np.array([1.0, 2.0, 3.0, 4.0], dtypenp.float32) np_b np.array([10.0, 20.0, 30.0, 40.0], dtypenp.float32) # 2. Transfer to NPU npu_a ap.ndarray.from_numpy(np_a) npu_b ap.ndarray.from_numpy(np_b) # 3. Run operations on NPU npu_sum ap.add(npu_a, npu_b) npu_prod ap.multiply(npu_a, npu_b) npu_total ap.sum(npu_prod) # 4. Transfer results back to CPU print(Sum: , npu_sum.to_numpy()) # [11. 22. 33. 44.] print(Prod: , npu_prod.to_numpy()) # [ 10. 40. 90. 160.] print(Total: , npu_total.to_numpy()) # 300.0 # 5. Verify against NumPy assert np.allclose(npu_sum.to_numpy(), np.add(np_a, np_b)) assert np.allclose(npu_prod.to_numpy(), np.multiply(np_a, np_b)) print(Verification passed.)这个例子展示了 AsNumpy 迁移的完整流水CPU 造数 → 搬运上卡 → NPU 计算链 → 回读结果 → 与 NumPy 交叉验证。最后一步验证非常重要它用np.allclose确认 NPU 计算结果与 CPU 基准在浮点容差内一致这正是保证迁移正确性的防线。仓库 examples/ 下的全部脚本都遵循这一模式用asnumpy与numpy的同名函数分别计算再用numpy.allclose对比结果见 examples/README.md。检查与切换 NPU 设备在多卡场景下需要查询可用设备并指定计算设备原文档给出如下用法import asnumpy as ap # Query available NPU devices print(ap.get_device_count()) # e.g. 8 # Switch to a specific NPU (default is 0) ap.set_device(1)需要说明的是当前仓库源码中检索到的设备管理 API 为set_device、reset_device、reset_device_force、init、finalize五个get_device_count出现在 docs/quick_start.md 文档示例中属于文档示例接口使用时请以实际安装版本导出的 API 为准。其中set_device(device_id)直接绑定 CANN 运行时接口aclrtSetDevice实现在 bindings/python/bind_cann.cppcann.def(set_device, aclrtSetDevice, pybind11::arg(device_id)); cann.def(reset_device, aclrtResetDevice, pybind11::arg(device_id)); cann.def(reset_device_force, aclrtResetDeviceForce, pybind11::arg(device_id)); cann.def(init, asnumpy::cann::init); cann.def(finalize, asnumpy::cann::finalize);对应的 Python 层封装在 src/asnumpy/cann.py五个函数均带有logger.catch装饰出现异常时会通过 loguru 记录日志后抛出便于定位设备侧问题。自动资源管理无需手写 init / finalize与 CUDA 生态常见的init()/finalize()手动管理不同AsNumpy 在模块导入时自动完成设备初始化进程退出时自动完成资源释放。这一机制在 src/asnumpy/init.py 中有明确实现atexit.register def reset(): reset_device(0) finalize() init() set_device(0)展开来看它做了三件事导入即初始化模块加载末尾依次调用init()与set_device(0)对应 CANN 层的aclInit(nullptr)与aclrtSetDevice(0)见 csrc/cann/driver.cpp因此用户 import asnumpy 后即可直接使用无需任何前置初始化代码退出即释放通过atexit.register注册的reset()会在解释器退出时自动执行reset_device(0)与finalize()对应aclrtResetDevice与aclFinalize保证 NPU 设备资源与 CANN 上下文被有序回收避免设备句柄泄漏RAII 内存管理在更细粒度上NPUArray析构时自动释放设备显存详见 README.md 特性说明AclWorkspace等临时工作区同样遵循 RAII析构时调用aclrtFree释放设备内存见 csrc/utils/acl_resource.cpp。因此原文档注释中强调的AsNumpy auto-initializes the NPU device on import and releases it on exit正是这一机制的准确描述普通用户完全不需要手动管理设备生命周期。更多可运行示例原文档给出了 examples/ 目录下的可运行脚本清单这里完整继承并补充脚本能力说明脚本操作examples/01_add.py逐元素加法ap.add对比np.add结果与运行时间examples/02_exp2.py以 2 为底的指数运算ap.exp2输入限制在 [-10, 10] 以避免数值溢出examples/03_multiply.py逐元素乘法ap.multiply带基准测试examples/04_all.py逻辑与归约ap.all判断数组所有元素是否均为 Trueexamples/05_divide.py逐元素除法ap.divide这些脚本的共性模式以 examples/01_add.py 为例是先用 examples/utils.py 的create_arrays生成 NumPy 随机数据并转换为 AsNumpy 数组然后分别跑ap.add与np.add并统计耗时最后计算相对误差relative_diff 1e-4视为验证通过。基准统计采用中段最快速度策略排序后剔除最慢 10% 再取最小值见 examples/utils.py以降低系统调度抖动的影响并做了显存优化及时del中间结果、gc.collect()避免大矩阵下显存溢出。运行方式# 在 examples 目录下先安装 asnumpy python 01_add.py仓库 README.md 提供的基准数据显示3000×3000 的 float32 数据上ap.mean()相对np.mean()加速约 35.70×且数据规模越大 NPU 并行优势越明显完整复现方法见 docs/benchmarks.md。不过需要注意该数据是特定软硬件环境下的示例结果实际加速比会随算子类型归约 vs 逐元素、数据规模与机器配置而变化。常见坑位与注意事项结合源码梳理几条实操中值得注意的点dtype 建议使用 float32基准与示例均以np.float32为主NPU 算子对 float32 的支持最完整实际迁移时优先保证数据是 float32或显式astype转换结果回读是必须的NPU 上的计算结果无法直接交给纯 NumPy 函数消费务必先to_numpy()回 CPU 内存再交给下游 CPU 代码反过来混用场景下 NumPy 的 ufunc/函数协议会自动把调用转发给 AsNumpy见前文__array_ufunc__/__array_function__不支持的关键字会安全回退__array_ufunc__实现中where、casting、subok、order、signature、extobj等 NumPy ufunc 关键字不被支持时会返回NotImplemented让 NumPy 用自己的实现兜底而不是静默忽略src/asnumpy/utils.py因此混用出错时优先检查是否传入了这些关键字设备资源自动管理普通脚本无需手动 init/finalize但多进程或多线程场景下需注意进程退出时的自动清理逻辑与显式设备切换的配合。深入阅读完整安装与构建说明README.md架构与内部设计docs/architecture.md性能基准与复现步骤docs/benchmarks.md常见问题解答docs/faq.md新增算子与二次开发指引docs/developer_guide.md可运行示例脚本examples/测试用例验证协议兼容与结果一致性tests/asnumpy_tests/creation_tests/test_move_wrapper.py【免费下载链接】asnumpy哈尔滨工业大学计算学部苏统华、王甜甜老师团队联合华为CANN团队开发的华为昇腾NPU原生Numpy仓库项目地址: https://gitcode.com/cann/asnumpy创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考