已合并
feat: Python 自定义算子原型执行样例 #4519
lfz2812创建于 14 天前
feat: Python 自定义算子原型执行样例 #4519
已合并
lfz2812创建于 14 天前
lfz2812成员
14 天前

Pull Request

描述

完善 Python 自定义算子公开资料,并将现有 args_refresh_add_custom/python 样例升级为 Python 原型注册、infer-meta 和 PyPTO 执行的端到端样例。

主要变更:

  • src/ge/add_custom.py 中使用 @register_op 注册 AddPythonCustomOp 原型和 infer-meta。
  • infer-meta 校验输入 shape 和 data type,并通过 TensorDesc 返回输出元信息。
  • 使用 @register_op_impl 注册 schema-bound execute(x, y) 实现。
  • 使用 PyPTO 实现与 add_custom.asc 等价的 z = x + y NPU kernel。
  • 通过 DLPack 将 GE device Tensor 零拷贝转换为 torch_npu Tensor,并在 GE 当前 stream 上执行 PyPTO kernel。
  • 删除样例原有的 Bisheng 编译、AI Core binary 提取、_load_kernel()、ACL binary 加载和 kernel launch 流程。
  • 更新 run.sh,仅生成和安装 ES Python API,并运行 Python Session 样例。
  • 移除构图阶段的 _out_shape_locked 属性和手工设置输出 shape、format、data type 的逻辑,输出元信息由 Python infer-meta 提供。
  • C++ REG_OP 仅作为 gen_esb 生成 ge.es.custom.AddPythonCustomOp 构造器的输入,不参与运行时原型注册。
  • 将 C++ schema 输出名改为 output0,与 Python register_op 对普通非别名输出生成的默认名称保持一致,避免 canonical IR 兼容性检查失败。
  • 增加 ge.runtime.TensorDesc 中文 API 文档,并说明其与 ge.graph.TensorDesc 的区别。
  • 同步更新中英文设计文档和自定义算子样例说明。

变更类型

请选择本次引入的变更类型:

关联的Issue

#299

如何测试

已执行以下静态检查:

python3 -m py_compile \
  examples/custom_op/args_refresh_add_custom/python/src/ge/add_custom.py \
  examples/custom_op/args_refresh_add_custom/python/src/run.py

bash -n examples/custom_op/args_refresh_add_custom/python/run.sh

git diff --check

已针对全部变更文件执行 pre-commit,以下检查均通过:

  • trim trailing whitespace
  • fix end of files
  • clang-format
  • ruff check
  • ruff format
  • codespell
  • OAT Compliance Check

NPU 环境验收步骤:

source ${ASCEND_HOME_PATH}/set_env.sh
cd examples/custom_op/args_refresh_add_custom/python
bash run.sh

核对清单

其他信息

  • args_refresh_add_custom/cpp/add_custom_kernel/add_custom.asc 继续供 C++ 样例使用,并作为 PyPTO 加法实现的逻辑参考;Python 样例运行时不再编译或加载该文件。
  • 当前 PyPTO 样例支持 DT_FLOAT 和连续的确定 Storage Shape。
  • GE 保持输入和输出 device memory 的所有权,DLPack deleter 仅释放 Python 侧临时元数据。
  • C++ REG_OP 输出名必须与 Python canonical proto 保持一致。当前普通非别名输出使用默认名称 output0。
  • 当前环境未配置完整的 CANN、PyPTO 和 torch_npu NPU 运行环境,真机执行需要在对应 NPU 环境完成。1. -
likedislike
Pull Request已成功合入, 合并人@CANN-robot
(感谢 lfz2812 的贡献)
Llfz2812成员
14 天前 添加了label:enhancement
Llfz2812成员
14 天前 创建了 pull request,commit 56f8e1bb
Llfz2812成员
14 天前 virtual merging failed, update merge request[project_id: 8824148, iid: 4519, target_commit_sha: 56f8e1bb90b1c89bcbbffe3dd952361bfa449121], message: Conflict detected
Llfz2812成员
14 天前 强制推送  2 个提交:f4521dea-feat: add Python custom op proto contracts,1cfec12f-feat: add Python custom op prototype infer-meta sample
Llfz2812成员
14 天前 virtual merging failed, update merge request[project_id: 8824148, iid: 4519, target_commit_sha: 56f8e1bb90b1c89bcbbffe3dd952361bfa449121], message: Conflict detected
此处折叠了572条消息 查看更多
yangyongqiang
yangyongqiang成员
7 天前 评论:

/lgtm

likedislike
CANN-robotCANN-robot成员
7 天前 添加了label:lgtm
张德鹏成员
7 天前 评论:

/approve

likedislike
CANN-robotCANN-robot成员
7 天前 添加了label:approved
CANN-robotCANN-robot成员
7 天前 合入了pull request