Ascend C 自定义算子算子运行时编译并下沉 om离线模型样例【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge样例概述构图入口GE构图 ATC离线编译算子编成语言Ascend C编译方式RTC算子运行时编译模型下沉能力支持下沉om离线模型核心链路Ascend C kernel 源码 RTC算子运行时编译 - GE 交付件 - AIR - ATC离线编译 - om离线模型 - ACL 执行与其他 sample 的区别本样例在Compile回调中完成 Ascend C kernel 的算子运行时编译并把编译产物序列化进 om离线模型后续可在模型执行阶段反序列化使用。多 shape 说明本样例主要提供按 shape key 管理和下沉多份 kernel binary 的多 shape / 多 kernel 处理框架。编译期输入说明样例在Compile阶段按Tensor布局读取输入因此可以同时拿到Shape、DataType、Format等元信息当前 key 仍然只按 shape 生成。本样例展示一条最小可运行的 GE 自定义算子下沉链路先构图生成AIR再通过ATC离线编译生成om离线模型最后由ACL程序加载并执行。样例中的 AddCustom 同时继承CompilableOp、PortableOp和ShapeInferOp因此既能在Compile回调中通过 RTC 完成 Ascend C kernel 的算子运行时编译也能把编译结果序列化到最终模型中并在同一个实现类中完成 Shape 和 DataType 推导。适用场景想看GE构图 ATC离线编译 om离线模型链路中的 Ascend C 自定义算子样例。想参考自定义算子 RTC算子运行时编译、模型下沉和最终模型执行的完整流程。想看AIR - ATC离线编译 - om离线模型 - ACL这条最小可运行链路。不适合用于了解PyTorch/TorchAir或TensorFlow/Triton场景。前置依赖CANN已正确安装并配置 CANN 环境例如执行过source ${ASCEND_HOME_PATH}/set_env.sh。atc命令可用。当前环境具备ACL、GE、Graph相关头文件与库。参考 安装指导 完成 toolkit 和 ops 包安装。框架与插件本样例不依赖 PyTorch 或 TensorFlow。环境变量ASCEND_HOME_PATHASCEND_CUSTOM_OPP_PATH会在run.sh中自动追加output/而自定义算子交付件会放到output/op_graph/lib/linux/x86_64/或output/op_graph/lib/linux/aarch64/额外依赖cmakeg快速运行在examples/custom_op/compilable_add_custom目录下执行推荐方式bash run.shrun.sh在生成output/op_graph/lib/os/arch/libcust_opapi.so后会自动将output/追加到ASCEND_CUSTOM_OPP_PATH。若运行成功会生成output/single_add.air、output/single_add.om并在终端打印Model executed successfully!与First element of output: 2.000000。分步方式如果想手动观察各阶段产物可执行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 output ../build/compilable_add_graph_build cd .. atc \ --model$(pwd)/output/single_add.air \ --framework1 \ --output$(pwd)/output/single_add \ --soc_versionAscend910B1 cd build ./compilable_add_model_exec ../output/single_add.om cd ..其中export ASCEND_CUSTOM_OPP_PATH$(pwd)/output:$ASCEND_CUSTOM_OPP_PATH用于将自定义算子包根目录加入环境变量随后 GE/ATC 会按output/op_graph/lib/os/arch/libcust_opapi.so规则加载交付件。目录结构与关键文件compilable_add_custom ├── CMakeLists.txt ├── README.md ├── run.sh ├── ge │ ├── add_custom_kernel.asc // 算子运行时编译时读取的 Ascend C kernel 源码 │ ├── add_custom.h // AddCustom 的唯一 proto 定义供交付件和构图侧共用 │ ├── custom_op.cpp // CompilableOp/PortableOp/ShapeInferOp 主流程实现 │ └── utils │ ├── compile_utils.cpp // 编译期输入 Tensor 读取、平台信息、kernel 路径和 shape key 等辅助逻辑 │ └── kernel_binary_map_utils.cpp // 多 bin 序列化/反序列化逻辑 ├── graph_build │ └── main.cc // 构图并导出 single_add.air └── model_exec └── main.cc // 加载 om离线模型并执行重点文件ge/custom_op.cpp自定义算子的核心主流程负责串起算子运行时编译、序列化、反序列化、Shape/DataType 推导和执行编译期按输入 Tensor 读取Shape/DataType/Format元信息。ge/add_custom.hAddCustom的唯一 proto 定义构图侧使用生成到output/op_graph/include/下的同名头文件。ge/add_custom_kernel.asc算子运行时编译阶段读取的 Ascend C kernel 源码。ge/utils/compile_utils.cpp封装编译期输入 Tensor 读取、平台信息读取、kernel 源码路径定位以及基于 shape 的 binary key 生成逻辑用于演示多 shape 场景下的 binary 管理方式。ge/utils/kernel_binary_map_utils.cpp负责多份 kernel 二进制的序列化与反序列化重点体现多 key / 多 binary 的落盘与恢复。graph_build/main.cc构建最小 Add 图并导出single_add.air。model_exec/main.cc加载 om离线模型准备输入输出并执行模型。run.sh串起 configure、build、install、ATC离线编译和模型执行。核心链路ge/custom_op.cpp在Compile回调中先读取输入 Tensor 的Shape/DataType/Format等元信息再读取ge/add_custom_kernel.asc通过aclrtc完成 Add kernel 的算子运行时编译。graph_build/main.cc构造最小 Add 图并导出single_add.air。atc将single_add.air离线编译为single_add.om。model_exec/main.cc通过 ACL 加载single_add.om并执行输出第一项结果用于校验。说明样例中编译期输入已经按 Tensor 方式组织因此可以直接拿到 shape、data type、format当前仍然只用输入 shape 生成 key并据此缓存、序列化和反序列化 kernel binary整体上提供的是多 shape / 多 kernel 的处理框架。构建产物output/op_graph/lib/linux/x86_64/libcust_opapi.soLinux x86_64 环境下 GE/ATC离线编译使用的自定义算子交付件aarch64 环境对应output/op_graph/lib/linux/aarch64/libcust_opapi.so。output/op_graph/include/add_custom.h构图侧可直接使用的 AddCustom proto 头文件由构建阶段从ge/add_custom.h复制生成。output/op_graph/lib/linux/x86_64/add_custom_kernel.asc随交付件一起输出的 kernel 源码文件供算子运行时编译阶段读取aarch64 环境对应output/op_graph/lib/linux/aarch64/add_custom_kernel.asc。output/single_add.air构图程序导出的 AIR 文件作为 ATC离线编译输入。output/single_add.omATC离线编译生成的.om离线模型文件。build/compilable_add_graph_buildgraph_build 程序。build/compilable_add_model_execom离线模型执行程序。结果校验成功时可观察到output/single_add.air已生成。output/single_add.om已生成。终端输出包含Model executed successfully!。终端输出包含First element of output: 2.000000。若失败优先检查ASCEND_HOME_PATH是否已设置并已正确sourceCANN 环境。atc是否可用。output/op_graph/lib/os/arch/libcust_opapi.so、output/op_graph/lib/os/arch/add_custom_kernel.asc以及output/op_graph/include/add_custom.h是否已生成。当前soc_version是否与实际环境一致run.sh默认使用Ascend910B1。注意事项 / 限制run.sh当前固定soc_versionAscend910B1如需适配其他产品请按实际环境调整。build/compilable_add_model_exec需要以位置参数传入模型路径。run.sh会自动将ASCEND_CUSTOM_OPP_PATH追加为output/以便ATC离线编译按op_graph/lib/os/arch/目录规范加载libcust_opapi.so。当前样例按 shape key 组织多份 kernel binary主要提供多 shape / 多 kernel 的处理框架。当前 sample 依赖已安装的 CANN toolkit 头文件并直接通过OpCompileContext::GetInputTensor读取编译期输入。附录参考文档RTC 算子运行时编译参考Ascend C 算子即时编译ATC 离线编译参考ATC 工具使用指南算子运行时编译与下沉机制CompilableOp负责在Compile回调中读取 kernel 源码并调用aclrtc完成算子运行时编译。Compile阶段的输入读取已经按Tensor组织因此除 shape 外还可以直接访问 data type 和 format。PortableOp负责把编译产物序列化进最终模型并在执行阶段反序列化恢复。ShapeInferOp负责在InferShape和InferDataType回调中完成输出 shape 与 data type 推导。这几部分配合后ATC离线编译生成的 om离线模型中会包含样例自定义算子运行所需的设备侧二进制若后续扩展为 key 对应不同 kernel 的场景本样例已给出多份 binary 的管理与下沉方式。关键实现职责Compile读取源码并触发算子运行时编译。InferShape / InferDataType基于输入 shape 和 data type 推导输出描述。Serialize / Deserialize把编译产物写入模型并在执行阶段恢复。Execute在模型执行阶段加载设备侧二进制并发起 kernel 调用。【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
