Star722
475
代码
代码
Issues30
Pull Requests117
Actions
项目讨论
Wiki
分析
Star722
475
已合并
feat: complete Python compilable custom-op online and offline example #4477
shangdf创建于 8月20日
feat: complete Python compilable custom-op online and offline example #4477
已合并
shangdf创建于 8月20日
从shangdf/gefeature/python-compilable-example-develop合入到cann/gedevelop
共 15 个文件变更+1394-0
Mexamples/custom_op/README.md+1-0
@@ -9,6 +9,7 @@
9| `ascendc_add_custom` | Ascend C 算子通过 GE 入图 | PyTorch + TorchAir | Ascend C | CMake编译 | 不涉及 | [README](./ascendc_add_custom/README.md) |9| `ascendc_add_custom` | Ascend C 算子通过 GE 入图 | PyTorch + TorchAir | Ascend C | CMake编译 | 不涉及 | [README](./ascendc_add_custom/README.md) |
10| `triton_add_custom` | Triton 算子通过 GE 入图 | TensorFlow | Triton | 预编译为 `npubin` | 不涉及 | [README](./triton_add_custom/README.md) |10| `triton_add_custom` | Triton 算子通过 GE 入图 | TensorFlow | Triton | 预编译为 `npubin` | 不涉及 | [README](./triton_add_custom/README.md) |
11| `compilable_add_custom` | Ascend C 算子通过 GE 入图并生成 om离线模型 | GE + ATC离线编译 | Ascend C | RTC算子运行时编译 | 支持模型下沉到 om离线模型 | [README](./compilable_add_custom/README.md) |11| `compilable_add_custom` | Ascend C 算子通过 GE 入图并生成 om离线模型 | GE + ATC离线编译 | Ascend C | RTC算子运行时编译 | 支持模型下沉到 om离线模型 | [README](./compilable_add_custom/README.md) |
12+| `python_compilable_add_custom` | Python 算子在在线图编译和 ATC 离线编译阶段生成 kernel，并验证在线执行与 OM 执行 | GE 在线执行 + ATC 离线编译 | Python + Ascend C | Python compile 回调中调用 BiSheng | 支持 OM 脱离 Python 插件执行 | [README](./python_compilable_add_custom/README.md) |
12| `data_dependent_shape_custom` | 数据依赖 shape 算子 | GE | Ascend C | CMake编译 | 不涉及 | [README](data_dependent_shape_custom/README.md) |13| `data_dependent_shape_custom` | 数据依赖 shape 算子 | GE | Ascend C | CMake编译 | 不涉及 | [README](data_dependent_shape_custom/README.md) |
13| `args_refresh_add_custom` | ArgsUpdater 地址刷新 + MallocReadOnlyDevArgs + 性能对比 | GE 在线执行 | Ascend C | RTC 运行时编译 | 在线地址刷新性能对比 | [README](./args_refresh_add_custom/cpp/README.md) |14| `args_refresh_add_custom` | ArgsUpdater 地址刷新 + MallocReadOnlyDevArgs + 性能对比 | GE 在线执行 | Ascend C | RTC 运行时编译 | 在线地址刷新性能对比 | [README](./args_refresh_add_custom/cpp/README.md) |
14| `annotated_args_refresh_add_custom` | AnnotatedArgsOp 声明式地址刷新在线场景性能对比+离线场景 | GE 在线执行 + ATC 离线编译 | Ascend C | RTC 运行时编译 | 支持在线性能对比和 OM 模型下沉 | [README](./annotated_args_refresh_add_custom/README.md) |15| `annotated_args_refresh_add_custom` | AnnotatedArgsOp 声明式地址刷新在线场景性能对比+离线场景 | GE 在线执行 + ATC 离线编译 | Ascend C | RTC 运行时编译 | 支持在线性能对比和 OM 模型下沉 | [README](./annotated_args_refresh_add_custom/README.md) |
Mexamples/custom_op/README_en.md+1-0
@@ -9,6 +9,7 @@ This directory provides samples related to custom operator graph integration, co
9| `ascendc_add_custom` | Ascend C operator enters graph through GE | PyTorch + TorchAir | Ascend C | CMake compilation | Not involved | [README](./ascendc_add_custom/README_en.md) |9| `ascendc_add_custom` | Ascend C operator enters graph through GE | PyTorch + TorchAir | Ascend C | CMake compilation | Not involved | [README](./ascendc_add_custom/README_en.md) |
10| `triton_add_custom` | Triton operator enters graph through GE | TensorFlow | Triton | Pre-compiled as `npubin` | Not involved | [README](./triton_add_custom/README_en.md) |10| `triton_add_custom` | Triton operator enters graph through GE | TensorFlow | Triton | Pre-compiled as `npubin` | Not involved | [README](./triton_add_custom/README_en.md) |
11| `compilable_add_custom` | Ascend C operator enters graph through GE and generates om offline model | GE + ATC offline compilation | Ascend C | RTC operator runtime compilation | Supports model sink to om offline model | [README](./compilable_add_custom/README_en.md) |11| `compilable_add_custom` | Ascend C operator enters graph through GE and generates om offline model | GE + ATC offline compilation | Ascend C | RTC operator runtime compilation | Supports model sink to om offline model | [README](./compilable_add_custom/README_en.md) |
12+| `python_compilable_add_custom` | Python operator compiles a kernel in online graph compilation and offline ATC compilation, then verifies online and OM execution | GE online execution + ATC offline compilation | Python + Ascend C | BiSheng invoked by the Python compile callback | OM executes without the Python plugin | [README](./python_compilable_add_custom/README_en.md) |
12| `data_dependent_shape_custom` | Data dependent shape operator | GE | Ascend C | CMake compilation | Not involved | [README](data_dependent_shape_custom/README_en.md) |13| `data_dependent_shape_custom` | Data dependent shape operator | GE | Ascend C | CMake compilation | Not involved | [README](data_dependent_shape_custom/README_en.md) |
13| `args_refresh_add_custom` | ArgsUpdater address refresh + MallocReadOnlyDevArgs + performance comparison | GE online execution | Ascend C | RTC runtime compilation | Online address refresh performance comparison | [README](./args_refresh_add_custom/cpp/README_en.md) |14| `args_refresh_add_custom` | ArgsUpdater address refresh + MallocReadOnlyDevArgs + performance comparison | GE online execution | Ascend C | RTC runtime compilation | Online address refresh performance comparison | [README](./args_refresh_add_custom/cpp/README_en.md) |
14| `annotated_args_refresh_add_custom` | AnnotatedArgsOp declarative address refresh for online performance comparison and offline model | GE online execution + ATC offline compilation | Ascend C | RTC runtime compilation | Supports online performance comparison and OM model sink | [README](./annotated_args_refresh_add_custom/README_en.md) |15| `annotated_args_refresh_add_custom` | AnnotatedArgsOp declarative address refresh for online performance comparison and offline model | GE online execution + ATC offline compilation | Ascend C | RTC runtime compilation | Supports online performance comparison and OM model sink | [README](./annotated_args_refresh_add_custom/README_en.md) |
Aexamples/custom_op/python_compilable_add_custom/CMakeLists.txt+114-0
@@ -0,0 +1,114 @@
1+# -----------------------------------------------------------------------------------------------------------
2+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+# CANN Open Software License Agreement Version 2.0 (the "License").
5+# Please refer to the License for details. You may not use this file except in compliance with the License.
6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+# See LICENSE in the root of the software repository for the full text of the License.
9+ 
10+cmake_minimum_required(VERSION 3.16)
11+project(python_compilable_add_custom LANGUAGES CXX)
12+ 
13+option(PYTHON_COMPILABLE_ADD_BUILD_OFFLINE_TOOLS
14+ "Build the AIR exporter and ACL OM runner used by run_offline.sh" ON)
15+option(PYTHON_COMPILABLE_ADD_BUILD_ES
16+ "Build the ES Python graph-builder wrapper used by run_online.sh" ON)
17+ 
18+set(CMAKE_CXX_STANDARD 17)
19+set(CMAKE_CXX_STANDARD_REQUIRED ON)
20+set(CMAKE_CXX_EXTENSIONS OFF)
21+ 
22+if(NOT CMAKE_BUILD_TYPE)
23+ set(CMAKE_BUILD_TYPE Release CACHE STRING "Build type" FORCE)
24+endif()
25+ 
26+set(ES_OUTPUT_PATH "${CMAKE_BINARY_DIR}/es_output")
27+set(ASCEND_HOME_PATH_OVERRIDE "" CACHE PATH "Optional ASCEND_HOME_PATH override")
28+if(ASCEND_HOME_PATH_OVERRIDE)
29+ set(ASCEND_HOME_PATH "${ASCEND_HOME_PATH_OVERRIDE}")
30+else()
31+ set(ASCEND_HOME_PATH "$ENV{ASCEND_HOME_PATH}")
32+endif()
33+ 
34+if(NOT ASCEND_HOME_PATH)
35+ message(FATAL_ERROR "ASCEND_HOME_PATH is empty. Please source CANN set_env.sh first.")
36+endif()
37+ 
38+add_compile_definitions(_GLIBCXX_USE_CXX11_ABI=0 google=ascend_private)
39+ 
40+add_library(python_compilable_add_custom_op_proto SHARED
41+ proto/add_custom.cc
42+)
43+target_compile_definitions(python_compilable_add_custom_op_proto PRIVATE OP_PROTO_LIB)
44+target_compile_options(python_compilable_add_custom_op_proto PRIVATE -fvisibility=hidden)
45+target_include_directories(python_compilable_add_custom_op_proto PRIVATE
46+ "${ASCEND_HOME_PATH}/include"
47+ "${ASCEND_HOME_PATH}/include/external"
48+)
49+target_link_libraries(python_compilable_add_custom_op_proto PRIVATE
50+ "${ASCEND_HOME_PATH}/lib64/libopp_registry.so"
51+ "${ASCEND_HOME_PATH}/lib64/libgraph.so"
52+ "${ASCEND_HOME_PATH}/lib64/libgraph_base.so"
53+)
54+ 
55+if(PYTHON_COMPILABLE_ADD_BUILD_ES)
56+ list(APPEND CMAKE_MODULE_PATH "${ASCEND_HOME_PATH}/include/ge/cmake")
57+ find_package(GenerateEsPackage REQUIRED)
58+ add_es_library(
59+ ES_LINKABLE_AND_ALL_TARGET es_custom
60+ OPP_PROTO_TARGET python_compilable_add_custom_op_proto
61+ OUTPUT_PATH ${ES_OUTPUT_PATH}
62+ )
63+endif()
64+ 
65+if(PYTHON_COMPILABLE_ADD_BUILD_OFFLINE_TOOLS)
66+ add_executable(python_compilable_add_graph_build
67+ src/offline_graph_build.cc
68+ )
69+ target_compile_options(python_compilable_add_graph_build PRIVATE
70+ -Wall
71+ -Wextra
72+ -Wno-unused-parameter
73+ )
74+ target_compile_definitions(python_compilable_add_graph_build PRIVATE
75+ _GLIBCXX_USE_CXX11_ABI=0
76+ )
77+ target_include_directories(python_compilable_add_graph_build PRIVATE
78+ "${CMAKE_CURRENT_SOURCE_DIR}/proto"
79+ "${ASCEND_HOME_PATH}/include"
80+ "${ASCEND_HOME_PATH}/include/graph"
81+ "${ASCEND_HOME_PATH}/include/ge"
82+ "${ASCEND_HOME_PATH}/opp/built-in/op_proto/inc"
83+ )
84+ target_link_directories(python_compilable_add_graph_build PRIVATE
85+ "${ASCEND_HOME_PATH}/lib64"
86+ )
87+ target_link_libraries(python_compilable_add_graph_build PRIVATE
88+ graph
89+ ge_compiler
90+ ascendcl
91+ graph_base
92+ )
93+ 
94+ add_executable(python_compilable_add_model_exec
95+ src/offline_model_exec.cc
96+ )
97+ target_compile_options(python_compilable_add_model_exec PRIVATE
98+ -Wall
99+ -Wextra
100+ -Wno-unused-parameter
101+ )
102+ target_compile_definitions(python_compilable_add_model_exec PRIVATE
103+ _GLIBCXX_USE_CXX11_ABI=0
104+ )
105+ target_include_directories(python_compilable_add_model_exec PRIVATE
106+ "${ASCEND_HOME_PATH}/include"
107+ )
108+ target_link_directories(python_compilable_add_model_exec PRIVATE
109+ "${ASCEND_HOME_PATH}/lib64"
110+ )
111+ target_link_libraries(python_compilable_add_model_exec PRIVATE
112+ ascendcl
113+ )
114+endif()
Aexamples/custom_op/python_compilable_add_custom/README.md+109-0
@@ -0,0 +1,109 @@
1+# Python CompilableAddCustom 在线与离线编译样例
2+ 
3+本样例对应 GE Python `CompilableOp` Python 化接口，使用同一个 Python
4+`compile`/`declare_launch_args` 实现覆盖两条编译链路：
5+ 
6+- 在线：GE `Session` 构图时由 `CustomGraphOptimizer` 调用 `compile`，随后直接执行图。
7+- 离线：C++ 构图程序生成 AIR，ATC 调用同一个 Python `compile` 回调生成 OM；随后只保留
8+ C++ OPP 交付件加载 OM，验证执行阶段不再导入 Python 插件。
9+ 
10+## 样例链路
11+ 
12+```text
13+Python plugin 加载
14+ -> register_op_impl 识别 compile / declare_launch_args
15+ -> CustomGraphOptimizer 调用 compile(x, y, z)
16+ -> get_compile_platform_info() 查询 NpuArch/SoC
17+ -> BiSheng + llvm-objcopy 生成并拥有 kernel bytes
18+ -> declare_launch_args 发布 kernel launch 描述
19+ ├── run_online.sh -> GE Session 在线执行
20+ └── run_offline.sh -> AIR -> ATC -> OM -> 脱离 Python 插件的 ACL 执行
21+```
22+ 
23+## 前置条件
24+ 
25+- 已安装并配置与 GE 版本匹配的 CANN，先执行
26+ `source /path/to/cann/set_env.sh`。
27+- `cmake`、`atc`、`bisheng`、`llvm-objcopy`、Python 3；离线执行还需要 ACL 开发库。
28+- 在线脚本和离线 OM 执行都需要 NPU；AIR→OM 的 ATC 编译本身可以在 host 上完成。
29+- 当前 kernel 固定使用 float32，输入元素个数必须是 1024 的整数倍。
30+ 
31+## 运行
32+ 
33+### 在线
34+ 
35+```bash
36+source /path/to/cann/set_env.sh
37+cd examples/custom_op/python_compilable_add_custom
38+bash run_online.sh
39+```
40+ 
41+脚本构建 proto/ES wrapper，设置同时包含 OPP 包根和 Python plugin 目录的
42+`ASCEND_CUSTOM_OPP_PATH`，然后运行 Python `Session`。成功时会看到：
43+ 
44+```text
45+PY_COMPILE_MODULE_LOADED=1
46+PY_COMPILE_CALLBACK_ENTER=1 mode=online ...
47+PY_COMPILE_ONLINE_NPU=PASS
48+```
49+ 
50+### 离线
51+ 
52+```bash
53+source /path/to/cann/set_env.sh
54+cd examples/custom_op/python_compilable_add_custom
55+bash run_offline.sh
56+```
57+ 
58+脚本默认使用 `Ascend910B1` 生成 OM；其他芯片可先设置
59+`PYTHON_COMPILABLE_ADD_SOC_VERSION` 为目标 `soc_version`。
60+ 
61+脚本依次执行：
62+ 
63+1. 构建 C++ `REG_OP` 交付件、AIR 构图程序和 ACL OM 执行程序；
64+2. 设置 `ASCEND_CUSTOM_OPP_PATH=<OPP根目录>:<Python插件目录>`，生成 AIR；
65+3. 调用 ATC。ATC 日志必须包含 `PY_COMPILE_CALLBACK_ENTER=1 mode=offline`；
66+4. 清除 Python 插件路径，仅保留 OPP 根目录加载 OM，并检查 `x + y = 3`。
67+ 
68+最后一步由 C++ 程序打印：
69+ 
70+```text
71+PY_COMPILE_OFFLINE_OM=PASS
72+```
73+ 
74+这一步证明 Python compile 只参与模型编译，执行 OM 时不依赖 Python callback。
75+ 
76+## 目录结构
77+ 
78+```text
79+python_compilable_add_custom
80+├── CMakeLists.txt # proto、ES wrapper、离线工具
81+├── run_online.sh # GE Session 在线编译与执行
82+├── run_offline.sh # AIR -> ATC -> OM -> ACL
83+├── kernel/add_custom.asc # compile callback 使用的 Ascend C 源码
84+├── python/es_custom/__init__.py # 生成 ES wrapper 的包入口模板
85+├── proto/add_custom.h # PythonCompilableAddCustom 原型
86+├── proto/add_custom.cc # shape/data type 推导
87+├── src/ge/python_compilable_add_custom.py # compile + declare_launch_args
88+├── src/run.py # 在线构图和执行
89+├── src/offline_graph_build.cc # 离线 AIR 构图
90+└── src/offline_model_exec.cc # 脱离 Python 的 OM 执行
91+```
92+ 
93+## 关键实现说明
94+ 
95+- `compile` 使用 `get_compile_platform_info()` 查询
96+ `get_platform_resource("version", "NpuArch")` 和 `get_soc_version()`；编译失败直接
97+ 传播为图编译失败。
98+- 生成的 `.aicore.o` 只在 Python holder 实例内按 shape/dtype key 缓存；
99+ 当编译目标的 SoC 或 NPU 架构变化时会清空该缓存，避免复用错误平台的二进制。
100+ 本地文件缓存同时包含 SoC、源码内容和 Ascend C 头文件路径，源码或目标平台更新后会重新编译。
101+ `declare_launch_args` 发生 cache miss 时显式报错，不会偷偷重复编译。
102+- `ASCEND_CUSTOM_OPP_PATH` 必须同时包含 OPP 包根目录和 Python plugin 目录：前者供
103+ C++ `REG_OP` 原型发现，后者供 Python custom-op loader 发现。
104+- 在线和离线都复用 GE 已有 `AnnotatedArgs` 下发路径；不新增执行期 Python callback，
105+ 也不把 Python 状态写入 OM 或模型缓存。
106+ 
107+## 验证
108+ 
109+真实链路分别使用 `bash run_online.sh` 和 `bash run_offline.sh` 验证。
Aexamples/custom_op/python_compilable_add_custom/README_en.md+120-0
@@ -0,0 +1,120 @@
1+# Python CompilableAddCustom online and offline compilation sample
2+ 
3+This sample exercises the GE Python `CompilableOp` interface through both
4+compilation paths with one shared Python `compile`/`declare_launch_args`
5+implementation:
6+ 
7+- Online: GE `CustomGraphOptimizer` calls `compile` while a Python `Session`
8+ graph is being built, then executes the graph immediately.
9+- Offline: a C++ graph builder emits AIR, ATC invokes the same Python callback
10+ to produce an OM, and a C++ ACL runner loads that OM after the Python plugin
11+ has been removed from `ASCEND_CUSTOM_OPP_PATH`.
12+ 
13+## Pipeline
14+ 
15+```text
16+Python plugin load
17+ -> register_op_impl discovers compile / declare_launch_args
18+ -> CustomGraphOptimizer calls compile(x, y, z)
19+ -> get_compile_platform_info() reads NpuArch/SoC
20+ -> BiSheng + llvm-objcopy produce owned kernel bytes
21+ -> declare_launch_args publishes the launch descriptor
22+ ├── run_online.sh -> GE Session online execution
23+ └── run_offline.sh -> AIR -> ATC -> OM -> ACL without the Python plugin
24+```
25+ 
26+## Prerequisites
27+ 
28+- A CANN installation matching the GE build (`source /path/to/cann/set_env.sh`).
29+- `cmake`, `atc`, `bisheng`, `llvm-objcopy`, and Python 3; offline execution also
30+ needs the ACL development libraries.
31+- An NPU for online execution and OM execution. AIR-to-OM compilation itself
32+ can run on the host.
33+- The sample kernel is float32-only and requires an element count divisible by
34+ 1024.
35+ 
36+## Run
37+ 
38+### Online
39+ 
40+```bash
41+source /path/to/cann/set_env.sh
42+cd examples/custom_op/python_compilable_add_custom
43+bash run_online.sh
44+```
45+ 
46+The script builds the proto/ES wrapper, sets `ASCEND_CUSTOM_OPP_PATH` to both
47+the OPP root and the Python plugin directory, and runs the Python `Session`. A
48+successful run prints:
49+ 
50+```text
51+PY_COMPILE_MODULE_LOADED=1
52+PY_COMPILE_CALLBACK_ENTER=1 mode=online ...
53+PY_COMPILE_ONLINE_NPU=PASS
54+```
55+ 
56+### Offline
57+ 
58+```bash
59+source /path/to/cann/set_env.sh
60+cd examples/custom_op/python_compilable_add_custom
61+bash run_offline.sh
62+```
63+ 
64+The script defaults to `Ascend910B1`; set `PYTHON_COMPILABLE_ADD_SOC_VERSION`
65+before running it when compiling for another SoC.
66+ 
67+The script:
68+ 
69+1. Builds the C++ `REG_OP` deliverable, AIR exporter, and ACL OM runner.
70+2. Sets `ASCEND_CUSTOM_OPP_PATH=<OPP-root>:<Python-plugin-dir>` and generates AIR.
71+3. Runs ATC; its log must contain `PY_COMPILE_CALLBACK_ENTER=1 mode=offline`.
72+4. Removes the Python plugin path, keeps only the OPP root, loads the OM, and
73+ checks `x + y = 3`.
74+ 
75+The final C++ runner prints:
76+ 
77+```text
78+PY_COMPILE_OFFLINE_OM=PASS
79+```
80+ 
81+This demonstrates that Python `compile` participates in model compilation but
82+is not needed to execute the resulting OM.
83+ 
84+## Layout
85+ 
86+```text
87+python_compilable_add_custom
88+├── CMakeLists.txt # proto, ES wrapper, offline tools
89+├── run_online.sh # online compile and execution
90+├── run_offline.sh # AIR -> ATC -> OM -> ACL
91+├── kernel/add_custom.asc # Ascend C source compiled by callback
92+├── python/es_custom/__init__.py # package entry template for generated ES wrapper
93+├── proto/add_custom.h # PythonCompilableAddCustom prototype
94+├── proto/add_custom.cc # shape/data-type inference
95+├── src/ge/python_compilable_add_custom.py # compile + declare_launch_args
96+├── src/run.py # online graph and execution
97+├── src/offline_graph_build.cc # offline AIR exporter
98+└── src/offline_model_exec.cc # OM execution without Python
99+```
100+ 
101+## Implementation notes
102+ 
103+- `compile` uses `get_compile_platform_info()` to obtain
104+ `get_platform_resource("version", "NpuArch")` and `get_soc_version()`;
105+ compiler failures propagate as graph-compilation failures.
106+- The generated `.aicore.o` is cached by shape/dtype key in the Python holder.
107+ The cache is cleared when the SoC or NPU architecture changes, so a binary
108+ from another target cannot be reused. The on-disk key also includes the
109+ SoC, source contents, and Ascend C include path, so changing the source or
110+ target platform triggers a rebuild. A cache miss in `declare_launch_args` is
111+ an explicit error; launch
112+ declaration never silently recompiles.
113+- `ASCEND_CUSTOM_OPP_PATH` must contain both the OPP root (for the C++ `REG_OP`
114+ prototype) and the Python plugin directory (for the Python custom-op loader).
115+- Both paths reuse GE's existing `AnnotatedArgs` launch mechanism. No execution-
116+ time Python callback is added and no Python state is written into the OM.
117+ 
118+## Verification
119+ 
120+Verify the complete paths with `bash run_online.sh` and `bash run_offline.sh`.
Aexamples/custom_op/python_compilable_add_custom/kernel/add_custom.asc+86-0
@@ -0,0 +1,86 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#include "kernel_operator.h"
12+ 
13+namespace {
14+constexpr int32_t kBlockLength = 1024;
15+constexpr int32_t kTileNum = 8;
16+constexpr int32_t kBufferNum = 2;
17+}
18+ 
19+class KernelAdd {
20+ public:
21+ __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR z) {
22+ x_gm_.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(x) +
23+ kBlockLength * AscendC::GetBlockIdx(),
24+ kBlockLength);
25+ y_gm_.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(y) +
26+ kBlockLength * AscendC::GetBlockIdx(),
27+ kBlockLength);
28+ z_gm_.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(z) +
29+ kBlockLength * AscendC::GetBlockIdx(),
30+ kBlockLength);
31+ tile_length_ = kBlockLength / kTileNum / kBufferNum;
32+ pipe_.InitBuffer(input_x_, kBufferNum, tile_length_ * sizeof(float));
33+ pipe_.InitBuffer(input_y_, kBufferNum, tile_length_ * sizeof(float));
34+ pipe_.InitBuffer(output_z_, kBufferNum, tile_length_ * sizeof(float));
35+ }
36+ 
37+ __aicore__ inline void Process() {
38+ for (int32_t index = 0; index < kTileNum * kBufferNum; ++index) {
39+ CopyIn(index);
40+ Compute();
41+ CopyOut(index);
42+ }
43+ }
44+ 
45+ private:
46+ __aicore__ inline void CopyIn(int32_t index) {
47+ auto x_local = input_x_.AllocTensor<float>();
48+ auto y_local = input_y_.AllocTensor<float>();
49+ AscendC::DataCopy(x_local, x_gm_[index * tile_length_], tile_length_);
50+ AscendC::DataCopy(y_local, y_gm_[index * tile_length_], tile_length_);
51+ input_x_.EnQue(x_local);
52+ input_y_.EnQue(y_local);
53+ }
54+ 
55+ __aicore__ inline void Compute() {
56+ auto x_local = input_x_.DeQue<float>();
57+ auto y_local = input_y_.DeQue<float>();
58+ auto z_local = output_z_.AllocTensor<float>();
59+ AscendC::Add(z_local, x_local, y_local, tile_length_);
60+ output_z_.EnQue(z_local);
61+ input_x_.FreeTensor(x_local);
62+ input_y_.FreeTensor(y_local);
63+ }
64+ 
65+ __aicore__ inline void CopyOut(int32_t index) {
66+ auto z_local = output_z_.DeQue<float>();
67+ AscendC::DataCopy(z_gm_[index * tile_length_], z_local, tile_length_);
68+ output_z_.FreeTensor(z_local);
69+ }
70+ 
71+ AscendC::TPipe pipe_;
72+ AscendC::TQue<AscendC::TPosition::VECIN, kBufferNum> input_x_;
73+ AscendC::TQue<AscendC::TPosition::VECIN, kBufferNum> input_y_;
74+ AscendC::TQue<AscendC::TPosition::VECOUT, kBufferNum> output_z_;
75+ AscendC::GlobalTensor<float> x_gm_;
76+ AscendC::GlobalTensor<float> y_gm_;
77+ AscendC::GlobalTensor<float> z_gm_;
78+ int32_t tile_length_{0};
79+};
80+ 
81+extern "C" __global__ __aicore__ void add_custom(GM_ADDR x, GM_ADDR y, GM_ADDR z) {
82+ KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY);
83+ KernelAdd op;
84+ op.Init(x, y, z);
85+ op.Process();
86+}
Aexamples/custom_op/python_compilable_add_custom/proto/add_custom.cc+29-0
@@ -0,0 +1,29 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#include "add_custom.h"
12+ 
13+namespace {
14+ge::graphStatus InferShape(gert::InferShapeContext *ctx) {
15+ const auto *input_shape = ctx->GetInputShape(0U);
16+ auto *output_shape = ctx->GetOutputShape(0U);
17+ if ((input_shape == nullptr) || (output_shape == nullptr)) {
18+ return ge::GRAPH_FAILED;
19+ }
20+ *output_shape = *input_shape;
21+ return ge::GRAPH_SUCCESS;
22+}
23+ 
24+ge::graphStatus InferDataType(gert::InferDataTypeContext *ctx) {
25+ return ctx->SetOutputDataType(0U, ctx->GetInputDataType(0U));
26+}
27+ 
28+IMPL_OP(PythonCompilableAddCustom).InferShape(InferShape).InferDataType(InferDataType);
29+} // namespace
Aexamples/custom_op/python_compilable_add_custom/proto/add_custom.h+26-0
@@ -0,0 +1,26 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#ifndef EXAMPLES_CUSTOM_OP_PYTHON_COMPILABLE_ADD_CUSTOM_PROTO_ADD_CUSTOM_H_
12+#define EXAMPLES_CUSTOM_OP_PYTHON_COMPILABLE_ADD_CUSTOM_PROTO_ADD_CUSTOM_H_
13+ 
14+#include "graph/operator_reg.h"
15+#include "register/op_impl_registry.h"
16+ 
17+namespace ge {
18+REG_OP(PythonCompilableAddCustom)
19+ .INPUT(x1, "T")
20+ .INPUT(x2, "T")
21+ .OUTPUT(y, "T")
22+ .DATATYPE(T, TensorType({DT_FLOAT}))
23+ .OP_END_FACTORY_REG(PythonCompilableAddCustom);
24+} // namespace ge
25+ 
26+#endif // EXAMPLES_CUSTOM_OP_PYTHON_COMPILABLE_ADD_CUSTOM_PROTO_ADD_CUSTOM_H_
Aexamples/custom_op/python_compilable_add_custom/python/es_custom/__init__.py+24-0
@@ -0,0 +1,24 @@
1+#!/usr/bin/env python3
2+# -*- coding: utf-8 -*-
3+# -----------------------------------------------------------------------------------------------------------
4+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
5+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
6+# CANN Open Software License Agreement Version 2.0 (the "License").
7+# Please refer to the License for details. You may not use this file except in compliance with the License.
8+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
9+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
10+# See LICENSE in the root of the software repository for the full text of the License.
11+# -----------------------------------------------------------------------------------------------------------
12+"""Package wrapper for the ES code generated by GE's gen_esb."""
13+ 
14+__all__ = [name for name in dir() if not name.startswith("_")]
15+ 
16+from .es_custom_ops import * # noqa: F403
17+ 
18+ 
19+def get_module():
20+ """Return this module for GE's ``ge.es.plugins`` entry point."""
21+ 
22+ import sys
23+ 
24+ return sys.modules[__name__]
Aexamples/custom_op/python_compilable_add_custom/run_offline.sh+119-0
@@ -0,0 +1,119 @@
1+#!/usr/bin/env bash
2+# -----------------------------------------------------------------------------------------------------------
3+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
4+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
5+# CANN Open Software License Agreement Version 2.0 (the "License").
6+# Please refer to the License for details. You may not use this file except in compliance with the License.
7+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
8+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
9+# See LICENSE in the root of the software repository for the full text of the License.
10+# -----------------------------------------------------------------------------------------------------------
11+ 
12+set -euo pipefail
13+ 
14+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
15+BUILD_DIR="${SCRIPT_DIR}/build/offline"
16+OUTPUT_DIR="${BUILD_DIR}/output"
17+AIR_PATH="${OUTPUT_DIR}/python_compilable_add.air"
18+OM_PATH="${OUTPUT_DIR}/python_compilable_add.om"
19+CUSTOM_OPP_ROOT="${BUILD_DIR}/custom_op_package"
20+PLUGIN_DIR="${SCRIPT_DIR}/src/ge"
21+LOG_FILE="${BUILD_DIR}/offline_compile.log"
22+ATC_BIN="${ASCEND_HOME_PATH:-}/bin/atc"
23+ 
24+detect_os() {
25+ case "$(uname -s | tr '[:upper:]' '[:lower:]')" in
26+ mingw*|msys*|cygwin*) echo "windows" ;;
27+ *) echo "linux" ;;
28+ esac
29+}
30+ 
31+detect_arch() {
32+ case "$(uname -m | tr '[:upper:]' '[:lower:]')" in
33+ arm64) echo "aarch64" ;;
34+ amd64) echo "x86_64" ;;
35+ *) uname -m | tr '[:upper:]' '[:lower:]' ;;
36+ esac
37+}
38+ 
39+require_command() {
40+ if ! command -v "$1" >/dev/null 2>&1; then
41+ echo "[ERROR] required command not found: $1" >&2
42+ exit 1
43+ fi
44+}
45+ 
46+if [[ -z "${ASCEND_HOME_PATH:-}" || ! -d "${ASCEND_HOME_PATH}" ]]; then
47+ echo "[ERROR] ASCEND_HOME_PATH is empty. Please source CANN set_env.sh first." >&2
48+ exit 1
49+fi
50+for command_name in cmake python3 bisheng llvm-objcopy; do
51+ require_command "${command_name}"
52+done
53+if [[ ! -x "${ATC_BIN}" ]]; then
54+ require_command atc
55+ ATC_BIN="$(command -v atc)"
56+fi
57+ 
58+HOST_OS="$(detect_os)"
59+HOST_ARCH="$(detect_arch)"
60+CUSTOM_OPP_LIB_DIR="${CUSTOM_OPP_ROOT}/op_graph/lib/${HOST_OS}/${HOST_ARCH}"
61+CUSTOM_OPP_PROTO="${BUILD_DIR}/libpython_compilable_add_custom_op_proto.so"
62+SOC_VERSION="${PYTHON_COMPILABLE_ADD_SOC_VERSION:-Ascend910B1}"
63+ 
64+mkdir -p "${OUTPUT_DIR}" "${CUSTOM_OPP_LIB_DIR}"
65+ 
66+echo "[INFO] configure and build the offline exporter, OM runner, and proto"
67+cmake -S "${SCRIPT_DIR}" -B "${BUILD_DIR}" \
68+ -DCMAKE_BUILD_TYPE=Release \
69+ -DPYTHON_COMPILABLE_ADD_BUILD_ES=OFF
70+cmake --build "${BUILD_DIR}" --target python_compilable_add_custom_op_proto \
71+ python_compilable_add_graph_build python_compilable_add_model_exec \
72+ -j"$(nproc 2>/dev/null || echo 8)"
73+ 
74+if [[ ! -s "${CUSTOM_OPP_PROTO}" ]]; then
75+ echo "[ERROR] proto library was not generated: ${CUSTOM_OPP_PROTO}" >&2
76+ exit 1
77+fi
78+cp -f "${CUSTOM_OPP_PROTO}" "${CUSTOM_OPP_LIB_DIR}/libcust_opapi.so"
79+ 
80+# ATC needs both halves of the package: the OPP root for the C++ proto and the
81+# Python directory for the compile/launch implementation.
82+export ASCEND_CUSTOM_OPP_PATH="${CUSTOM_OPP_ROOT}:${PLUGIN_DIR}${ASCEND_CUSTOM_OPP_PATH:+:${ASCEND_CUSTOM_OPP_PATH}}"
83+export PYTHON_COMPILABLE_ADD_CUSTOM_BUILD_DIR="${BUILD_DIR}/python_compile"
84+export PYTHON_COMPILABLE_ADD_CUSTOM_MODE=offline
85+ 
86+echo "[INFO] generate AIR"
87+rm -f "${AIR_PATH}" "${OM_PATH}"
88+(
89+ cd "${OUTPUT_DIR}"
90+ "${BUILD_DIR}/python_compilable_add_graph_build"
91+)
92+if [[ ! -s "${AIR_PATH}" ]]; then
93+ echo "[ERROR] AIR was not generated: ${AIR_PATH}" >&2
94+ exit 1
95+fi
96+ 
97+echo "[INFO] compile AIR with ATC; Python compile callback is expected"
98+"${ATC_BIN}" \
99+ --model="${AIR_PATH}" \
100+ --framework=1 \
101+ --output="${OM_PATH%.*}" \
102+ --soc_version="${SOC_VERSION}" 2>&1 | tee "${LOG_FILE}"
103+if [[ ! -s "${OM_PATH}" ]]; then
104+ echo "[ERROR] OM was not generated: ${OM_PATH}" >&2
105+ exit 1
106+fi
107+grep -Fq "PY_COMPILE_MODULE_LOADED=1" "${LOG_FILE}"
108+grep -Fq "PY_COMPILE_CALLBACK_ENTER=1 mode=offline" "${LOG_FILE}"
109+ 
110+echo "[INFO] execute OM without the Python plugin"
111+unset ASCEND_CUSTOM_OPP_PATH
112+export ASCEND_CUSTOM_OPP_PATH="${CUSTOM_OPP_ROOT}"
113+unset PYTHON_COMPILABLE_ADD_CUSTOM_MODE
114+(
115+ cd "${OUTPUT_DIR}"
116+ "${BUILD_DIR}/python_compilable_add_model_exec" "${OM_PATH}"
117+) 2>&1 | tee "${BUILD_DIR}/offline_run.log"
118+grep -Fq "PY_COMPILE_OFFLINE_OM=PASS" "${BUILD_DIR}/offline_run.log"
119+echo "[INFO] Python compilable custom-op offline sample PASS"
Aexamples/custom_op/python_compilable_add_custom/run_online.sh+88-0
@@ -0,0 +1,88 @@
1+#!/usr/bin/env bash
2+# -----------------------------------------------------------------------------------------------------------
3+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
4+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
5+# CANN Open Software License Agreement Version 2.0 (the "License").
6+# Please refer to the License for details. You may not use this file except in compliance with the License.
7+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
8+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
9+# See LICENSE in the root of the software repository for the full text of the License.
10+# -----------------------------------------------------------------------------------------------------------
11+ 
12+set -euo pipefail
13+ 
14+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
15+BUILD_DIR="${SCRIPT_DIR}/build"
16+PLUGIN_DIR="${SCRIPT_DIR}/src/ge"
17+ES_LIB_DIR="${BUILD_DIR}/es_output/lib64"
18+ES_PYTHON_PACKAGE_DIR="${BUILD_DIR}/python_package"
19+ES_GENERATED_CODE_DIR="${BUILD_DIR}/es_custom_build/generated_code"
20+CUSTOM_OPP_ROOT="${BUILD_DIR}/custom_op_package"
21+LOG_FILE="${BUILD_DIR}/online_compile.log"
22+ 
23+detect_os() {
24+ case "$(uname -s | tr '[:upper:]' '[:lower:]')" in
25+ mingw*|msys*|cygwin*) echo "windows" ;;
26+ *) echo "linux" ;;
27+ esac
28+}
29+ 
30+detect_arch() {
31+ case "$(uname -m | tr '[:upper:]' '[:lower:]')" in
32+ arm64) echo "aarch64" ;;
33+ amd64) echo "x86_64" ;;
34+ *) uname -m | tr '[:upper:]' '[:lower:]' ;;
35+ esac
36+}
37+ 
38+require_command() {
39+ if ! command -v "$1" >/dev/null 2>&1; then
40+ echo "[ERROR] required command not found: $1" >&2
41+ exit 1
42+ fi
43+}
44+ 
45+if [[ -z "${ASCEND_HOME_PATH:-}" || ! -d "${ASCEND_HOME_PATH}" ]]; then
46+ echo "[ERROR] ASCEND_HOME_PATH is empty. Please source CANN set_env.sh first." >&2
47+ exit 1
48+fi
49+for command_name in cmake python3 bisheng llvm-objcopy; do
50+ require_command "${command_name}"
51+done
52+ 
53+HOST_OS="$(detect_os)"
54+HOST_ARCH="$(detect_arch)"
55+CUSTOM_OPP_LIB_DIR="${CUSTOM_OPP_ROOT}/op_graph/lib/${HOST_OS}/${HOST_ARCH}"
56+CUSTOM_OPP_PROTO="${BUILD_DIR}/libpython_compilable_add_custom_op_proto.so"
57+ 
58+mkdir -p "${BUILD_DIR}" "${CUSTOM_OPP_LIB_DIR}"
59+ 
60+echo "[INFO] configure and build the GE Python ES package"
61+cmake -S "${SCRIPT_DIR}" -B "${BUILD_DIR}" -DCMAKE_BUILD_TYPE=Release
62+cmake --build "${BUILD_DIR}" --target install_es_custom -j"$(nproc 2>/dev/null || echo 8)"
63+ 
64+if [[ ! -s "${CUSTOM_OPP_PROTO}" || ! -s "${ES_LIB_DIR}/libes_custom.so" ]]; then
65+ echo "[ERROR] C++ proto library or ES shared library was not generated" >&2
66+ exit 1
67+fi
68+cp -f "${CUSTOM_OPP_PROTO}" "${CUSTOM_OPP_LIB_DIR}/libcust_opapi.so"
69+ 
70+echo "[INFO] stage the generated ES Python wrapper without network access"
71+rm -rf "${ES_PYTHON_PACKAGE_DIR}"
72+mkdir -p "${ES_PYTHON_PACKAGE_DIR}/es_custom"
73+cp "${SCRIPT_DIR}/python/es_custom/__init__.py" "${ES_PYTHON_PACKAGE_DIR}/es_custom/"
74+cp "${ES_GENERATED_CODE_DIR}"/*.py "${ES_PYTHON_PACKAGE_DIR}/es_custom/"
75+export PYTHONPATH="${ES_PYTHON_PACKAGE_DIR}${PYTHONPATH:+:${PYTHONPATH}}"
76+export LD_LIBRARY_PATH="${ES_LIB_DIR}${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}"
77+# The OPP root is for the C++ proto; the plugin directory is for the Python
78+# custom-op loader. Both are required by GE's two loaders.
79+export ASCEND_CUSTOM_OPP_PATH="${CUSTOM_OPP_ROOT}:${PLUGIN_DIR}${ASCEND_CUSTOM_OPP_PATH:+:${ASCEND_CUSTOM_OPP_PATH}}"
80+export PYTHON_COMPILABLE_ADD_CUSTOM_BUILD_DIR="${BUILD_DIR}/online_compile"
81+export PYTHON_COMPILABLE_ADD_CUSTOM_MODE=online
82+ 
83+echo "[INFO] run the online GE graph; Python compile is invoked by CustomGraphOptimizer"
84+python3 "${SCRIPT_DIR}/src/run.py" 2>&1 | tee "${LOG_FILE}"
85+grep -Fq "PY_COMPILE_MODULE_LOADED=1" "${LOG_FILE}"
86+grep -Fq "PY_COMPILE_CALLBACK_ENTER=1 mode=online" "${LOG_FILE}"
87+grep -Fq "PY_COMPILE_ONLINE_NPU=PASS" "${LOG_FILE}"
88+echo "[INFO] Python compilable custom-op online sample PASS"
Aexamples/custom_op/python_compilable_add_custom/src/ge/python_compilable_add_custom.py+254-0
@@ -0,0 +1,254 @@
1+#!/usr/bin/env python3
2+# -*- coding: utf-8 -*-
3+# -----------------------------------------------------------------------------------------------------------
4+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
5+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
6+# CANN Open Software License Agreement Version 2.0 (the "License").
7+# Please refer to the License for details. You may not use this file except in compliance with the License.
8+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
9+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
10+# See LICENSE in the root of the software repository for the full text of the License.
11+# -----------------------------------------------------------------------------------------------------------
12+"""Python compile/launch implementation for the GE Add sample.
13+ 
14+The callback deliberately keeps only owned values (the metadata key and the
15+kernel bytes). Borrowed Tensor and OpCompileContext objects never escape the
16+callback that created them.
17+"""
18+ 
19+from __future__ import annotations
20+ 
21+import hashlib
22+import os
23+import subprocess
24+import tempfile
25+import threading
26+from dataclasses import dataclass
27+from pathlib import Path
28+ 
29+from ge.custom_op import (
30+ AnnotatedKernelLaunchInfo,
31+ get_compile_platform_info,
32+ get_declare_launch_args_ctx,
33+ register_op_impl,
34+)
35+from ge.runtime import Tensor
36+ 
37+ 
38+_KERNEL_NAME = "add_custom"
39+_KERNEL_BLOCK_SIZE = 1024
40+_KERNEL_SOURCE = Path(__file__).resolve().parents[2] / "kernel" / "add_custom.asc"
41+ 
42+ 
43+@dataclass(frozen=True)
44+class KernelArtifact:
45+ """Owned launch data produced by the online compile callback."""
46+ 
47+ kernel_name: str
48+ kernel_bin: bytes
49+ block_dim: int
50+ 
51+ 
52+def _shape_dims(tensor: Tensor) -> tuple[int, ...]:
53+ return tuple(int(dim) for dim in tensor.storage_shape.dims)
54+ 
55+ 
56+def _tensor_key(x: Tensor, y: Tensor, z: Tensor) -> tuple:
57+ return (
58+ _shape_dims(x),
59+ _shape_dims(y),
60+ _shape_dims(z),
61+ str(x.data_type),
62+ str(y.data_type),
63+ str(z.data_type),
64+ )
65+ 
66+ 
67+def _normalise_npu_arch(value: str) -> str:
68+ arch = str(value).strip()
69+ if arch.startswith("dav-"):
70+ arch = arch[4:]
71+ if not arch.isdigit():
72+ raise ValueError(
73+ "NpuArch must be a numeric architecture, got {!r}".format(value)
74+ )
75+ return arch
76+ 
77+ 
78+def _build_output_dir() -> Path:
79+ configured = os.environ.get("PYTHON_COMPILABLE_ADD_CUSTOM_BUILD_DIR")
80+ if configured:
81+ output_dir = Path(configured)
82+ else:
83+ output_dir = Path(tempfile.gettempdir()) / "python_compilable_add_custom"
84+ output_dir.mkdir(parents=True, exist_ok=True)
85+ return output_dir
86+ 
87+ 
88+def _find_ascendc_include(ascend_home: str) -> Path:
89+ root = Path(ascend_home)
90+ candidates = (
91+ root / "asc" / "include",
92+ root / "aarch64-linux" / "asc" / "include",
93+ root / "x86_64-linux" / "asc" / "include",
94+ )
95+ for candidate in candidates:
96+ if candidate.is_dir():
97+ return candidate
98+ raise FileNotFoundError(
99+ "Ascend C include directory not found under: {}".format(root)
100+ )
101+ 
102+ 
103+def _compile_kernel(
104+ source_path: Path, npu_arch: str, soc_version: str
105+) -> KernelArtifact:
106+ """Compile the sample Ascend C source and own its device binary bytes."""
107+ 
108+ if not source_path.is_file():
109+ raise FileNotFoundError("kernel source not found: {}".format(source_path))
110+ ascend_home = os.environ.get("ASCEND_HOME_PATH")
111+ if not ascend_home:
112+ raise RuntimeError("ASCEND_HOME_PATH is required to locate Ascend C headers")
113+ ascendc_include = _find_ascendc_include(ascend_home)
114+ output_dir = _build_output_dir()
115+ digest_input = b"\0".join(
116+ (
117+ str(source_path).encode("utf-8"),
118+ str(ascendc_include).encode("utf-8"),
119+ npu_arch.encode("utf-8"),
120+ soc_version.encode("utf-8"),
121+ source_path.read_bytes(),
122+ )
123+ )
124+ digest = hashlib.sha256(digest_input).hexdigest()[:16]
125+ host_object = output_dir / ("add_custom_{}.host.o".format(digest))
126+ device_binary = output_dir / ("add_custom_{}.aicore.o".format(digest))
127+ if not device_binary.is_file() or device_binary.stat().st_size == 0:
128+ subprocess.run(
129+ [
130+ "bisheng",
131+ "-c",
132+ str(source_path),
133+ "-o",
134+ str(host_object),
135+ "--npu-arch=dav-{}".format(npu_arch),
136+ "-I{}".format(ascendc_include),
137+ ],
138+ check=True,
139+ )
140+ subprocess.run(
141+ [
142+ "llvm-objcopy",
143+ "-O",
144+ "binary",
145+ "--only-section=.aicore_binary",
146+ str(host_object),
147+ str(device_binary),
148+ ],
149+ check=True,
150+ )
151+ kernel_bin = device_binary.read_bytes()
152+ if not kernel_bin:
153+ raise RuntimeError("empty Ascend C kernel binary: {}".format(device_binary))
154+ return KernelArtifact(
155+ kernel_name=_KERNEL_NAME,
156+ kernel_bin=bytes(kernel_bin),
157+ block_dim=1,
158+ )
159+ 
160+ 
161+def _validate_tensors(x: Tensor, y: Tensor, z: Tensor) -> None:
162+ if _shape_dims(x) != _shape_dims(y) or _shape_dims(x) != _shape_dims(z):
163+ raise ValueError("PythonCompilableAddCustom requires matching tensor shapes")
164+ if str(x.data_type) != str(y.data_type) or str(x.data_type) != str(z.data_type):
165+ raise ValueError("PythonCompilableAddCustom requires matching data types")
166+ try:
167+ dtype_value = int(x.data_type)
168+ except (TypeError, ValueError):
169+ dtype_value = str(x.data_type)
170+ if dtype_value not in (0, "0", "DT_FLOAT", "DataType.DT_FLOAT", "float32"):
171+ raise ValueError("the sample kernel supports only float32 tensors")
172+ if not _shape_dims(x) or any(dim <= 0 for dim in _shape_dims(x)):
173+ raise ValueError("PythonCompilableAddCustom requires a concrete positive shape")
174+ element_count = 1
175+ for dim in _shape_dims(x):
176+ element_count *= dim
177+ if element_count % _KERNEL_BLOCK_SIZE != 0:
178+ raise ValueError(
179+ "the sample kernel requires an element count divisible by {}".format(
180+ _KERNEL_BLOCK_SIZE
181+ )
182+ )
183+ 
184+ 
185+@register_op_impl(op_type="PythonCompilableAddCustom")
186+class PythonCompilableAddCustom:
187+ """Compile an Add kernel and publish it to the AnnotatedArgs path."""
188+ 
189+ def __init__(self) -> None:
190+ self._artifacts: dict[tuple, KernelArtifact] = {}
191+ self._platform_key: tuple[str, str] | None = None
192+ self._lock = threading.RLock()
193+ 
194+ def compile(self, x: Tensor, y: Tensor, z: Tensor) -> None:
195+ _validate_tensors(x, y, z)
196+ platform_info = get_compile_platform_info()
197+ npu_arch = _normalise_npu_arch(
198+ platform_info.get_platform_resource("version", "NpuArch")
199+ )
200+ soc_version = str(platform_info.get_soc_version()).strip()
201+ platform_key = (soc_version, npu_arch)
202+ key = _tensor_key(x, y, z)
203+ with self._lock:
204+ if self._platform_key != platform_key:
205+ self._artifacts.clear()
206+ self._platform_key = platform_key
207+ if key in self._artifacts:
208+ return None
209+ print(
210+ "PY_COMPILE_CALLBACK_ENTER=1 mode={} soc={} arch={} shape={}".format(
211+ os.environ.get("PYTHON_COMPILABLE_ADD_CUSTOM_MODE", "unknown"),
212+ soc_version,
213+ npu_arch,
214+ _shape_dims(x),
215+ ),
216+ flush=True,
217+ )
218+ artifact = _compile_kernel(_KERNEL_SOURCE, npu_arch, soc_version)
219+ element_count = 1
220+ for dim in _shape_dims(x):
221+ element_count *= dim
222+ self._artifacts[key] = KernelArtifact(
223+ kernel_name=artifact.kernel_name,
224+ kernel_bin=artifact.kernel_bin,
225+ block_dim=element_count // _KERNEL_BLOCK_SIZE,
226+ )
227+ return None
228+ 
229+ def declare_launch_args(self, x: Tensor, y: Tensor, z: Tensor) -> None:
230+ _validate_tensors(x, y, z)
231+ key = _tensor_key(x, y, z)
232+ with self._lock:
233+ artifact = self._artifacts.get(key)
234+ if artifact is None:
235+ raise RuntimeError("PythonCompilableAddCustom compile cache miss")
236+ 
237+ ctx = get_declare_launch_args_ctx()
238+ args = ctx.create_kernel_args()
239+ args.append_input(0, x)
240+ args.append_input(1, y)
241+ args.append_output(0, z)
242+ ctx.add_launch(
243+ AnnotatedKernelLaunchInfo(
244+ kernel_name=artifact.kernel_name,
245+ kernel_bin=artifact.kernel_bin,
246+ block_dim=artifact.block_dim,
247+ stream_id=ctx.get_stream_id(),
248+ ),
249+ args,
250+ )
251+ return None
252+ 
253+ 
254+print("PY_COMPILE_MODULE_LOADED=1", flush=True)
Aexamples/custom_op/python_compilable_add_custom/src/offline_graph_build.cc+58-0
@@ -0,0 +1,58 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#include <cstdint>
12+#include <iostream>
13+#include <vector>
14+ 
15+#include "add_custom.h"
16+#include "graph.h"
17+#include "ops_proto_legacy.h"
18+#include "tensor.h"
19+#include "types.h"
20+ 
21+namespace ge {
22+namespace {
23+ 
24+constexpr int64_t kNumElements = 1024;
25+ 
26+bool BuildGraph(Graph &graph) {
27+ const std::vector<int64_t> shape = {kNumElements};
28+ TensorDesc data_desc(Shape(shape), FORMAT_ND, DT_FLOAT);
29+ 
30+ auto data_x = op::Data("data_x");
31+ data_x.update_input_desc_x(data_desc);
32+ data_x.update_output_desc_y(data_desc);
33+ 
34+ auto data_y = op::Data("data_y");
35+ data_y.update_input_desc_x(data_desc);
36+ data_y.update_output_desc_y(data_desc);
37+ 
38+ auto add = op::PythonCompilableAddCustom("python_compilable_add").set_input_x1(data_x).set_input_x2(data_y);
39+ add.update_output_desc_y(data_desc);
40+ 
41+ std::vector<Operator> inputs{data_x, data_y};
42+ std::vector<Operator> outputs{add};
43+ graph.SetInputs(inputs).SetOutputs(outputs);
44+ return graph.SaveToFile("./python_compilable_add.air") == GRAPH_SUCCESS;
45+}
46+ 
47+} // namespace
48+} // namespace ge
49+ 
50+int main() {
51+ ge::Graph graph("PythonCompilableAddOfflineGraph");
52+ if (!ge::BuildGraph(graph)) {
53+ std::cerr << "failed to generate AIR" << std::endl;
54+ return 1;
55+ }
56+ std::cout << "PY_COMPILE_AIR_BUILD=PASS" << std::endl;
57+ return 0;
58+}
Aexamples/custom_op/python_compilable_add_custom/src/offline_model_exec.cc+235-0
@@ -0,0 +1,235 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#include <cmath>
12+#include <cstddef>
13+#include <iostream>
14+#include <vector>
15+ 
16+#include "acl/acl.h"
17+ 
18+namespace {
19+ 
20+constexpr int kExpectedArgc = 2;
21+constexpr int32_t kDeviceId = 0;
22+constexpr float kExpectedValue = 3.0F;
23+ 
24+struct DeviceBuffer {
25+ void *address = nullptr;
26+};
27+ 
28+void DestroyDataset(aclmdlDataset *dataset, std::vector<DeviceBuffer> &buffers) {
29+ if (dataset != nullptr) {
30+ for (size_t i = 0U; i < aclmdlGetDatasetNumBuffers(dataset); ++i) {
31+ aclDataBuffer *data_buffer = aclmdlGetDatasetBuffer(dataset, i);
32+ if (data_buffer != nullptr) {
33+ (void)aclDestroyDataBuffer(data_buffer);
34+ }
35+ }
36+ (void)aclmdlDestroyDataset(dataset);
37+ }
38+ for (auto &buffer : buffers) {
39+ if (buffer.address != nullptr) {
40+ (void)aclrtFree(buffer.address);
41+ buffer.address = nullptr;
42+ }
43+ }
44+ buffers.clear();
45+}
46+ 
47+bool AddInput(aclmdlDataset *dataset, size_t size, size_t index, std::vector<DeviceBuffer> &buffers) {
48+ if (size == 0U || (size % sizeof(float)) != 0U) {
49+ return false;
50+ }
51+ void *address = nullptr;
52+ if (aclrtMalloc(&address, size, ACL_MEM_MALLOC_NORMAL_ONLY) != ACL_SUCCESS) {
53+ return false;
54+ }
55+ std::vector<float> host_data(size / sizeof(float), static_cast<float>(index + 1U));
56+ if (aclrtMemcpy(address, size, host_data.data(), size, ACL_MEMCPY_HOST_TO_DEVICE) != ACL_SUCCESS) {
57+ (void)aclrtFree(address);
58+ return false;
59+ }
60+ aclDataBuffer *data_buffer = aclCreateDataBuffer(address, size);
61+ if (data_buffer == nullptr || aclmdlAddDatasetBuffer(dataset, data_buffer) != ACL_SUCCESS) {
62+ if (data_buffer != nullptr) {
63+ (void)aclDestroyDataBuffer(data_buffer);
64+ }
65+ (void)aclrtFree(address);
66+ return false;
67+ }
68+ buffers.push_back({address});
69+ return true;
70+}
71+ 
72+bool CheckOutput(aclDataBuffer *data_buffer, size_t size) {
73+ if (data_buffer == nullptr || size == 0U || (size % sizeof(float)) != 0U) {
74+ return false;
75+ }
76+ std::vector<float> host_data(size / sizeof(float));
77+ if (aclrtMemcpy(host_data.data(), size, aclGetDataBufferAddr(data_buffer), size, ACL_MEMCPY_DEVICE_TO_HOST) !=
78+ ACL_SUCCESS) {
79+ return false;
80+ }
81+ for (const float value : host_data) {
82+ if (std::fabs(value - kExpectedValue) > 1.0e-5F) {
83+ return false;
84+ }
85+ }
86+ return true;
87+}
88+ 
89+class AclResourceGuard {
90+ public:
91+ ~AclResourceGuard() {
92+ DestroyDataset(input_dataset_, input_buffers_);
93+ DestroyDataset(output_dataset_, output_buffers_);
94+ if (model_desc_ != nullptr) {
95+ (void)aclmdlDestroyDesc(model_desc_);
96+ }
97+ if (model_id_ != 0U) {
98+ (void)aclmdlUnload(model_id_);
99+ }
100+ if (device_set_) {
101+ (void)aclrtResetDevice(kDeviceId);
102+ }
103+ if (acl_initialized_) {
104+ (void)aclFinalize();
105+ }
106+ }
107+ 
108+ bool Initialize() {
109+ if (aclInit(nullptr) != ACL_SUCCESS) {
110+ return false;
111+ }
112+ acl_initialized_ = true;
113+ if (aclrtSetDevice(kDeviceId) != ACL_SUCCESS) {
114+ return false;
115+ }
116+ device_set_ = true;
117+ return true;
118+ }
119+ 
120+ void SetModelId(uint32_t model_id) {
121+ model_id_ = model_id;
122+ }
123+ void SetModelDesc(aclmdlDesc *model_desc) {
124+ model_desc_ = model_desc;
125+ }
126+ void SetInputDataset(aclmdlDataset *dataset) {
127+ input_dataset_ = dataset;
128+ }
129+ void SetOutputDataset(aclmdlDataset *dataset) {
130+ output_dataset_ = dataset;
131+ }
132+ std::vector<DeviceBuffer> &InputBuffers() {
133+ return input_buffers_;
134+ }
135+ std::vector<DeviceBuffer> &OutputBuffers() {
136+ return output_buffers_;
137+ }
138+ 
139+ private:
140+ bool acl_initialized_ = false;
141+ bool device_set_ = false;
142+ uint32_t model_id_ = 0U;
143+ aclmdlDesc *model_desc_ = nullptr;
144+ aclmdlDataset *input_dataset_ = nullptr;
145+ aclmdlDataset *output_dataset_ = nullptr;
146+ std::vector<DeviceBuffer> input_buffers_;
147+ std::vector<DeviceBuffer> output_buffers_;
148+};
149+ 
150+bool LoadModel(const char *model_path, AclResourceGuard &resources, uint32_t &model_id, aclmdlDesc *&model_desc) {
151+ if (aclmdlLoadFromFile(model_path, &model_id) != ACL_SUCCESS) {
152+ return false;
153+ }
154+ resources.SetModelId(model_id);
155+ model_desc = aclmdlCreateDesc();
156+ resources.SetModelDesc(model_desc);
157+ return model_desc != nullptr && aclmdlGetDesc(model_desc, model_id) == ACL_SUCCESS;
158+}
159+ 
160+bool FillInputDataset(aclmdlDesc *model_desc, aclmdlDataset *input_dataset, std::vector<DeviceBuffer> &input_buffers) {
161+ const size_t input_count = aclmdlGetNumInputs(model_desc);
162+ if (input_count != 2U) {
163+ return false;
164+ }
165+ for (size_t i = 0U; i < input_count; ++i) {
166+ if (!AddInput(input_dataset, aclmdlGetInputSizeByIndex(model_desc, i), i, input_buffers)) {
167+ return false;
168+ }
169+ }
170+ return input_buffers.size() == input_count;
171+}
172+ 
173+bool CreateOutputDataset(aclmdlDesc *model_desc, aclmdlDataset *output_dataset,
174+ std::vector<DeviceBuffer> &output_buffers, aclDataBuffer *&output_buffer,
175+ size_t &output_size) {
176+ output_size = aclmdlGetOutputSizeByIndex(model_desc, 0U);
177+ void *output_address = nullptr;
178+ if (aclrtMalloc(&output_address, output_size, ACL_MEM_MALLOC_NORMAL_ONLY) != ACL_SUCCESS) {
179+ return false;
180+ }
181+ output_buffer = aclCreateDataBuffer(output_address, output_size);
182+ if (output_buffer == nullptr || aclmdlAddDatasetBuffer(output_dataset, output_buffer) != ACL_SUCCESS) {
183+ if (output_buffer != nullptr) {
184+ (void)aclDestroyDataBuffer(output_buffer);
185+ }
186+ (void)aclrtFree(output_address);
187+ return false;
188+ }
189+ output_buffers.push_back({output_address});
190+ return true;
191+}
192+ 
193+bool ExecuteModel(uint32_t model_id, aclmdlDataset *input_dataset, aclmdlDataset *output_dataset,
194+ aclDataBuffer *output_buffer, size_t output_size) {
195+ return aclmdlExecute(model_id, input_dataset, output_dataset) == ACL_SUCCESS &&
196+ CheckOutput(output_buffer, output_size);
197+}
198+ 
199+} // namespace
200+ 
201+int RunModel(const char *model_path) {
202+ AclResourceGuard resources;
203+ if (!resources.Initialize()) {
204+ return 1;
205+ }
206+ uint32_t model_id = 0U;
207+ aclmdlDesc *model_desc = nullptr;
208+ if (!LoadModel(model_path, resources, model_id, model_desc)) {
209+ return 1;
210+ }
211+ aclmdlDataset *input_dataset = aclmdlCreateDataset();
212+ resources.SetInputDataset(input_dataset);
213+ aclmdlDataset *output_dataset = aclmdlCreateDataset();
214+ resources.SetOutputDataset(output_dataset);
215+ if (input_dataset == nullptr || output_dataset == nullptr ||
216+ !FillInputDataset(model_desc, input_dataset, resources.InputBuffers())) {
217+ return 1;
218+ }
219+ aclDataBuffer *output_buffer = nullptr;
220+ size_t output_size = 0U;
221+ if (!CreateOutputDataset(model_desc, output_dataset, resources.OutputBuffers(), output_buffer, output_size) ||
222+ !ExecuteModel(model_id, input_dataset, output_dataset, output_buffer, output_size)) {
223+ return 1;
224+ }
225+ std::cout << "PY_COMPILE_OFFLINE_OM=PASS" << std::endl;
226+ return 0;
227+}
228+ 
229+int main(int argc, char *argv[]) {
230+ if (argc != kExpectedArgc) {
231+ std::cerr << "usage: " << argv[0] << " <model_path>" << std::endl;
232+ return 1;
233+ }
234+ return RunModel(argv[1]);
235+}
Aexamples/custom_op/python_compilable_add_custom/src/run.py+130-0
@@ -0,0 +1,130 @@
1+#!/usr/bin/env python3
2+# -*- coding: utf-8 -*-
3+# -----------------------------------------------------------------------------------------------------------
4+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
5+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
6+# CANN Open Software License Agreement Version 2.0 (the "License").
7+# Please refer to the License for details. You may not use this file except in compliance with the License.
8+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
9+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
10+# See LICENSE in the root of the software repository for the full text of the License.
11+# -----------------------------------------------------------------------------------------------------------
12+"""Build and execute the Python-compilable Add graph in online GE mode."""
13+ 
14+from __future__ import annotations
15+ 
16+import traceback
17+from typing import List
18+ 
19+from ge.es.graph_builder import GraphBuilder
20+from ge.ge_global import GeApi
21+from ge.graph import Tensor
22+from ge.graph.types import DataType, Format
23+from ge.session import Session
24+ 
25+ 
26+GRAPH_ID = 0
27+DEVICE_ID = 0
28+NUM_ELEMENTS = 1024
29+ 
30+ 
31+def build_graph():
32+ builder = GraphBuilder("python_compilable_add_graph")
33+ input_x = builder.create_input(
34+ index=0,
35+ name="data_x",
36+ data_type=DataType.DT_FLOAT,
37+ format=Format.FORMAT_ND,
38+ shape=[NUM_ELEMENTS],
39+ )
40+ input_y = builder.create_input(
41+ index=1,
42+ name="data_y",
43+ data_type=DataType.DT_FLOAT,
44+ format=Format.FORMAT_ND,
45+ shape=[NUM_ELEMENTS],
46+ )
47+ # The run scripts stage the generated ES module directly. Importing it
48+ # by package name keeps the sample independent of an installed wheel and
49+ # still exercises the same GraphBuilder API.
50+ from es_custom import PythonCompilableAddCustom
51+ 
52+ output_z = PythonCompilableAddCustom(input_x, input_y)
53+ output_z.set_shape([NUM_ELEMENTS]).set_format(Format.FORMAT_ND).set_data_type(
54+ DataType.DT_FLOAT
55+ )
56+ builder.set_graph_output(output_z, 0)
57+ return builder.build_and_reset()
58+ 
59+ 
60+def build_input_data(start: float) -> List[float]:
61+ return [start + float(index) for index in range(NUM_ELEMENTS)]
62+ 
63+ 
64+def _check_output(output: Tensor) -> None:
65+ output_data = output.get_data()
66+ expected = [float(index) * 2.0 + 3.0 for index in range(NUM_ELEMENTS)]
67+ if output_data != expected:
68+ raise RuntimeError(
69+ "output mismatch: first={}, expected_first={}".format(
70+ output_data[0], expected[0]
71+ )
72+ )
73+ 
74+ 
75+def run_graph() -> int:
76+ options = {
77+ "ge.exec.deviceId": str(DEVICE_ID),
78+ "ge.graphRunMode": "1",
79+ }
80+ ge_api = GeApi()
81+ session = None
82+ ge_initialized = False
83+ graph_added = False
84+ try:
85+ ge_api.ge_initialize(options)
86+ ge_initialized = True
87+ session = Session(options)
88+ session.add_graph(GRAPH_ID, build_graph())
89+ graph_added = True
90+ print("PY_COMPILE_GRAPH_ADDED=1", flush=True)
91+ outputs = session.run_graph(
92+ GRAPH_ID,
93+ [
94+ Tensor(
95+ build_input_data(1.0),
96+ None,
97+ DataType.DT_FLOAT,
98+ Format.FORMAT_ND,
99+ [NUM_ELEMENTS],
100+ ),
101+ Tensor(
102+ build_input_data(2.0),
103+ None,
104+ DataType.DT_FLOAT,
105+ Format.FORMAT_ND,
106+ [NUM_ELEMENTS],
107+ ),
108+ ],
109+ )
110+ if len(outputs) != 1:
111+ raise RuntimeError("expected one output, got {}".format(len(outputs)))
112+ _check_output(outputs[0])
113+ print("PY_COMPILE_ONLINE_NPU=PASS", flush=True)
114+ return 0
115+ except Exception as error:
116+ print("PY_COMPILE_ONLINE_NPU=FAIL: {}".format(error), flush=True)
117+ traceback.print_exc()
118+ return 1
119+ finally:
120+ try:
121+ if graph_added and session is not None:
122+ session.remove_graph(GRAPH_ID)
123+ finally:
124+ session = None
125+ if ge_initialized:
126+ ge_api.ge_finalize()
127+ 
128+ 
129+if __name__ == "__main__":
130+ raise SystemExit(run_graph())