ArgsUpdater Address Refresh Custom Operator Sample【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/geSample OverviewGraph construction entry:GEOperator programming language:Ascend C(RTC runtime compilation)Compilation method: The.cppfile compiles the host-side custom op, while kernel source code is compiled to device binary through RTC at runtimeCore pipeline:Ascend C kernel source code - RTC runtime compilation - GE deliverables - In-process graph construction - Session::ExecuteGraphWithStreamAsync online executionDifference from other samples: This sample focuses on theArgsUpdaterinterface working withMallocReadOnlyDevArgsto implement address refresh. The GE framework manages args synchronization, avoiding extra D2D copies (MEMCPY_ASYNC) and improving repeated execution performance.This sample demonstrates the complete pipeline of theArgsUpdateraddress refresh mechanism: using an Ascend C Add operator with input shape[4096, 4096]float32 (16M elements, 64MB), defining two functionally identical operators—AddRefreshOp(implements theArgsUpdaterinterface) andAddNoRefreshOp(does not implement it). The performance difference is compared throughSession::ExecuteGraphWithStreamAsynconline execution.The core concept ofArgsUpdater: During model loading,MallocReadOnlyDevArgsallocates a read-only kernel args memory region on the device side. For subsequent repeated executions, theUpdateHostArgscallback only refreshes address fields in args (input/output tensor pointers). This eliminates the extra D2D copies (MEMCPY_ASYNC) that the GE framework inserts to synchronize operator input/output tensor content to the device side, achieving approximately 1.17x performance improvement in high-frequency execution scenarios.Applicable ScenariosUnderstanding the implementation and performance benefits of theArgsUpdaterMallocReadOnlyDevArgsaddress refresh mechanism.Viewing the complete process of Ascend C kernel compilation through RTC and invocation in GE custom operators.Comparing performance differences between implementations with and without address refresh in high-frequency execution scenarios.PrerequisitesCANNThe CANN environment is properly installed and configured, for example, by executingsource ${ASCEND_HOME_PATH}/set_env.sh.The current environment hasACL,GE, andGraphrelated header files and libraries.Refer to the Installation Guide to complete toolkit and ops package installation.Framework and PluginsThis sample does not depend on PyTorch, TensorFlow, or TorchAir.The kernel source codeadd_custom_kernel/add_custom.ascis compiled through RTC at runtime and does not require pre-compilation.Environment VariablesASCEND_HOME_PATHASCEND_CUSTOM_OPP_PATHwill be automatically appended to the current samplesoutput/inrun.shAdditional DependenciescmakegQuick RunExecute in theexamples/custom_op/args_refresh_add_customdirectory:Recommended Methodsource ${ASCEND_HOME_PATH}/set_env.sh bash run.shrun.shautomatically completes configure, build, and install, and appendsoutput/toASCEND_CUSTOM_OPP_PATH. The script executes the following 2 steps sequentially:Compile custom operator deliverables and executable programsRunsession_run(online performance comparison)If successful, the terminal will print output similar to:[Perf] input shape: [4096, 4096], float32, 64MB [Perf] iters: 100 [Perf] With ArgsUpdater: xxx us (avg xxx us/iter) [Perf] Without ArgsUpdater: xxx us (avg xxx us/iter) [Perf] Speedup: xxx xStep-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 # Online execution: performance comparison cd build ./args_refresh_session_run cd ..The commandexport ASCEND_CUSTOM_OPP_PATH$(pwd)/output:$ASCEND_CUSTOM_OPP_PATHadds the custom operator package root directory to the environment variable. Then GE loads deliverables according to the ruleoutput/op_graph/lib/os/arch/libcust_opapi.so.Directory Structure and Key Filesargs_refresh_add_custom ├── CMakeLists.txt ├── README.md ├── run.sh ├── add_custom_kernel │ ├── add_custom.asc // Ascend C Add kernel source code (RTC runtime compilation) │ └── add_custom_kernel.h // kernel header file ├── ge │ ├── add_custom.h // AddRefreshOp / AddNoRefreshOp proto definition │ ├── custom_op.cpp // Implementation of Execute, ArgsUpdater, InferShape, etc. for both operators │ └── utils │ ├── log.h // Unified log macros (LOG_ERROR/LOG_WARNING/LOG_INFO) │ ├── rtc_kernel_loader.h // RTC kernel loader interface │ └── rtc_kernel_loader.cpp // RTC compilation and loading implementation └── session_run └── main.cc // In-process graph construction, online performance comparisonKey files:ge/custom_op.cppThe core main process of custom operators.AddRefreshOpsimultaneously implementsEagerExecuteOp,ArgsUpdater, andShapeInferOp;AddNoRefreshOponly implementsEagerExecuteOpandShapeInferOp. Both load kernels throughRtcKernelLoader, allocate output tensors, and invokeaclrtLaunchKernelV2to launch kernels. The difference is thatAddRefreshOpregisters args throughMallocReadOnlyDevArgsand implements theUpdateHostArgscallback, with the GE framework managing args synchronization;AddNoRefreshOpdoes not register args, so the GE framework cannot perceive address changes and must insert extra D2D copies (MEMCPY_ASYNC) to synchronize operator input/output tensor content to the device side during each execution.ge/utils/rtc_kernel_loader.cppRTC kernel loader, encapsulating the complete pipeline from source code compilation to loading: read kernel source code →aclrtcCreateProg→aclrtcCompileProg→aclrtcGetBinData→aclrtBinaryLoadFromData→aclrtBinaryGetFunction. Supports dynamically obtaining NPU architecture to generate compilation options.ge/utils/log.hUnified log macros supporting three levels:LOG_ERROR,LOG_WARNING, andLOG_INFO, automatically appending filename and line number.ge/add_custom.hGraph construction operator proto definition, registeringAddRefreshOpandAddNoRefreshOp.add_custom_kernel/add_custom.ascAscend C Add kernel source code, performing element-wise addition withBLOCK_SIZE1024, compiled through RTC at runtime.session_run/main.ccConstructs two graphs (usingAddRefreshOpandAddNoRefreshOprespectively), executes throughSession::ExecuteGraphWithStreamAsyncand performs 100 rounds of performance comparison. Uses two sets of memory to alternately triggerUpdateHostArgsaddress changes.run.shConnects the complete pipeline of compilation and online execution.Core PipelineOnline Execution (Session::ExecuteGraphWithStreamAsync)session_run/main.ccconstructs two graphs:refresh_graph(usingAddRefreshOp) andno_refresh_graph(usingAddNoRefreshOp), both with input shape[4096, 4096]float32.Inge/custom_op.cpp, theExecutecallback compiles and loads the kernel throughRtcKernelLoaderduring model loading, and allocates output tensors throughctx-MallocOutputTensor(...). Note:Executeis called only once during model loading and will not be invoked again when the model sinks to device execution.AddRefreshOpregisters theAddArgsstructure to the GE framework throughctx-MallocReadOnlyDevArgs(...)and additionally implements theUpdateHostArgscallback: during subsequent executions, the GE framework calls this callback to refresh input/output tensor addresses in host-side args, then the GE framework efficiently synchronizes changes to the device side.AddNoRefreshOpdoes not implementArgsUpdaterand does not useMallocReadOnlyDevArgsto register args. AlthoughaclrtMallocaclrtMemcpyin Execute only occurs once during model loading, since args are not registered, the GE framework cannot perceive address changes and must insert extra Identity operators in the graph to transfer data, generating extra D2D copies (MEMCPY_ASYNC) during each execution to synchronize operator input/output tensor content to the device side.Both launch Ascend C kernels throughaclrtLaunchKernelV2.After execution completes,session_run/main.cccalculates and prints the total time and speedup ratio for both.ArgsUpdater MallocReadOnlyDevArgs MechanismDuring model loading (Execute called only once): Execute() ├─ RtcKernelLoader::Load() → RTC compiles and loads kernel ├─ MallocReadOnlyDevArgs(args, sizeof(args)) → Allocates device-side read-only args memory ├─ Fill AddArgs { x_ptr, y_ptr, z_ptr } └─ aclrtLaunchKernelV2(registered_args) → Kernel launch Subsequent executions (AddRefreshOp): UpdateHostArgs(ctx) ├─ GetKernelArgs(kPlacementHost, 0) → Get host-side args pointer └─ Only refresh args-x_ptr / y_ptr / z_ptr → Update tensor addresses GE framework automatically synchronizes address changes to device side, no need to re-copy argsMallocReadOnlyDevArgscopies the args structure to the device side and caches it during model loading; for subsequent executions,UpdateHostArgsonly updates address fields in host-side args, and the GE framework synchronizes changes to the device side, avoiding the extra D2D copies (MEMCPY_ASYNC) that the GE framework inserts to synchronize operator input/output tensor content to the device side.RTC Runtime CompilationRtcKernelLoader::Load() ├─ GetCurrentLibraryDir() → Get dynamic library directory ├─ LoadTextFromFile(source_path) → Read kernel source code ├─ GetRtcCompileOption() → Dynamically obtain NPU architecture (such as dav-2201) ├─ aclrtcCreateProg() → Create compilation program ├─ aclrtcCompileProg() → Compile kernel ├─ aclrtcGetBinData() → Get compiled binary ├─ aclrtBinaryLoadFromData() → Load binary └─ aclrtBinaryGetFunction() → Get function handleRTC compilation completes during model loading, and subsequent executions directly reuse the compiled kernel without re-compilation.Build Productsoutput/op_graph/lib/linux/x86_64/libcust_opapi.soCustom operator deliverable used by GE in Linux x86_64 environment; aarch64 environment corresponds tooutput/op_graph/lib/linux/aarch64/libcust_opapi.so.output/op_graph/lib/os/arch/add_custom.ascKernel source file, copied by CMake fromadd_custom_kernel/for RTC compilation use.output/op_graph/include/add_custom.hOperator proto header file that can be directly used for graph construction.build/args_refresh_session_runOnline execution performance comparison program (Session::ExecuteGraphWithStreamAsync).Result ValidationWhen successful, you can observe:output/op_graph/lib/os/arch/libcust_opapi.sois generated.output/op_graph/include/add_custom.his generated.session_runterminal output contains[Perf] Speedup: xxx x, andWith ArgsUpdatertime is lower thanWithout ArgsUpdater.If failed, prioritize checking:WhetherASCEND_HOME_PATHis set and the CANN environment is properly sourced.WhetherASCEND_CUSTOM_OPP_PATHincludes the current samplesoutput/.Whetheroutput/op_graph/lib/os/arch/libcust_opapi.soandoutput/op_graph/include/add_custom.hare generated.Whether the current environment has an available NPU.Precautions / LimitationsThe kernel is compiled through RTC at runtime, with compilation overhead during model loading, and subsequent executions directly reuse it.RTC compilation options dynamically obtain NPU architecture throughaclrtGetDeviceInfo, automatically adapting to different chip models.Performance comparison results are affected by NPU model, system load, and other factors; the speedup ratio is for reference only.Insession_run,ge.graphRunModeis set to1(that is,PRIORITY_GRAPHmode), ensuring the online execution pipeline.The kernel logic ofAddRefreshOpandAddNoRefreshOpis completely identical; the performance difference comes from the D2D copies (MEMCPY_ASYNC) that the GE framework inserts to synchronize operator input/output tensor content to the device side.Performance testing uses two sets of memory for alternate execution, triggeringUpdateHostArgsaddress changes, more realistically reflecting optimization effects.AppendixOperator SpecificationsItemContentOperator typeAddRefreshOp/AddNoRefreshOpInputx,yOutputzInput shape[4096, 4096]Output shape[4096, 4096]Input data typefloat32Output data typefloat32FormatNDkernel nameadd_custom(Ascend C, RTC runtime compilation)BLOCK_SIZE1024ArgsUpdater Interface DescriptionInterfaceClassPurposeEagerExecuteOp::ExecuteAddRefreshOp/AddNoRefreshOpDuring model loading: load kernel, allocate output, allocate device args, launch kernelArgsUpdater::UpdateHostArgsAddRefreshOpSubsequent executions: get host-side args, refresh tensor address fieldsShapeInferOp::InferShapeAddRefreshOp/AddNoRefreshOpCompile-time output shape inference (same as input)ShapeInferOp::InferDataTypeAddRefreshOp/AddNoRefreshOpCompile-time output dtype inference (same as input)Performance AnalysisThrough profiling data analysis,AddNoRefreshOphas approximately 323 us extra overhead per round, mainly from the following sources (data below is for reference only; actual time may vary depending on NPU model, system load, and other factors):Overhead sourceTimePercentageMEMCPY_ASYNC (D2D copy)~308 us95%Identity operator scheduling overhead~15 us5%SinceAddNoRefreshOpdoes not register args throughMallocReadOnlyDevArgs, the GE framework inserts extra Identity operators to transfer data during graph compilation. These Identity operators generate MEMCPY_ASYNC (D2D copy) during device-side execution, approximately 102 us each time, 3 times per round. In contrast,AddRefreshOprefreshes addresses through theUpdateHostArgscallback, and the GE framework efficiently synchronizes without needing to insert Identity operators.【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考