已合并
feat:声明式自定义算子的辅流申请功能 #4934
feat:声明式自定义算子的辅流申请功能 #4934
已合并
Chang-an-HW创建于 9月13日
Chang-an-HW成员
9月13日

Pull Request

描述

为声明式自定义算子(AnnotatedArgsOp / DeclareLaunchArgs)新增编译期辅流申请与任务依赖声明能力,使算子可以在 DeclareLaunchArgs 回调内把多个 kernel launch 下发到不同逻辑流,并由 GE 自动完成流分配与事件同步插入。

对外接口(gert::AnnotatedArgsContext,@since 9.3.0(2026-09))

  • 新增 uint32_t RequestAttachedStream(const ge::AscendString &key):按模型级 key 申请逻辑辅流,相同 key 复用同一 ID;key 为空或申请失败返回 UINT32_MAX。
  • 新增 graphStatus AddLaunchWithDependencies(launch_info, args, dependencies, dependency_num, token):以 AnnotatedLaunchToken 表达 launch 间前驱依赖,成功后输出本次 launch 的 token。
  • AddLaunch 的 ABI 与行为保持不变(不产生 token,也不能作为依赖源)。

graph_metadef 侧

  • AnnotatedArgsHandler:LaunchRecord 增加 token 与依赖列表;新增 SetAttachedStreamRequestFunc / RequestAttachedStream / GetAttachedStreamIds;AddLaunchWithDependencies 校验 token 有效性(依赖必须指向更早 launch、不可重复、不可溢出)。
  • stub_offline_launch_runtime.cc 同步补齐 stub 实现;stub_lowering 增加 runtime_headers、slog 链接。

compiler 侧

  • 新增 DeclarativeStreamRegistry:挂到 ComputeGraph ext attr(declarative_stream_registry),从当前逻辑流数之后按 key 去重分配辅流 ID,并做溢出保护。
  • 新增 AnnotatedArgsTaskPlan 序列化/反序列化(带 version、长度与依赖边合法性校验):缓存属性由进程内 ext attr 改为可序列化的 _custom_annotated_args_task_plan_bytes,使计划在 ReGetTaskInfo / 二次生成场景下可无回调恢复;计划中新增 launch_stream_ids、dependencies、attached_stream_ids。
  • ModelBuilder:PrepareDeclarativeAttachedStreamRegistry(AssignLogicalStreams 之后建注册表)、FinalizeDeclarativeAttachedStreams(把计划辅流写回 op_desc->SetAttachedStreamIds 并扩展 stream_num)、MaterializeAnnotatedArgsTaskDependencies。
  • StreamAllocator:InsertAnnotatedArgsEventTasks 为跨流依赖插入 EVENT_RECORD/EVENT_WAIT 任务;ReconcileAnnotatedArgsEventTasks 在拆流后重映射 event id、剪除同流事件对,并回写 stream/event/notify 总数(RTS 与 GE 两条拆流路径均覆盖)。
  • CustomOpsKernelBuilder:物化时按计划流 ID 下发任务(无依赖且无辅流的 legacy 计划仍刷新为节点当前主流,兼容跨流复用);CalcOpRunningParam 改为 DelAttr 清理序列化计划。

约束校验

  • 计划中至少存在一个主流 launch;辅流任务是否汇合回主流由算子自行保证。
  • 依赖 token 必须指向更早的 launch,不可重复、不可自依赖。
  • mobile OMC 单任务模式不支持依赖与辅流;端侧仍只允许一次 AddLaunch。

文档:新增 RequestAttachedStream.md、AddLaunchWithDependencies.md,更新 overview.md、README.md 目录、AddLaunch.md 与 DeclareLaunchArgs.md 约束描述。

变更类型

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

关联的Issue

无

如何测试

描述测试此变更的步骤和前提条件:

  1. 编译并运行 UT:
    • graph_metadef UT:AnnotatedArgsContextUT.AddLaunchWithDependenciesValidatesAndCopiesTokens
    • ge UT:
      • stream_allocator_unittest:DeclarativeStreamRegistryTest.*、AnnotatedArgsTaskPlanTest.*(序列化往返、跨流事件插入/重映射/同流剪枝)、ModelBuilderTest.FinalizeDeclarativeAttachedStreams*
      • custom_ops_kernel_info_store_unittest:计划无回调恢复、辅流不汇合主流、缺主流 launch 报错
      • task_generator_unittest:GenerateTask_AnnotatedKernelRecordWait
      • custom_task_info_unittest:AnnotatedEventWaitUsesExistingWaitAndResetSemantics
  2. 编译并运行 ST:graph_engine_test 中 DeclarativeCustomOpAttachedStreamSt.*(按 key 复用辅流 ID、回调缺失返回 UINT32_MAX)。
  3. 手工验证:自定义算子在 DeclareLaunchArgs 中调用 RequestAttachedStream + AddLaunchWithDependencies 下发主流/辅流任务,出 OM 后确认 stream_num 扩展、辅流 launch 与 EVENT_RECORD/EVENT_WAIT 任务正确生成且执行结果正确。

核对清单

其他信息

  • PR 当前带 stat/needs-squash 标签,38 个 commit 中含多个 CI 规避/revert 提交,合入前需 squash。
  • 标题 feat: 使用了全角冒号,建议改为半角 feat: 。
likedislike
Pull Request已成功合入, 合并人@CANN-robot
(感谢 Chang-an-HW 的贡献)
CChang-an-HW成员
9月13日 添加了label:enhancement
CChang-an-HW成员
9月13日 创建了 pull request,commit a41980ea
CANN-robotCANN-robot成员
9月13日 添加了label:cann-cla/yes
CANN-robot
CANN-robot成员
9月13日 评论:

Thanks for your pull-request.
The full list of commands accepted by me can be found at here.
You can get sig-info at here.
You can self-configure the PR merge rules for this repository. For more details, please refer to Here.
For more, you also can visit HICANN.


PR Approval Progress

✅ Congratulations! All modules have met the lgtm and approve requirements.

Module Approval Details

module lgtm status approve status
inc/graph_metadef/external ✅ yangyongqiang0606, 王涛 (2/2) ✅ 王涛 (1/1)
repo-cann/ge ✅ 王涛, yangyongqiang0606 (2/2) ✅ 王涛 (1/1)

💡 Tip:

  • Committer can comment /approve or /lgtm
  • Commenting /approve implies both code review (lgtm) and intent to merge (approve)

CLA Signature Pass

Chang-an-HW, thanks for your pull request. All authors of the commits have signed the CLA. 👍

likedislike
CANN-robotCANN-robot成员
9月13日 将yangyongqiang0606,kobemini,ClarkXie,tangqunzhang,xchu42,hugo111,lining23666,depeng1994,wqtshg_wt,peiyang,shengnan666,zhangfan_hanq,zhanj,zhujingjing,Mexyy,wangxiaotian995,stevenaw0,yuht9设为评审人
此处折叠了434条消息 查看更多
王涛成员
24 天前 评论:

/lgtm

likedislike
CANN-robotCANN-robot成员
24 天前 添加了label:lgtm
王涛成员
24 天前 评论:

/approve

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