已合并
feat: Eager自定义算子RT2路径辅流申请功能 #5139
feat: Eager自定义算子RT2路径辅流申请功能 #5139
已合并
Chang-an-HW创建于 4 天前
Chang-an-HW成员
4 天前

Pull Request

描述

2026829评审通过

为 Eager 自定义算子(EagerExecuteOp)在 RT2 动态图执行链路 补齐按 key 申请框架托管物理辅流的能力。此前该能力只在 V1 静态图链路可用,RT2 下 EagerOpExecutionContext::RequestAttachedStream 因 EagerArgsHandler 未提供 provider 而恒返回 nullptr。

公共 API 与 ABI 零改动,用户算子代码零改动:RequestAttachedStream 签名、AdditionalInputIndex/AdditionalOutputIndex 布局、ArgsHandler::GetAttachedStreamProvider() 均不变,仅扩展 v2 内部枚举 CustomOpInput。

核心链路

LoweringCustomNode(Main 帧)
  └─ bg::GetAttachedStreamProvider → Init 图 CreateAttachedStreamProvider(每模型一份)
     └─ 作为 ExecuteCustomOp 附加输入 CustomOpInput::kAttachedStreamProvider
ExecuteCustomOpImpl → 首轮 !IsInitialized() 时注入 EagerArgsHandler
用户 Execute(ctx) → RequestAttachedStream(key) → Rt2AttachedStreamCollection

关键设计

  • 归属层级 per-ModelV2Executor:容器由 Init 图节点的输出 Chain 持有。StreamExecutor 为每条 aclrtStream 各建一个执行器,故不同执行流即使用相同 key 也拿到不同辅流;换来天然隔离且无需加锁(ExecuteCustomOp 已注册 kKernelUseMemory,固定在唯一 memory worker 上串行执行)。
  • 与 V1 的差异:RT2 无 rtModel_t,不做 aclmdlRIBindStream;StreamActive 在 RT2 已退化为透传,不存在 V1 的 WAIT_ACTIVE 激活坑,按默认 flag 建流即可。
  • 生命周期:执行器析构时 Chain deleter 触发 Destroy()(先 aclrtSynchronizeStream 再 aclrtDestroyStream,幂等)。现有卸载路径(UnloadRt2Model → DeleteExecutor、StreamExecutor::Erase)都在 UnLoad() 后立即销毁执行器,故无需额外 DeInit 节点。
  • 执行期开销:RT2 下 Execute 每轮都会被调用(V1 只在下沉期一次),故 key 查找用透明比较容器(std::less<> + string_view),命中路径零堆分配,host 微基准约 25.3 ns/次。
  • 不做数量限制:单执行实例的辅流数量不设上限(与 V1 实现一致);建流失败时打印 error 日志并返回 nullptr,由算子自行降级。
  • 内存回收契约:辅流不是框架逻辑流,RT2 多流 L2 回收不感知它,故算子必须在 Execute 返回前用 event 把辅流 join 回 ctx->GetStream(),已写入 API 文档约束说明。

顺带修复的既存缺陷

RT2 内嵌 V1 静态子图场景下 Eager 算子申请的辅流确定性泄漏:DavinciModelFinalizer 未清理辅流,而 DestroyResources() 会置 has_finalized_ = true,导致 ~DavinciModel() 中含 UnbindAndDestroy() 的兜底清理块被整体跳过。修复:新增幂等的 DavinciModel::UnbindAndDestroyAttachedStreams(),由 Finalizer 在 DestroyStream() 之前调用。

行为变化(需评审确认)

  1. RT2 下 RequestAttachedStream 由恒 nullptr 变为可返回有效句柄——对以 nullptr 判定"链路不支持"的算子属行为变化。
  2. 取消 CANN-FMK- 前缀的硬拦截(后续 HCCL 等框架组件也要通过本接口建流),改为 API 文档中的命名约定,key 校验统一为"仅判空"。属放宽型变化,原被拒绝的 key 现在会成功。

对既有动态图流程的影响

  • 不含 device 自定义算子的图:零影响。bg::GetAttachedStreamProvider 全仓唯一调用点在 LoweringCustomNode,而它只注册给 kCustomOpKernelLibName + kOnDeviceHbm,这类图不会新增任何节点或输入。
  • 含 device 自定义算子的图:Init 图 +1 个 no-op 节点(每模型一份)、每个算子节点 +1 条输入边、Main 帧 +1 个 InnerData、每执行器 +1 个约 64B 容器对象;每轮执行开销为 0(provider 只在首轮读取一次)。
  • Host CPU 自定义算子路径未触碰;OM/OM2 序列化不受影响(SerializeComputeNodeInfo 只遍历 IR 计算节点);profiling/dump/trace 只按 IR 输入数遍历,不受附加输入影响。

本 PR 不包含

Python 绑定(因 ST fallback 编译问题临时移除)、动态 shape 端到端样例、OM2 路径(其 context 由 codegen SO 内部构造且跨 OM2 C ABI)。

变更类型

关联的Issue

无

如何测试

  1. 以动态 shape 图验证:Execute 每轮都被调用(RT2 链路生效)、辅流句柄多轮恒定不变(跨轮次跨 shape 复用)、异 key 隔离、每轮精度全部通过。

核对清单

其他信息

无

likedislike
Pull Request已成功合入, 合并人@CANN-robot
(感谢 Chang-an-HW 的贡献)
CChang-an-HW成员
4 天前 添加了label:enhancement
CChang-an-HW成员
4 天前 创建了 pull request,commit 97c21b98
atomgit-bot
atomgit-bot
4 天前 评论:

变更摘要

本 PR 为 Eager 自定义算子在 RT2 动态图链路下补齐辅流申请能力。核心是在 EagerOpExecutionContext 上新增 request_attached_stream(key) 接口,并在 RT2 侧引入按执行器隔离、按 key 缓存的辅流容器 Rt2AttachedStreamCollection,通过 AttachedStreamProvider 接口和 CreateAttachedStreamProvider 初始化图节点注入到 EagerArgsHandler;同时调整 V1 侧 key 校验与辅流清理时机,并补充 RT2 动态 shape 样例与单元测试。

主要改动

  • 新增 Python/绑定层辅流接口: 在 _ge_custom_op_native.pyi 中新增 EagerOpExecutionContext.request_attached_stream(key),并在 context_binding.cc 中实现 BorrowedEagerOpExecutionContext::RequestAttachedStream,空 key 抛 ValueError,链路不支持/建流失败/超配额时返回 None。

  • RT2 辅流容器与生命周期管理: 新增 Rt2AttachedStreamCollection(实现 AttachedStreamProvider),按 key 惰性建流并复用、单实例辅流上限 kMaxAttachedStreamNum = 8、Destroy() 先 aclrtSynchronizeStream 再销毁且幂等;配套新增 attached_stream_provider_kernel.cc 在 Init 图创建容器并由 Chain deleter 持有,执行器析构时释放,实现 per-executor 隔离。

  • RT2 图构建与执行注入链路: 新增 bg_attached_stream_provider 获取当前执行器辅流容器;custom_node_converter.cc 在 custom_executor_func 之后追加该附加输入,custom_op_kernel.cc 扩展 CustomOpInput::kAttachedStreamProvider 并通过 EagerArgsHandler::SetAttachedStreamProvider 注入;Release() 与析构不清理 provider,保证跨轮次复用。

  • V1 侧 key 校验与辅流清理调整: AttachedStreamCollection::IsValidKey 改为仅校验空 key,不再拦截 CANN-FMK- 前缀(保留为软约束);DavinciModel 新增 UnbindAndDestroyAttachedStreams 并在 DavinciModelFinalizer 中于 DestroyStream 之前调用,重复调用无副作用。

  • 样例与测试补充: 新增 eager_attached_stream_dynamic_shape 动态 shape 样例(kernel、算子、session_run、run.sh 日志断言、RtcKernelLoader 等),并更新 eager_attached_stream_add 注释;新增/更新 RT2 辅流容器、CustomNodeKernel、EagerArgsHandler、custom_node_converter、DavinciModel 等单元测试。

likedislike
不准确?
atomgit-bot
atomgit-bot
4 天前 评论:

代码审查

✅ 未发现问题

likedislike
不准确?
CChang-an-HW成员
4 天前 修改了pull request 的描述
此处折叠了95条消息 查看更多
CChang-an-HW成员
3 天前 解决了最后一个问题
CChang-an-HW成员
3 天前 修改了pull request 的描述
yangyongqiang
yangyongqiang成员
3 天前 评论:

/approve

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