已合并
feat: run Python custom op infer_meta at compile time and RT2 #4414
lfz2812创建于 7 天前
feat: run Python custom op infer_meta at compile time and RT2 #4414
已合并
lfz2812创建于 7 天前
lfz2812成员
7 天前

Pull Request

描述

在 Python 自定义算子原型注册、Bridge、Adapter 和 canonical IR 链路的基础上,接通 Python infer_meta 回调,使注册的 Python 原型可以参与编译期和 RT2 动态 Shape 推导。

本次主要包含:

  1. 接通 Python infer_meta 回调:

    • register_op 装饰函数作为 Python infer_meta 函数使用。
    • 根据 canonical IR 组装 required、optional、dynamic 输入。
    • 根据属性 IR 类型读取运行时属性。
    • 将输入转换为 TensorDesc,将属性转换为 Python 参数。
    • 支持 required、optional、dynamic 输入和 required、dynamic 输出。
    • 支持多输出和动态输出实例数量校验。
    • 校验 Python 返回值类型、结构和输出数量。
  2. 新增 InferShape native context:

    • 提供输入 shape 和 data type 读取能力。
    • 提供动态输入实例数量查询。
    • 提供运行时属性读取能力。
    • 提供动态输出实例数量查询。
    • 提供输出 shape 写回能力。
    • callback 结束后使 borrowed context 及其派生对象失效。
  3. 接通编译期 Meta 推导:

    • 在 Graph Metadef 编译期识别 CustomOpInferMetaProvider
    • 每个算子节点只调用一次 Python infer_meta
    • 推导成功后统一更新输出 shape、data type 和 origin data type。
    • 使用临时 OpDesc 承载推导结果,所有校验成功后再提交到原始 OpDesc
    • 任意推导、返回值或格式转换失败时,不提交部分输出结果。
  4. 接通 RT2 动态 Shape 推导:

    • RT2 的 custom op shape kernel 优先复用 infer_meta 回调。
    • Python callback 读取当前运行时输入元信息和属性。
    • 运行期只更新输出 shape,已编译的 data type 不被覆盖。
    • 保留普通 C++ custom op 的 ShapeInferOp fallback 路径。
  5. 扩展 Bridge、Adapter 和注册结构:

    • 扩展 Bridge C ABI,传递 Python infer_meta callback。
    • 扩展 Python custom op Adapter 的 infer_meta 能力。
    • 支持 Python proto-only、Python proto + impl 和 infer-only 场景。
    • 保持现有 Python execute、schema-bound execute 和 legacy execute(ctx) 行为。
    • 保持 callback、holder、runtime registry 和卸载生命周期管理方式不变。
  6. 补充测试和设计说明:

    • 补充 Python infer_meta 参数组装、属性读取、返回值校验和动态输出测试。
    • 补充 InferShapeContext native binding 测试。
    • 补充 C++ Adapter、Bridge callback 和 infer-only adapter 测试。
    • 补充 Graph Metadef 编译期 shape/dtype/origin dtype 回写及失败原子性测试。
    • 补充 RT2 Python infer_meta 动态 Shape ST。
    • 同步更新中英文 Python 自定义算子设计文档。

变更类型

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

关联的Issue

无。

如何测试

建议执行:

cmake --build build_ut \
  --target install_ge_py_for_tests \
           ut_exe_graph \
           ut_libge_multiparts_utest \
           ut_fast_runtime2_test \
  -j4

python3 -m pytest \
  tests/ge/ut/ge/graph/pyge_tests/python_custom_op_infer_meta_test.py \
  -q

ctest --test-dir build_ut \
  -R 'ut_exe_graph|ut_libge_multiparts_utest|ut_fast_runtime2_test' \
  --output-on-failure

cmake --build build_st --target graph_engine_test -j4

ctest --test-dir build_st \
  -R graph_engine_test \
  --output-on-failure

git diff --check

核对清单

其他信息

  • 本次仅接通 Python infer_meta 的编译期和 RT2 动态 Shape 链路。
  • 本次不新增公开 API,不修改 register_op API 文档,不修改 user_guides。
  • 不支持 data-dependent infer、InferShapeRange、format inference、symbolic inference 和 shape rule generation。
  • 不修改 OM 格式、DavinciModel 接口或 RT2 到 V1 的边界数据。
  • 不修改 build_fwk.sh 等公共构建脚本。
  • 保持 Python native、Bridge 和 Adapter 按同一 ABI 版本整体构建和替换。
likedislike
Pull Request已成功合入, 合并人@CANN-robot
(感谢 lfz2812 的贡献)
Llfz2812成员
7 天前 添加了label:enhancement
Llfz2812成员
7 天前 创建了 pull request,commit e2068f56
atomgit-bot
atomgit-bot
7 天前 评论:

🤖 正在生成合并请求摘要,请稍候…

likedislike
atomgit-bot
atomgit-bot
7 天前 评论:

🤖 AI 代码检视正在进行中,请稍候…

likedislike
CANN-robotCANN-robot成员
7 天前 添加了label:cann-cla/no
此处折叠了486条消息 查看更多
黄桂军
黄桂军成员
3 天前 评论:

/lgtm

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

/approve

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