Triton Custom Operator TensorFlow Graph Integration Sample【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/geSample OverviewGraph construction entry:TensorFlowOperator programming language:TritonCompilation method:Pre-compiled to npubinModel sinking capability:Not applicableCore pipeline:Triton kernel - TensorFlow deliverables GE deliverables - TensorFlow graph integration executionDifference from other samples: This sample focuses on Triton kernel and TensorFlow/GE deliverables cooperation, not involving TorchAir, norATC offline compilation - om offline modelmodel sinking.This sample demonstrates how to integrate Triton Add kernel into TensorFlow and complete NPU execution through GE custom operator deliverables. The pipeline includesnpubingeneration, TensorFlow custom operator so building, GE deliverables building, and final TensorFlow script result verification.Applicable ScenariosWant to see custom operator graph integration example inTensorFlow Tritonscenario.Want to reference cooperation betweennpubin, TensorFlow deliverables, and GE deliverables.Want to verify numerical correctness of Triton kernel in TensorFlow graph.Not suitable for understandingTorchAirorATC offline compilation - om offline modelpipeline.PrerequisitesCANNCANN toolkit and ops packages installed and configured.ASCEND_HOME_PATHset.Follow installation guide to complete toolkit and ops package installation.Framework and PluginsInstalledTriton-Ascend.InstalledTensorFlow 1.15orTensorFlow 2.6.5and corresponding framework plugin packages.References:Triton-Ascend Installation GuideTriton-Ascend Quick StartTensorFlow 1.15 Migration GuideTensorFlow 2.6.5 Migration GuideEnvironment VariablesASCEND_HOME_PATHASCEND_CUSTOM_OPP_PATHIf need to customizenpubinoutput directory, can setTRITON_CACHE_DIRAdditional Dependenciespython3gcmakeQuick RunExecute shortest path inexamples/custom_op/triton_add_customdirectory:Recommended Methodcd add_custom_kernel python3 add_custom_kernel.py cd ../tensorflow bash build_tf.sh cd ../ge bash run.sh export ASCEND_CUSTOM_OPP_PATH$(pwd)/build_out:$ASCEND_CUSTOM_OPP_PATH cd ../tensorflow python3 ../script/run_add_custom_tf_1.15.pyBefore execution, confirmbin_pathinge/src/custom_op.cppcan accessadd_kernel.npubingenerated in step 1. If successful, terminal will printThe result of tf and ac is the same.Step-by-step MethodFirst executepython3 add_custom_kernel.pyinadd_custom_kernel/to generatenpubin.Then executebash build_tf.shintensorflow/to generateoutputs/libcustom_ops.so.Next executebash run.shinge/to generatebuild_out/libcust_opapi.soandframework/tensorflow/npu_supported_ops.json.Executeexport ASCEND_CUSTOM_OPP_PATH$(pwd)/build_out:$ASCEND_CUSTOM_OPP_PATHto addlibcust_opapi.sodirectory to environment variables.Finally return totensorflow/directory and executepython3 ../script/run_add_custom_tf_1.15.py.If need to adjustnpubinoutput location, can first setTRITON_CACHE_DIR; if need to re-specifynpubinpath, can modifybin_pathinge/src/custom_op.cppaccording to actual cache location.Directory Structure and Key Filestriton_add_custom ├── README.md ├── add_custom_kernel │ └── add_custom_kernel.py // Triton Add kernel implementation and compilation entry ├── ge │ ├── CMakeLists.txt // GE deliverables build script │ ├── gen_npu_supported_ops_json.sh │ ├── run.sh │ └── src │ └── custom_op.cpp // Triton operator GE deliverables ├── tensorflow │ ├── add_custom_triton_tf.cc // TensorFlow-side custom operator declaration │ └── build_tf.sh // TensorFlow deliverables build script ├── script │ └── run_add_custom_tf_1.15.py // TensorFlow test script ├── static_shape.png └── dynamic_shape.pngKey files:add_custom_kernel/add_custom_kernel.pyTriton Add kernel implementation and compilation entry.tensorflow/add_custom_triton_tf.ccTensorFlow-side custom operator declaration.tensorflow/build_tf.shCompilelibcustom_ops.so.ge/src/custom_op.cppGE custom operator deliverables, implement loadingnpubinand initiating kernel execution.ge/CMakeLists.txtBuildlibcust_opapi.soand generatenpu_supported_ops.json.ge/run.shOne-click script for GE substeps, responsible for configure, build, and install.script/run_add_custom_tf_1.15.pyTensorFlow graph execution script, responsible for numerical consistency verification.Core Pipelineadd_custom_kernel/add_custom_kernel.pygenerates Triton kernel correspondingnpubin.tensorflow/build_tf.shbuilds TensorFlow-sidelibcustom_ops.so.ge/CMakeLists.txtbuilds GE deliverablelibcust_opapi.so, and generatesnpu_supported_ops.jsonfor TensorFlow use.script/run_add_custom_tf_1.15.pyloads TensorFlow-side so, and executes custom AddCustom operator through TensorFlow graph.Build Artifactstensorflow/outputs/libcustom_ops.soTensorFlow-side custom operator deliverable.ge/build_out/libcust_opapi.soGE-side custom operator deliverable.ge/build_out/framework/tensorflow/npu_supported_ops.jsonDescription file used by TensorFlow Adapter when loading custom operator support information.add_kernel.npubinTriton kernel compiled binary file, default written by Triton to cache directory, path can be controlled throughTRITON_CACHE_DIR.Result VerificationWhen successful, can observe:tensorflow/outputs/libcustom_ops.sogenerated.ge/build_out/libcust_opapi.sogenerated.ge/build_out/framework/tensorflow/npu_supported_ops.jsongenerated.Terminal output includesThe result of tf and ac is the same..If failed, prioritize checking:Whetherbin_pathinge/src/custom_op.cppcan correctly access generatednpubin.WhetherASCEND_CUSTOM_OPP_PATHalready includesge/build_out.Whether./outputs/libcustom_ops.socan be accessed under test script startup directory.Whether TensorFlow, framework plugin package and Triton-Ascend installation match.Notes / LimitationsCurrent sample depends on Triton-Ascend target hardware support scope, please refer to Triton-Ascend project description.bin_pathinge/src/custom_op.cppneeds to be consistent with actual generatednpubinpath.Test script defaults to loading TensorFlow-side so from./outputs/libcustom_ops.soin current working directory, recommend running test script intensorflow/directory.Sample defaults to using static shape; if need dynamic shape, can enablecompile_dynamic_modeconfiguration in test script.AppendixDynamic / Static Shape Debug SwitchInscript/run_add_custom_tf_1.15.pycan enable dynamic shape through following configuration:custom_op.parameter_map[compile_dynamic_mode].b TrueStatic Shape Dump Data ConfigurationInscript/run_add_custom_tf_1.15.pycan enable dump data capability through following configuration:custom_op.parameter_map[enable_dump].b True custom_op.parameter_map[dump_mode].s tf.compat.as_bytes(all) custom_op.parameter_map[dump_path].s tf.compat.as_bytes(/home/test/output)Profiling Switchcustom_op.parameter_map[profiling_mode].b True custom_op.parameter_map[profiling_options].s tf.compat.as_bytes( {output:.,training_trace:on,task_trace:on,hccl:on,aicpu:on,aic_metrics:PipeUtilization,msproftx:off} )After exporting profiling data, can use following command to view result directory:msprof --exporton --output${profiling_path} cd ${profiling_path}/mindstudio_profiler_outputThen findmsprof_*.jsonand use tools likechrome://tracing/to load.Profiling Phenomenon DescriptionStatic shape result illustration:Dynamic shape result illustration:Under static shape can observeMODEL_EXECUTE,EVENT_WAIT,MEMCPY_ASYNCand other stream synchronization and data transfer processes; under dynamic shape these phenomena will be different.【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考