CANN/GE数据依赖形状自定义算子样例

Data Dependent Shape 自定义算子(三类算子)样例

【免费下载链接】geGE（Graph Engine）是面向昇腾的图编译器和执行器，提供了计算图优化、多流并行、内存复用和模型下沉等技术手段，加速模型执行效率，减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力，并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge

样例概述

• 构图入口：GE
• 算子编成语言：Ascend C
• 编译方式：.asc与.cpp同 target 编译
• 模型下沉能力：不涉及
• 核心链路：Ascend C kernel 与 host 侧 custom op 同库编译 -> GE 交付件 -> 进程内构图 -> Session::RunGraph -> hybrid/RT2 动态执行
• 与其他 sample 的区别：本样例聚焦三类自定义算子在unknown graph动态执行场景下的 shape buffer 用法。

本样例展示一条最小可运行的三类自定义算子执行链路：在进程内构图后直接调用Session::RunGraph，由 CMake 将 Ascend C kernel 和 host 侧 custom op 编译到同一个libwhere_like_custom_op.so，再由WhereLikeCustom::Execute()申请最大输出和 shape buffer，最终在 device 侧写回实际输出 shape。

WhereLikeCustom输入为bool[N]，用于返回输入中值为true的元素下标，输出为int64[true_count, rank]。由于true_count依赖运行时输入数据，编译期只能确定输出上界，实际 shape 需要通过三类 custom op 的 shape buffer 协议在执行后回传。

适用场景

• 想看三类自定义算子在Session::RunGraph动态执行链路中的最小实现方式。
• 想参考用户在Execute内自行同步 stream、回读 shape buffer 并更新输出 shape 的方式。
• 想看.asc和.cpp直接一起编译，并在Execute中发起 kernel 调用的完整流程。
• 不适合用于了解ATC离线编译 -> om离线模型下沉链路。

前置依赖

CANN

• 已正确安装并配置 CANN 环境，例如执行过source ${ASCEND_HOME_PATH}/set_env.sh。
• 当前环境具备ACL、GE、Graph、Ascend C相关头文件与库。
• 参考 安装指导 完成 toolkit 和 ops 包安装。

框架与插件

• 本样例不依赖 PyTorch、TensorFlow 或 TorchAir。

环境变量

• ASCEND_HOME_PATH
• ASCEND_CUSTOM_OPP_PATH会在run.sh中自动追加为当前 sample 的output/

额外依赖

• cmake
• g++

快速运行

在examples/custom_op/data_dependent_shape_custom目录下执行：

推荐方式

source ${ASCEND_HOME_PATH}/set_env.sh bash run.sh

run.sh会自动完成 configure、build、install，并把output/追加到ASCEND_CUSTOM_OPP_PATH。若运行成功，终端会打印：

output shape: [4, 1] output values: 0 2 4 7

分步方式

source ${ASCEND_HOME_PATH}/set_env.sh cmake -S . -B build -DCMAKE_BUILD_TYPE=Release cmake --build build -j"$(nproc)" cmake --install build export ASCEND_CUSTOM_OPP_PATH="$(pwd)/output:$ASCEND_CUSTOM_OPP_PATH" cd build ./data_dependent_shape_custom_session_run cd ..

其中export ASCEND_CUSTOM_OPP_PATH="$(pwd)/output:$ASCEND_CUSTOM_OPP_PATH"用于将自定义算子包根目录加入环境变量，随后 GE 会按output/op_graph/lib/<os>/<arch>/libwhere_like_custom_op.so规则加载交付件。

目录结构与关键文件

data_dependent_shape_custom ├── CMakeLists.txt ├── README.md ├── run.sh ├── ge │ ├── custom_op.cpp // EagerExecuteOp + ShapeInferOp 主流程实现 │ ├── where_like_custom.h // WhereLikeCustom proto 定义 │ └── where_like_custom_kernel.asc // Ascend C kernel 源码 └── session_run └── main.cc // 进程内构图并直接调用 Session::RunGraph

重点文件：

• ge/custom_op.cpp自定义算子的核心主流程，实现Execute、InferShape和InferDataType；其中Execute负责申请最大输出、shape buffer、调用.asc中导出的 launch wrapper，并在 kernel 完成后自行回读 shape buffer 更新输出 shape，InferShape/InferDataType负责编译期输出 shape / dtype 推导。
• ge/where_like_custom_kernel.ascdevice kernel 和 host 侧 launch wrapper 实现，负责写输出数据和 shape buffer。
• session_run/main.cc构建最小图并直接通过Session::AddGraph + Session::RunGraph执行。
• run.sh串起 configure、build、install 和运行过程。

核心链路

1. session_run/main.cc构建包含Data -> WhereLikeCustom的最小图，并把输入/输出描述设置为动态 shape，使整图走unknown graph执行链路。
2. ge/custom_op.cpp中的InferShape/InferDataType在构图阶段给出输出 shape / dtype；本样例不依赖框架 lowering 插入 shape 回写节点。
3. CMakeLists.txt将ge/custom_op.cpp和ge/where_like_custom_kernel.asc一起编译到libwhere_like_custom_op.so。
4. ge/custom_op.cpp在Execute回调中按最大 shape 调用ctx->MallocOutputTensor(...)申请输出，再通过第一次ctx->MallocWorkSpace(...)申请 shape buffer。
5. ge/where_like_custom_kernel.asc在 device 侧写入输出索引和 shape buffer，其中 shape buffer 用于回传实际输出 shape。
6. ge/custom_op.cpp的Execute在 launch 后同步当前 stream，将 shape buffer 拷回 host，解析真实 shape 并更新输出 tensor 的 logical shape 和有效 size。
7. 执行完成后，session_run/main.cc读取输出 tensor，打印实际 shape 和输出值。

构建产物

• output/op_graph/lib/linux/x86_64/libwhere_like_custom_op.soLinux x86_64 环境下 GE 使用的自定义算子交付件；aarch64 环境对应output/op_graph/lib/linux/aarch64/libwhere_like_custom_op.so。
• output/op_graph/include/where_like_custom.h构图侧可直接使用的算子 proto 头文件。
• build/data_dependent_shape_custom_session_run直接Session::RunGraph的执行程序。

结果校验

成功时可观察到：

• output/op_graph/lib/<os>/<arch>/libwhere_like_custom_op.so已生成。
• output/op_graph/include/where_like_custom.h已生成。
• 终端输出包含output shape: [4, 1]。
• 终端输出包含output values: 0 2 4 7。

若失败，优先检查：

• ASCEND_HOME_PATH是否已设置并已正确sourceCANN 环境。
• ASCEND_CUSTOM_OPP_PATH是否已包含当前 sample 的output/。
• output/op_graph/lib/<os>/<arch>/libwhere_like_custom_op.so和output/op_graph/include/where_like_custom.h是否已生成。
• 当前环境是否具备可用 NPU 和可用 Ascend C 编译环境。

注意事项 / 限制

• WhereLikeCustom当前样例输入为一维bool[8]，因此实际输出为匹配位置索引，shape 为[true_count, 1]。
• .asc编译参数当前固定为--npu-arch=dav-2201，如目标芯片不同需调整CMakeLists.txt。

附录

算子规格

shape buffer 约定

本样例由用户自定义 shape buffer 协议，并在Execute中自行解析回写实际输出 shape。当前 kernel 会写入：

• shape[0] = 2U
• shape[1] = true_count
• shape[2] = rank

对当前一维输入场景，rank = 1，因此输入[true, false, true, false, true, false, false, true]的实际输出 shape 为[4, 1]，输出数据为0 2 4 7。

【免费下载链接】geGE（Graph Engine）是面向昇腾的图编译器和执行器，提供了计算图优化、多流并行、内存复用和模型下沉等技术手段，加速模型执行效率，减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力，并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge