昇腾CANN/GE数据依赖形状自定义算子示例

昇腾CANN/GE数据依赖形状自定义算子示例 Data Dependent Shape Custom Operator (Type III Operator) Sample【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/geSample OverviewGraph construction entry:GEOperator programming language:Ascend CCompilation method:.ascand.cppcompiled in same targetModel sinking capability:Not applicableCore pipeline:Ascend C kernel and host-side custom op compiled in same library - GE deliverables - In-process graph construction - Session::RunGraph - hybrid/RT2 dynamic executionDifference from other samples: This sample focuses on shape buffer usage in unknown graph dynamic execution scenarios for type III custom operators.This sample demonstrates a minimum runnable type III custom operator execution pipeline: after in-process graph construction, directly callSession::RunGraph, compile Ascend C kernel and host-side custom op to samelibwhere_like_custom_op.soby CMake, thenWhereLikeCustom::Execute()allocates maximum output and shape buffer, finally writes back actual output shape on device side.WhereLikeCustominput isbool[N], used to return indices of elements with valuetruein input, output isint64[true_count, rank]. Sincetrue_countdepends on runtime input data, only output upper bound can be determined at compile time, actual shape needs to be passed back through type III custom op shape buffer protocol after execution.Applicable ScenariosWant to see minimum implementation of type III custom operators inSession::RunGraphdynamic execution pipeline.Want to reference how user synchronizes stream, reads back shape buffer and updates output shape inExecute.Want to see complete process of.ascand.cppcompiled directly together and initiating kernel calls inExecute.Not suitable for understandingATC offline compilation - om offline modelsinking pipeline.PrerequisitesCANNCANN environment has been correctly installed and configured, e.g., executedsource ${ASCEND_HOME_PATH}/set_env.sh.Current environment hasACL,GE,Graph,Ascend Crelated headers and libraries.Follow installation guide to complete toolkit and ops package installation.Framework and PluginsThis sample does not depend on PyTorch, TensorFlow or TorchAir.Environment VariablesASCEND_HOME_PATHASCEND_CUSTOM_OPP_PATHwill be automatically appended as current samplesoutput/inrun.shAdditional DependenciescmakegQuick RunExecute inexamples/custom_op/data_dependent_shape_customdirectory:Recommended Methodsource ${ASCEND_HOME_PATH}/set_env.sh bash run.shrun.shwill automatically complete configure, build, install and appendoutput/toASCEND_CUSTOM_OPP_PATH. If successful, terminal will print:output shape: [4, 1] output values: 0 2 4 7Step-by-step Methodsource ${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 ..Whereexport ASCEND_CUSTOM_OPP_PATH$(pwd)/output:$ASCEND_CUSTOM_OPP_PATHis used to add custom operator package root directory to environment variable, then GE will load deliverables according tooutput/op_graph/lib/os/arch/libwhere_like_custom_op.sorule.Directory Structure and Key Filesdata_dependent_shape_custom ├── CMakeLists.txt ├── README.md ├── run.sh ├── ge │ ├── custom_op.cpp // EagerExecuteOp ShapeInferOp main process implementation │ ├── where_like_custom.h // WhereLikeCustom proto definition │ └── where_like_custom_kernel.asc // Ascend C kernel source code └── session_run └── main.cc // In-process graph construction and direct Session::RunGraph callKey files:ge/custom_op.cppCustom operator core main process, implementsExecute,InferShapeandInferDataType; whereExecuteis responsible for allocating maximum output, shape buffer, calling.ascexported launch wrapper, and reading back shape buffer to update output shape after kernel completion,InferShape/InferDataTypeis responsible for compile-time output shape / dtype derivation.ge/where_like_custom_kernel.ascDevice kernel and host-side launch wrapper implementation, responsible for writing output data and shape buffer.session_run/main.ccBuild minimum graph and execute directly viaSession::AddGraph Session::RunGraph.run.shChain configure, build, install and run process.Core Pipelinesession_run/main.ccbuilds minimum graph containingData - WhereLikeCustomand sets input/output description to dynamic shape, making the whole graph gounknown graphexecution pipeline.InferShape/InferDataTypeinge/custom_op.cppgives output shape / dtype at graph construction phase; this sample does not depend on framework lowering to insert shape write-back nodes.CMakeLists.txtcompilesge/custom_op.cppandge/where_like_custom_kernel.asctogether tolibwhere_like_custom_op.so.ge/custom_op.cppallocates output viactx-MallocOutputTensor(...)according to maximum shape inExecutecallback, then allocates shape buffer via firstctx-MallocWorkSpace(...).ge/where_like_custom_kernel.ascwrites output indices and shape buffer on device side, where shape buffer is used to pass back actual output shape.Executeinge/custom_op.cppsynchronizes current stream after launch, copies shape buffer back to host, parses real shape and updates output tensors logical shape and valid size.After execution,session_run/main.ccreads output tensor and prints actual shape and output values.Build Artifactsoutput/op_graph/lib/linux/x86_64/libwhere_like_custom_op.soCustom operator deliverables used by GE in Linux x86_64 environment; aarch64 environment corresponds tooutput/op_graph/lib/linux/aarch64/libwhere_like_custom_op.so.output/op_graph/include/where_like_custom.hOperator proto header file that can be directly used by graph construction side.build/data_dependent_shape_custom_session_runDirectSession::RunGraphexecution program.Result VerificationWhen successful, you can observe:output/op_graph/lib/os/arch/libwhere_like_custom_op.sohas been generated.output/op_graph/include/where_like_custom.hhas been generated.Terminal output containsoutput shape: [4, 1].Terminal output containsoutput values: 0 2 4 7.If failed, priority checks:WhetherASCEND_HOME_PATHhas been set and CANN environment has been correctlysourced.WhetherASCEND_CUSTOM_OPP_PATHalready contains current samplesoutput/.Whetheroutput/op_graph/lib/os/arch/libwhere_like_custom_op.soandoutput/op_graph/include/where_like_custom.hhave been generated.Whether current environment has available NPU and available Ascend C compilation environment.Notes / LimitationsWhereLikeCustomcurrent sample input is one-dimensionalbool[8], so actual output is matching position indices, shape is[true_count, 1]..asccompilation parameter currently fixed as--npu-archdav-2201, if target chip is different need to adjustCMakeLists.txt.AppendixOperator SpecificationItemContentOperator typeWhereLikeCustomInputxOutputyInput shapeNOutput shape upper bound[N, rank(x)]Output actual shape[true_count, rank(x)]Input data typeboolOutput data typeint64FormatNDKernel namewhere_like_customShape Buffer ConventionThis sample uses user-defined shape buffer protocol and user parses and writes back actual output shape inExecute. Current kernel will write:shape[0] 2Ushape[1] true_countshape[2] rankFor current one-dimensional input scenario,rank 1, so input[true, false, true, false, true, false, false, true]actual output shape is[4, 1], output data is0 2 4 7.【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考