Merged
feat: 新增 PyPTO 算子 GE 入图样例 #4797
lfz2812Add on Sep 8
feat: 新增 PyPTO 算子 GE 入图样例 #4797
Merged
Merged fromlfz2812/gepypto_tryintocann/gedevelop
Total: 18 files changed+1616-139
| @@ -14,6 +14,7 @@ | |||
| 14 | | `data_dependent_shape_custom` | 数据依赖 shape 算子 | GE | Ascend C | CMake编译 | 不涉及 | [README](data_dependent_shape_custom/README.md) | | 14 | | `data_dependent_shape_custom` | 数据依赖 shape 算子 | GE | Ascend C | CMake编译 | 不涉及 | [README](data_dependent_shape_custom/README.md) | |
| 15 | | `args_refresh_add_custom` | ArgsUpdater 地址刷新 + MallocReadOnlyDevArgs + 性能对比 | GE 在线执行 | Ascend C | RTC 运行时编译 | 在线地址刷新性能对比 | [README](./args_refresh_add_custom/cpp/README.md) | | 15 | | `args_refresh_add_custom` | ArgsUpdater 地址刷新 + MallocReadOnlyDevArgs + 性能对比 | GE 在线执行 | Ascend C | RTC 运行时编译 | 在线地址刷新性能对比 | [README](./args_refresh_add_custom/cpp/README.md) | |
| 16 | | `annotated_args_refresh_add_custom` | AnnotatedArgsOp 声明式地址刷新在线场景性能对比+离线场景 | GE 在线执行 + ATC 离线编译 | Ascend C | RTC 运行时编译 | 支持在线性能对比和 OM 模型下沉 | [README](./annotated_args_refresh_add_custom/README.md) | | 16 | | `annotated_args_refresh_add_custom` | AnnotatedArgsOp 声明式地址刷新在线场景性能对比+离线场景 | GE 在线执行 + ATC 离线编译 | Ascend C | RTC 运行时编译 | 支持在线性能对比和 OM 模型下沉 | [README](./annotated_args_refresh_add_custom/README.md) | |
| 17 | +| `pypto_add_custom` | PyPTO 算子通过 Python 自定义算子入图执行 | GE 在线执行（RT2 动态执行） | PyPTO | JIT 主线程预热编译 + 进程内缓存 | 不涉及 | [README](./pypto_add_custom/README.md) | | ||
| 17 | | `tilelang_add_custom` | TileLang 算子通过 GE 入图 | GE 原生 (Session API) | TileLang | TileLang 预编译产出 `.so` | 不涉及 | [README](./tilelang_add_custom/README.md) | | 18 | | `tilelang_add_custom` | TileLang 算子通过 GE 入图 | GE 原生 (Session API) | TileLang | TileLang 预编译产出 `.so` | 不涉及 | [README](./tilelang_add_custom/README.md) | |
| 18 | | `tilelang_add_custom_online` | TileLang 算子在线编译 + 在线执行 | GE 原生 (Session API) | TileLang | GE 编译阶段 `CompilableOp::Compile` subprocess 调用 Python 编译器 | 不涉及 | [README](./tilelang_add_custom_online/README.md) | | 19 | | `tilelang_add_custom_online` | TileLang 算子在线编译 + 在线执行 | GE 原生 (Session API) | TileLang | GE 编译阶段 `CompilableOp::Compile` subprocess 调用 Python 编译器 | 不涉及 | [README](./tilelang_add_custom_online/README.md) | |
| 19 | | `tilelang_add_custom_offline` | TileLang 算子离线 OM 模型下沉 | GE 原生 (`aclgrphBuildModel`) | TileLang | `CompilableOp::Compile` + `PortableOp::Serialize` 序列化到 OM | 支持 OM 模型下沉 | [README](./tilelang_add_custom_offline/README.md) | | 20 | | `tilelang_add_custom_offline` | TileLang 算子离线 OM 模型下沉 | GE 原生 (`aclgrphBuildModel`) | TileLang | `CompilableOp::Compile` + `PortableOp::Serialize` 序列化到 OM | 支持 OM 模型下沉 | [README](./tilelang_add_custom_offline/README.md) | |
| @@ -14,6 +14,7 @@ This directory provides samples related to custom operator graph integration, co | |||
| 14 | | `data_dependent_shape_custom` | Data dependent shape operator | GE | Ascend C | CMake compilation | Not involved | [README](data_dependent_shape_custom/README_en.md) | | 14 | | `data_dependent_shape_custom` | Data dependent shape operator | GE | Ascend C | CMake compilation | Not involved | [README](data_dependent_shape_custom/README_en.md) | |
| 15 | | `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) | | 15 | | `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) | |
| 16 | | `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) | | 16 | | `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) | |
| 17 | +| `pypto_add_custom` | PyPTO operator enters graph through Python custom operator | GE online execution (RT2 dynamic execution) | PyPTO | JIT warm-up compiled on main thread + in-process cache | Not involved | [README](./pypto_add_custom/README_en.md) | | ||
| 17 | | `tilelang_add_custom` | TileLang operator enters graph through GE | GE native (Session API) | TileLang | TileLang pre-compiled `.so` | Not involved | [README](./tilelang_add_custom/README_en.md) | | 18 | | `tilelang_add_custom` | TileLang operator enters graph through GE | GE native (Session API) | TileLang | TileLang pre-compiled `.so` | Not involved | [README](./tilelang_add_custom/README_en.md) | |
| 18 | | `tilelang_add_custom_online` | TileLang operator online compilation + online execution | GE native (Session API) | TileLang | `CompilableOp::Compile` subprocess invokes Python compiler during GE compile phase | Not involved | [README](./tilelang_add_custom_online/README_en.md) | | 19 | | `tilelang_add_custom_online` | TileLang operator online compilation + online execution | GE native (Session API) | TileLang | `CompilableOp::Compile` subprocess invokes Python compiler during GE compile phase | Not involved | [README](./tilelang_add_custom_online/README_en.md) | |
| 19 | | `tilelang_add_custom_offline` | TileLang operator offline OM model sinking | GE native (`aclgrphBuildModel`) | TileLang | `CompilableOp::Compile` + `PortableOp::Serialize` to OM | Supports OM model sinking | [README](./tilelang_add_custom_offline/README_en.md) | | 20 | | `tilelang_add_custom_offline` | TileLang operator offline OM model sinking | GE native (`aclgrphBuildModel`) | TileLang | `CompilableOp::Compile` + `PortableOp::Serialize` to OM | Supports OM model sinking | [README](./tilelang_add_custom_offline/README_en.md) | |
| @@ -26,6 +26,7 @@ Compile 回调 RTC 编译 Ascend C kernel | |||
| 26 | ## 前置依赖 | 26 | ## 前置依赖 |
| 27 | 27 | ||
| 28 | - 已安装并配置 CANN，例如执行过 `source /usr/local/Ascend/cann/set_env.sh`（路径按实际安装位置调整）。 | 28 | - 已安装并配置 CANN，例如执行过 `source /usr/local/Ascend/cann/set_env.sh`（路径按实际安装位置调整）。 |
| 29 | +- 参考 [安装指导](../../../../../docs/zh/quick_install.md) 完成 toolkit 和 ops 包安装。 | ||
| 29 | - CANN 头文件和库需要包含 `AnnotatedArgsOp`、`AnnotatedKernelLaunchInfo`、`AnnotatedKernelArgs` 等声明式地址刷新接口。 | 30 | - CANN 头文件和库需要包含 `AnnotatedArgsOp`、`AnnotatedKernelLaunchInfo`、`AnnotatedKernelArgs` 等声明式地址刷新接口。 |
| 30 | - `atc` 可用。 | 31 | - `atc` 可用。 |
| 31 | - 当前环境具备 `ACL`、`GE`、`Graph` 相关头文件与库。 | 32 | - 当前环境具备 `ACL`、`GE`、`Graph` 相关头文件与库。 |
| @@ -26,6 +26,7 @@ Key implementation points: | |||
| 26 | ## Prerequisites | 26 | ## Prerequisites |
| 27 | 27 | ||
| 28 | - CANN has been installed and configured, for example by running `source /usr/local/Ascend/cann/set_env.sh` (adjust the path for your installation). | 28 | - CANN has been installed and configured, for example by running `source /usr/local/Ascend/cann/set_env.sh` (adjust the path for your installation). |
| 29 | +- Refer to the [Installation Guide](../../../../../docs/en/quick_install.md) to install the toolkit and ops packages. | ||
| 29 | - The CANN headers and libraries must provide the declarative address refresh APIs, including `AnnotatedArgsOp`, `AnnotatedKernelLaunchInfo`, and `AnnotatedKernelArgs`. | 30 | - The CANN headers and libraries must provide the declarative address refresh APIs, including `AnnotatedArgsOp`, `AnnotatedKernelLaunchInfo`, and `AnnotatedKernelArgs`. |
| 30 | - `atc` is available. | 31 | - `atc` is available. |
| 31 | - The environment provides ACL, GE, and Graph headers and libraries. | 32 | - The environment provides ACL, GE, and Graph headers and libraries. |
| @@ -1,47 +1,129 @@ | |||
| 1 | -# Python 离线自定义算子样例 | 1 | +# Python 声明式参数地址刷新 Add 自定义算子离线样例 |
| 2 | 2 | ||
| 3 | -本样例演示使用 Python 装饰器完成 `AnnotatedAddCustom` 的原型注册、`infer_meta`、编译和 `declare_launch_args`，再生成 AIR/OM 并通过 ACL 在 NPU 上执行。`declare_launch_args` 只在 ATC 编译期运行，离线运行期只消费已经生成的任务描述。Python 回调复用 Ascend C kernel 源码并生成 device binary。 | 3 | +## 样例概述 |
| 4 | 4 | ||
| 5 | -## 目录 | 5 | +本离线样例展示使用 Python 装饰器完成自定义算子原型注册、`infer_meta`、`compile` 和 |
| 6 | +`declare_launch_args`，生成 AIR/OM 并由 ACL 在 NPU 上执行的最小链路。`declare_launch_args` | ||
| 7 | +只在 ATC 编译期运行，离线运行期只消费已生成的任务描述，不再导入 Python 模块。 | ||
| 8 | + | ||
| 9 | +核心链路： | ||
| 6 | 10 | ||
| 7 | ```text | 11 | ```text |
| 8 | -offline/python | 12 | +Python 装饰器注册原型和 infer_meta |
| 9 | -├── CMakeLists.txt # 构建 custom OPP 注册库和 ES wheel | 13 | + -> build_graph.py 生成 AIR |
| 10 | -├── run.sh # kernel、ES、AIR、ATC 和 ACL 验证入口 | 14 | + -> ATC 导入 Python 模块执行 compile 回调（bisheng 编译 Ascend C kernel） |
| 11 | -├── proto/ # gen_esb 使用的 C++ 构图原型 | 15 | + -> DeclareLaunchArgs 声明 kernel/bin/block_dim 和地址槽位 |
| 12 | -└── src/ | 16 | + -> ATC 生成 TaskDef.kernel().args 和 context.args_format |
| 13 | - ├── build_graph.py # Python 构图并生成 AIR | 17 | + -> OM 加载期 GE 按 TaskDef 下发 custom kernel |
| 14 | - ├── run_model.py # ACL 两轮离线 NPU 执行 | 18 | + -> 执行前刷新 INPUT/OUTPUT 地址槽位 |
| 15 | - └── ge/annotated_add_custom.py # Python 原型和 infer_meta | ||
| 16 | ``` | 19 | ``` |
| 17 | 20 | ||
| 18 | -## 依赖与运行 | 21 | +主要实现点： |
| 19 | 22 | ||
| 20 | -需要已安装并配置的 CANN、ATC、CMake、Python 3/pip、numpy 和可用 NPU。 | 23 | +- 本样例的算子类型是 `AnnotatedAddCustom`，原型、`infer_meta`、`compile` 和 |
| 24 | + `declare_launch_args` 均由 `src/ge/annotated_add_custom.py` 的 Python 装饰器提供。 | ||
| 25 | +- 离线模型路径不依赖 Python `execute` 回调；OM 运行期无需 Python 环境。 | ||
| 26 | +- `declare_launch_args` 使用 `append_input`/`append_output` 标注两个输入地址槽和一个输出地址槽， | ||
| 27 | + GE 根据 args 布局完成地址刷新。 | ||
| 28 | +- kernel binary 由 Python `compile` 回调调用 `bisheng` 编译 `add_custom_kernel.cpp` 生成， | ||
| 29 | + 并随地址布局一起写入 OM。 | ||
| 21 | 30 | ||
| 22 | -`proto/add_custom.h` 中的 C++ `REG_OP` 仅用于 `gen_esb` 生成 `ge.es.custom.AnnotatedAddCustom` 构图接口；运行时原型、`infer_meta`、`compile` 和 `declare_launch_args` 均由 `src/ge/annotated_add_custom.py` 的 Python 装饰器提供。 | 31 | +## 前置依赖 |
| 32 | + | ||
| 33 | +- 已正确安装并配置 CANN 环境，例如执行过 `source ${ASCEND_HOME_PATH}/set_env.sh`。 | ||
| 34 | +- 当前环境具备 `ACL`、`GE`、`Graph` 相关头文件与库。 | ||
| 35 | +- 参考 [安装指导](../../../../../docs/zh/quick_install.md) 完成 toolkit 和 ops 包安装。 | ||
| 36 | +- 需要 ATC、CMake、Python 3/pip 和可用 NPU。 | ||
| 37 | +- `bisheng`、`llvm-objcopy`（CANN 工具链自带，由 Python `compile` 回调调用）。 | ||
| 38 | + | ||
| 39 | +## 快速运行 | ||
| 40 | + | ||
| 41 | +在 `examples/custom_op/annotated_args_refresh_add_custom/offline/python` 目录执行： | ||
| 23 | 42 | ||
| 24 | ```bash | 43 | ```bash |
| 25 | -source /path/to/cann/set_env.sh | ||
| 26 | -cd examples/custom_op/annotated_args_refresh_add_custom/offline/python | ||
| 27 | bash run.sh | 44 | bash run.sh |
| 28 | ``` | 45 | ``` |
| 29 | 46 | ||
| 47 | +脚本会完成以下步骤： | ||
| 48 | + | ||
| 49 | +1. 构建 custom OPP 注册库和 Python ES wheel 并安装。 | ||
| 50 | +2. 运行 `build_graph.py` 生成 `build/annotated_add.air`。 | ||
| 51 | +3. 调用 `atc` 导入 Python 模块，执行 `infer_meta`、`compile` 和 `declare_launch_args`， | ||
| 52 | + 生成 `build/annotated_add.om`。 | ||
| 53 | +4. 运行 `run_model.py` 通过 ACL 加载 OM，分两轮独立数据集执行。 | ||
| 54 | + | ||
| 55 | +运行成功时，终端应打印： | ||
| 56 | + | ||
| 57 | +```text | ||
| 58 | +ROUND_1_FIRST=3 | ||
| 59 | +ROUND_2_FIRST=9 | ||
| 60 | +NPU_TWO_ROUND_VALIDATION=PASS | ||
| 61 | +``` | ||
| 62 | + | ||
| 30 | 可覆盖 `SOC_VERSION`（默认 `Ascend910B1`）和 `DEVICE_ID`（默认 `0`）。 | 63 | 可覆盖 `SOC_VERSION`（默认 `Ascend910B1`）和 `DEVICE_ID`（默认 `0`）。 |
| 31 | 64 | ||
| 32 | -## 执行链路 | 65 | +## 关键文件 |
| 33 | 66 | ||
| 34 | -1. 构建 custom OPP 和 ES wheel，`build_graph.py` 生成 AIR。 | 67 | +```text |
| 35 | -2. ATC 导入 Python 模块执行 `infer_meta`、`compile` 和 `declare_launch_args`，将 binary 和地址布局写入 | 68 | +annotated_args_refresh_add_custom |
| 36 | - `build/annotated_add.om`。 | 69 | +└── offline/python |
| 37 | -3. `run_model.py` 使用 ACL 加载 OM，分两轮创建独立数据集并执行，验证 AnnotatedArgs 地址刷新。 | 70 | + ├── CMakeLists.txt # 构建 custom OPP 注册库和 ES wheel |
| 71 | + ├── README.md | ||
| 72 | + ├── README_en.md | ||
| 73 | + ├── run.sh # ES、AIR、ATC 和 ACL 验证入口 | ||
| 74 | + ├── proto | ||
| 75 | + │ └── add_custom.h # gen_esb 使用的 C++ 构图原型 | ||
| 76 | + └── src | ||
| 77 | + ├── build_graph.py # Python 构图并生成 AIR | ||
| 78 | + ├── run_model.py # ACL 两轮离线 NPU 执行 | ||
| 79 | + └── ge | ||
| 80 | + └── annotated_add_custom.py # Python 原型、infer_meta 和编译回调 | ||
| 81 | +``` | ||
| 38 | 82 | ||
| 39 | -运行日志中的 `NPU_TWO_ROUND_VALIDATION=PASS` 表示两轮输出均通过校验。ATC 日志应包含 Python 模块加载、 | 83 | +- `proto/add_custom.h` 中的 C++ `REG_OP` 仅用于 `gen_esb` 生成 `ge.es.custom.AnnotatedAddCustom` |
| 40 | -`infer_meta`、Python compile 和地址声明日志；OM 运行期不应再次导入 Python 模块。 | 84 | + 构图接口。 |
| 85 | +- `src/ge/annotated_add_custom.py` 提供运行时原型、`infer_meta`、`compile` 和 | ||
| 86 | + `declare_launch_args`。 | ||
| 87 | +- `src/build_graph.py` 构建 `AnnotatedAddCustom` 图并导出 AIR。 | ||
| 88 | +- `src/run_model.py` 加载 OM 并分两轮校验 AnnotatedArgs 地址刷新。 | ||
| 41 | 89 | ||
| 42 | -Add kernel 使用 Ascend C 编写，源码位于 `offline/cpp/ge/add_custom_kernel.cpp`。开发者可参考 CANN | 90 | +## 声明式地址刷新实现点 |
| 43 | -Ascend C 文档编写 kernel，Python `compile` 回调负责生成并交付 binary。 | ||
| 44 | 91 | ||
| 45 | -## callback 约束 | 92 | +```text |
| 93 | +ATC 编译期: | ||
| 94 | + infer_meta / compile / declare_launch_args | ||
| 95 | + ├─ compile: bisheng 编译 add_custom_kernel.cpp 并提取 .aicore_binary | ||
| 96 | + ├─ KernelArgs: InputAddr{0}, InputAddr{1}, OutputAddr{0} | ||
| 97 | + ├─ AnnotatedKernelLaunchInfo { kernel_name, kernel_bin, block_dim } | ||
| 98 | + └─ binary 和地址布局写入 annotated_add.om | ||
| 46 | 99 | ||
| 47 | -`append_input` 和 `append_output` 使用当前节点 input/output 的平铺 index。`AnnotatedArgsContext`、Tensor、workspace 和 args builder 都是 callback 期间的 borrowed 对象，不能逃逸；args builder 在 `add_launch` 后已 consumed，不能复用。 | 100 | +OM 执行期: |
| 101 | + ACL 加载 OM | ||
| 102 | + ├─ GE 按 TaskDef 下发 custom kernel | ||
| 103 | + └─ 执行前刷新 INPUT/OUTPUT 地址槽位 | ||
| 104 | + （运行期不导入 Python 模块） | ||
| 105 | +``` | ||
| 106 | + | ||
| 107 | +callback 约束： | ||
| 108 | + | ||
| 109 | +- `append_input` 和 `append_output` 使用当前节点 input/output 的平铺 index。 | ||
| 110 | +- `AnnotatedArgsContext`、Tensor、workspace 和 args builder 都是 callback 期间的 borrowed 对象， | ||
| 111 | + 不能逃逸；args builder 在 `add_launch` 后已 consumed，不能复用。 | ||
| 112 | + | ||
| 113 | +## 构建产物 | ||
| 114 | + | ||
| 115 | +- `build/annotated_add.air` | ||
| 116 | + `build_graph.py` 生成的 AIR 中间图。 | ||
| 117 | +- `build/annotated_add.om` | ||
| 118 | + ATC 编译生成的离线模型，内嵌 kernel binary 和 AnnotatedArgs 地址布局，运行期无需 Python 环境。 | ||
| 119 | +- `build/opp/` | ||
| 120 | + 自定义算子 OPP 注册库（`op_proto/custom/libcust_opapi.so` 及平台拷贝），通过 `ASCEND_CUSTOM_OPP_PATH` 暴露给 GE/ATC。 | ||
| 121 | +- `build/es_output/` | ||
| 122 | + ES 构图接口交付件：`lib64/libes_custom.so` 和 `whl/es_custom-1.0.0-py3-none-any.whl`（已安装到 `build/whl_package/`）。 | ||
| 123 | + | ||
| 124 | +## 注意事项 | ||
| 125 | + | ||
| 126 | +- ATC 日志应包含 Python 模块加载、`infer_meta`、Python compile 和地址声明标记；OM 运行期 | ||
| 127 | + 不应再次出现这些标记，`run.sh` 会自动校验。 | ||
| 128 | +- `run.sh` 在 ATC 阶段后 unset `ASCEND_CUSTOM_OPP_PATH`，验证 OM 运行不依赖 Python 插件路径。 | ||
| 129 | +- OM 内嵌 kernel binary 和地址布局，可脱离本样例目录分发执行。 | ||
| @@ -1,55 +1,135 @@ | |||
| 1 | -# Python Offline Custom Operator | 1 | +# Python Declarative Address Refresh Add Custom Operator Offline Sample |
| 2 | 2 | ||
| 3 | -This sample uses Python decorators to register the `AnnotatedAddCustom` prototype, `infer_meta`, `compile`, and | 3 | +## Sample Overview |
| 4 | -`declare_launch_args`. It generates AIR/OM and runs the model on an NPU through ACL. The declaration callback runs | ||
| 5 | -only during ATC compilation; offline execution consumes the generated task description. The Python callback compiles | ||
| 6 | -and serializes the Ascend C kernel binary. | ||
| 7 | 4 | ||
| 8 | -## Directory | 5 | +This offline sample demonstrates the minimal chain of registering a custom operator prototype with Python |
| 6 | +decorators, running `infer_meta`, `compile`, and `declare_launch_args`, generating AIR/OM, and executing | ||
| 7 | +the model on an NPU through ACL. `declare_launch_args` runs only during ATC compilation; offline execution | ||
| 8 | +consumes the generated task description and never imports the Python module again. | ||
| 9 | + | ||
| 10 | +Core path: | ||
| 9 | 11 | ||
| 10 | ```text | 12 | ```text |
| 11 | -offline/python | 13 | +Register prototype and infer_meta with Python decorators |
| 12 | -├── CMakeLists.txt # Build the custom OPP library and ES wheel | 14 | + -> build_graph.py generates AIR |
| 13 | -├── run.sh # Kernel, ES, AIR, ATC, and ACL validation entry point | 15 | + -> ATC imports the Python module and invokes the compile callback (bisheng compiles the Ascend C kernel) |
| 14 | -├── proto/ # C++ graph prototype consumed by gen_esb | 16 | + -> DeclareLaunchArgs declares kernel/bin/block_dim and address slots |
| 15 | -└── src/ | 17 | + -> ATC generates TaskDef.kernel().args and context.args_format |
| 16 | - ├── build_graph.py # Build the Python graph and save AIR | 18 | + -> GE dispatches the custom kernel per TaskDef during OM loading |
| 17 | - ├── run_model.py # Two-round ACL execution on an NPU | 19 | + -> INPUT/OUTPUT address slots are refreshed before execution |
| 18 | - └── ge/annotated_add_custom.py # Python prototype and infer_meta | ||
| 19 | ``` | 20 | ``` |
| 20 | 21 | ||
| 21 | -## Requirements and usage | 22 | +Main implementation points: |
| 22 | 23 | ||
| 23 | -Install and configure CANN, ATC, CMake, Python 3/pip, numpy, and an available NPU. | 24 | +- The operator type is `AnnotatedAddCustom`. The prototype, `infer_meta`, `compile`, and |
| 25 | + `declare_launch_args` are all supplied by the Python decorators in `src/ge/annotated_add_custom.py`. | ||
| 26 | +- The offline model path does not depend on the Python `execute` callback; no Python environment is | ||
| 27 | + required at OM runtime. | ||
| 28 | +- `declare_launch_args` uses `append_input`/`append_output` to annotate two input address slots and one | ||
| 29 | + output address slot; GE refreshes addresses according to the args layout. | ||
| 30 | +- The kernel binary is produced by the Python `compile` callback invoking `bisheng` on | ||
| 31 | + `add_custom_kernel.cpp` and is written into the OM together with the address layout. | ||
| 24 | 32 | ||
| 25 | -The C++ `REG_OP` in `proto/add_custom.h` is used only by `gen_esb` to generate the | 33 | +## Prerequisites |
| 26 | -`ge.es.custom.AnnotatedAddCustom` graph-building interface. The runtime prototype, `infer_meta`, `compile`, and | 34 | + |
| 27 | -`declare_launch_args` are supplied by the Python decorators in `src/ge/annotated_add_custom.py`. | 35 | +- The CANN environment is installed and configured, for example by running `source ${ASCEND_HOME_PATH}/set_env.sh`. |
| 36 | +- The environment provides the required `ACL`, `GE`, and `Graph` headers and libraries. | ||
| 37 | +- Refer to the [Installation Guide](../../../../../docs/en/quick_install.md) to install the toolkit and ops packages. | ||
| 38 | +- ATC, CMake, Python 3/pip, and an available NPU are required. | ||
| 39 | +- `bisheng` and `llvm-objcopy` (shipped with the CANN toolchain, invoked by the Python `compile` callback). | ||
| 40 | + | ||
| 41 | +## Quick Run | ||
| 42 | + | ||
| 43 | +Run the following command in the `examples/custom_op/annotated_args_refresh_add_custom/offline/python` directory: | ||
| 28 | 44 | ||
| 29 | ```bash | 45 | ```bash |
| 30 | -source /path/to/cann/set_env.sh | ||
| 31 | -cd examples/custom_op/annotated_args_refresh_add_custom/offline/python | ||
| 32 | bash run.sh | 46 | bash run.sh |
| 33 | ``` | 47 | ``` |
| 34 | 48 | ||
| 49 | +The script performs the following steps: | ||
| 50 | + | ||
| 51 | +1. Builds the custom OPP registration library and the Python ES wheel, then installs it. | ||
| 52 | +2. Runs `build_graph.py` to generate `build/annotated_add.air`. | ||
| 53 | +3. Invokes `atc`, which imports the Python module, executes `infer_meta`, `compile`, and | ||
| 54 | + `declare_launch_args`, and generates `build/annotated_add.om`. | ||
| 55 | +4. Runs `run_model.py` to load the OM through ACL and execute two rounds with independent datasets. | ||
| 56 | + | ||
| 57 | +On success the terminal prints: | ||
| 58 | + | ||
| 59 | +```text | ||
| 60 | +ROUND_1_FIRST=3 | ||
| 61 | +ROUND_2_FIRST=9 | ||
| 62 | +NPU_TWO_ROUND_VALIDATION=PASS | ||
| 63 | +``` | ||
| 64 | + | ||
| 35 | The script accepts `SOC_VERSION` (default `Ascend910B1`) and `DEVICE_ID` (default `0`). | 65 | The script accepts `SOC_VERSION` (default `Ascend910B1`) and `DEVICE_ID` (default `0`). |
| 36 | 66 | ||
| 37 | -## Execution flow | 67 | +## Key Files |
| 38 | 68 | ||
| 39 | -1. Build the custom OPP and ES wheel; `build_graph.py` generates AIR. | 69 | +```text |
| 40 | -2. ATC imports the Python module and invokes `infer_meta`, `compile`, and `declare_launch_args` to write the Ascend C | 70 | +annotated_args_refresh_add_custom |
| 41 | - binary and address layout into `build/annotated_add.om`. | 71 | +└── offline/python |
| 42 | -3. `run_model.py` loads the OM with ACL, creates independent datasets for two rounds, and validates address refresh. | 72 | + ├── CMakeLists.txt # Build the custom OPP library and ES wheel |
| 73 | + ├── README.md | ||
| 74 | + ├── README_en.md | ||
| 75 | + ├── run.sh # ES, AIR, ATC, and ACL validation entry point | ||
| 76 | + ├── proto | ||
| 77 | + │ └── add_custom.h # C++ graph prototype consumed by gen_esb | ||
| 78 | + └── src | ||
| 79 | + ├── build_graph.py # Build the Python graph and save AIR | ||
| 80 | + ├── run_model.py # Two-round ACL execution on an NPU | ||
| 81 | + └── ge | ||
| 82 | + └── annotated_add_custom.py # Python prototype, infer_meta, and compile callback | ||
| 83 | +``` | ||
| 43 | 84 | ||
| 44 | -`NPU_TWO_ROUND_VALIDATION=PASS` in the runtime log indicates that both rounds passed. The ATC log should contain | 85 | +- The C++ `REG_OP` in `proto/add_custom.h` is used only by `gen_esb` to generate the |
| 45 | -the Python module-load, `infer_meta`, compile, and address-declaration markers. These | 86 | + `ge.es.custom.AnnotatedAddCustom` graph-building interface. |
| 46 | -compile-time markers must not appear during OM execution. | 87 | +- `src/ge/annotated_add_custom.py` supplies the runtime prototype, `infer_meta`, `compile`, and |
| 88 | + `declare_launch_args`. | ||
| 89 | +- `src/build_graph.py` builds the `AnnotatedAddCustom` graph and exports AIR. | ||
| 90 | +- `src/run_model.py` loads the OM and validates AnnotatedArgs address refresh over two rounds. | ||
| 47 | 91 | ||
| 48 | -The Add kernel is written with Ascend C in `offline/cpp/ge/add_custom_kernel.cpp`. Use the CANN Ascend C | 92 | +## Declarative Address Refresh Implementation |
| 49 | -documentation when developing kernels and ACL RTC when delivering their binaries. | ||
| 50 | 93 | ||
| 51 | -## Callback constraints | 94 | +```text |
| 95 | +ATC compile phase: | ||
| 96 | + infer_meta / compile / declare_launch_args | ||
| 97 | + ├─ compile: bisheng compiles add_custom_kernel.cpp and extracts .aicore_binary | ||
| 98 | + ├─ KernelArgs: InputAddr{0}, InputAddr{1}, OutputAddr{0} | ||
| 99 | + ├─ AnnotatedKernelLaunchInfo { kernel_name, kernel_bin, block_dim } | ||
| 100 | + └─ The binary and address layout are written into annotated_add.om | ||
| 52 | 101 | ||
| 53 | -Use the flattened input/output index of the current node with `append_input` and `append_output`. | 102 | +OM execution phase: |
| 54 | -`AnnotatedArgsContext`, tensors, workspace, and the argument builder are borrowed objects valid only during the | 103 | + ACL loads the OM |
| 55 | -callback and must not escape. The argument builder is consumed by `add_launch` and cannot be reused. | 104 | + ├─ GE dispatches the custom kernel per TaskDef |
| 105 | + └─ INPUT/OUTPUT address slots are refreshed before execution | ||
| 106 | + (The Python module is not imported at runtime) | ||
| 107 | +``` | ||
| 108 | + | ||
| 109 | +Callback constraints: | ||
| 110 | + | ||
| 111 | +- `append_input` and `append_output` use the flattened input/output index of the current node. | ||
| 112 | +- `AnnotatedArgsContext`, tensors, workspace, and the argument builder are borrowed objects valid only | ||
| 113 | + during the callback and must not escape. The argument builder is consumed by `add_launch` and cannot | ||
| 114 | + be reused. | ||
| 115 | + | ||
| 116 | +## Build Artifacts | ||
| 117 | + | ||
| 118 | +- `build/annotated_add.air` | ||
| 119 | + The AIR intermediate graph generated by `build_graph.py`. | ||
| 120 | +- `build/annotated_add.om` | ||
| 121 | + The offline model compiled by ATC. It embeds the kernel binary and the AnnotatedArgs address layout, so no Python environment is required at runtime. | ||
| 122 | +- `build/opp/` | ||
| 123 | + The custom operator OPP registration library (`op_proto/custom/libcust_opapi.so` plus platform copies), exposed to GE/ATC through `ASCEND_CUSTOM_OPP_PATH`. | ||
| 124 | +- `build/es_output/` | ||
| 125 | + ES graph-building deliverables: `lib64/libes_custom.so` and `whl/es_custom-1.0.0-py3-none-any.whl` (installed into `build/whl_package/`). | ||
| 126 | + | ||
| 127 | +## Notes | ||
| 128 | + | ||
| 129 | +- The ATC log should contain the Python module-load, `infer_meta`, compile, and address-declaration | ||
| 130 | + markers. These compile-time markers must not appear during OM execution; `run.sh` verifies this | ||
| 131 | + automatically. | ||
| 132 | +- `run.sh` unsets `ASCEND_CUSTOM_OPP_PATH` after the ATC phase to prove that OM execution does not | ||
| 133 | + depend on the Python plugin path. | ||
| 134 | +- The OM embeds the kernel binary and address layout and can be distributed for execution outside this | ||
| 135 | + sample directory. | ||
| @@ -1,55 +1,241 @@ | |||
| 1 | -# Python 在线自定义算子性能对比样例 | 1 | +# Python 声明式参数地址刷新 Add 自定义算子在线样例 |
| 2 | 2 | ||
| 3 | -本样例使用 Python 装饰器注册 `AnnotatedAddCustom` 和 `NoRefreshAddCustom` 的原型与 `infer_meta`， | 3 | +## 样例概述 |
| 4 | -在两张在线 GE 图中对比声明式地址刷新和不声明地址刷新的执行性能。两条链路使用同一份 Ascend C Add kernel， | ||
| 5 | -脚本以两组独立的 device Tensor 交替运行，完成精度校验、预热和计时，不生成 AIR/OM。 | ||
| 6 | 4 | ||
| 7 | -## 目录 | 5 | +- 构图入口：`GE`（Python ES 构图接口） |
| 6 | +- 算子编程语言：`Python`（原型与回调）+ `Ascend C`（kernel） | ||
| 7 | +- 编译方式：`bisheng` 编译 Ascend C kernel，Python `compile` 回调生成并按 shape 缓存 binary | ||
| 8 | +- 核心链路：`Python 原型注册 -> ES 构图 -> Session 在线执行 -> compile/declare_launch_args/execute 回调` | ||
| 9 | +- 对比目标：声明式地址刷新 `AnnotatedAddCustom` 与无地址刷新 `NoRefreshAddCustom` | ||
| 8 | 10 | ||
| 9 | -```text | 11 | +本样例定义两个功能相同的 Add 自定义算子。输入 shape 均为 `[8192]` float32，两个算子下发同一个 |
| 10 | -online/python | 12 | +`add_custom` Ascend C kernel： |
| 11 | -├── CMakeLists.txt # 构建 custom OPP 注册库和 ES wheel | ||
| 12 | -├── run.sh # 构建并执行在线性能对比 | ||
| 13 | -├── proto/ # gen_esb 使用的 C++ 构图原型 | ||
| 14 | -└── src/ | ||
| 15 | - ├── run.py # 两张在线 GE 图的精度和性能对比 | ||
| 16 | - └── ge/ | ||
| 17 | - └── annotated_add_custom.py # Python 原型和 infer_meta | ||
| 18 | -``` | ||
| 19 | 13 | ||
| 20 | -## 依赖与运行 | 14 | +- `AnnotatedAddCustom` 通过 `register_op_impl` 提供 `compile` 和 `declare_launch_args`。`compile` |
| 21 | - | 15 | + 调用 `bisheng` 编译 `add_custom.asc` 并按输入 shape/dtype 缓存 binary；`declare_launch_args` |
| 22 | -需要已安装并配置的 CANN、CMake、Python 3/pip 和可用 NPU。 | 16 | + 声明两个输入地址槽和一个输出地址槽，GE 在重复执行时按该布局刷新地址。 |
| 17 | +- `NoRefreshAddCustom` 通过 `register_op_impl` 提供 `execute`。它在每次图执行时申请输出 Tensor， | ||
| 18 | + 通过 ACL Python API 加载 binary 并重新组装 kernel 参数，用作性能对比基线。 | ||
| 23 | 19 | ||
| 24 | `proto/add_custom.h` 中的 C++ `REG_OP` 仅用于 `gen_esb` 生成构图接口；运行时原型和 `infer_meta` | 20 | `proto/add_custom.h` 中的 C++ `REG_OP` 仅用于 `gen_esb` 生成构图接口；运行时原型和 `infer_meta` |
| 25 | -由 Python 装饰器提供；`AnnotatedAddCustom` 的 `compile`、`declare_launch_args` 和 `NoRefreshAddCustom` 的 | 21 | +由 Python 装饰器提供，算子执行全部由 Python 回调完成，不编译 `online/cpp/ge/custom_op.cpp`。 |
| 26 | -`execute` 均由 Python `register_op_impl` 提供。两个实现复用 `cpp/add_custom_kernel/add_custom.asc` 生成的 | ||
| 27 | -Ascend C binary。 | ||
| 28 | 22 | ||
| 29 | -本目录不编译 `online/cpp/ge/custom_op.cpp`；该文件仅作为 C++ 对照样例，Python 样例的算子执行全部由 Python 回调完成。 | 23 | +`run.py` 构建上述两张图。每张图分别使用两套设备内存交替执行，完成预热、100 轮性能统计和精度校验， |
| 24 | +最后打印两者耗时及 speedup。 | ||
| 25 | + | ||
| 26 | +## 适用场景 | ||
| 27 | + | ||
| 28 | +- 了解 Python `register_op`/`register_op_impl` 如何声明 kernel task 和 args 地址布局。 | ||
| 29 | +- 对比声明式地址刷新与无地址刷新在重复在线执行中的性能差异。 | ||
| 30 | +- 想在 Python 侧完成自定义算子原型、编译回调与执行回调的完整链路验证。 | ||
| 31 | + | ||
| 32 | +## 前置依赖 | ||
| 33 | + | ||
| 34 | +### CANN | ||
| 35 | + | ||
| 36 | +- 已正确安装并配置 CANN 环境，例如执行过 `source ${ASCEND_HOME_PATH}/set_env.sh`。 | ||
| 37 | +- 当前环境具备 `ACL`、`GE`、`Graph` 相关头文件与库。 | ||
| 38 | +- 参考 [安装指导](../../../../../docs/zh/quick_install.md) 完成 toolkit 和 ops 包安装。 | ||
| 39 | + | ||
| 40 | +### 框架与插件 | ||
| 41 | + | ||
| 42 | +- Python 3 和 `pip`。 | ||
| 43 | +- ES 构图接口由本样例构建的 `es_custom` wheel 提供，无需额外安装框架插件。 | ||
| 44 | + | ||
| 45 | +### 环境变量 | ||
| 46 | + | ||
| 47 | +- `ASCEND_HOME_PATH` | ||
| 48 | +- `run.sh` 会自动设置 `PYTHONPATH`、`LD_LIBRARY_PATH`、`ASCEND_CUSTOM_OPP_PATH`、 | ||
| 49 | + `GE_PYTHON_CUSTOM_OP_SOURCE` 和 `GE_PYTHON_CUSTOM_OP_BINARY`。 | ||
| 50 | + | ||
| 51 | +### 额外依赖 | ||
| 52 | + | ||
| 53 | +- `cmake` | ||
| 54 | +- `bisheng`、`llvm-objcopy`（CANN 工具链自带，用于编译和提取 kernel binary） | ||
| 55 | + | ||
| 56 | +## 快速运行 | ||
| 57 | + | ||
| 58 | +在 `examples/custom_op/annotated_args_refresh_add_custom/online/python` 目录下执行： | ||
| 59 | + | ||
| 60 | +### 推荐方式 | ||
| 30 | 61 | ||
| 31 | ```bash | 62 | ```bash |
| 32 | -source /path/to/cann/set_env.sh | 63 | +source ${ASCEND_HOME_PATH}/set_env.sh |
| 33 | -cd examples/custom_op/annotated_args_refresh_add_custom/online/python | ||
| 34 | bash run.sh | 64 | bash run.sh |
| 35 | ``` | 65 | ``` |
| 36 | 66 | ||
| 37 | -可通过 `DEVICE_ID`（默认 `0`）选择 NPU： | 67 | +`run.sh` 会自动完成 kernel 编译、OPP/ES wheel 构建、安装和在线对比： |
| 38 | 68 | ||
| 39 | -```bash | 69 | +1. `bisheng` 编译 `add_custom.asc`，`llvm-objcopy` 提取 `.aicore_binary` 生成 kernel binary。 |
| 40 | -DEVICE_ID=1 bash run.sh | 70 | +2. 构建 custom OPP 注册库和 Python ES wheel 并安装。 |
| 71 | +3. 运行两张图的在线精度与性能对比。 | ||
| 72 | + | ||
| 73 | +若运行成功，终端会打印类似： | ||
| 74 | + | ||
| 75 | +```text | ||
| 76 | +[OnlinePython] AnnotatedAddCustom precision check PASS | ||
| 77 | +[OnlinePython] NoRefreshAddCustom precision check PASS | ||
| 78 | +[Perf] input shape: [8192], dtype: float32 | ||
| 79 | +[Perf] iters: 100 | ||
| 80 | +[Perf] AnnotatedAddCustom: xxx us (avg xxx us/iter) | ||
| 81 | +[Perf] NoRefreshAddCustom: xxx us (avg xxx us/iter) | ||
| 82 | +[Perf] Annotated speedup: xxx x | ||
| 83 | +[OnlinePython] NPU_EXECUTION=PASS | ||
| 41 | ``` | 84 | ``` |
| 42 | 85 | ||
| 43 | -## 执行链路 | 86 | +### 分步方式 |
| 44 | 87 | ||
| 45 | -1. 构建 custom OPP 和 ES wheel，生成两个构图接口。 | 88 | +```bash |
| 46 | -2. Python `compile` 回调在 GE 编译图时编译并保存 Ascend C Add kernel binary。 | 89 | +source ${ASCEND_HOME_PATH}/set_env.sh |
| 47 | -3. Python `declare_launch_args` 回调声明输入、输出地址槽和 kernel launch；GE 在重复执行时刷新地址。 | 90 | +mkdir -p build |
| 48 | -4. Python `execute` 为 schema 输出申请 Tensor，并通过 ACL Python API 下发同一 Ascend C kernel，作为无地址声明基线。 | 91 | +# 1. 编译 Ascend C kernel 并提取 binary |
| 49 | -5. `run.py` 为两张图交替传入两组 device Tensor，分别完成精度校验、5 次预热和 100 次计时。 | 92 | +bisheng -c ../cpp/add_custom_kernel/add_custom.asc -o build/add_custom.host.o --npu-arch=dav-2201 |
| 93 | +llvm-objcopy -O binary --only-section=.aicore_binary build/add_custom.host.o build/add_custom.aicore.o | ||
| 94 | +# 2. 构建 custom OPP 和 ES wheel | ||
| 95 | +cmake -S . -B build -DCMAKE_BUILD_TYPE=Release | ||
| 96 | +cmake --build build --target build_es_custom -j"$(nproc)" | ||
| 97 | +# 3. 安装 ES wheel | ||
| 98 | +python3 -m pip install --target build/whl_package build/es_output/whl/es_custom-1.0.0-py3-none-any.whl | ||
| 99 | +# 4. 设置环境变量并运行对比 | ||
| 100 | +export PYTHONPATH="$(pwd)/build/whl_package:$(pwd)/src:$PYTHONPATH" | ||
| 101 | +export LD_LIBRARY_PATH="$(pwd)/build/es_output/lib64:$LD_LIBRARY_PATH" | ||
| 102 | +export ASCEND_CUSTOM_OPP_PATH="$(pwd)/build/opp:$(pwd)/src/ge" | ||
| 103 | +export GE_PYTHON_CUSTOM_OP_SOURCE="$(pwd)/../cpp/add_custom_kernel/add_custom.asc" | ||
| 104 | +export GE_PYTHON_CUSTOM_OP_BINARY="$(pwd)/build/add_custom.aicore.o" | ||
| 105 | +DEVICE_ID=0 python3 src/run.py | ||
| 106 | +``` | ||
| 50 | 107 | ||
| 51 | -日志会输出 `AnnotatedAddCustom`、`NoRefreshAddCustom` 的总耗时、平均耗时和 `Annotated speedup`； | 108 | +NPU 架构默认 `dav-2201`，可通过 `ADD_CUSTOM_NPU_ARCH` 覆盖；`DEVICE_ID` 默认 `0`。 |
| 52 | -`NPU_EXECUTION=PASS` 表示两条链路均通过精度和执行校验。 | ||
| 53 | 109 | ||
| 54 | -Add kernel 使用 Ascend C 编写，源码位于 `online/cpp/add_custom_kernel/add_custom.asc`。开发者可参考 | 110 | +## 目录结构与关键文件 |
| 55 | -CANN Ascend C 文档编写并通过 ACL RTC 编译自定义 kernel。 | 111 | + |
| 112 | +```text | ||
| 113 | +annotated_args_refresh_add_custom | ||
| 114 | +└── online/python | ||
| 115 | + ├── CMakeLists.txt # 构建 custom OPP 注册库和 ES wheel | ||
| 116 | + ├── README.md | ||
| 117 | + ├── README_en.md | ||
| 118 | + ├── run.sh # 构建并执行在线性能对比 | ||
| 119 | + ├── proto | ||
| 120 | + │ └── add_custom.h # gen_esb 使用的 C++ 构图原型 | ||
| 121 | + └── src | ||
| 122 | + ├── run.py # 两张在线 GE 图的精度和性能对比 | ||
| 123 | + └── ge | ||
| 124 | + └── annotated_add_custom.py # Python 原型、infer_meta 和算子实现 | ||
| 125 | +``` | ||
| 126 | + | ||
| 127 | +重点文件： | ||
| 128 | + | ||
| 129 | +- `src/ge/annotated_add_custom.py` | ||
| 130 | + 实现 `AnnotatedAddCustom` 和 `NoRefreshAddCustom`。前者在编译回调生成 binary，并在 task 生成 | ||
| 131 | + 阶段声明地址槽位；后者使用 eager 执行链路直接下发同一 kernel。 | ||
| 132 | +- `proto/add_custom.h` | ||
| 133 | + 注册 `AnnotatedAddCustom` 和 `NoRefreshAddCustom` 两个构图侧算子类型，供 `gen_esb` 生成 | ||
| 134 | + Python ES 构图接口。 | ||
| 135 | +- `src/run.py` | ||
| 136 | + 构建两张图，交替两套输入/输出设备地址，分别预热 5 次、统计 100 轮并校验精度。 | ||
| 137 | +- `CMakeLists.txt` | ||
| 138 | + 构建 custom OPP 注册库（`libcust_opapi.so`）和 Python ES wheel。 | ||
| 139 | + | ||
| 140 | +## 核心链路 | ||
| 141 | + | ||
| 142 | +### 在线执行 | ||
| 143 | + | ||
| 144 | +1. `run.py` 构建 `python_annotated_graph` 和 `python_no_refresh_graph`，输入均为 `[8192]` float32。 | ||
| 145 | +2. `AnnotatedAddCustom` 的 `compile` 回调读取 `add_custom.asc`，调用 `bisheng` 编译并提取 | ||
| 146 | + `.aicore_binary`，按输入 shape/dtype 缓存 binary。 | ||
| 147 | +3. `AnnotatedAddCustom` 的 `declare_launch_args` 回调设置 kernel 名称、binary、block dim， | ||
| 148 | + 并声明 `InputAddr{0}`、`InputAddr{1}` 和 `OutputAddr{0}`。 | ||
| 149 | +4. `NoRefreshAddCustom` 的 `execute` 回调申请输出 Tensor，通过 ACL Python API 加载 binary、 | ||
| 150 | + 组装 args 并下发 kernel。 | ||
| 151 | +5. 两张图分别交替使用两套设备地址执行，输出总耗时、平均耗时和 speedup。 | ||
| 152 | + | ||
| 153 | +### 声明式地址刷新 | ||
| 154 | + | ||
| 155 | +```text | ||
| 156 | +编译期: | ||
| 157 | + compile(x, y, z) | ||
| 158 | + ├─ 读取 add_custom.asc | ||
| 159 | + ├─ bisheng 编译 + llvm-objcopy 提取 .aicore_binary -> kernel binary | ||
| 160 | + └─ 按输入 shape/dtype 缓存 binary | ||
| 161 | + | ||
| 162 | +Task 生成期: | ||
| 163 | + declare_launch_args(x, y, z) | ||
| 164 | + ├─ KernelArgs: InputAddr{0}, InputAddr{1}, OutputAddr{0} | ||
| 165 | + ├─ AnnotatedKernelLaunchInfo { kernel_name, kernel_bin, block_dim, stream_id } | ||
| 166 | + └─ ctx.add_launch(launch_info, args) | ||
| 167 | + | ||
| 168 | +执行期: | ||
| 169 | + GE 根据保存的 args 布局刷新当前输入/输出地址 | ||
| 170 | +``` | ||
| 171 | + | ||
| 172 | +### 无地址刷新基线 | ||
| 173 | + | ||
| 174 | +```text | ||
| 175 | +每次图执行: | ||
| 176 | + NoRefreshAddCustom.execute(x, y) | ||
| 177 | + ├─ ctx.malloc_output_tensor() | ||
| 178 | + ├─ acl.rt.binary_load_from_file() / binary_get_function() | ||
| 179 | + ├─ kernel_args_init / append(x, y, z) / finalize | ||
| 180 | + └─ acl.rt.launch_kernel_with_config() | ||
| 181 | +``` | ||
| 182 | + | ||
| 183 | +该实现不声明 args 中的地址槽位。两张图的计算逻辑和 kernel 完全一致，因此性能差异用于反映 | ||
| 184 | +声明式地址刷新的效果。 | ||
| 185 | + | ||
| 186 | +## 构建产物 | ||
| 187 | + | ||
| 188 | +- `build/add_custom.aicore.o` | ||
| 189 | + Ascend C Add kernel 的设备侧 binary，由 `bisheng` 编译并提取 `.aicore_binary` 段生成。 | ||
| 190 | +- `build/opp/` | ||
| 191 | + 自定义算子 OPP 注册库（`op_proto/custom/libcust_opapi.so` 及平台拷贝），通过 `ASCEND_CUSTOM_OPP_PATH` 暴露给 GE。 | ||
| 192 | +- `build/es_output/` | ||
| 193 | + ES 构图接口交付件：`lib64/libes_custom.so` 和 `whl/es_custom-1.0.0-py3-none-any.whl`（已安装到 `build/whl_package/`）。 | ||
| 194 | + | ||
| 195 | +## 结果校验 | ||
| 196 | + | ||
| 197 | +成功时可观察到： | ||
| 198 | + | ||
| 199 | +- 终端输出包含 `AnnotatedAddCustom precision check PASS` 和 `NoRefreshAddCustom precision check PASS`。 | ||
| 200 | +- 终端输出包含 `AnnotatedAddCustom`、`NoRefreshAddCustom` 和 `Annotated speedup`。 | ||
| 201 | +- 终端输出包含 `NPU_EXECUTION=PASS` 和 `Online Python custom-op pipeline PASS`。 | ||
| 202 | + | ||
| 203 | +若失败，优先检查： | ||
| 204 | + | ||
| 205 | +- `ASCEND_HOME_PATH` 是否已设置并正确加载 CANN 环境。 | ||
| 206 | +- `bisheng`、`llvm-objcopy`、`cmake` 是否可用。 | ||
| 207 | +- `build/es_output/whl/es_custom-1.0.0-py3-none-any.whl` 和 `build/add_custom.aicore.o` 是否生成。 | ||
| 208 | +- 当前环境是否具备可用 NPU。 | ||
| 209 | + | ||
| 210 | +## 注意事项 / 限制 | ||
| 211 | + | ||
| 212 | +- kernel binary 由 `compile` 回调调用 `bisheng` 编译，图编译阶段包含该开销；100 轮计时不包含这些阶段。 | ||
| 213 | +- `ge.graphRunMode` 设置为 `1`（`PRIORITY_GRAPH`），确保使用在线执行链路。 | ||
| 214 | +- 性能测试使用两套设备内存交替执行，以触发输入/输出地址变化。 | ||
| 215 | +- 加速比受 NPU 型号、系统负载等因素影响，仅供参考。 | ||
| 216 | +- NPU 架构默认 `dav-2201`（Atlas A2），其他型号需通过 `ADD_CUSTOM_NPU_ARCH` 覆盖。 | ||
| 217 | + | ||
| 218 | +## 附录 | ||
| 219 | + | ||
| 220 | +### 算子规格 | ||
| 221 | + | ||
| 222 | +| 项目 | 内容 | | ||
| 223 | +| --- | --- | | ||
| 224 | +| 算子类型 | `AnnotatedAddCustom` / `NoRefreshAddCustom` | | ||
| 225 | +| 输入 | `x`, `y` | | ||
| 226 | +| 输出 | `z` | | ||
| 227 | +| 输入/输出 shape | `[8192]` | | ||
| 228 | +| 输入/输出数据类型 | `float32` | | ||
| 229 | +| 格式 | `ND` | | ||
| 230 | +| kernel 名称 | `add_custom`（Ascend C，`bisheng` 编译） | | ||
| 231 | +| block 大小 | `1024` | | ||
| 232 | + | ||
| 233 | +### 关键接口 | ||
| 234 | + | ||
| 235 | +| 接口 | 算子 | 用途 | | ||
| 236 | +| --- | --- | --- | | ||
| 237 | +| `ge.custom_op.register_op` + `infer_meta` | 两个算子 | 原型注册与输出 shape/dtype 推导 | | ||
| 238 | +| `register_op_impl` + `compile` | `AnnotatedAddCustom` | `bisheng` 编译 kernel 并按 shape 缓存 binary | | ||
| 239 | +| `register_op_impl` + `declare_launch_args` | `AnnotatedAddCustom` | 声明 kernel launch 和输入/输出地址槽位 | | ||
| 240 | +| `register_op_impl` + `execute` | `NoRefreshAddCustom` | 申请输出 Tensor 并经 ACL API 下发 kernel | | ||
| 241 | +| `acl.rt.launch_kernel_with_config` | `NoRefreshAddCustom` | ACL Python API 直接下发 kernel | | ||
| @@ -1,58 +1,254 @@ | |||
| 1 | -# Python Online Custom Operator Performance Comparison | 1 | +# Python Declarative Address Refresh Add Custom Operator Online Sample |
| 2 | 2 | ||
| 3 | -This sample uses Python decorators to register the prototypes and `infer_meta` functions of `AnnotatedAddCustom` and `NoRefreshAddCustom`. Two online GE | 3 | +## Sample Overview |
| 4 | -graphs compare declarative address refresh with a path that does not declare | ||
| 5 | -refreshable addresses. Both paths use the same Ascend C Add kernel. The script validates results and measures | ||
| 6 | -execution without generating AIR/OM. | ||
| 7 | 4 | ||
| 8 | -## Directory | 5 | +- Graph building entry: `GE` (Python ES graph-building APIs) |
| 6 | +- Operator language: `Python` (prototype and callbacks) + `Ascend C` (kernel) | ||
| 7 | +- Compilation mode: `bisheng` compiles the Ascend C kernel; the Python `compile` callback generates and caches the binary by shape | ||
| 8 | +- Core path: `Python prototype registration -> ES graph building -> Session online execution -> compile/declare_launch_args/execute callbacks` | ||
| 9 | +- Comparison target: declarative address refresh `AnnotatedAddCustom` versus no-refresh `NoRefreshAddCustom` | ||
| 9 | 10 | ||
| 10 | -```text | 11 | +This sample defines two functionally identical Add custom operators. Both take `[8192]` float32 inputs and |
| 11 | -online/python | 12 | +dispatch the same `add_custom` Ascend C kernel: |
| 12 | -├── CMakeLists.txt # Build the custom OPP library and ES wheel | ||
| 13 | -├── run.sh # Build and run the online comparison | ||
| 14 | -├── proto/ # C++ graph prototypes consumed by gen_esb | ||
| 15 | -└── src/ | ||
| 16 | - ├── run.py # Validate and benchmark the two online graphs | ||
| 17 | - └── ge/annotated_add_custom.py # Python prototypes and infer_meta | ||
| 18 | -``` | ||
| 19 | 13 | ||
| 20 | -## Requirements and usage | 14 | +- `AnnotatedAddCustom` provides `compile` and `declare_launch_args` through `register_op_impl`. `compile` |
| 15 | + invokes `bisheng` to compile `add_custom.asc` and caches the binary by input shape/dtype; | ||
| 16 | + `declare_launch_args` declares two input address slots and one output address slot, and GE refreshes | ||
| 17 | + addresses according to that layout during repeated execution. | ||
| 18 | +- `NoRefreshAddCustom` provides `execute` through `register_op_impl`. On every graph execution it allocates | ||
| 19 | + the output tensor, loads the binary through the ACL Python API, reassembles the kernel arguments, and | ||
| 20 | + serves as the performance baseline. | ||
| 21 | 21 | ||
| 22 | -Install and configure CANN, CMake, Python 3/pip, and an available NPU. | 22 | +The C++ `REG_OP` declarations in `proto/add_custom.h` are used only by `gen_esb` to generate the |
| 23 | +graph-building APIs. The runtime prototypes and `infer_meta` functions are supplied by Python decorators, | ||
| 24 | +and all operator execution is implemented by Python callbacks. `online/cpp/ge/custom_op.cpp` is not compiled. | ||
| 23 | 25 | ||
| 24 | -The C++ `REG_OP` declarations in `proto/add_custom.h` are used only by `gen_esb` to generate the graph APIs. | 26 | +`run.py` builds the two graphs above. Each graph alternates between two sets of device buffers, performs |
| 25 | -The runtime prototypes and `infer_meta` functions of both operators are supplied by Python decorators. Python `register_op_impl` | 27 | +warm-up, 100 benchmark iterations, and precision checks, then prints the latency of both paths and the |
| 26 | -supplies `compile`, `declare_launch_args`, and `execute` for both operators. They reuse the Ascend C binary generated | 28 | +speedup. |
| 27 | -from `cpp/add_custom_kernel/add_custom.asc`. | ||
| 28 | 29 | ||
| 29 | -This directory does not compile `online/cpp/ge/custom_op.cpp`; that file is provided only as the C++ comparison sample. | 30 | +## Applicable Scenarios |
| 30 | -All operator execution in this Python sample is implemented by Python callbacks. | 31 | + |
| 32 | +- Understanding how Python `register_op`/`register_op_impl` declare kernel tasks and args address layouts. | ||
| 33 | +- Comparing declarative address refresh with no-refresh paths in repeated online execution. | ||
| 34 | +- Verifying the complete chain of Python-side prototype, compile callback, and execute callback. | ||
| 35 | + | ||
| 36 | +## Prerequisites | ||
| 37 | + | ||
| 38 | +### CANN | ||
| 39 | + | ||
| 40 | +- The CANN environment is installed and configured, for example by running `source ${ASCEND_HOME_PATH}/set_env.sh`. | ||
| 41 | +- The environment provides the required `ACL`, `GE`, and `Graph` headers and libraries. | ||
| 42 | +- Refer to the [Installation Guide](../../../../../docs/en/quick_install.md) to install the toolkit and ops packages. | ||
| 43 | + | ||
| 44 | +### Frameworks and Plugins | ||
| 45 | + | ||
| 46 | +- Python 3 and `pip`. | ||
| 47 | +- The ES graph-building APIs are provided by the `es_custom` wheel built by this sample. No extra | ||
| 48 | + framework plugin is required. | ||
| 49 | + | ||
| 50 | +### Environment Variables | ||
| 51 | + | ||
| 52 | +- `ASCEND_HOME_PATH` | ||
| 53 | +- `run.sh` automatically sets `PYTHONPATH`, `LD_LIBRARY_PATH`, `ASCEND_CUSTOM_OPP_PATH`, | ||
| 54 | + `GE_PYTHON_CUSTOM_OP_SOURCE`, and `GE_PYTHON_CUSTOM_OP_BINARY`. | ||
| 55 | + | ||
| 56 | +### Additional Dependencies | ||
| 57 | + | ||
| 58 | +- `cmake` | ||
| 59 | +- `bisheng` and `llvm-objcopy` (shipped with the CANN toolchain, used to compile and extract the kernel binary) | ||
| 60 | + | ||
| 61 | +## Quick Run | ||
| 62 | + | ||
| 63 | +Run the following commands in the `examples/custom_op/annotated_args_refresh_add_custom/online/python` directory: | ||
| 64 | + | ||
| 65 | +### Recommended | ||
| 31 | 66 | ||
| 32 | ```bash | 67 | ```bash |
| 33 | -source /path/to/cann/set_env.sh | 68 | +source ${ASCEND_HOME_PATH}/set_env.sh |
| 34 | -cd examples/custom_op/annotated_args_refresh_add_custom/online/python | ||
| 35 | bash run.sh | 69 | bash run.sh |
| 36 | ``` | 70 | ``` |
| 37 | 71 | ||
| 38 | -Use `DEVICE_ID` (default `0`) to select the NPU: | 72 | +`run.sh` automatically compiles the kernel, builds and installs the OPP/ES wheel, and runs the online |
| 73 | +comparison: | ||
| 39 | 74 | ||
| 40 | -```bash | 75 | +1. `bisheng` compiles `add_custom.asc` and `llvm-objcopy` extracts the `.aicore_binary` section into the |
| 41 | -DEVICE_ID=1 bash run.sh | 76 | + kernel binary. |
| 77 | +2. Builds the custom OPP registration library and the Python ES wheel, then installs it. | ||
| 78 | +3. Runs the online precision and performance comparison of the two graphs. | ||
| 79 | + | ||
| 80 | +On success the terminal prints something like: | ||
| 81 | + | ||
| 82 | +```text | ||
| 83 | +[OnlinePython] AnnotatedAddCustom precision check PASS | ||
| 84 | +[OnlinePython] NoRefreshAddCustom precision check PASS | ||
| 85 | +[Perf] input shape: [8192], dtype: float32 | ||
| 86 | +[Perf] iters: 100 | ||
| 87 | +[Perf] AnnotatedAddCustom: xxx us (avg xxx us/iter) | ||
| 88 | +[Perf] NoRefreshAddCustom: xxx us (avg xxx us/iter) | ||
| 89 | +[Perf] Annotated speedup: xxx x | ||
| 90 | +[OnlinePython] NPU_EXECUTION=PASS | ||
| 42 | ``` | 91 | ``` |
| 43 | 92 | ||
| 44 | -## Execution flow | 93 | +### Step by Step |
| 45 | 94 | ||
| 46 | -1. The custom OPP and ES wheel are built to generate both Python graph APIs. | 95 | +```bash |
| 47 | -2. The Python `compile` callback compiles and retains the Ascend C Add kernel binary while GE compiles the graph. | 96 | +source ${ASCEND_HOME_PATH}/set_env.sh |
| 48 | -3. The Python `declare_launch_args` callback declares the input/output address slots and kernel launch, allowing GE to refresh addresses during | 97 | +mkdir -p build |
| 49 | - repeated execution. | 98 | +# 1. Compile the Ascend C kernel and extract the binary |
| 50 | -4. Python `execute` allocates the schema output and launches the same Ascend C kernel through the ACL Python API as the | 99 | +bisheng -c ../cpp/add_custom_kernel/add_custom.asc -o build/add_custom.host.o --npu-arch=dav-2201 |
| 51 | - no-declaration baseline. | 100 | +llvm-objcopy -O binary --only-section=.aicore_binary build/add_custom.host.o build/add_custom.aicore.o |
| 52 | -5. `run.py` validates both graphs, performs five warm-up iterations, and measures 100 iterations. | 101 | +# 2. Build the custom OPP and ES wheel |
| 102 | +cmake -S . -B build -DCMAKE_BUILD_TYPE=Release | ||
| 103 | +cmake --build build --target build_es_custom -j"$(nproc)" | ||
| 104 | +# 3. Install the ES wheel | ||
| 105 | +python3 -m pip install --target build/whl_package build/es_output/whl/es_custom-1.0.0-py3-none-any.whl | ||
| 106 | +# 4. Set environment variables and run the comparison | ||
| 107 | +export PYTHONPATH="$(pwd)/build/whl_package:$(pwd)/src:$PYTHONPATH" | ||
| 108 | +export LD_LIBRARY_PATH="$(pwd)/build/es_output/lib64:$LD_LIBRARY_PATH" | ||
| 109 | +export ASCEND_CUSTOM_OPP_PATH="$(pwd)/build/opp:$(pwd)/src/ge" | ||
| 110 | +export GE_PYTHON_CUSTOM_OP_SOURCE="$(pwd)/../cpp/add_custom_kernel/add_custom.asc" | ||
| 111 | +export GE_PYTHON_CUSTOM_OP_BINARY="$(pwd)/build/add_custom.aicore.o" | ||
| 112 | +DEVICE_ID=0 python3 src/run.py | ||
| 113 | +``` | ||
| 53 | 114 | ||
| 54 | -The log reports total and average times for both operators plus `Annotated speedup`. | 115 | +The NPU architecture defaults to `dav-2201` and can be overridden with `ADD_CUSTOM_NPU_ARCH`; `DEVICE_ID` |
| 55 | -`NPU_EXECUTION=PASS` indicates that both paths passed validation and execution. | 116 | +defaults to `0`. |
| 56 | 117 | ||
| 57 | -The Add kernel is written with Ascend C in `online/cpp/add_custom_kernel/add_custom.asc`. Use the CANN Ascend C | 118 | +## Directory Structure and Key Files |
| 58 | -documentation when developing kernels for the target hardware. | 119 | + |
| 120 | +```text | ||
| 121 | +annotated_args_refresh_add_custom | ||
| 122 | +└── online/python | ||
| 123 | + ├── CMakeLists.txt # Build the custom OPP library and ES wheel | ||
| 124 | + ├── README.md | ||
| 125 | + ├── README_en.md | ||
| 126 | + ├── run.sh # Build and run the online comparison | ||
| 127 | + ├── proto | ||
| 128 | + │ └── add_custom.h # C++ graph prototype consumed by gen_esb | ||
| 129 | + └── src | ||
| 130 | + ├── run.py # Validate and benchmark the two online graphs | ||
| 131 | + └── ge | ||
| 132 | + └── annotated_add_custom.py # Python prototypes, infer_meta, and implementations | ||
| 133 | +``` | ||
| 134 | + | ||
| 135 | +Key files: | ||
| 136 | + | ||
| 137 | +- `src/ge/annotated_add_custom.py` | ||
| 138 | + Implements `AnnotatedAddCustom` and `NoRefreshAddCustom`. The former generates the binary in the compile | ||
| 139 | + callback and declares the address slots during task generation; the latter dispatches the same kernel | ||
| 140 | + through the eager execution path. | ||
| 141 | +- `proto/add_custom.h` | ||
| 142 | + Registers the `AnnotatedAddCustom` and `NoRefreshAddCustom` graph operator types for `gen_esb` to | ||
| 143 | + generate the Python ES graph-building APIs. | ||
| 144 | +- `src/run.py` | ||
| 145 | + Builds the two graphs, alternates two sets of input/output device addresses, performs 5 warm-up | ||
| 146 | + iterations, 100 benchmark iterations, and validates precision. | ||
| 147 | +- `CMakeLists.txt` | ||
| 148 | + Builds the custom OPP registration library (`libcust_opapi.so`) and the Python ES wheel. | ||
| 149 | + | ||
| 150 | +## Core Path | ||
| 151 | + | ||
| 152 | +### Online Execution | ||
| 153 | + | ||
| 154 | +1. `run.py` builds `python_annotated_graph` and `python_no_refresh_graph`, both with `[8192]` float32 inputs. | ||
| 155 | +2. The `compile` callback of `AnnotatedAddCustom` reads `add_custom.asc`, invokes `bisheng`, extracts the | ||
| 156 | + `.aicore_binary` section, and caches the binary by input shape/dtype. | ||
| 157 | +3. The `declare_launch_args` callback of `AnnotatedAddCustom` sets the kernel name, binary, and block | ||
| 158 | + dimension, then declares `InputAddr{0}`, `InputAddr{1}`, and `OutputAddr{0}`. | ||
| 159 | +4. The `execute` callback of `NoRefreshAddCustom` allocates the output tensor, loads the binary through | ||
| 160 | + the ACL Python API, assembles the args, and dispatches the kernel. | ||
| 161 | +5. The two graphs alternate between two sets of device addresses and report total time, average time, | ||
| 162 | + and speedup. | ||
| 163 | + | ||
| 164 | +### Declarative Address Refresh | ||
| 165 | + | ||
| 166 | +```text | ||
| 167 | +Compile phase: | ||
| 168 | + compile(x, y, z) | ||
| 169 | + ├─ Read add_custom.asc | ||
| 170 | + ├─ bisheng + llvm-objcopy extract .aicore_binary -> kernel binary | ||
| 171 | + └─ Cache the binary by input shape/dtype | ||
| 172 | + | ||
| 173 | +Task generation phase: | ||
| 174 | + declare_launch_args(x, y, z) | ||
| 175 | + ├─ KernelArgs: InputAddr{0}, InputAddr{1}, OutputAddr{0} | ||
| 176 | + ├─ AnnotatedKernelLaunchInfo { kernel_name, kernel_bin, block_dim, stream_id } | ||
| 177 | + └─ ctx.add_launch(launch_info, args) | ||
| 178 | + | ||
| 179 | +Execution phase: | ||
| 180 | + GE refreshes the current input/output addresses according to the saved args layout | ||
| 181 | +``` | ||
| 182 | + | ||
| 183 | +### No-Refresh Baseline | ||
| 184 | + | ||
| 185 | +```text | ||
| 186 | +Every graph execution: | ||
| 187 | + NoRefreshAddCustom.execute(x, y) | ||
| 188 | + ├─ ctx.malloc_output_tensor() | ||
| 189 | + ├─ acl.rt.binary_load_from_file() / binary_get_function() | ||
| 190 | + ├─ kernel_args_init / append(x, y, z) / finalize | ||
| 191 | + └─ acl.rt.launch_kernel_with_config() | ||
| 192 | +``` | ||
| 193 | + | ||
| 194 | +This implementation declares no address slots in the args. Both graphs share identical computation logic | ||
| 195 | +and kernel, so the performance difference reflects the effect of declarative address refresh. | ||
| 196 | + | ||
| 197 | +## Build Artifacts | ||
| 198 | + | ||
| 199 | +- `build/add_custom.aicore.o` | ||
| 200 | + The device-side binary of the Ascend C Add kernel, compiled with `bisheng` and extracted from the `.aicore_binary` section. | ||
| 201 | +- `build/opp/` | ||
| 202 | + The custom operator OPP registration library (`op_proto/custom/libcust_opapi.so` plus platform copies), exposed to GE through `ASCEND_CUSTOM_OPP_PATH`. | ||
| 203 | +- `build/es_output/` | ||
| 204 | + ES graph-building deliverables: `lib64/libes_custom.so` and `whl/es_custom-1.0.0-py3-none-any.whl` (installed into `build/whl_package/`). | ||
| 205 | + | ||
| 206 | +## Result Verification | ||
| 207 | + | ||
| 208 | +On success you can observe: | ||
| 209 | + | ||
| 210 | +- The terminal output contains `AnnotatedAddCustom precision check PASS` and `NoRefreshAddCustom precision check PASS`. | ||
| 211 | +- The terminal output contains `AnnotatedAddCustom`, `NoRefreshAddCustom`, and `Annotated speedup`. | ||
| 212 | +- The terminal output contains `NPU_EXECUTION=PASS` and `Online Python custom-op pipeline PASS`. | ||
| 213 | + | ||
| 214 | +If it fails, check first: | ||
| 215 | + | ||
| 216 | +- Whether `ASCEND_HOME_PATH` is set and the CANN environment is loaded correctly. | ||
| 217 | +- Whether `bisheng`, `llvm-objcopy`, and `cmake` are available. | ||
| 218 | +- Whether `build/es_output/whl/es_custom-1.0.0-py3-none-any.whl` and `build/add_custom.aicore.o` were generated. | ||
| 219 | +- Whether an NPU is available in the current environment. | ||
| 220 | + | ||
| 221 | +## Notes / Limitations | ||
| 222 | + | ||
| 223 | +- The kernel binary is compiled by `bisheng` inside the `compile` callback; the graph compilation phase | ||
| 224 | + includes this overhead, which is excluded from the 100 benchmark iterations. | ||
| 225 | +- `ge.graphRunMode` is set to `1` (`PRIORITY_GRAPH`) to ensure the online execution path. | ||
| 226 | +- The benchmark alternates two sets of device buffers to trigger input/output address changes. | ||
| 227 | +- The speedup is affected by the NPU model and system load, and is for reference only. | ||
| 228 | +- The NPU architecture defaults to `dav-2201` (Atlas A2); override it with `ADD_CUSTOM_NPU_ARCH` for | ||
| 229 | + other models. | ||
| 230 | + | ||
| 231 | +## Appendix | ||
| 232 | + | ||
| 233 | +### Operator Specification | ||
| 234 | + | ||
| 235 | +| Item | Content | | ||
| 236 | +| --- | --- | | ||
| 237 | +| Operator type | `AnnotatedAddCustom` / `NoRefreshAddCustom` | | ||
| 238 | +| Inputs | `x`, `y` | | ||
| 239 | +| Output | `z` | | ||
| 240 | +| Input/output shape | `[8192]` | | ||
| 241 | +| Input/output data type | `float32` | | ||
| 242 | +| Format | `ND` | | ||
| 243 | +| Kernel name | `add_custom` (Ascend C, compiled with `bisheng`) | | ||
| 244 | +| Block size | `1024` | | ||
| 245 | + | ||
| 246 | +### Key Interfaces | ||
| 247 | + | ||
| 248 | +| Interface | Operator | Purpose | | ||
| 249 | +| --- | --- | --- | | ||
| 250 | +| `ge.custom_op.register_op` + `infer_meta` | Both operators | Prototype registration and output shape/dtype inference | | ||
| 251 | +| `register_op_impl` + `compile` | `AnnotatedAddCustom` | Compile the kernel with `bisheng` and cache the binary by shape | | ||
| 252 | +| `register_op_impl` + `declare_launch_args` | `AnnotatedAddCustom` | Declare the kernel launch and input/output address slots | | ||
| 253 | +| `register_op_impl` + `execute` | `NoRefreshAddCustom` | Allocate the output tensor and dispatch the kernel through the ACL API | | ||
| 254 | +| `acl.rt.launch_kernel_with_config` | `NoRefreshAddCustom` | Dispatch the kernel directly through the ACL Python API | | ||
| @@ -0,0 +1,76 @@ | |||
| 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 | +cmake_minimum_required(VERSION 3.16) | ||
| 12 | +project(pypto_add_custom LANGUAGES CXX) | ||
| 13 | + | ||
| 14 | +set(CMAKE_CXX_STANDARD 17) | ||
| 15 | +set(CMAKE_CXX_STANDARD_REQUIRED ON) | ||
| 16 | +set(CMAKE_CXX_EXTENSIONS OFF) | ||
| 17 | + | ||
| 18 | +if(NOT CMAKE_BUILD_TYPE) | ||
| 19 | + set(CMAKE_BUILD_TYPE Release CACHE STRING "Build type" FORCE) | ||
| 20 | +endif() | ||
| 21 | + | ||
| 22 | +set(ES_OUTPUT_PATH "${CMAKE_BINARY_DIR}/es_output") | ||
| 23 | +set(OPP_OUTPUT_PATH "${CMAKE_BINARY_DIR}/opp") | ||
| 24 | +set(ASCEND_HOME_PATH "$ENV{ASCEND_HOME_PATH}") | ||
| 25 | +if(NOT ASCEND_HOME_PATH) | ||
| 26 | + message(FATAL_ERROR "ASCEND_HOME_PATH is empty. Please source CANN set_env.sh first.") | ||
| 27 | +endif() | ||
| 28 | + | ||
| 29 | +list(APPEND CMAKE_MODULE_PATH "${ASCEND_HOME_PATH}/include/ge/cmake") | ||
| 30 | +find_package(GenerateEsPackage REQUIRED) | ||
| 31 | +add_compile_definitions(_GLIBCXX_USE_CXX11_ABI=0 google=ascend_private) | ||
| 32 | + | ||
| 33 | +add_library(pypto_add_custom_op_proto SHARED proto/add_custom.cc) | ||
| 34 | +target_compile_definitions(pypto_add_custom_op_proto PRIVATE OP_PROTO_LIB) | ||
| 35 | +target_compile_options(pypto_add_custom_op_proto PRIVATE -fvisibility=hidden) | ||
| 36 | +target_include_directories(pypto_add_custom_op_proto PRIVATE | ||
| 37 | + "${ASCEND_HOME_PATH}/include" | ||
| 38 | + "${ASCEND_HOME_PATH}/include/external" | ||
| 39 | +) | ||
| 40 | +target_link_libraries(pypto_add_custom_op_proto PRIVATE | ||
| 41 | + "${ASCEND_HOME_PATH}/lib64/libopp_registry.so" | ||
| 42 | + "${ASCEND_HOME_PATH}/lib64/libgraph.so" | ||
| 43 | + "${ASCEND_HOME_PATH}/lib64/libgraph_base.so" | ||
| 44 | +) | ||
| 45 | + | ||
| 46 | +if(CMAKE_SYSTEM_NAME STREQUAL "Windows") | ||
| 47 | + set(OPP_OS_TYPE "windows") | ||
| 48 | +else() | ||
| 49 | + set(OPP_OS_TYPE "linux") | ||
| 50 | +endif() | ||
| 51 | +string(TOLOWER "${CMAKE_SYSTEM_PROCESSOR}" PROCESSOR_LOWER) | ||
| 52 | +if(PROCESSOR_LOWER MATCHES "^(aarch64|arm64)$") | ||
| 53 | + set(OPP_CPU_TYPE "aarch64") | ||
| 54 | +elseif(PROCESSOR_LOWER MATCHES "^(x86_64|amd64)$") | ||
| 55 | + set(OPP_CPU_TYPE "x86_64") | ||
| 56 | +else() | ||
| 57 | + set(OPP_CPU_TYPE "${PROCESSOR_LOWER}") | ||
| 58 | +endif() | ||
| 59 | +set(OPP_PROTO_OUTPUT_PATH "${OPP_OUTPUT_PATH}/op_proto/custom") | ||
| 60 | +set(CUSTOM_OP_OUTPUT_PATH "${OPP_OUTPUT_PATH}/op_graph/lib/${OPP_OS_TYPE}/${OPP_CPU_TYPE}") | ||
| 61 | +set_target_properties(pypto_add_custom_op_proto PROPERTIES | ||
| 62 | + OUTPUT_NAME cust_opapi | ||
| 63 | + LIBRARY_OUTPUT_DIRECTORY "${OPP_PROTO_OUTPUT_PATH}" | ||
| 64 | + RUNTIME_OUTPUT_DIRECTORY "${OPP_PROTO_OUTPUT_PATH}" | ||
| 65 | + ARCHIVE_OUTPUT_DIRECTORY "${OPP_PROTO_OUTPUT_PATH}" | ||
| 66 | +) | ||
| 67 | +add_custom_command(TARGET pypto_add_custom_op_proto POST_BUILD | ||
| 68 | + COMMAND ${CMAKE_COMMAND} -E make_directory "${CUSTOM_OP_OUTPUT_PATH}" | ||
| 69 | + COMMAND ${CMAKE_COMMAND} -E copy "$<TARGET_FILE:pypto_add_custom_op_proto>" "${CUSTOM_OP_OUTPUT_PATH}" | ||
| 70 | +) | ||
| 71 | + | ||
| 72 | +add_es_library_and_whl( | ||
| 73 | + ES_LINKABLE_AND_ALL_TARGET es_custom | ||
| 74 | + OPP_PROTO_TARGET pypto_add_custom_op_proto | ||
| 75 | + OUTPUT_PATH ${ES_OUTPUT_PATH} | ||
| 76 | +) | ||
| @@ -0,0 +1,199 @@ | |||
| 1 | +# PyPTO 算子通过 GE 自定义算子入图样例 | ||
| 2 | + | ||
| 3 | +## 样例概述 | ||
| 4 | + | ||
| 5 | +- 构图入口：`GE`（Python ES 构图接口，RT2 动态执行路径） | ||
| 6 | +- 算子编程语言：`PyPTO`（Python Tensor 算子编程） | ||
| 7 | +- 编译方式：`JIT`（主线程预热编译，进程内缓存复用） | ||
| 8 | +- 模型下沉能力：`不涉及` | ||
| 9 | +- 核心链路：`PyPTO kernel -> Python 自定义算子 execute 回调 -> torchair 零拷贝桥接 -> GE 图执行` | ||
| 10 | +- 与其他 sample 的区别：本样例把 PyPTO JIT kernel 接入 GE 自定义算子的 Python `execute` 回调，不编译任何 `.so` 或 kernel binary，也不涉及模型下沉。 | ||
| 11 | + | ||
| 12 | +本样例注册 `PyptoAddCustom` 自定义算子：`@register_op` 注册原型和 `infer_meta`， | ||
| 13 | +`@register_op_impl` 注册 schema-bound `execute(x, y)` 实现。执行时 execute 回调通过 GE | ||
| 14 | +`EagerOpExecutionContext` 申请输出 Tensor，把输入/输出 device 地址经 torchair `as_torch_tensors` 零拷贝包装为 | ||
| 15 | +torch Tensor，再调用 `@pypto.jit` 定义的 Add kernel。PyPTO kernel 由 `src/run.py` 在 GE | ||
| 16 | +初始化前于主线程预热编译，execute 回调只消费进程内编译缓存。 | ||
| 17 | + | ||
| 18 | +## 适用场景 | ||
| 19 | + | ||
| 20 | +- 想了解 PyPTO kernel 如何通过 Python 自定义算子接入 GE 图执行。 | ||
| 21 | +- 想参考 GE Tensor 与 torch Tensor 之间的零拷贝桥接方式。 | ||
| 22 | +- 想验证 PyPTO JIT kernel 在 GE 在线执行链路下的数值正确性。 | ||
| 23 | + | ||
| 24 | +## 前置依赖 | ||
| 25 | + | ||
| 26 | +### CANN | ||
| 27 | + | ||
| 28 | +- 参考 [安装指导](../../../docs/zh/quick_install.md#1-环境准备) 完成 toolkit 和 ops 包安装。 | ||
| 29 | +- CANN 版本需 >= 9.2.0，`pypto` Python 包随 CANN toolkit 发布（位于 `python/site-packages/pypto`）。 | ||
| 30 | + | ||
| 31 | +### 框架与插件 | ||
| 32 | + | ||
| 33 | +- 已安装 `PyTorch` 和 `torch_npu`。 | ||
| 34 | + | ||
| 35 | +参考： | ||
| 36 | + | ||
| 37 | +- [Ascend Extension for PyTorch](https://gitcode.com/Ascend/pytorch) | ||
| 38 | +- [Ascend Extension for PyTorch 昇腾社区说明](https://hiascend.com/document/redirect/Pytorch-index) | ||
| 39 | + | ||
| 40 | +### 环境变量 | ||
| 41 | + | ||
| 42 | +- `ASCEND_HOME_PATH` | ||
| 43 | +- `LD_LIBRARY_PATH` 等 CANN 运行时相关变量 | ||
| 44 | +- `run.sh` 会自动设置 `PYTHONPATH`、`LD_LIBRARY_PATH` 和 `ASCEND_CUSTOM_OPP_PATH` | ||
| 45 | + | ||
| 46 | +### 额外依赖 | ||
| 47 | + | ||
| 48 | +- `cmake`、Python 3/`pip` | ||
| 49 | + | ||
| 50 | +## 快速运行 | ||
| 51 | + | ||
| 52 | +在 `examples/custom_op/pypto_add_custom` 目录下执行： | ||
| 53 | + | ||
| 54 | +### 推荐方式 | ||
| 55 | + | ||
| 56 | +```bash | ||
| 57 | +source ${ASCEND_HOME_PATH}/set_env.sh | ||
| 58 | +bash run.sh | ||
| 59 | +``` | ||
| 60 | + | ||
| 61 | +`run.sh` 会自动完成构建、安装和在线执行： | ||
| 62 | + | ||
| 63 | +1. 构建 custom OPP 注册库和 ES Python 构图接口（产物全部位于 `build/` 下）。 | ||
| 64 | +2. 运行 `src/run.py`：主线程预热编译 PyPTO kernel，构建 `PyptoAddCustom` 图并在线执行。 | ||
| 65 | + | ||
| 66 | +若运行成功，终端会打印类似： | ||
| 67 | + | ||
| 68 | +```text | ||
| 69 | +[OnlinePython] PyPTO kernel warm-up PASS | ||
| 70 | +[OnlinePython] PyptoAddCustom precision check PASS | ||
| 71 | +[Perf] input shape: [8192], dtype: float32 | ||
| 72 | +[Perf] iters: 100 | ||
| 73 | +[Perf] PyptoAddCustom: xxx us (avg xxx us/iter) | ||
| 74 | +[OnlinePython] NPU_EXECUTION=PASS | ||
| 75 | +``` | ||
| 76 | + | ||
| 77 | +## 目录结构与关键文件 | ||
| 78 | + | ||
| 79 | +```text | ||
| 80 | +pypto_add_custom | ||
| 81 | +├── CMakeLists.txt # 构建 custom OPP 注册库和 ES wheel | ||
| 82 | +├── README.md | ||
| 83 | +├── README_en.md | ||
| 84 | +├── run.sh # 构建并执行在线 GE 图 | ||
| 85 | +├── proto | ||
| 86 | +│ └── add_custom.h # gen_esb 使用的 C++ 构图原型 | ||
| 87 | +└── src | ||
| 88 | + ├── tensor_bridge.py # GE Tensor 到 torch Tensor 的零拷贝桥接（torchair） | ||
| 89 | + ├── pypto_add_kernel.py # @pypto.jit Add kernel（共享） | ||
| 90 | + ├── run.py # 预热 + GE 图构建与在线执行 | ||
| 91 | + └── ge | ||
| 92 | + └── add_custom.py # Python 原型、infer_meta 和 execute 回调 | ||
| 93 | +``` | ||
| 94 | + | ||
| 95 | +重点文件： | ||
| 96 | + | ||
| 97 | +- `src/ge/add_custom.py` | ||
| 98 | + 实现 `PyptoAddCustom`：`@register_op` 注册原型与 `infer_meta`；`@register_op_impl` 的 | ||
| 99 | + `execute` 申请输出 Tensor、经 torchair 桥接输入输出地址并在设备同步保护下执行 PyPTO kernel。 | ||
| 100 | +- `src/tensor_bridge.py` | ||
| 101 | + GE device 地址到 torch Tensor 的零拷贝桥接，复用 torchair `as_torch_tensors` | ||
| 102 | + （torchair Python 自定义算子回调的同款机制），设备内存始终归 GE 所有。 | ||
| 103 | +- `src/pypto_add_kernel.py` | ||
| 104 | + `@pypto.jit` 定义的 Add kernel，含 `pypto.set_vec_tile_shapes(1, 1024)` tile 配置， | ||
| 105 | + 供 execute 回调和直连测试共享。 | ||
| 106 | +- `src/run.py` | ||
| 107 | + 在 GE 初始化前主线程预热编译 kernel，随后构图、精度校验和 100 轮计时。 | ||
| 108 | +- `proto/add_custom.h` | ||
| 109 | + 注册 `PyptoAddCustom` 构图侧算子类型，供 `gen_esb` 生成 Python ES 构图接口。 | ||
| 110 | + | ||
| 111 | +## 核心链路 | ||
| 112 | + | ||
| 113 | +1. `run.py` 在主线程导入 torch/torch_npu/pypto 并预热编译 PyPTO kernel，填充进程内编译缓存。 | ||
| 114 | +2. `run.py` 以 `ge.exec.static_model_ops_lower_limit=-1` 初始化 GE，使图走 RT2 动态执行路径， | ||
| 115 | + Python execute 回调在每次图执行时被调用。 | ||
| 116 | +3. `infer_meta` 校验输入并返回输出元信息。 | ||
| 117 | +4. `execute` 通过 `ctx.malloc_output_tensor` 申请输出，把 x/y/z 的 device 地址经 torchair | ||
| 118 | + 零拷贝包装为 torch Tensor。 | ||
| 119 | +5. execute 在设备同步保护下调用 PyPTO kernel（torch 当前 stream），kernel 结果由 GE 图 | ||
| 120 | + 后续读取。 | ||
| 121 | + | ||
| 122 | +## 构建产物 | ||
| 123 | + | ||
| 124 | +- `build/opp/` | ||
| 125 | + 自定义算子 OPP 注册库（`op_proto/custom/libcust_opapi.so` 及 `op_graph/lib/<os>/<arch>/` 平台拷贝），通过 `ASCEND_CUSTOM_OPP_PATH` 暴露给 GE。 | ||
| 126 | +- `build/es_custom_build/python_package/es_custom/` | ||
| 127 | + 构建自动生成的 ES Python 构图接口（含 `__init__.py` 与 `gen_esb` 生成的 wrapper），通过 `PYTHONPATH` 加载。 | ||
| 128 | +- `build/es_output/` | ||
| 129 | + ES 库交付件：`lib64/libes_custom.so`（`LD_LIBRARY_PATH` 加载）和 `whl/es_custom-1.0.0-py3-none-any.whl`。 | ||
| 130 | + | ||
| 131 | +## 结果校验 | ||
| 132 | + | ||
| 133 | +成功时可观察到： | ||
| 134 | + | ||
| 135 | +- 终端输出包含 `PyPTO kernel warm-up PASS`。 | ||
| 136 | +- 终端输出包含 `PyptoAddCustom precision check PASS`。 | ||
| 137 | +- 终端输出包含 `NPU_EXECUTION=PASS` 和 `Online Python PyPTO custom-op pipeline PASS`。 | ||
| 138 | + | ||
| 139 | +若失败，优先检查： | ||
| 140 | + | ||
| 141 | +- `ASCEND_HOME_PATH` 是否已设置并正确加载 CANN 环境。 | ||
| 142 | +- torch/torch_npu/pypto 是否可导入（CANN >= 9.2.0）。 | ||
| 143 | +- `build/opp/`、`build/es_custom_build/python_package/es_custom/` 和 `build/es_output/lib64/` 是否生成。 | ||
| 144 | +- 当前环境是否具备可用 NPU。 | ||
| 145 | + | ||
| 146 | +## PyPTO 支持现状 | ||
| 147 | + | ||
| 148 | +- 当前 `pypto` Python 包仅支持 **JIT 直接执行**：首次调用 `@pypto.jit` kernel 时自动完成 | ||
| 149 | + 解析、编译和执行，编译结果缓存在进程内部。 | ||
| 150 | +- **尚未提供独立的编译接口**：无法把 PyPTO kernel 预编译为可独立发布的 kernel binary 文件 | ||
| 151 | + （例如 Ascend C 通过 `bisheng` 编译得到的 `.aicore.o`）。 | ||
| 152 | +- **尚未提供 binary 导出接口**：JIT 编译产物为 PyPTO 内部格式，无法转换为 GE `AnnotatedArgs` | ||
| 153 | + 声明式地址刷新所需的 `kernel_bin`，因此无法接入 `compile`/`declare_launch_args` 链路， | ||
| 154 | + 也无法通过 ATC 下沉到 om 离线模型。 | ||
| 155 | +- 本样例展示的 **Python `execute` 回调 + torchair 桥接** 是当前 PyPTO kernel 接入 GE 的可行路径。 | ||
| 156 | + | ||
| 157 | +## 注意事项 / 限制 | ||
| 158 | + | ||
| 159 | +- 已验证的产品：`Atlas A2 训练系列产品`。样例不包含硬编码的 NPU 架构，PyPTO JIT 在编译时自动探测实际设备架构；理论上 CANN 9.2 中 PyPTO 和 torch_npu 支持的其他平台（如 Atlas A3、Ascend 950 系列）也可运行，其他平台未实测。 | ||
| 160 | +- 样例当前使用 `float32` 的 `8192` 一维输入做结果校验，输入元素个数需能被 tile size（1024）整除。 | ||
| 161 | +- **kernel 必须在 GE 初始化前于主线程预热编译**：在 GE 回调内触发 JIT 编译会因编译器子进程 | ||
| 162 | + 与 GE 运行时冲突而挂起。 | ||
| 163 | +- `run.py` 通过 `ge.exec.static_model_ops_lower_limit=-1` 强制走 RT2 动态执行路径，使 execute | ||
| 164 | + 回调在每次图执行时被调用；默认静态路径下 execute 仅在任务生成阶段调用一次，与 PyPTO 的 | ||
| 165 | + 即时启动模型不兼容。 | ||
| 166 | +- **torchair 桥接要求 `import torchair.llm_datadist` 可用**（`src/tensor_bridge.py` 直接使用公开 | ||
| 167 | + `create_npu_tensors` 接口，不含兼容性降级逻辑）。部分 torch/torchair 版本组合下 torchair | ||
| 168 | + 包初始化会失败（例如 `torch.fx` 缺少 `hint_int`），此时需将 torch_npu 与 torchair 升级为 | ||
| 169 | + 相互匹配的版本后再运行本样例。 | ||
| 170 | +- Python 导入顺序要求 `import torch` / `import torch_npu` 在 `import pypto` 之前， | ||
| 171 | + 否则可能触发 `libc10.so` 静态 TLS 相关报错。 | ||
| 172 | +- execute 回调以设备同步保证数据可见性，单轮耗时（约 4 ms）高于 Ascend C + AnnotatedArgs | ||
| 173 | + 路径（参考 `annotated_args_refresh_add_custom`，约 0.4 ms），本样例聚焦链路打通而非性能。 | ||
| 174 | + | ||
| 175 | +## 附录 | ||
| 176 | + | ||
| 177 | +### 算子规格 | ||
| 178 | + | ||
| 179 | +| 项目 | 内容 | | ||
| 180 | +| --- | --- | | ||
| 181 | +| 算子类型 | `PyptoAddCustom` | | ||
| 182 | +| 输入 | `x`, `y` | | ||
| 183 | +| 输出 | `z` | | ||
| 184 | +| 输入 shape | `8192` | | ||
| 185 | +| 输出 shape | `8192` | | ||
| 186 | +| 数据类型 | `float32` | | ||
| 187 | +| 格式 | `ND` | | ||
| 188 | +| kernel 名称 | `pypto_add_kernel`（PyPTO JIT） | | ||
| 189 | +| tile 大小 | `1024` | | ||
| 190 | + | ||
| 191 | +### 关键接口 | ||
| 192 | + | ||
| 193 | +| 接口 | 用途 | | ||
| 194 | +| --- | --- | | ||
| 195 | +| `ge.custom_op.register_op` + `infer_meta` | 原型注册与输出 shape/dtype 推导 | | ||
| 196 | +| `ge.custom_op.register_op_impl` + `execute` | 申请输出 Tensor 并执行 PyPTO kernel | | ||
| 197 | +| `ge.custom_op.get_execute_ctx` | 获取 execute 回调上下文（输出申请） | | ||
| 198 | +| `pypto.jit` + `pypto.set_vec_tile_shapes` | PyPTO kernel 定义与 tile 配置 | | ||
| 199 | +| `torchair as_torch_tensors` | GE device 地址到 torch Tensor 的零拷贝桥接 | | ||
| @@ -0,0 +1,218 @@ | |||
| 1 | +# PyPTO Operator Enters GE Graph Through Custom Operator | ||
| 2 | + | ||
| 3 | +## Sample Overview | ||
| 4 | + | ||
| 5 | +- Graph building entry: `GE` (Python ES graph-building APIs, RT2 dynamic execution path) | ||
| 6 | +- Operator language: `PyPTO` (Python Tensor operator programming) | ||
| 7 | +- Compilation mode: `JIT` (warm-up compiled on the main thread, reused from the in-process cache) | ||
| 8 | +- Model sinking capability: `N/A` | ||
| 9 | +- Core path: `PyPTO kernel -> Python custom-op execute callback -> torchair zero-copy bridge -> GE graph execution` | ||
| 10 | +- Difference from other samples: this sample wires a PyPTO JIT kernel into the Python `execute` callback of a GE custom operator. No `.so` or kernel binary is compiled, and model sinking is not involved. | ||
| 11 | + | ||
| 12 | +This sample registers the `PyptoAddCustom` custom operator: `@register_op` registers the prototype | ||
| 13 | +and `infer_meta`, and `@register_op_impl` registers the schema-bound `execute(x, y)` implementation. | ||
| 14 | +At run time the execute callback allocates the output tensor through the GE | ||
| 15 | +`EagerOpExecutionContext`, wraps the input/output device memory as torch tensors through a | ||
| 16 | +zero-copy torchair bridge (`as_torch_tensors`), and invokes the Add kernel defined with `@pypto.jit`. The PyPTO kernel is | ||
| 17 | +warm-up compiled by `src/run.py` on the main thread before GE starts; the execute callback only | ||
| 18 | +consumes the in-process compilation cache. | ||
| 19 | + | ||
| 20 | +## Applicable Scenarios | ||
| 21 | + | ||
| 22 | +- Understanding how a PyPTO kernel enters GE graph execution through a Python custom operator. | ||
| 23 | +- Referencing the zero-copy bridge between GE tensors and torch tensors. | ||
| 24 | +- Validating the numerics of a PyPTO JIT kernel on the GE online execution path. | ||
| 25 | + | ||
| 26 | +## Prerequisites | ||
| 27 | + | ||
| 28 | +### CANN | ||
| 29 | + | ||
| 30 | +- Follow [installation guide](../../../docs/en/quick_install.md#1-environment-preparation) to complete toolkit and ops package installation. | ||
| 31 | +- CANN version must be >= 9.2.0. The `pypto` Python package is released with the CANN toolkit (located in `python/site-packages/pypto`). | ||
| 32 | + | ||
| 33 | +### Frameworks and Plugins | ||
| 34 | + | ||
| 35 | +- `PyTorch` and `torch_npu` have been installed. | ||
| 36 | + | ||
| 37 | +References: | ||
| 38 | + | ||
| 39 | +- [Ascend Extension for PyTorch](https://gitcode.com/Ascend/pytorch) | ||
| 40 | +- [Ascend Extension for PyTorch on Ascend community](https://hiascend.com/document/redirect/Pytorch-index) | ||
| 41 | + | ||
| 42 | +### Environment Variables | ||
| 43 | + | ||
| 44 | +- `ASCEND_HOME_PATH` | ||
| 45 | +- `LD_LIBRARY_PATH` and other CANN runtime related variables | ||
| 46 | +- `run.sh` automatically sets `PYTHONPATH`, `LD_LIBRARY_PATH`, and `ASCEND_CUSTOM_OPP_PATH` | ||
| 47 | + | ||
| 48 | +### Additional Dependencies | ||
| 49 | + | ||
| 50 | +- `cmake`, Python 3/`pip` | ||
| 51 | + | ||
| 52 | +## Quick Run | ||
| 53 | + | ||
| 54 | +Run the following commands in the `examples/custom_op/pypto_add_custom` directory: | ||
| 55 | + | ||
| 56 | +### Recommended | ||
| 57 | + | ||
| 58 | +```bash | ||
| 59 | +source ${ASCEND_HOME_PATH}/set_env.sh | ||
| 60 | +bash run.sh | ||
| 61 | +``` | ||
| 62 | + | ||
| 63 | +`run.sh` automatically builds, installs, and runs the online graph: | ||
| 64 | + | ||
| 65 | +1. Builds the custom OPP registration library and the ES Python graph-building APIs | ||
| 66 | + (all artifacts are placed under `build/`). | ||
| 67 | +2. Runs `src/run.py`: warm-up compiles the PyPTO kernel on the main thread, builds the | ||
| 68 | + `PyptoAddCustom` graph, and executes it online. | ||
| 69 | + | ||
| 70 | +On success the terminal prints something like: | ||
| 71 | + | ||
| 72 | +```text | ||
| 73 | +[OnlinePython] PyPTO kernel warm-up PASS | ||
| 74 | +[OnlinePython] PyptoAddCustom precision check PASS | ||
| 75 | +[Perf] input shape: [8192], dtype: float32 | ||
| 76 | +[Perf] iters: 100 | ||
| 77 | +[Perf] PyptoAddCustom: xxx us (avg xxx us/iter) | ||
| 78 | +[OnlinePython] NPU_EXECUTION=PASS | ||
| 79 | +``` | ||
| 80 | + | ||
| 81 | +## Directory Structure and Key Files | ||
| 82 | + | ||
| 83 | +```text | ||
| 84 | +pypto_add_custom | ||
| 85 | +├── CMakeLists.txt # Build the custom OPP library and ES wheel | ||
| 86 | +├── README.md | ||
| 87 | +├── README_en.md | ||
| 88 | +├── run.sh # Build and run the online GE graph | ||
| 89 | +├── proto | ||
| 90 | +│ └── add_custom.h # C++ graph prototype consumed by gen_esb | ||
| 91 | +└── src | ||
| 92 | + ├── tensor_bridge.py # Zero-copy bridge from GE tensors to torch tensors (torchair) | ||
| 93 | + ├── pypto_add_kernel.py # @pypto.jit Add kernel (shared) | ||
| 94 | + ├── run.py # Warm-up + GE graph building and online execution | ||
| 95 | + └── ge | ||
| 96 | + └── add_custom.py # Python prototype, infer_meta, and execute callback | ||
| 97 | +``` | ||
| 98 | + | ||
| 99 | +Key files: | ||
| 100 | + | ||
| 101 | +- `src/ge/add_custom.py` | ||
| 102 | + Implements `PyptoAddCustom`: `@register_op` registers the prototype and `infer_meta`; the | ||
| 103 | + `@register_op_impl` `execute` allocates the output tensor, bridges the input/output device memory | ||
| 104 | + through the torchair bridge, and runs the PyPTO kernel under device synchronization. | ||
| 105 | +- `src/tensor_bridge.py` | ||
| 106 | + The zero-copy bridge from GE device memory to torch tensors. It reuses the torchair | ||
| 107 | + ``as_torch_tensors`` entry (the same machinery behind torchair's Python custom-op | ||
| 108 | + callbacks), while the device memory stays owned by GE. | ||
| 109 | +- `src/pypto_add_kernel.py` | ||
| 110 | + The Add kernel defined with `@pypto.jit`, including the `pypto.set_vec_tile_shapes(1, 1024)` | ||
| 111 | + tile configuration; shared by the execute callback and the direct test. | ||
| 112 | +- `src/run.py` | ||
| 113 | + Warm-up compiles the kernel on the main thread before GE starts, then builds the graph, validates | ||
| 114 | + precision, and measures 100 iterations. | ||
| 115 | +- `proto/add_custom.h` | ||
| 116 | + Registers the `PyptoAddCustom` graph operator type for `gen_esb` to generate the Python ES | ||
| 117 | + graph-building APIs. | ||
| 118 | + | ||
| 119 | +## Core Path | ||
| 120 | + | ||
| 121 | +1. `run.py` imports torch/torch_npu/pypto on the main thread and warm-up compiles the PyPTO kernel | ||
| 122 | + to populate the in-process compilation cache. | ||
| 123 | +2. `run.py` initializes GE with `ge.exec.static_model_ops_lower_limit=-1` so the graph takes the | ||
| 124 | + RT2 dynamic execution path, where the Python execute callback runs for every graph execution. | ||
| 125 | +3. `infer_meta` validates the inputs and returns the output metadata. | ||
| 126 | +4. `execute` allocates the output through `ctx.malloc_output_tensor` and wraps the x/y/z device | ||
| 127 | + device memory as torch tensors through the zero-copy torchair bridge. | ||
| 128 | +5. `execute` invokes the PyPTO kernel (current torch stream) under device synchronization; the | ||
| 129 | + result is consumed by the subsequent tasks of the GE graph. | ||
| 130 | + | ||
| 131 | +## Build Artifacts | ||
| 132 | + | ||
| 133 | +- `build/opp/` | ||
| 134 | + The custom operator OPP registration library (`op_proto/custom/libcust_opapi.so` plus `op_graph/lib/<os>/<arch>/` platform copies), exposed to GE through `ASCEND_CUSTOM_OPP_PATH`. | ||
| 135 | +- `build/es_custom_build/python_package/es_custom/` | ||
| 136 | + The build-generated ES Python graph-building APIs (including the auto-generated `__init__.py` and the wrappers generated by `gen_esb`), loaded through `PYTHONPATH`. | ||
| 137 | +- `build/es_output/` | ||
| 138 | + ES deliverables: `lib64/libes_custom.so` (loaded through `LD_LIBRARY_PATH`) and `whl/es_custom-1.0.0-py3-none-any.whl`. | ||
| 139 | + | ||
| 140 | +## Result Verification | ||
| 141 | + | ||
| 142 | +On success you can observe: | ||
| 143 | + | ||
| 144 | +- The terminal output contains `PyPTO kernel warm-up PASS`. | ||
| 145 | +- The terminal output contains `PyptoAddCustom precision check PASS`. | ||
| 146 | +- The terminal output contains `NPU_EXECUTION=PASS` and `Online Python PyPTO custom-op pipeline PASS`. | ||
| 147 | + | ||
| 148 | +If it fails, check first: | ||
| 149 | + | ||
| 150 | +- Whether `ASCEND_HOME_PATH` is set and the CANN environment is loaded correctly. | ||
| 151 | +- Whether torch/torch_npu/pypto can be imported (CANN >= 9.2.0). | ||
| 152 | +- Whether `build/opp/`, `build/es_custom_build/python_package/es_custom/`, and `build/es_output/lib64/` were generated. | ||
| 153 | +- Whether an NPU is available in the current environment. | ||
| 154 | + | ||
| 155 | +## PyPTO Support Status | ||
| 156 | + | ||
| 157 | +- The current `pypto` Python package only supports **JIT direct execution**: the first call of a | ||
| 158 | + `@pypto.jit` kernel automatically parses, compiles, and executes it, and the compiled result is | ||
| 159 | + cached in the process. | ||
| 160 | +- **No standalone compilation interface yet**: a PyPTO kernel cannot be pre-compiled into an | ||
| 161 | + independently publishable kernel binary file (for example the `.aicore.o` produced by compiling | ||
| 162 | + Ascend C with `bisheng`). | ||
| 163 | +- **No binary export interface yet**: the JIT compilation artifact is an internal PyPTO format and | ||
| 164 | + cannot be converted into the `kernel_bin` required by the GE `AnnotatedArgs` declarative refresh | ||
| 165 | + mechanism. Therefore it cannot be wired into the `compile`/`declare_launch_args` path, nor sunk | ||
| 166 | + into an offline om model through ATC. | ||
| 167 | +- The **Python `execute` callback + torchair bridge** shown in this sample is the currently viable | ||
| 168 | + path to integrate PyPTO kernels with GE. | ||
| 169 | + | ||
| 170 | +## Notes / Limitations | ||
| 171 | + | ||
| 172 | +- Verified products: `Atlas A2 training series`. The sample contains no hardcoded NPU architecture; PyPTO JIT detects the actual device architecture at compile time. Other platforms supported by PyPTO and torch_npu in CANN 9.2 (such as Atlas A3 or the Ascend 950 series) should also work, but are untested. | ||
| 173 | +- The sample validates results with a `float32` one-dimensional input of `8192` elements. The number | ||
| 174 | + of input elements must be divisible by the tile size (1024). | ||
| 175 | +- **The kernel must be warm-up compiled on the main thread before GE starts**: triggering JIT | ||
| 176 | + compilation inside a GE callback hangs because the compiler subprocess conflicts with the GE | ||
| 177 | + runtime. | ||
| 178 | +- `run.py` forces the RT2 dynamic execution path through `ge.exec.static_model_ops_lower_limit=-1` | ||
| 179 | + so the execute callback runs for every graph execution. On the default static path the callback | ||
| 180 | + runs only once at task generation time, which is incompatible with the immediate launch model of | ||
| 181 | + PyPTO. | ||
| 182 | +- **The torchair bridge requires a working ``import torchair.llm_datadist``** (``src/tensor_bridge.py`` | ||
| 183 | + uses the public ``create_npu_tensors`` entry directly, without compatibility fallbacks). Some | ||
| 184 | + torch/torchair version combinations fail during the torchair package init (for example a missing | ||
| 185 | + ``hint_int`` in ``torch.fx``); upgrade torch_npu and torchair to matching versions before running | ||
| 186 | + this sample in that case. | ||
| 187 | +- Python import order requires `import torch` / `import torch_npu` before `import pypto`; otherwise a | ||
| 188 | + `libc10.so` static TLS error may occur. | ||
| 189 | +- The execute callback guarantees data visibility with device synchronization, so the per-iteration | ||
| 190 | + latency (about 4 ms) is higher than the Ascend C + AnnotatedArgs path (see | ||
| 191 | + `annotated_args_refresh_add_custom`, about 0.4 ms). This sample focuses on wiring the path, not | ||
| 192 | + on performance. | ||
| 193 | + | ||
| 194 | +## Appendix | ||
| 195 | + | ||
| 196 | +### Operator Specification | ||
| 197 | + | ||
| 198 | +| Item | Content | | ||
| 199 | +| --- | --- | | ||
| 200 | +| Operator type | `PyptoAddCustom` | | ||
| 201 | +| Inputs | `x`, `y` | | ||
| 202 | +| Output | `z` | | ||
| 203 | +| Input shape | `8192` | | ||
| 204 | +| Output shape | `8192` | | ||
| 205 | +| Data type | `float32` | | ||
| 206 | +| Format | `ND` | | ||
| 207 | +| Kernel name | `pypto_add_kernel` (PyPTO JIT) | | ||
| 208 | +| Tile size | `1024` | | ||
| 209 | + | ||
| 210 | +### Key Interfaces | ||
| 211 | + | ||
| 212 | +| Interface | Purpose | | ||
| 213 | +| --- | --- | | ||
| 214 | +| `ge.custom_op.register_op` + `infer_meta` | Prototype registration and output shape/dtype inference | | ||
| 215 | +| `ge.custom_op.register_op_impl` + `execute` | Allocate the output tensor and run the PyPTO kernel | | ||
| 216 | +| `ge.custom_op.get_execute_ctx` | Access the execute callback context (output allocation) | | ||
| 217 | +| `pypto.jit` + `pypto.set_vec_tile_shapes` | PyPTO kernel definition and tile configuration | | ||
| 218 | +| `torchair as_torch_tensors` | Zero-copy bridge from GE device memory to torch tensors | | ||
| @@ -0,0 +1,15 @@ | |||
| 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 | +// Keep this translation unit as the OPP proto library input for gen_esb. | ||
| 14 | +// Runtime prototype and infer-meta implementations are registered by the | ||
| 15 | +// Python decorators in src/ge/pypto_add_custom.py. | ||
| @@ -0,0 +1,25 @@ | |||
| 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_PYPTO_ADD_CUSTOM_PROTO_ADD_CUSTOM_H_ | ||
| 12 | +#define EXAMPLES_CUSTOM_OP_PYPTO_ADD_CUSTOM_PROTO_ADD_CUSTOM_H_ | ||
| 13 | + | ||
| 14 | +#include "graph/operator_reg.h" | ||
| 15 | + | ||
| 16 | +namespace ge { | ||
| 17 | +REG_OP(PyptoAddCustom) | ||
| 18 | + .INPUT(x1, "T") | ||
| 19 | + .INPUT(x2, "T") | ||
| 20 | + .OUTPUT(output0, "T") | ||
| 21 | + .DATATYPE(T, TensorType({DT_FLOAT})) | ||
| 22 | + .OP_END_FACTORY_REG(PyptoAddCustom); | ||
| 23 | +} // namespace ge | ||
| 24 | + | ||
| 25 | +#endif // EXAMPLES_CUSTOM_OP_PYPTO_ADD_CUSTOM_PROTO_ADD_CUSTOM_H_ | ||
| @@ -0,0 +1,62 @@ | |||
| 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 | +OPP_ROOT="${BUILD_DIR}/opp" | ||
| 18 | +ES_LIB_DIR="${BUILD_DIR}/es_output/lib64" | ||
| 19 | +ES_PYTHON_PACKAGE_DIR="${BUILD_DIR}/es_custom_build/python_package" | ||
| 20 | +LOG_FILE="${BUILD_DIR}/online_pypto.log" | ||
| 21 | + | ||
| 22 | +require_command() { | ||
| 23 | + if ! command -v "$1" >/dev/null 2>&1; then | ||
| 24 | + echo "[ERROR] required command not found: $1" >&2 | ||
| 25 | + exit 1 | ||
| 26 | + fi | ||
| 27 | +} | ||
| 28 | + | ||
| 29 | +if [[ -z "${ASCEND_HOME_PATH:-}" || ! -d "${ASCEND_HOME_PATH}" ]]; then | ||
| 30 | + echo "[ERROR] ASCEND_HOME_PATH is empty. Please source CANN set_env.sh first." >&2 | ||
| 31 | + exit 1 | ||
| 32 | +fi | ||
| 33 | +for command_name in cmake python3; do | ||
| 34 | + require_command "${command_name}" | ||
| 35 | +done | ||
| 36 | + | ||
| 37 | +mkdir -p "${BUILD_DIR}" | ||
| 38 | + | ||
| 39 | +echo "[INFO] configure and build the custom OPP package and the ES Python API" | ||
| 40 | +cmake -S "${SCRIPT_DIR}" -B "${BUILD_DIR}" -DCMAKE_BUILD_TYPE=Release | ||
| 41 | +cmake --build "${BUILD_DIR}" --target build_es_custom -j"$(nproc 2>/dev/null || echo 8)" | ||
| 42 | + | ||
| 43 | +if ! find "${OPP_ROOT}/op_graph/lib" -name libcust_opapi.so -print -quit 2>/dev/null | grep -q .; then | ||
| 44 | + echo "[ERROR] OPP library was not generated under ${OPP_ROOT}/op_graph/lib" >&2 | ||
| 45 | + exit 1 | ||
| 46 | +fi | ||
| 47 | +if [[ ! -s "${ES_PYTHON_PACKAGE_DIR}/es_custom/__init__.py" || ! -s "${ES_LIB_DIR}/libes_custom.so" ]]; then | ||
| 48 | + echo "[ERROR] ES Python package or ES shared library was not generated" >&2 | ||
| 49 | + exit 1 | ||
| 50 | +fi | ||
| 51 | + | ||
| 52 | +# The OPP root is for the C++ proto; the plugin directory is for the Python | ||
| 53 | +# custom-op loader. Both are required by GE's two loaders. | ||
| 54 | +export ASCEND_CUSTOM_OPP_PATH="${OPP_ROOT}:${PLUGIN_DIR}${ASCEND_CUSTOM_OPP_PATH:+:${ASCEND_CUSTOM_OPP_PATH}}" | ||
| 55 | +export PYTHONPATH="${ES_PYTHON_PACKAGE_DIR}:${SCRIPT_DIR}/src${PYTHONPATH:+:${PYTHONPATH}}" | ||
| 56 | +export LD_LIBRARY_PATH="${ES_LIB_DIR}${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" | ||
| 57 | + | ||
| 58 | +echo "[INFO] run the PyPTO Add custom operator on NPU through the online GE graph" | ||
| 59 | +DEVICE_ID="${DEVICE_ID:-0}" python3 "${SCRIPT_DIR}/src/run.py" 2>&1 | tee "${LOG_FILE}" | ||
| 60 | +grep -Fq "PyptoAddCustom precision check PASS" "${LOG_FILE}" | ||
| 61 | +grep -Fq "NPU_EXECUTION=PASS" "${LOG_FILE}" | ||
| 62 | +echo "[INFO] Online Python PyPTO custom-op pipeline PASS" | ||
| @@ -0,0 +1,70 @@ | |||
| 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 | + | ||
| 13 | +"""GE custom operator whose execute callback runs a PyPTO JIT kernel. | ||
| 14 | + | ||
| 15 | +The GE tensors are exposed to PyPTO as zero-copy torch tensors through | ||
| 16 | +the torchair bridge in src/tensor_bridge.py. PyPTO always launches on the | ||
| 17 | +current torch stream, so the callback synchronizes the device around the | ||
| 18 | +launch to keep the kernel ordered with the other tasks of the same GE | ||
| 19 | +graph. | ||
| 20 | + | ||
| 21 | +The kernel itself is compiled by the runner before GE starts (see | ||
| 22 | +src/run.py); the execute callback only consumes the in-process PyPTO | ||
| 23 | +compilation cache, because compiling inside the GE runtime is unsafe. | ||
| 24 | +""" | ||
| 25 | + | ||
| 26 | +import torch_npu | ||
| 27 | + | ||
| 28 | +from tensor_bridge import to_torch | ||
| 29 | +from pypto_add_kernel import pypto_add_kernel | ||
| 30 | + | ||
| 31 | +from ge.custom_op import get_execute_ctx, register_op, register_op_impl | ||
| 32 | +from ge.runtime import Tensor, TensorDesc | ||
| 33 | + | ||
| 34 | + | ||
| 35 | +def _launch_kernel(kernel, tensors): | ||
| 36 | + """Launch the compiled PyPTO kernel with device-wide ordering. | ||
| 37 | + | ||
| 38 | + PyPTO always launches on the current torch stream, so the callback | ||
| 39 | + synchronizes the device before the launch (graph inputs become | ||
| 40 | + visible) and after the launch (the kernel finishes before the graph | ||
| 41 | + reads the output). The kernel itself must already be compiled; see | ||
| 42 | + the warm-up in src/run.py. | ||
| 43 | + """ | ||
| 44 | + torch_npu.npu.synchronize() | ||
| 45 | + kernel(*tensors) | ||
| 46 | + torch_npu.npu.synchronize() | ||
| 47 | + | ||
| 48 | + | ||
| 49 | +@register_op(op_type="PyptoAddCustom") | ||
| 50 | +def pypto_add_infer_meta(x1: TensorDesc, x2: TensorDesc) -> TensorDesc: | ||
| 51 | + """Infer output metadata for the PyPTO Add operator.""" | ||
| 52 | + | ||
| 53 | + if x1.shape.origin_shape.dims != x2.shape.origin_shape.dims: | ||
| 54 | + raise ValueError("PyptoAddCustom inputs must have the same shape") | ||
| 55 | + if x1.data_type != x2.data_type: | ||
| 56 | + raise ValueError("PyptoAddCustom inputs must have the same data type") | ||
| 57 | + return TensorDesc(x1.shape, x1.data_type) | ||
| 58 | + | ||
| 59 | + | ||
| 60 | +@register_op_impl(op_type="PyptoAddCustom") | ||
| 61 | +class PyptoAddCustom: | ||
| 62 | + """Allocate the output tensor and run the PyPTO kernel for GE.""" | ||
| 63 | + | ||
| 64 | + def execute(self, x: Tensor, y: Tensor) -> None: | ||
| 65 | + ctx = get_execute_ctx() | ||
| 66 | + z = ctx.malloc_output_tensor(0, x.shape, x.format, x.data_type) | ||
| 67 | + x_torch = to_torch(x) | ||
| 68 | + y_torch = to_torch(y) | ||
| 69 | + z_torch = to_torch(z) | ||
| 70 | + _launch_kernel(pypto_add_kernel, (x_torch, y_torch, z_torch)) | ||
| @@ -0,0 +1,30 @@ | |||
| 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 | + | ||
| 13 | +"""Shared PyPTO Add kernel used by the direct test and the GE custom op.""" | ||
| 14 | + | ||
| 15 | +import pypto | ||
| 16 | + | ||
| 17 | +NUM_ELEMENTS = 8 * 1024 | ||
| 18 | +TILE_SIZE = 1024 | ||
| 19 | + | ||
| 20 | + | ||
| 21 | +@pypto.jit(new_ir=False) | ||
| 22 | +def pypto_add_kernel( | ||
| 23 | + x: pypto.Tensor([NUM_ELEMENTS], pypto.DT_FP32), | ||
| 24 | + y: pypto.Tensor([NUM_ELEMENTS], pypto.DT_FP32), | ||
| 25 | + out: pypto.Tensor([NUM_ELEMENTS], pypto.DT_FP32), | ||
| 26 | +): | ||
| 27 | + # A one-dimensional tensor needs an explicit vector tile shape whose last | ||
| 28 | + # dimension is 32-byte aligned. | ||
| 29 | + pypto.set_vec_tile_shapes(1, TILE_SIZE) | ||
| 30 | + out.move(pypto.add(x, y)) | ||
| @@ -0,0 +1,184 @@ | |||
| 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 | + | ||
| 13 | +"""Run the PyPTO Add custom operator inside an online GE graph.""" | ||
| 14 | + | ||
| 15 | +import os | ||
| 16 | +import time | ||
| 17 | +import traceback | ||
| 18 | + | ||
| 19 | +# torch must be imported before the GE libraries so that its dynamic | ||
| 20 | +# libraries are loaded with enough static TLS slots. Importing torch_npu, | ||
| 21 | +# pypto, and the torchair bridge in the main thread as well keeps the | ||
| 22 | +# plugin load inside GEInitialize lightweight (sys.modules cache hits) and | ||
| 23 | +# loads the torchair extension before GE starts, which is required because | ||
| 24 | +# loading the shared library in the middle of GE initialization disturbs | ||
| 25 | +# the ACL context binding of the loading thread. | ||
| 26 | +import torch | ||
| 27 | +import torch_npu | ||
| 28 | + | ||
| 29 | +from pypto_add_kernel import NUM_ELEMENTS, pypto_add_kernel | ||
| 30 | +import tensor_bridge | ||
| 31 | + | ||
| 32 | +from es_custom import PyptoAddCustom | ||
| 33 | +from ge.es.graph_builder import GraphBuilder | ||
| 34 | +from ge.ge_global import GeApi | ||
| 35 | +from ge.graph import Tensor | ||
| 36 | +from ge.graph.types import DataType, Format, Placement | ||
| 37 | +from ge.session import Session | ||
| 38 | + | ||
| 39 | +GRAPH_ID = 0 | ||
| 40 | +DEVICE_ID = int(os.environ.get("DEVICE_ID", "0")) | ||
| 41 | +WARMUP_ITERS = 5 | ||
| 42 | +BENCHMARK_ITERS = 100 | ||
| 43 | + | ||
| 44 | + | ||
| 45 | +def warmup_kernel(): | ||
| 46 | + """Compile the PyPTO kernel before GE starts. | ||
| 47 | + | ||
| 48 | + The PyPTO compilation cache is process-wide and keyed by kernel source, | ||
| 49 | + so the execute callback hits this cache and never compiles inside the | ||
| 50 | + GE runtime, where spawning the compiler subprocess is unsafe. | ||
| 51 | + """ | ||
| 52 | + if not torch_npu.npu.is_available(): | ||
| 53 | + raise RuntimeError("NPU is not available!") | ||
| 54 | + # tensor_bridge is imported at module level so that the torchair | ||
| 55 | + # extension is loaded before GE starts. | ||
| 56 | + assert hasattr(tensor_bridge, "to_torch"), "tensor bridge is unavailable" | ||
| 57 | + x = torch.empty(NUM_ELEMENTS, dtype=torch.float32, device="npu") | ||
| 58 | + y = torch.empty_like(x) | ||
| 59 | + out = torch.empty_like(x) | ||
| 60 | + pypto_add_kernel(x, y, out) | ||
| 61 | + torch.npu.synchronize() | ||
| 62 | + print("[OnlinePython] PyPTO kernel warm-up PASS") | ||
| 63 | + | ||
| 64 | + | ||
| 65 | +def build_graph(): | ||
| 66 | + builder = GraphBuilder("python_pypto_add_graph") | ||
| 67 | + input_x = builder.create_input( | ||
| 68 | + index=0, | ||
| 69 | + name="data_x", | ||
| 70 | + data_type=DataType.DT_FLOAT, | ||
| 71 | + format=Format.FORMAT_ND, | ||
| 72 | + shape=[NUM_ELEMENTS], | ||
| 73 | + ) | ||
| 74 | + input_y = builder.create_input( | ||
| 75 | + index=1, | ||
| 76 | + name="data_y", | ||
| 77 | + data_type=DataType.DT_FLOAT, | ||
| 78 | + format=Format.FORMAT_ND, | ||
| 79 | + shape=[NUM_ELEMENTS], | ||
| 80 | + ) | ||
| 81 | + output0 = PyptoAddCustom(input_x, input_y) | ||
| 82 | + builder.set_graph_output(output0, 0) | ||
| 83 | + return builder.build_and_reset() | ||
| 84 | + | ||
| 85 | + | ||
| 86 | +def build_input_data(start, step): | ||
| 87 | + return [start + float(index) * step for index in range(NUM_ELEMENTS)] | ||
| 88 | + | ||
| 89 | + | ||
| 90 | +def build_device_tensor(data): | ||
| 91 | + return Tensor( | ||
| 92 | + data, | ||
| 93 | + None, | ||
| 94 | + DataType.DT_FLOAT, | ||
| 95 | + Format.FORMAT_ND, | ||
| 96 | + [NUM_ELEMENTS], | ||
| 97 | + Placement.PLACEMENT_DEVICE, | ||
| 98 | + ) | ||
| 99 | + | ||
| 100 | + | ||
| 101 | +def build_input_sets(): | ||
| 102 | + values = ( | ||
| 103 | + (build_input_data(1.0, 1.0), build_input_data(10.0, 0.5)), | ||
| 104 | + (build_input_data(3.0, 2.0), build_input_data(20.0, 0.25)), | ||
| 105 | + ) | ||
| 106 | + return [([build_device_tensor(x), build_device_tensor(y)], x, y) for x, y in values] | ||
| 107 | + | ||
| 108 | + | ||
| 109 | +def validate_graph(session, input_sets): | ||
| 110 | + for inputs, values_x, values_y in input_sets: | ||
| 111 | + outputs = session.run_graph(GRAPH_ID, inputs) | ||
| 112 | + if len(outputs) != 1: | ||
| 113 | + raise RuntimeError( | ||
| 114 | + "PyptoAddCustom returned {} outputs".format(len(outputs)) | ||
| 115 | + ) | ||
| 116 | + actual = outputs[0].data | ||
| 117 | + expected = [x + y for x, y in zip(values_x, values_y)] | ||
| 118 | + max_error = max(abs(value - golden) for value, golden in zip(actual, expected)) | ||
| 119 | + if max_error > 1.0e-5: | ||
| 120 | + raise RuntimeError( | ||
| 121 | + "PyptoAddCustom precision check failed, max_error={}".format(max_error) | ||
| 122 | + ) | ||
| 123 | + print("[OnlinePython] PyptoAddCustom precision check PASS") | ||
| 124 | + | ||
| 125 | + | ||
| 126 | +def benchmark_graph(session, input_sets): | ||
| 127 | + for iteration in range(WARMUP_ITERS): | ||
| 128 | + session.run_graph(GRAPH_ID, input_sets[iteration % len(input_sets)][0]) | ||
| 129 | + start = time.perf_counter() | ||
| 130 | + for iteration in range(BENCHMARK_ITERS): | ||
| 131 | + session.run_graph(GRAPH_ID, input_sets[iteration % len(input_sets)][0]) | ||
| 132 | + return (time.perf_counter() - start) * 1.0e6 | ||
| 133 | + | ||
| 134 | + | ||
| 135 | +def run_graph(): | ||
| 136 | + options = { | ||
| 137 | + "ge.exec.deviceId": str(DEVICE_ID), | ||
| 138 | + "ge.graphRunMode": "1", | ||
| 139 | + # Force the dynamic (RT2) execution path so the Python execute | ||
| 140 | + # callback runs for every graph execution with real input data. | ||
| 141 | + # The default known-shape path records the callback once at task | ||
| 142 | + # generation time, which is incompatible with a PyPTO launch. | ||
| 143 | + "ge.exec.static_model_ops_lower_limit": "-1", | ||
| 144 | + } | ||
| 145 | + ge_api = GeApi() | ||
| 146 | + session = None | ||
| 147 | + ge_initialized = False | ||
| 148 | + graph_ids = [] | ||
| 149 | + | ||
| 150 | + try: | ||
| 151 | + ge_api.ge_initialize(options) | ||
| 152 | + ge_initialized = True | ||
| 153 | + session = Session(options) | ||
| 154 | + session.add_graph(GRAPH_ID, build_graph()) | ||
| 155 | + graph_ids.append(GRAPH_ID) | ||
| 156 | + input_sets = build_input_sets() | ||
| 157 | + | ||
| 158 | + validate_graph(session, input_sets) | ||
| 159 | + total_us = benchmark_graph(session, input_sets) | ||
| 160 | + avg_us = total_us / BENCHMARK_ITERS | ||
| 161 | + print("[Perf] input shape: [{}], dtype: float32".format(NUM_ELEMENTS)) | ||
| 162 | + print("[Perf] iters: {}".format(BENCHMARK_ITERS)) | ||
| 163 | + print( | ||
| 164 | + "[Perf] PyptoAddCustom: {:.3f} us (avg {:.3f} us/iter)".format( | ||
| 165 | + total_us, avg_us | ||
| 166 | + ) | ||
| 167 | + ) | ||
| 168 | + print("[OnlinePython] NPU_EXECUTION=PASS") | ||
| 169 | + return 0 | ||
| 170 | + except Exception as exc: | ||
| 171 | + print("[OnlinePython] run_graph failed: {}".format(exc)) | ||
| 172 | + traceback.print_exc() | ||
| 173 | + return 1 | ||
| 174 | + finally: | ||
| 175 | + if session is not None: | ||
| 176 | + for graph_id in graph_ids: | ||
| 177 | + session.remove_graph(graph_id) | ||
| 178 | + session = None | ||
| 179 | + if ge_initialized: | ||
| 180 | + ge_api.ge_finalize() | ||
| 181 | + | ||
| 182 | + | ||
| 183 | +if __name__ == "__main__": | ||
| 184 | + raise SystemExit(run_graph()) | ||
| @@ -0,0 +1,50 @@ | |||
| 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 | + | ||
| 13 | +"""Zero-copy bridge from GE device tensors to torch tensors. | ||
| 14 | + | ||
| 15 | +Reuses the torchair ``create_npu_tensors`` entry (the same machinery | ||
| 16 | +behind torchair's Python custom-op callbacks) to wrap raw NPU addresses | ||
| 17 | +as torch tensors. The device memory stays owned by GE. | ||
| 18 | + | ||
| 19 | +The host must import torch and torch_npu before importing this module | ||
| 20 | +so that the npu backend is registered with enough static TLS slots. | ||
| 21 | +""" | ||
| 22 | + | ||
| 23 | +import torch | ||
| 24 | + | ||
| 25 | +from torchair.llm_datadist import create_npu_tensors | ||
| 26 | + | ||
| 27 | +from ge.graph.types import DataType | ||
| 28 | +from ge.runtime import Tensor | ||
| 29 | + | ||
| 30 | +# torch dtypes keyed by GE data types. | ||
| 31 | +_GE_TO_TORCH_DTYPE = {DataType.DT_FLOAT: torch.float32} | ||
| 32 | + | ||
| 33 | + | ||
| 34 | +def to_torch(tensor: Tensor): | ||
| 35 | + """Wrap a GE device tensor as a zero-copy torch tensor.""" | ||
| 36 | + dtype = _GE_TO_TORCH_DTYPE.get(tensor.data_type) | ||
| 37 | + if dtype is None: | ||
| 38 | + raise ValueError( | ||
| 39 | + "unsupported GE data type for the tensor bridge: {}".format( | ||
| 40 | + tensor.data_type | ||
| 41 | + ) | ||
| 42 | + ) | ||
| 43 | + addr = tensor.addr | ||
| 44 | + if addr == 0: | ||
| 45 | + raise RuntimeError("to_torch: tensor.addr is 0 (null)") | ||
| 46 | + shape = list(tensor.storage_shape.dims) | ||
| 47 | + tensors = create_npu_tensors(shape, dtype, [int(addr)]) | ||
| 48 | + if not tensors: | ||
| 49 | + raise RuntimeError("to_torch: create_npu_tensors returned empty") | ||
| 50 | + return tensors[0] | ||