Data Dependent Shape 自定义算子(三类算子)样例【免费下载链接】geGEGraph 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_PATHASCEND_CUSTOM_OPP_PATH会在run.sh中自动追加为当前 sample 的output/额外依赖cmakeg快速运行在examples/custom_op/data_dependent_shape_custom目录下执行推荐方式source ${ASCEND_HOME_PATH}/set_env.sh bash run.shrun.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_TYPERelease 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 更新输出 shapeInferShape/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 和运行过程。核心链路session_run/main.cc构建包含Data - WhereLikeCustom的最小图并把输入/输出描述设置为动态 shape使整图走unknown graph执行链路。ge/custom_op.cpp中的InferShape/InferDataType在构图阶段给出输出 shape / dtype本样例不依赖框架 lowering 插入 shape 回写节点。CMakeLists.txt将ge/custom_op.cpp和ge/where_like_custom_kernel.asc一起编译到libwhere_like_custom_op.so。ge/custom_op.cpp在Execute回调中按最大 shape 调用ctx-MallocOutputTensor(...)申请输出再通过第一次ctx-MallocWorkSpace(...)申请 shape buffer。ge/where_like_custom_kernel.asc在 device 侧写入输出索引和 shape buffer其中 shape buffer 用于回传实际输出 shape。ge/custom_op.cpp的Execute在 launch 后同步当前 stream将 shape buffer 拷回 host解析真实 shape 并更新输出 tensor 的 logical shape 和有效 size。执行完成后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-archdav-2201如目标芯片不同需调整CMakeLists.txt。附录算子规格项目内容算子类型WhereLikeCustom输入x输出y输入 shapeN输出 shape 上界[N, rank(x)]输出实际 shape[true_count, rank(x)]输入数据类型bool输出数据类型int64格式NDkernel 名称where_like_customshape buffer 约定本样例由用户自定义 shape buffer 协议并在Execute中自行解析回写实际输出 shape。当前 kernel 会写入shape[0] 2Ushape[1] true_countshape[2] rank对当前一维输入场景rank 1因此输入[true, false, true, false, true, false, false, true]的实际输出 shape 为[4, 1]输出数据为0 2 4 7。【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考