Merged
feat: 新增 PyPTO 算子 GE 入图样例 #4797
lfz2812Add on Sep 8
feat: 新增 PyPTO 算子 GE 入图样例 #4797
Merged
lfz2812Add on Sep 8
lfz2812Collaborators
Sep 8

Pull Request

描述

新增 pypto_add_custom 样例,把 PyPTO JIT kernel 通过 Python 自定义算子接入 GE 图执行。

主要变更:

  • 使用 @register_op 注册 PyptoAddCustom 原型与 infer_meta,@register_op_impl 注册 schema-bound execute(x, y) 实现。
  • execute 回调通过 GE EagerOpExecutionContext 申请输出 Tensor,把输入/输出 device 地址经 torchair create_npu_tensors 零拷贝包装为 torch Tensor,再执行 @pypto.jit Add kernel。
  • PyPTO kernel 由 src/run.py 在 GE 初始化前于主线程预热编译,execute 回调只消费进程内编译缓存(GE 回调内触发 JIT 编译会因编译器子进程与 GE 运行时冲突而挂起)。
  • 通过 ge.exec.static_model_ops_lower_limit=-1 强制 RT2 动态执行路径,使 execute 回调在每次图执行时被调用(默认静态路径下回调仅在任务生成阶段执行一次,与 PyPTO 即时启动模型不兼容)。
  • execute 以设备同步保证输入/输出数据可见性,100 轮计时约 3-4 ms/iter。
  • README 中英文双语说明 PyPTO 支持现状:仅支持 JIT 直接执行,无独立编译接口和 binary 导出接口,无法预编译为 kernel binary 或转换为 AnnotatedArgs 的 kernel_bin;Python execute 回调 + torchair 桥接是当前 PyPTO 接入 GE 的可行路径。
  • examples/custom_op/README.md 与 README_en.md 总览表同步登记新样例。

变更类型

关联的Issue

无。

如何测试

静态检查:

python3 -m py_compile examples/custom_op/pypto_add_custom/src/*.py examples/custom_op/pypto_add_custom/src/ge/*.py
bash -n examples/custom_op/pypto_add_custom/run.sh
git diff --check

已对全部变更文件执行 pre-commit(ruff/codespell/clang-format/OAT 等),全部通过。

NPU 环境验收(CANN 9.2.0 + Atlas A2):

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

实测输出:

[OnlinePython] PyPTO kernel warm-up PASS
[OnlinePython] PyptoAddCustom precision check PASS
[Perf] input shape: [8192], dtype: float32
[Perf] iters: 100
[Perf] PyptoAddCustom: 371476 us (avg 3714 us/iter)
[OnlinePython] NPU_EXECUTION=PASS
[INFO] Online Python PyPTO custom-op pipeline PASS

核对清单

其他信息

  • PyPTO JIT 编译产物(output/ 下 program.json、dy_kernel_*.o 等)为进程内部缓存,未纳入版本管理。
  • Python 导入顺序要求 import torch / import torch_npu 在 import pypto 之前,否则可能触发 libc10.so 静态 TLS 报错,README 已注明。
  • torchair 桥接优先使用公开 torchair.llm_datadist.create_npu_tensors;torchair 包初始化失败的环境自动降级为直接加载其 C++ 扩展。
  • execute 回调的单轮耗时(约 3-4 ms)高于 Ascend C + AnnotatedArgs 路径(约 0.4 ms,参考 annotated_args_refresh_add_custom):开销来自设备同步与 Python 回调,本样例聚焦链路打通而非性能。
  • 如需声明式地址刷新或 om 模型下沉,请参考 annotated_args_refresh_add_custom 等基于 Ascend C 的样例。
likedislike
Pull Request has been successfully merged by CANN-robot.
(Thank you for lfz2812's contribution.)
Llfz2812Collaborators
Sep 8 Added label: enhancement
Llfz2812Collaborators
Sep 8 Created pull request, commit 896b9f3d
atomgit-bot
atomgit-bot
Sep 8 Comments:

变更摘要

该 PR 的标题与描述均为模板占位内容(未填写具体说明),从实际补丁看,本次改动集中在示例文件 examples/custom_op/annotated_args_refresh_add_custom/online/python/src/ge/annotated_add_custom.py,属于一次实验性的实现替换:将 _compile_kernel(platform) 中基于 bisheng/llvm-objcopy 子进程编译 Ascend C kernel 的完整流程,改写为基于 pypto.jit 的 kernel 编译尝试。由于 PR 描述未提供业务背景,可推断此改动主要为验证 PyPTO 编译路径(源分支名 pypto_try 与之吻合),尚未接通最终产物导出。

主要改动

  • _compile_kernel 实现方式替换:删除了通过 subprocess 调用 bisheng 与 llvm-objcopy 完成 host 与 .aicore 二进制编译的整套逻辑,改为调用 @pypto.jit(new_ir=False) 装饰的 add_kernel 完成 kernel 编译。
  • 删除编译缓存与依赖导入:移除了基于 _KERNEL_SOURCE 内容 sha256 摘要生成 .o/.aicore.o 并判断产物是否已存在的缓存机制,同时清理了 hashlib、subprocess、tempfile 等不再使用的 import。
  • kernel 逻辑改用 PyPTO 描述:通过 pypto.Tensor([8192], pypto.DT_FP32) 声明 x/y/out 张量,以 pypto.set_vec_tile_shapes(1, 1024) 设置分块并执行 out.move(pypto.add(x, y)) 表达加法计算。
  • 移除环境变量相关逻辑:删除了对 ASCEND_HOME_PATH、GE_PYTHON_CUSTOM_OP_BUILD_DIR 的读取及 include/输出目录构建逻辑,新实现改为直接构造 torch 的 npu 张量驱动执行,并以 torch.npu.synchronize() 等待完成。
  • 结尾改为占位异常:新实现末尾抛出 RuntimeError("PyPTO kernel compiled, but GE kernel binary export is not implemented"),明确表示 PyPTO 编译已跑通、但 GE kernel 二进制导出尚未实现,属于阶段性的占位变更。
likedislike
Incorrect?
atomgit-bot
atomgit-bot
Sep 8 Comments:

代码审查

审查结论

审查文件(共 1 个):

  • examples/custom_op/annotated_args_refresh_add_custom/online/python/src/ge/annotated_add_custom.py — 发现 2 个问题:
# 严重度 置信度 行号 摘要
1 P1 0.8 123–125 _compile_kernel 无条件抛 RuntimeError,AnnotatedAddCustom.compile() 必然失败,样例的声明式地址刷新算子整体不可用(确定性功能回归)
2 P3 0.9 69–79 _KERNEL_SOURCE 模块常量随旧编译实现删除而成为死代码(可选项)

总体风险评估:

本次 diff 把 _compile_kernel 从"可产出真实 device binary 的 bisheng 编译实现"替换为"pypto JIT 编译并在真实 NPU 上启动 kernel 后无条件抛错"的桩实现。除渲染样例必然无法运行外,该实现还在抛错前执行了完整的设备 kernel 启动与同步副作用,把编译回调从纯离线编译变成对 NPU/torch_npu 的硬性运行期依赖,并引入硬编码 [8192]/FP32 形状、dtype 与 _KERNEL_NAME 符号名与剩余未改代码契约的不一致(这些不一致当前因必然抛错而不可达,一旦补上导出即会暴露)。就当前代码形态而言,该 PR 若目标为使样例跑通 pypto 链路则不可合入;建议补全导出实现或明确标记为 WIP。整体风险等级:中(样例级功能完全回归,不涉生产数据安全)。

类型 数量
🔴 阻塞 1
🟡 建议 0

⛔ 需要修改

likedislike
Incorrect?
CANN-robotCANN-robotCollaborators
Sep 8 Added label: cann-cla/yes
256 messages are collapsed here。 See more
黄桂军
黄桂军Collaborators
Sep 16 Comments:

/lgtm

likedislike
CANN-robotCANN-robotCollaborators
Sep 16 Added label: lgtm
yangyongqiang
yangyongqiangCollaborators
Sep 16 Comments:

/approve

likedislike
CANN-robotCANN-robotCollaborators
Sep 16 Added label: approved
CANN-robotCANN-robotCollaborators
Sep 16 Merged pull request