Ascend C Operator Graph Integration Sample via TorchAir【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/geSample OverviewGraph construction entry:PyTorch TorchAirOperator programming language:Ascend CCompilation method:Pre-compiled/Direct invocationModel sinking capability:Not applicableCore pipeline:Ascend C kernel - GE deliverables - TorchAir graph integration - Graph mode executionDifference from other samples: This sample focuses onPyTorch TorchAirscenario, does not involveATC offline compilation - om offline modelmodel sinking, nor Triton kernel.This sample demonstrates how to register Ascend C custom operators and integrate them into GE graph viaTorchAir. Using Add operator as an example, the deliverable side completes GE graph integration execution viaEagerExecuteOp, while the Python side covers both eager invocation and graph mode compiled invocation.Applicable ScenariosWant to see the minimum graph integration pipeline for Ascend C custom operators inPyTorch TorchAirscenario.Want to reference the collaboration betweenTORCH_LIBRARYregistration, custom converter and GE custom operator.Want to verify numerical results of Ascend C kernel in eager and graph modes.PrerequisitesCANNFollow installation guide to complete toolkit and ops package installation.Framework and PluginsAlready installedPyTorch,torch_npuandTorchAir.TorchAirdependent Python version needs to meet its release requirements.Reference:Ascend Extension for PyTorchAscend Extension for PyTorch Ascend Community DocumentationEnvironment VariablesASCEND_HOME_PATHLD_LIBRARY_PATHand other CANN runtime related variablesAdditional Dependenciespip3 install expecttestQuick RunExecute inexamples/custom_op/ascendc_add_customdirectory:Recommended Methodmkdir -p build cd build cmake .. make -j$(nproc) export ASCEND_CUSTOM_OPP_PATH$(pwd):$ASCEND_CUSTOM_OPP_PATH python3 ../script/add_custom_test.pyIf the terminal showsRan 1 testandOK, the sample runs successfully.Step-by-step MethodFirst executecmake ..andmake -j$(nproc)to generatebuild/libcust_opapi.so.Then executeexport ASCEND_CUSTOM_OPP_PATH$(pwd):$ASCEND_CUSTOM_OPP_PATHto add the directory wherelibcust_opapi.sois located to environment variable.Finally executepython3 ../script/add_custom_test.pyto run eager and TorchAir graph mode tests.If you need to execute installation target, you can additionally runcmake --install build.Directory Structure and Key Filesascendc_add_custom ├── CMakeLists.txt ├── README.md ├── add_custom_kernel │ ├── add_custom.asc // Ascend C kernel implementation and PyTorch registration code │ ├── add_custom_kernel.h // kernel declaration │ └── custom_op.cpp // GE deliverables, implementing Execute graph integration logic └── script └── add_custom_test.py // Python test script, covering eager and graph modesKey files:add_custom_kernel/add_custom.ascAscend C Add kernel implementation file, also contains PyTorch custom operator registration logic.add_custom_kernel/custom_op.cppCustom operator GE deliverables, implementingExecutepath and completing GE graph integration.script/add_custom_test.pyPython test script, covering eager execution and TorchAir graph mode compiled execution.CMakeLists.txtCompileslibcust_opapi.soand configures PyTorch, TorchAir, CANN dependencies.Core PipelineImplement Ascend C Add kernel inadd_custom_kernel/add_custom.ascand complete PyTorch custom operator registration.ImplementEagerExecuteOpinadd_custom_kernel/custom_op.cppto register custom operator to GE.Loadlibcust_opapi.soviascript/add_custom_test.py, execute and verify results in eager and TorchAir graph modes.Build Artifactsbuild/libcust_opapi.soCustom operator deliverables used by GE/TorchAir.Result VerificationWhen successful, you can observe:Terminal output containsRan 1 test.Terminal output containsOK.NoNumerical values do not match within toleranceerror during testing.If failed, priority checks:WhetherASCEND_HOME_PATH,LD_LIBRARY_PATHand other environment variables are correct.Whethertorch,torch_npu,TorchAirversions match.Whetherbuild/libcust_opapi.sohas been successfully generated and can be loaded by Python script.Whether current environment has available NPU, script will directly verifytorch.npu.is_available().Notes / LimitationsSupported products:Atlas A2 training series products / Atlas A2 inference series products.Sample currently usesfloat168 x 2048input for result verification.script/add_custom_test.pydepends on available NPU environment, cannot pass in CPU environment.This sample demonstratesPyTorch TorchAirgraph integration path, does not involve model sinking to om offline model.AppendixOperator SpecificationItemContentOperator typeAddCustomInputx,yOutputzInput shape8 x 2048Output shape8 x 2048Data typefloat16FormatNDKernel nameadd_customOperator Implementation SupplementAscend C Add kernel execution process can be summarized in three steps:CopyIn: Move input from Global Memory to on-chip Local Memory.Compute: Perform addition operation on two LocalTensors.CopyOut: Move computation result back to Global Memory corresponding to output Tensor.Registration Mechanism SupplementTORCH_LIBRARY_FRAGMENTis used to define Python-side visible operator signatures.TORCH_LIBRARY_IMPLis used to bind operator implementation toPrivateUse1,XLA,Metaand otherDispatchKeys.torchair.register_fx_node_ge_converter(...)inscript/add_custom_test.pyis responsible for mapping PyTorch nodes to GE custom operatorAddCustom.【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考