asnumpy：让你的 NumPy 代码零改动跑在昇腾 NPU 上

有个同事之前跟我说，他写了一套数据预处理的 pipeline，全是 NumPy 写的，后来要迁移到昇腾 NPU 上跑，“感觉天都要塌了”——几千行 NumPy 代码，难道要全部重写？

后来他发现了 asnumpy，这个东西基本上解决了这个问题。不需要改你的 NumPy 代码，直接换一下 import，数据搬上 NPU 跑，原来的语法一个不动。

---

先搞清楚 asnumpy 是什么

asnumpy 不是 NumPy 的替代品，它是昇腾 NPU 对 NumPy API 的一层兼容封装。你写的是 NumPy 代码，但底层跑在了昇腾 NPU 的向量计算单元上。

这里有个容易混淆的地方：asnumpy 不等于 NumPy 本身。NumPy 是 CPU 上的 Python 数值计算库，asnumpy 是昇腾 NPU 的同名实现，API 签名几乎一致，但底层硬件完全不同。如果你看到有人说"用 asnumpy 替换 NumPy"，这个说法是不准确的——正确理解是"用 asnumpy 的接口来写代码，享受 NPU 的加速"。

asnumpy 属于 ops-tensor 仓库的一部分，和 ops-math、ops-nn 这些算子库并列。同一个父项目下，不同的算子能力，asnumpy 主要负责向量级别的数值计算支持。

---

为什么要用它

标准 NumPy 在 CPU 上跑，数据要先从 NPU 的显存里搬到 CPU 内存，算完再搬回去。这个来回搬运的过程，在数据量大的时候开销非常可观。

举个例子，你有一段图像预处理的代码：

import numpy as np

# 原始 NumPy 版本（CPU）
def preprocess(image):
    image = image.astype(np.float32) / 255.0
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    image = (image - mean) / std
    return image

如果你想在昇腾 NPU 上跑这个预处理，通常的做法是把数据从 NPU 搬到 CPU，用 NumPy 算完再搬回去——等于白加速了预处理部分。

asnumpy 的思路是：让这部分计算直接在 NPU 上完成，不用搬运。

import ascendnumpy as np  # 直接替换 import

# 同一个函数，一个字不用改，数据全程在 NPU 上
def preprocess(image):
    image = image.astype(np.float32) / 255.0
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    image = (image - mean) / std
    return image

这样预处理和模型推理都在 NPU 上，数据零拷贝。这个"零改动"的体验是 asnumpy 最核心的价值。

---

常用操作速查

asnumpy 的 API 覆盖了 NumPy 的常用子集，不是 100% 全覆盖，但主流操作基本都有。我列几个最常用的：

import ascendnumpy as np

# 张量创建
a = np.zeros((1024, 1024), dtype=np.float32)
b = np.ones((512, 512))
c = np.random.randn(256, 256).astype(np.float16)

# 数学运算
d = np.matmul(a, b)         # 矩阵乘法，昇腾 NPU 硬件加速
e = np.sum(c, axis=0)       # 按列求和
f = np.mean(c, axis=1)      # 按行求均值
g = np.clip(a, 0, 1)         # 截断

# 数据类型转换，asnumpy 特有
h = a.asnpu()   # 转为 NPU 显存格式，提交给模型推理用
i = h.ascpu()   # 从 NPU 拿回 CPU，结果导出用

这里要特别说一说 asnpu() 和 ascpu() 这一对方法。这是 asnumpy 和标准 NumPy 最大的差异——你写的代码可以全程跑在 NPU 上，但总有需要把结果拿回 CPU 的时候，比如保存文件、打印日志、写磁盘。这两个方法就是连接 NPU 和 CPU 数据的桥。

# 典型场景：预处理在 NPU，算完拿回 CPU 存图
result = preprocess(image_npu)
result_cpu = result.ascpu()  # 这步才是真正有开销的，但只跑一次

# 相反方向：CPU 数据批量喂给 NPU
images = [load_image(f) for f in image_paths]  # CPU 上读取
images_npu = [img.asnpu() for img in images]  # 批量上传 NPU

---

asnumpy 和 catlass 的区别

有人会问：asnumpy 也能算矩阵乘法，catlass 也是做矩阵运算的，它们是什么关系？

这个问题问得好。简单说：asnumpy 是给数据处理用的，catlass 是给算子开发用的。

asnumpy 的矩阵乘法是封装好的高层接口，你调一下 np.matmul，它内部帮你把数据切块、调度硬件资源，你不需要关心底层怎么实现的。catlass 是昇腾算子模板库，提供了 GEMM（通用矩阵乘法）的底层开发框架，你需要自己写 tiling 逻辑、SRAM 使用策略，然后编译成可调用的算子。

一个面向应用开发者，一个面向算子开发者。日常写预处理脚本、做数据增强，asnumpy 够用了；如果你要开发新的融合算子、优化特定算子的性能上限，那得看 catlass。

昇腾异构计算架构里，这两层的定位是不同的：

AscendCL（应用层接口）
  └─ asnumpy（数据处理，API 友好）
catlass（算子开发层，底层灵活）
  └─ opbase（基础组件，上面两个都依赖它）

---

几个容易踩的坑

坑一：不是所有 NumPy API 都支持。 asnumpy 是 NumPy 的子集实现，不支持 FFT（用 ops-fft）、稀疏矩阵、某些特殊的线性代数操作。如果你的代码里用到了，先跑一遍看看报不报错。昇腾 CANN 的生态覆盖是逐步完善的，缺的功能可以提 Issue 到社区。

坑二：ascpu() 是隐性的性能杀手。 这个方法会触发一次完整的数据从 NPU 到 CPU 的搬运，有些场景下这一步会把前面的加速全部抵消。如果你的 pipeline 里频繁调用 ascpu()，比如每处理一张图就拿回 CPU 存一次，实际上并没有省到搬运时间。正确的做法是批量处理完再统一导出，或者干脆不导出直接喂给下游模型。

# 低效写法：每张图都搬运一次
for img in images:
    result = preprocess(img)
    save(result.ascpu())   # 每次都搬运，慢了

# 高效写法：批量处理，最后统一导出
batch = np.stack([img.asnpu() for img in images])
results = preprocess(batch)       # 全程 NPU
final = results.ascpu()           # 只搬运一次

坑三：dtype 精度问题。 asnumpy 的计算精度和 NumPy 有细微差异，主要体现在 float16 和 bfloat16 上。昇腾 NPU 的向量单元对 float16 有原生支持，asnumpy 默认会使用硬件加速；如果你明确需要 float32 的精度，需要显式指定 dtype，但这样可能会触发额外的精度转换开销。这个坑在图像处理里不明显，在科学计算场景里需要留意。

---

结尾

asnumpy 这个东西上手门槛很低，你甚至不需要懂昇腾 NPU 的编程模型，会 NumPy 就能用。但它解决的问题很实在——数据预处理不再需要绕道 CPU，NPU 加速从头到尾覆盖到。

如果你在昇腾 NPU 上跑 PyTorch 推理，可以把数据预处理部分用 asnumpy 重写，preprocess + 推理统一在 NPU 上完成，吞吐会比"CPU 预处理 + NPU 推理"的串行方式高出不少。

ops-tensor 仓库里除了 asnumpy，还有 tensorapi 和 Blaze 等其他张量工具，可以一起看看。源码在 https://atomgit.com/cann/ops-tensor。