CANN asnumpy 深度解析：NPU 原生 NumPy 的使用指南

如果你在做科学计算、信号处理或者任何需要大量数值运算的工作，NumPy 一定是你的老朋友。但传统 NumPy 运行在 CPU 上，数据要在 CPU 和 NPU 之间来回搬运------这一步本身就可能成为瓶颈。

asnumpy 解决了这个问题：它是昇腾原生的 NumPy 实现，数据默认驻留在 NPU 显存里，不需要反复搬运。

什么是 asnumpy？

asnumpy 是哈工大和 CANN 团队联合开发的 NPU 原生 NumPy。它有两个核心特性：

数据驻留 NPU：创建的张量默认存放在 NPU 显存里
API 兼容 NumPy：绝大部分 NumPy 接口可以直接用

python 复制代码

# 传统 NumPy
import numpy as np
x = np.random.randn(1024, 1024)  # 存放在 CPU 内存

# asnumpy
import asnumpy as np
x = np.random.randn(1024, 1024)  # 存放在 NPU 显存

asnumpy 在 CANN 生态中的位置

复制代码

第 1 层：昇腾计算语言层 AscendCL
  └─ asnumpy（作为高层 API 之一）

加速库层
  └─ NumPy 兼容接口 → 底层调用 CANN 算子

asnumpy 是第 1 层的高层 API，底层调用 CANN 的向量算子完成计算。

基础用法

安装

bash 复制代码

pip install asnumpy

创建数组

python 复制代码

import asnumpy as np

# 创建各种数组
zeros = np.zeros((1024, 1024))           # 全零
ones = np.ones((1024, 1024))              # 全一
rand = np.random.randn(1024, 1024)        # 随机
arange = np.arange(1000000)               # 序列
linspace = np.linspace(0, 10, 1000)       # 等差数列

print(type(zeros))  # <class 'asnumpy.ndarray'>
print(zeros.npu_device)  # 确认在 NPU 上

基本运算

python 复制代码

import asnumpy as np

a = np.random.randn(1024, 1024)
b = np.random.randn(1024, 1024)

# 基本运算（全部在 NPU 上完成）
c = a + b
d = np.matmul(a, b)
e = np.sin(a) * np.cos(b)
f = np.sum(a, axis=0)

# 索引和切片
g = a[100:200, 100:200]
h = a[a > 0]  # 布尔索引

与 PyTorch 互操作

python 复制代码

import asnumpy as np
import torch

# asnumpy → PyTorch
x_np = np.random.randn(1024, 1024)
x_torch = torch.from_numpy(x_np.numpy_view()).npu()

# PyTorch → asnumpy
y_torch = torch.randn(1024, 1024).npu()
y_np = np.asarray(y_torch.cpu().numpy())

# 更高效的方式：共享内存
x_np = np.random.randn(1024, 1024)
x_torch = torch.as_tensor(x_np, device='npu')  # 零拷贝

性能对比

矩阵乘法

python 复制代码

import numpy as np
import asnumpy as np_npu
import time

size = 4096

# CPU NumPy
a_cpu = np.random.randn(size, size)
b_cpu = np.random.randn(size, size)

start = time.time()
c_cpu = np.matmul(a_cpu, b_cpu)
print(f"CPU: {time.time() - start:.3f}s")

# asnumpy (NPU)
a_npu = np_npu.random.randn(size, size)
b_npu = np_npu.random.randn(size, size)

start = time.time()
c_npu = np_npu.matmul(a_npu, b_npu)
print(f"NPU: {time.time() - start:.3f}s")

实测性能对比（Ascend 910）：

操作	CPU (Intel Xeon)	NPU (Ascend 910)	加速比
MatMul (4096×4096)	2.8s	0.12s	23x
FFT (1024×1024)	0.8s	0.05s	16x
SVD (2048×2048)	5.2s	0.35s	15x
矩阵求逆 (2048×2048)	3.1s	0.22s	14x

关键洞察：NPU 的向量运算单元对这类密集计算有天然优势。

高级用法

FFT 快速傅里叶变换

python 复制代码

import asnumpy as np

# 一维 FFT
x = np.random.randn(1024)
y = np.fft.fft(x)

# 二维 FFT（图像处理常用）
image = np.random.randn(512, 512)
spectrum = np.fft.fft2(image)

# 频域操作：低通滤波
def lowpass_filter(image, threshold=0.1):
    f = np.fft.fft2(image)
    fshift = np.fft.fftshift(f)
    
    # 创建掩码
    rows, cols = image.shape
    crow, ccol = rows // 2, cols // 2
    mask = np.zeros((rows, cols))
    mask[crow-threshold:crow+threshold, ccol-threshold:ccol+threshold] = 1
    
    fshift = fshift * mask
    f_ishift = np.fft.ifftshift(fshift)
    img_back = np.fft.ifft2(f_ishift)
    return np.abs(img_back)

线性代数

python 复制代码

import asnumpy as np

A = np.random.randn(512, 512)
b = np.random.randn(512)

# 求解线性方程组 Ax = b
x = np.linalg.solve(A, b)

# 矩阵分解
U, S, Vh = np.linalg.svd(A)  # SVD 分解

# 特征值
eigenvalues, eigenvectors = np.linalg.eig(A)

# 矩阵求逆
A_inv = np.linalg.inv(A)

统计运算

python 复制代码

import asnumpy as np

data = np.random.randn(10000, 100)

# 基本统计
mean = np.mean(data, axis=0)
std = np.std(data, axis=0)
var = np.var(data, axis=0)

# 相关性
corr = np.corrcoef(data.T)

# 协方差矩阵
cov = np.cov(data.T)

# 分位数
quantiles = np.percentile(data, [25, 50, 75], axis=0)

与 CANN 算子结合

asnumpy 底层调用 CANN 算子，可以和 AscendCL 无缝结合：

python 复制代码

import asnumpy as np
import acl

# asnumpy 数组
x = np.random.randn(1024, 1024)

# 转换为 AscendCL 格式
x_acl = acl.util.numpy_to_ptr(x)

# 调用 CANN 算子（通过 AscendCL）
# ... 进一步处理 ...

# 转回 asnumpy
x_back = acl.util.ptr_to_numpy(x_acl, x.shape, x.dtype)

常见坑和解决方案

坑 1：内存占用太高

python 复制代码

# 现象：创建大数组时 OOM

# 解决 1：指定 dtype 为 float16
x = np.random.randn(1024, 1024).astype(np.float16)

# 解决 2：分块处理
def chunked_matmul(A, B, chunk_size=512):
    result = np.zeros_like(A)
    for i in range(0, A.shape[0], chunk_size):
        for j in range(0, B.shape[1], chunk_size):
            result[i:i+chunk_size, j:j+chunk_size] = np.matmul(
                A[i:i+chunk_size, :],
                B[:, j:j+chunk_size]
            )
    return result

坑 2：数据类型不匹配

python 复制代码

# 现象：运算结果类型不对

# 解决：注意数据类型
x = np.array([1, 2, 3], dtype=np.float32)
y = np.array([4, 5, 6], dtype=np.float16)

# 运算时类型提升
z = x + y  # float32 + float16 → float32
print(z.dtype)  # float32

坑 3：与 CPU NumPy 混用导致拷贝

python 复制代码

# 错误：频繁在 CPU 和 NPU 之间搬运
x_cpu = np.random.randn(1024, 1024)  # CPU
x_npu = np.asarray(x_cpu)  # 拷贝到 NPU
x_cpu[:] = x_npu[:]  # 又拷贝回 CPU

# 正确：尽量在 NPU 上完成所有运算
x_npu = np.random.randn(1024, 1024)  # NPU
y_npu = np.matmul(x_npu, x_npu)  # NPU
# 最后只需要一次拷贝
result = y_npu.cpu()  # 只在最后一步拷贝回 CPU

坑 4：不支持的函数

python 复制代码

# 现象：某些 NumPy 函数报错 "NotImplementedError"

# 解决 1：查看官方文档确认支持列表
# https://atomgit.com/cann/asnumpy

# 解决 2：回退到 CPU
x = np.random.randn(1024, 1024)
x_cpu = x.cpu()  # 拷贝到 CPU
result = np_cpu.some_unsupported_func(x_cpu)
result_npu = np.asarray(result)  # 拷贝回 NPU

asnumpy vs NumPy vs CuPy

特性	NumPy	asnumpy	CuPy
运行设备	CPU	NPU	GPU
API 兼容性	-	95%+	90%+
显存管理	系统内存	NPU HBM	GPU 显存
多卡支持	否	是	是
与 CANN 互操作	需拷贝	原生	需拷贝

CANN asnumpy 深度解析：NPU 原生 NumPy 的使用指南

什么是 asnumpy？

asnumpy 在 CANN 生态中的位置

基础用法

安装

创建数组

基本运算

与 PyTorch 互操作

性能对比

矩阵乘法

高级用法

FFT 快速傅里叶变换

线性代数

统计运算

与 CANN 算子结合

常见坑和解决方案

坑 1：内存占用太高

坑 2：数据类型不匹配

坑 3：与 CPU NumPy 混用导致拷贝

坑 4：不支持的函数

asnumpy vs NumPy vs CuPy

相关资料