已合并
更新自定义算子python callback生成的python代码的格式 #3554
更新自定义算子python callback生成的python代码的格式 #3554
已合并
Chang-an-HW创建于 6 天前
Chang-an-HW
6 天前

更新自定义算子python callback生成的python代码的格式

likedislike
Pull Request已成功合入, 合并人@ascend-robot
(感谢 Chang-an-HW 的贡献)
CChang-an-HW
6 天前 创建了 pull request,commit a35a182c
atomgit-bot
atomgit-bot
6 天前 评论:

变更摘要

该 PR 更新了自定义算子 Python callback 生成代码的格式,将原先基于 EagerExecuteOp/EagerOpExecutionContext 的执行模型,改为基于 GE schema-bound 执行框架的格式:生成的类不再继承 EagerExecuteOpexecute 方法改为显式声明位置参数和 keyword-only 属性参数,并在函数体内通过 get_execute_ctx() 获取上下文。为保持 REG_OP 原型与 Python 执行签名一致,代码生成侧抽象出了统一的 RegOpProto/RegOpInput/RegOpAttr 描述结构,由 build_reg_op_proto 作为单一事实来源同时供 C++ REG_OP 文本和 Python execute 签名生成使用,并同步更新了依赖检测与相关单元测试。

主要改动

  • 引入统一的 REG_OP 原型描述结构:在 custom_op_codegen_common.py 中新增 RegOpInputRegOpAttrRegOpProto 三个 frozen dataclass 及 build_reg_op_proto,将输入端口和属性信息从 SchemaInfo 统一推导出来,供 C++ REG_OP 生成与 Python execute 签名共同使用。
  • 重构 generate_reg_op 并新增 Triton 输入迭代器generate_reg_op 改为基于 build_reg_op_proto 生成 .INPUT/.OPTIONAL_INPUT/.DYNAMIC_INPUT.ATTR;同时新增 _iter_triton_hop_reg_op_inputs,统一枚举 Triton HOP 的 tensor/scalar/TMA/grid 输入,避免 schema 与 Python 参数列表漂移。
  • Python execute 改为 schema-bound 签名:在 custom_op_codegen_python.py 中新增 _safe_param_name_execute_param_names_generate_execute_signature,生成的类为普通类,execute 签名形如 def execute(self, 位置参数, *, 属性参数),并在函数内使用 ctx = get_execute_ctx() 获取执行上下文。
  • _generate_buildstack 与 Triton execute 生成逻辑改用参数变量:原先通过 ctx.get_input_tensor/get_optional_input_tensor/get_dynamic_input_tensor 提取输入的逻辑,改为直接从 execute 参数读取并包装/解析;generate_triton_hop_execute 也改为按 ordered_inputs 生成的参数处理 tensor、scalar、TMA 和 grid 输入,并移除 return True
  • 依赖检测与测试同步更新custom_op_pre_scan.py、相关 smoke/st 测试将 ge.custom_op 的导入从 EagerExecuteOp, register_op_impl 调整为 get_execute_ctx, register_op_impl,并新增/修改用例以验证 schema-bound 签名、无 return True、Triton HOP 执行参数绑定等行为。
likedislike
atomgit-bot
atomgit-bot
6 天前 评论:

代码审查

✅ 未发现问题

likedislike
ascend-robotascend-robot成员
6 天前 添加了label:smoke-pipeline-running
ascend-robotascend-robot成员
6 天前 添加了label:stat/needs-squash
此处折叠了85条消息 查看更多
ascend-robot
ascend-robot成员
3 天前 评论:
流水线 PR-pipeline_torchair#3731 [ commitID:018f6602 ] 已完成
>>>代码风格自动修复执行成功(无修复内容)
阶段 任务名 状态 详情
编译构建 Build_x86 >>>
Build_ARM >>>
恶意代码检查 Antipoison >>>
编码安全与规范检查 codecheck_pre-commit >>>
pre-commit >>>
开源片段检查 SCA >>>
流水线 PR-pipeline_torchair >>>
此流水线已支持下列评论快捷指令,仅PR创建者和白名单成员评论有效
  • compile : 运行流水线
  • retry : 重试流水线所有失败子任务
  • retry <任务名> : 仅重试指定失败子任务
  • stop : 停止流水线
likedislike
dingdairong成员
3 天前 评论:

/lgtm

likedislike
薛鹏
薛鹏成员
3 天前 评论:

/approve

likedislike
ascend-robotascend-robot成员
3 天前 添加了label:approvedlgtm
ascend-robotascend-robot成员
3 天前 合入了pull request