CANN 系列深度篇：基于 ge 图引擎构建高效 AI 执行图

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 构建一个"加法+指数"计算图

✅ 场景描述

我们要构建如下计算流程：

output = exp(input_A + input_B)

其中 input_A 和 input_B 是长度为 1024 的 float 向量。

我们将不使用任何预训练模型 ，而是通过 ge 的 API 手动构建图，然后编译并执行。

---

✅ 代码实现（C++）

#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;
}

---

🔧 关键说明

1. OperatorFactory::CreateOperator

   动态创建内置算子，名称需与 CANN 支持的 Op 列表一致（如 "Add", "Exp"）。

2. 图不可变性

   一旦调用 CreateSession，图结构即被冻结，不能再修改。

3. 内存生命周期

   绑定的设备内存必须在 Session 存活期间保持有效。

---

四、ge 的优化能力示例

假设我们将上述图扩展为：

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 插件将计算图转为 ge IR；
• 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（共享内存管理）？欢迎指定方向！