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 设备;
  • 运行时调度 :协调 hcllops-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_Ainput_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;
}

🔧 关键说明

  1. OperatorFactory::CreateOperator

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

  2. 图不可变性

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

  3. 内存生命周期

    绑定的设备内存必须在 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 插件将计算图转为 ge IR;
  • atc 编译器将 ONNX 模型转换为 OM 文件,由 ge 加载执行。

因此,理解 ge 有助于深入掌握 CANN 的执行机制。


七、结语

ge 图引擎是 CANN 的"大脑"------它不仅执行计算,更通过智能优化释放硬件潜能。对于追求极致性能或需要深度定制的开发者而言,掌握 ge 的使用方法,意味着你拥有了直接操控 NPU 计算流的能力

ops-mathhcll,再到 ge,我们看到 CANN 构建了一个层次清晰、协同高效的软件栈:
算子(Ops)→ 通信(HCLL)→ 调度(GE)→ 应用(AI Model)

下一步,你或许可以尝试:

  • ge 中注册自定义算子(结合 tbe);
  • 使用 ge 的 profiling API 分析瓶颈;
  • 构建包含控制流(如 If/While)的复杂图。

🔗 探索更多:https://gitcode.com/cann

📂 建议查看 ge/samples/ 目录下的官方示例。

是否希望继续解读 tbe(自定义算子开发框架)shmem(共享内存管理)?欢迎指定方向!

相关推荐
qq_12498707537 小时前
基于Hadoop的信贷风险评估的数据可视化分析与预测系统的设计与实现(源码+论文+部署+安装)
大数据·人工智能·hadoop·分布式·信息可视化·毕业设计·计算机毕业设计
Coder_Boy_7 小时前
TensorFlow小白科普
人工智能·深度学习·tensorflow·neo4j
L、2187 小时前
CANN 中的图优化技术详解:如何让 AI 模型跑得更快、更省
人工智能
大模型玩家七七7 小时前
梯度累积真的省显存吗?它换走的是什么成本
java·javascript·数据库·人工智能·深度学习
新缸中之脑7 小时前
像画家一样编程
人工智能
tq10867 小时前
心主神明:传统智慧如何启示AI的可靠之道
人工智能
珠海西格电力科技7 小时前
微电网能量平衡理论的实现条件在不同场景下有哪些差异?
运维·服务器·网络·人工智能·云计算·智慧城市
新缸中之脑7 小时前
“AI 裁员“神话
人工智能
零售ERP菜鸟7 小时前
范式革命:从“信息化”到“数字化”的本质跃迁
大数据·人工智能·职场和发展·创业创新·学习方法·业界资讯