昇腾 CANN 的五层架构，到底分了哪五层

前言

CANN 是昇腾的软件栈，从底层的硬件抽象到上层的应用框架，一共分五层。理解这五层的关系，就知道模型是怎么跑在 NPU 上的。

一、五层架构总览

复制代码

┌─────────────────────────────────────┐
│      应用层      │  ← PyTorch、TensorFlow、ONNX
├─────────────────────────────────────┤
│      框架适配层     │  ← torch_npu、TensorFlow插件
├─────────────────────────────────────┤
│      计算图引擎层          │  ← 图优化、算子调度
├─────────────────────────────────────┤
│      算子层                │  ← Conv、MatMul、FlashAttention
├─────────────────────────────────────┤
│      运行时层           │  ← 内存管理、设备调度
└─────────────────────────────────────┘

每一层只和相邻层交互，职责清晰。

二、应用层：模型怎么写

最上层是用户代码，用 PyTorch 或 ONNX 定义模型。

PyTorch 模型示例

python 复制代码

import torch
import torch.nn as nn

class ResNetBlock(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.conv1 = nn.Conv2d(channels, channels, 3, padding=1)
        self.bn1 = nn.BatchNorm2d(channels)
        self.relu = nn.ReLU()
    
    def forward(self, x):
        identity = x
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        return x + identity

# 创建模型并移到 NPU
model = ResNetBlock(64).npu()

这一层不关心底层硬件，只管模型结构。

三、框架适配层：打通 PyTorch 和 NPU

torch_npu 是 PyTorch 的昇腾后端，把 PyTorch 的调用转发给 CANN。

注册 NPU 设备

python 复制代码

import torch
import torch_npu

# 检查 NPU 是否可用
print(f"NPU available: {torch.npu.is_available()}")

# 查看设备数量
print(f"NPU count: {torch.npu.device_count()}")

# 设置当前设备
torch.npu.set_device(0)

# 把模型移到 NPU
model = model.npu()

自定义算子注册

如果模型里有自定义算子，需要在框架适配层注册：

python 复制代码

import torch_npu

# 注册自定义算子
torch_npu.npu.register_op(
    "custom_op",
    "CustomOp",
    inputs=["Tensor"],
    outputs=["Tensor"],
    attrs={"param": "float"}
)

# 使用自定义算子
output = torch.ops.npu.custom_op(input_tensor, param=0.5)

四、计算图引擎层：优化和调度

GE 负责计算图的优化和算子调度。

导出计算图

python 复制代码

import torch
import torch_npu

model = ResNetBlock(64).npu().eval（)
traced = torch.jit.trace(model, torch.randn(1, 64, 56, 56).npu())

# 保存计算图
torch.jit.save(traced, "model.pt")

# 导出 ONNX
torch.onnx.export(model, torch.randn(1, 64, 56, 56).npu(), "model.onnx")

ATC 编译

bash 复制代码

# 用 ATC 编译成 .om
atc --model=model.onnx \
    --framework=5 \
    --output=model \
    --enable_fusion=true

GE 会在编译时做算子融合、内存规划、算子选型。

五、算子层：具体的计算实现

算子层提供各种算子的实现，包括标准算子和优化算子。

标准算子

python 复制代码

import torch
import torch_npu

# 卷积
conv = torch.nn.Conv2d(3, 64, 3, padding=1).npu()
output = conv(input.npu())

# 矩阵乘法
a = torch.randn(1024, 1024).npu()
b = torch.randn(1024, 1024).npu()
c = torch.matmul(a, b)

优化算子（ops-adv）

python 复制代码

from ops_adv import flash_attention

# FlashAttention
Q = torch.randn(1, 32, 1024, 64).npu()
K = torch.randn(1, 32, 1024, 64).npu()
V = torch.randn(1, 32, 1024, 64).npu()

output = flash_attention(Q, K, V)

自定义算子开发（Ascend C）

cpp 复制代码

#include "kernel_operator.h"

class MatMulKernel {
public:
    __aicore__ inline void process(GM_ADDR x, GM_ADDR y, GM_ADDR z) {
        // 自定义矩阵乘法实现
        // 使用 Cube Unit 或 Vector Unit
    }
};

// 注册算子
extern "C" __global__ void matmul_custom(GM_ADDR x, GM_ADDR y, GM_ADDR z) {
    MatMulKernel op;
    op.process(x, y, z);
}

六、运行时层：内存和设备管理

最底层是 Ascend CL，管理 NPU 设备、内存、Stream。

初始化运行时

cpp 复制代码

#include <acl/acl.h>

// 初始化
aclError ret = aclInit(nullptr);
ret = aclrtSetDevice(0);

// 创建 Context 和 Stream
aclrtContext context;
aclrtStream stream;
aclrtCreateContext(&context, 0);
aclrtCreateStream(&stream);

内存管理

cpp 复制代码

// 分配 Device 内存
void* devicePtr = nullptr;
size_t size = 1024 * 1024;  // 1MB
aclrtMalloc(&devicePtr, size, ACL_MEM_MALLOC_HUGE_FIRST);

// Host 到 Device 拷贝
void* hostPtr = malloc(size);
aclrtMemcpy(devicePtr, size, hostPtr, size, ACL_MEMCPY_HOST_TO_DEVICE);

// Device 到 Host 拷贝
aclrtMemcpy(hostPtr, size, devicePtr, size, ACL_MEMCPY_DEVICE_TO_HOST);

// 释放内存
aclrtFree(devicePtr);
free(hostPtr);

执行推理

cpp 复制代码

// 加载模型
uint32_t modelId;
aclmdlLoadFromFile("model.om", &modelId);

// 创建输入输出 Dataset
aclmdlDataset* input = aclmdlCreateDataset();
aclmdlDataset* output = aclmdlCreateDataset();

// 执行推理
aclmdlExecuteAsync(modelId, input, output, stream);
aclrtSynchronizeStream(stream);

// 卸载模型
aclmdlUnload(modelId);

七、五层之间的调用链

从用户代码到硬件执行的完整流程：

python 复制代码

# 1. 应用层：用户代码
model = ResNetBlock(64).npu()
output = model(input)

# ↓ torch_npu 转发

# 2. 框架适配层：把 PyTorch 调用转成 GE 调用
# torch_npu 内部：
#   ge::Operator op = ConvertToGeOperator("Conv2d", ...)
#   ge::RunOp(op)

# ↓ GE 处理

# 3. 计算图引擎层：优化和调度
# GE 内部：
#   1. 算子融合：Conv + BN + ReLU → ConvBNReLU
#   2. 算子选型：选择 Cube 或 Vector 实现
#   3. 内存规划：分配 UB 和 HBM 空间

# ↓ 调用算子库

# 4. 算子层：执行具体计算
# 算子库内部：
#   ConvBNReLU kernel 调用 Cube Unit
#   中间数据在 UB 里流转

# ↓ 调用运行时

# 5. 运行时层：管理硬件资源
# ACL 内部：
#   分配 HBM 内存
#   启动 AI Core
#   同步 Stream

八、各层的性能影响

不同层对性能的影响不同：

层级	性能影响	优化手段
应用层	模型结构决定上限	换模型架构、量化
框架适配层	5-10%	用 torch.compile、避免频繁 CPU-NPU 同步
计算图引擎层	20-40%	算子融合、内存规划、AOE 调优
算子层	30-50%	选优化算子、自定义 Ascend C
运行时层	10-20%	Stream 并行、内存复用

优化示例

python 复制代码

# 框架适配层优化：避免频繁同步
for i, data in enumerate(dataloader):
    output = model(data.npu())
    
    # 不要每个 step 都同步
    # loss_val = loss.item()  # 这会触发同步
    
    # 改成每 100 步同步一次
    if i % 100 == 0:
        print(f"Step {i}, loss: {loss.item()}")

参考资源

CANN 架构文档：https://www.hiascend.com/document/detail/zh/CANN/
torch_npu API：https://gitee.com/ascend/pytorch
Ascend CL 编程指南：https://www.hiascend.com/document/detail/zh/CANN/
算子开发指南：https://www.hiascend.com/document/detail/zh/CANN/

总结

CANN 五层架构从上到下：应用层写模型、框架适配层打通 PyTorch、计算图引擎层做优化、算子层执行计算、运行时层管理硬件。每层职责清晰，通过标准接口交互。性能优化主要集中在计算图引擎层（算子融合、AOE 调优）和算子层（选优化算子、自定义实现）。理解五层关系，才能定位性能瓶颈在哪个层级，针对性优化。