Graph Engine 是什么，为什么需要它

前言

训练或推理一个神经网络，底层发生了什么？

框架层（PyTorch、MindSpore 等）定义好模型结构后，需要把计算图送到硬件上执行。早期的做法是逐个算子直接下发------卷积层调一次 kernel，BN 层再调一次 kernel，每层之间还要把中间结果写回显存。

这种做法的问题很明显：kernel 启动本身有开销，显存读写慢，而且相邻算子之间往往可以合并成更的计算。

Graph Engine（简称 GE）就是来解决这件事的。它位于 CANN 架构的第四层（计算执行层），职责是：接收上层（编译层）优化好的计算图，把它调度到底层硬件上高效执行。

GE 在 CANN 里的位置

CANN 的架构从上到下分为五层：

复制代码

应用层（PyTorch/MindSpore）
    ↓
第1层：AscendCL（统一编程接口）
    ↓
第2层：AOL算子库 + AOE调优引擎
    ↓
第3层：Graph Compiler（图编译）
    ↓
第4层：Graph Engine ← 这里
    ↓
第5层：Driver/Firmware（硬件驱动）
    ↓
NPU 硬件

GE 上面接的是编译层（已经把计算图优化好了），下面接的是驱动层（直接跟 NPU 对话）。GE 自己做的事情可以概括为三件：

图切分：把一张大图切成多个子图，分配到不同的计算单元
流调度：用多个 Stream 并发执行相互独立的子图
内存管理：复用显存，减少分配/释放的次数

快速上手：跑通第一个例子

环境准备

昇腾官方提供了预装好 CANN 的 Docker 镜像，直接拉取即可：

bash 复制代码

# 拉取镜像（单卡）
docker pull swr.cn-south-1.myhuaweicloud.com/ascend/ascend-pytoch:8.0.R1-cp38

# 启动容器
docker run -it --privileged \
    --device /dev/davinci0 \
    --network host \
    ascend-pytoch:8.0.R1-cp38 \
    /bin/bash

进入容器后，验证 NPU 是否可用：

bash 复制代码

npu-smi list

正常输出应该能看到 NPU 设备信息。如果报错，检查 --device 参数是否跟宿主机的 NPU 设备对应。

第一个 PyTorch + NPU 程序

python 复制代码

import torch
import torch.npu  # 导入即注册 NPU 后端

# 1. 检查 NPU 是否可用
print("NPU available:", torch.npu.is_available())
print("NPU count:", torch.npu.device_count())
print("NPU name:", torch.npu.get_device_name(0))

# 2. 创建简单计算
device = torch.device("npu:0")
x = torch.randn(1024, 1024, device=device)
w = torch.randn(1024, 1024, device=device)

# 3. 矩阵乘法（GE 在背后接管执行）
y = torch.matmul(x, w)
loss = y.sum()

# 4. 反向传播
loss.backward()

print("Output shape:", y.shape)
print("Loss:", loss.item())

运行这个程序，npu-smi 里对应进程的 NPU 利用率会跳上来，说明 GE 已经开始工作了。

GE 的核心概念

理解 GE，需要搞清楚三个概念：图（Graph） 、流（Stream） 、算子（Operator）。

图

计算任务在 GE 里表示为一张有向无环图（DAG）。节点是算子，边是数据依赖关系。

GE 拿到这张图之后，先做拓扑排序 ，确定算子的执行顺序；然后做内存规划 ，确定每个 tensor 存在显存的哪个位置；最后做流分配，把可以并行的算子放到不同的 Stream 上。

流

Stream 是 NPU 里的执行队列。同一个 Stream 里的任务按顺序执行，不同 Stream 里的任务可以并发。

默认情况下，所有算子都在 Stream 0 上执行。要手动控制多流并发，可以这样做：

python 复制代码

import torch.npu

# 创建两个流
stream0 = torch.npu.Stream(0)
stream1 = torch.npu.Stream(1)

# 在 stream0 上执行前一半层
with torch.npu.stream(stream0):
    x0 = model.layers[:6](x[:16])

# 在 stream1 上执行后一半层
with torch.npu.stream(stream1):
    x1 = model.layers[6:](x[16:])

# 等待两个流都完成
stream0.synchronize()
stream1.synchronize()

# 合并结果
x = torch.cat([x0, x1], dim=0)

算子

图里的节点。GE 内置了大量优化过的算子实现，覆盖 Conv、MatMul、Softmax、LayerNorm 等常见操作。

如果需要自定义算子，可以通过 Ascend C 编程接口实现，编译后注册到 GE 里即可调用。

进阶：查看 GE 的调度过程

想知道 GE 到底是怎么切图、怎么调度的？可以打开 GE 的日志：

bash 复制代码

# 设置 GE 日志级别为 DEBUG
export GE_LOG_LEVEL=3
export GENGINE_GRAPH_SAVE_PATH=./ge_graphs

# 运行程序，GE 会把计算图保存到指定目录
python train.py

GENGINE_GRAPH_SAVE_PATH 目录里会生成多张计算图的可视化文件（可以用 Netron 打开查看），包括：

原始计算图（框架层下发过来的）
优化后的计算图（算子融合、死代码消除之后）
最终执行的子图（分配好 Stream 和显存之后）

对比这几张图，能直观地看到 GE 做了哪些优化。

常见问题

Q：GE 和 ATC 的区别是什么？

ATC 是编译器，负责把计算图编译成 NPU 可执行的指令；GE 是执行引擎，负责在运行时调度这些指令。两者分工不同，但紧密协作。

Q：用 PyTorch 需要手动调用 GE 的 API 吗？

不需要。torch.npu 已经封装好了，模型放到 NPU 上之后，GE 自动接管。只有在做非常底层的优化（比如手动控制算子融合策略）时，才需要直接跟 GE 的接口打交道。

Q：多卡训练时 GE 怎么工作？

多卡场景下，GE 跟 HCCL（集合通信库）配合，在每个 NPU 上各跑一张子图，卡间通信通过 HCCL 完成。GE 负责切图，HCCL 负责通信，两者协同实现数据并行或模型并行。

总结

Graph Engine 是 CANN 里负责图执行的核心组件，位于架构的第四层。它的价值在于：把编译层优化好的计算图，高效地调度到 NPU 上执行------通过图切分、多流并发、显存复用等手段，让硬件利用率最大化。

对应用层开发者来说，GE 的存在是透明的（torch.npu 一行代码搞定）；但对性能优化来说，理解 GE 的调度逻辑，是定位瓶颈、做针对性优化的前提。