cann组织链接:https://atomgit.com/cann
ops-nn仓库链接:https://atomgit.com/cann/ops-nn
CANN 系列深度篇:基于 ge 图引擎构建高效 AI 执行图
在现代 AI 框架中,"计算图"(Computation Graph)是组织神经网络操作的核心抽象。无论是静态图(如 TensorFlow 1.x)还是动态图(如 PyTorch),最终都需要一个底层引擎来解析、优化并执行这些操作。CANN 提供的 ge(Graph Engine) 正是这样一个面向 NPU 的高性能图执行引擎。
本文将系统介绍 ge 的设计目标、关键模块,并通过一个从零构建计算图的完整 C++ 示例,展示如何利用 ge 实现端到端的 AI 推理流程------无需依赖高层框架,直接操作底层图结构。
📌 项目地址:https://gitcode.com/cann/ge
一、什么是 ge(Graph Engine)?
ge 是 CANN 中负责图构建、优化与调度执行的核心模块。它的主要功能包括:
- 图定义:支持以 IR(Intermediate Representation)形式描述计算流程;
- 图优化:执行常量折叠、算子融合、内存复用等优化策略;
- 设备映射:将图中的节点分配到合适的 NPU 设备;
- 运行时调度 :协调
hcll、ops-math等底层库完成实际执行; - Profiling 支持:提供性能分析接口,便于调优。
ge 的输入通常是 OM(Offline Model)文件 ,这是 CANN 工具链(如 atc 编译器)将 ONNX/TensorFlow 模型转换后的离线格式。但 ge 也支持程序化构建图,这为自定义算子集成、研究新型网络结构提供了极大灵活性。
二、ge 的核心概念
| 概念 | 说明 |
|---|---|
| Operator(Op) | 图的基本节点,如 Add, MatMul, Relu |
| TensorDesc | 描述张量的形状、数据类型、内存布局(NCHW / NHWC) |
| NodeBuilder | 用于构造图节点的辅助类 |
| Graph | 由多个 Op 节点组成的有向无环图(DAG) |
| Session | 图的执行上下文,管理资源与运行状态 |
| Input/Output Binding | 将 Host 内存绑定到图的输入/输出占位符 |
三、实战示例:用 ge 构建一个"加法+指数"计算图
✅ 场景描述
我们要构建如下计算流程:
text
output = exp(input_A + input_B)
其中 input_A 和 input_B 是长度为 1024 的 float 向量。
我们将不使用任何预训练模型 ,而是通过 ge 的 API 手动构建图,然后编译并执行。
✅ 代码实现(C++)
cpp
#include <iostream>
#include <vector>
#include "ge_api.h" // Graph Engine 主头文件
#include "hcll.h" // 用于内存管理
using namespace ge;
int main() {
const int size = 1024;
std::vector<float> input_A(size, 1.0f);
std::vector<float> input_B(size, 2.0f);
std::vector<float> output(size);
// Step 1: 创建计算图
Graph graph("exp_add_graph");
// 定义输入张量描述
TensorDesc desc(Dims({size}), FORMAT_ND, DT_FLOAT);
// 创建输入占位符
auto inputA = graph.AddInput("input_A", desc);
auto inputB = graph.AddInput("input_B", desc);
// 构建 Add 节点
auto add_op = ge::OperatorFactory::CreateOperator("Add", "Add");
add_op.SetInput("x", inputA).SetInput("y", inputB);
auto add_out = add_op.GetOutput("z");
// 构建 Exp 节点
auto exp_op = ge::OperatorFactory::CreateOperator("Exp", "Exp");
exp_op.SetInput("x", add_out);
auto exp_out = exp_op.GetOutput("y");
// 设置图输出
graph.SetOutput(exp_out, "output");
// Step 2: 编译图(生成可执行 Session)
SessionOptions opts;
opts.device_id = 0; // 使用 device 0
auto session = ge::CreateSession(graph, opts);
if (!session) {
std::cerr << "Failed to create session!" << std::endl;
return -1;
}
// Step 3: 绑定输入/输出内存
void* dev_inputA, *dev_inputB, *dev_output;
hcllMalloc(&dev_inputA, size * sizeof(float));
hcllMalloc(&dev_inputB, size * sizeof(float));
hcllMalloc(&dev_output, size * sizeof(float));
hcllMemcpy(dev_inputA, input_A.data(), size * sizeof(float), HCLL_MEMCPY_HOST_TO_DEVICE);
hcllMemcpy(dev_inputB, input_B.data(), size * sizeof(float), HCLL_MEMCPY_HOST_TO_DEVICE);
// 绑定
session->BindInput("input_A", dev_inputA);
session->BindInput("input_B", dev_inputB);
session->BindOutput("output", dev_output);
// Step 4: 执行图
if (!session->Run()) {
std::cerr << "Graph execution failed!" << std::endl;
return -1;
}
// Step 5: 拷回结果
hcllMemcpy(output.data(), dev_output, size * sizeof(float), HCLL_MEMCPY_DEVICE_TO_HOST);
// 验证结果(前5个)
std::cout << "Result (first 5): ";
for (int i = 0; i < 5; ++i) {
// input_A[i] + input_B[i] = 3.0 → exp(3.0) ≈ 20.0855
std::cout << output[i] << " ";
}
std::cout << std::endl;
// 清理
hcllFree(dev_inputA);
hcllFree(dev_inputB);
hcllFree(dev_output);
ge::DestroySession(session);
return 0;
}
🔧 关键说明
-
OperatorFactory::CreateOperator动态创建内置算子,名称需与 CANN 支持的 Op 列表一致(如
"Add","Exp")。 -
图不可变性
一旦调用
CreateSession,图结构即被冻结,不能再修改。 -
内存生命周期
绑定的设备内存必须在
Session存活期间保持有效。
四、ge 的优化能力示例
假设我们将上述图扩展为:
text
output = exp(A + B) + log(C)
ge 在编译阶段可能自动执行以下优化:
- 算子融合 :若硬件支持,将
Add + Exp融合为单个 kernel; - 内存复用 :中间结果
A+B不写回全局内存,直接传给Exp; - 常量传播 :若
B是常量,则提前计算部分结果。
这些优化对用户透明,但显著提升性能。
五、典型应用场景
- 自定义模型部署 :当标准框架不支持某类算子时,用
ge手动构建图; - 低延迟推理服务 :绕过 Python 层,直接 C++ 调用
ge,减少开销; - 科研实验:快速验证新型网络结构或算子组合;
- 边缘设备轻量化 :仅链接
ge+hcll+ops-math,构建极简推理引擎。
六、与高层框架的关系
虽然 ge 可独立使用,但它也是 CANN 对接 PyTorch/TensorFlow 的桥梁:
- PyTorch 通过
torch_npu插件将计算图转为geIR; atc编译器将 ONNX 模型转换为 OM 文件,由ge加载执行。
因此,理解 ge 有助于深入掌握 CANN 的执行机制。
七、结语
ge 图引擎是 CANN 的"大脑"------它不仅执行计算,更通过智能优化释放硬件潜能。对于追求极致性能或需要深度定制的开发者而言,掌握 ge 的使用方法,意味着你拥有了直接操控 NPU 计算流的能力。
从 ops-math 到 hcll,再到 ge,我们看到 CANN 构建了一个层次清晰、协同高效的软件栈:
算子(Ops)→ 通信(HCLL)→ 调度(GE)→ 应用(AI Model)。
下一步,你或许可以尝试:
- 在
ge中注册自定义算子(结合tbe); - 使用
ge的 profiling API 分析瓶颈; - 构建包含控制流(如 If/While)的复杂图。
🔗 探索更多:https://gitcode.com/cann
📂 建议查看
ge/samples/目录下的官方示例。
是否希望继续解读 tbe(自定义算子开发框架) 或 shmem(共享内存管理)?欢迎指定方向!