已合并
【PR】: broadcast 后移 优化 #1782
czways创建于 5 天前
【PR】: broadcast 后移 优化 #1782
已合并
Pull Request已成功合入, 合并人@CANN-robot
(感谢 czways 的贡献)5 天前 添加了label:enhancement
atomgit-bot
5 天前 评论:
5 天前 评论:
🤖 正在生成合并请求摘要,请稍候…


atomgit-bot
5 天前 评论:
5 天前 评论:
🤖 AI 代码检视正在进行中,请稍候…


5 天前 添加了label:cann-cla/yes
此处折叠了188条消息 查看更多
wangxiaotian995
3 天前 评论:
3 天前 评论:
/lgtm


3 天前 添加了label:lgtm
3 天前 添加了label:approved
3 天前 合入了pull request
Pull Request
描述
broadcast backward opt
变更类型
请选择本次引入的变更类型:
关联的Issue
如何测试
描述测试此变更的步骤和前提条件:
1.
2.
核对清单
其他信息
在此添加任何其他关于本次 PR 的说明。
BroadcastBackwardPass:GE 侧与 graph-autofusion PR 侧对比分析
简介
目的
本文档对比 GE 仓库与 graph-autofusion PR #1782 中
broadcast_backward_pass的实现差异与相同之处,为代码审查、后续维护和跨仓库同步提供参考。范围
ge_0728/compiler/graph/optimize/autofuse/autofuse/post_process/pass/broadcast_backward_pass.cpp(1515 行)graph-autofusion_0721/autofuse/optimize/graph_pass/broadcast_backward_pass.cpp(1324 行)+broadcast_backward_shared_split.cpp(388 行)optimize.cpp重构、frontend_shape_vars移除等 PR 中的其他变更。总体概述
PR 侧本质上是 GE 侧的移植 + 增强:
asc_adapt::适配层,直接使用af::GraphUtils、af::NodeUtils、ScheduleUtils::原生接口。TopologicalSorting。broadcast_backward_shared_split预处理模块(388 行)。GetBroAxisFromNode轴匹配逻辑,支持 axis 不对齐和隐式广播维度。IndirectLoad、RemovePad屏障算子检查。MoveBroadcastAfterMerge支持多后继。差异分析
一、架构层差异(适配层不同)
namespace genamespace optimizeasc_adapt::适配层(GetPeerInNodes、GetPeerOutNode、IsSingleInAndOutNode、GetOutputTensorAttr、CreateBroadcastNode等 20+ 函数)af::GraphUtils、af::NodeUtils、ScheduleUtils::原生接口,无中间层ge::NodePtr直接使用af::NodePtr基类指针,需ToAscNode()向下转型std::dynamic_pointer_cast<AscNode>kTransposeType、kBroadcastType等全局宏Transpose::Type、Broadcast::Type等类静态成员DataType(GE 命名空间)af::DataTypeRun(ComputeGraphPtr)→asc_adapt::ProcessAscBackendNodes遍历RunPass(AscGraph &)直接操作原因分析:
optimize命名空间隔离代码。asc_adapt::适配层桥接两种图模型;graph-autofusion 中 AscGraph 是原生图结构,直接调用af::GraphUtils即可,无需中间层。ge::Node下;graph-autofusion 中af::NodePtr是基类指针,AscNode是携带 tensor attr 的派生类,需dynamic_pointer_cast转型后才能访问outputs[i].attr等扩展字段。Broadcast::Type),比全局宏更符合面向对象设计且避免命名冲突。ProcessAscBackendNodes遍历找到 AscBackend 子图再执行;graph-autofusion 的 pass runner 直接在 AscGraph 层级注册和调用,无需遍历。二、PR 侧新增功能(GE 侧没有)
1.
broadcast_backward_shared_split模块(388 行,全新)PR 侧在
BroadcastBackward入口最先调用:GE_ASSERT_SUCCESS(broadcast_backward_shared_split::SplitSharedBroadcastBranches(graph)); GE_ASSERT_SUCCESS(broadcast_backward_shared_split::SplitSharedBroadcastConsumers(graph));功能:在执行后移前,先克隆拆分共享同一条 Broadcast 链的多分支/多消费者场景,避免改写时边冲突。GE 侧无此预处理步骤,直接在原图上改写。
原因分析:
graph-autofusion 的图结构中,一条 Broadcast 链的输出可能同时被多个分支或多个消费者引用。如果不先拆分,后移时
ReorderBroadcasts需要断开 Broadcast 输出边,但该边连接了多个消费者,无法同时断开所有边,会导致改写失败或图结构损坏。通过预处理阶段克隆 Broadcast 链,使每个分支/消费者拥有独立的 Broadcast 副本,后续后移可以独立操作每条链。GE 侧可能这类共享场景较少,或通过其他 pass 在更早阶段已经拆分。根本原因:GE 融合反推机制不可移植到 AF
GE 的 Broadcast 处理与子图融合深度绑定,采用 "消除 → 融合 → 恢复" 三段式流程:
UpdateBroadcastBeforeMerge):沿 Load → Broadcast 路径,把广播轴信息"折叠"进 Load 输出属性(repeats[i]=1, strides[i]=0),然后删除 Brc 节点。目的是让不同子图在融合前统一布局。MergeAscGraphByLoop):按循环轴合并两个 AscGraph。OptimizedFallback):遍历融合后图中每个计算节点的每个输入,从 Load 属性和图级AscGraphAttr倒推broadcast_info,按消费者实际需求重新插入 Brc。OptimizedFallback的核心是 从消费者需求倒推(bottom-up):对每个消费者输入,比较消费者所需布局 vs Load 当前布局,差异轴即为broadcast_info,若非空则在消费者输入前新建 Brc。一个共享源有多个消费者时,每个消费者独立判断需要哪些广播轴,只插需要的轴,不产生冗余:GE 这一机制依赖大量 GE 特有结构:
UpdateBroadcastInfoToLoadAscBackendSubGraphFusionDecider、MergeAscGraphByLoop、NodeFuseInfo、ViewOpAttrInfoTuningSubgraphBeforeMergeAscGraphAxisMappingOptimizedFallbackProAscGraphAttr、AutoFuseAttrs、FusedAscBackend节点、ProcessAscBackendNodes这些结构在 AF 中均不存在。AF 不走子图融合路径,直接在单个 AscGraph 上做图优化和 codegen。因此 GE 的融合反推机制 不可移植到 AF。
当前 GE 侧的
BroadcastBackwardPass已被注释停用(asc_graph_pass.cpp:20),AF PR #1782 是在 AF 侧重新实现,替代 GE 侧旧实现。broadcast_backward_shared_split正是为了弥补 GE 融合反推机制缺失后的共享链场景处理:GE 通过融合后OptimizedFallback按消费者需求精确重建 Brc,AF 无法走这条路径,只能在融合前先把共享 Brc 链拆分为独立副本,再逐条执行后移。2.
CheckBackwardCommon新增两个屏障算子PR 侧
broadcast_backward_pass.cpp:407-412:// IndirectLoad 的输出轴属于独立的物理视图,Broadcast 后移不能跨过该边界。 if (IsOps<af::ascir_op::IndirectLoad>(next_node)) { return false; } if (ScheduleUtils::IsRemovePad(next_node)) { return false; }GE 侧只检查 Store、ViewOp、DtypeNotSupport 三种。PR 侧额外屏蔽
IndirectLoad和RemovePad——这是 graph-autofusion 特有的算子类型。原因分析:
IndirectLoad:该算子的输出轴属于独立的物理视图(通过索引间接加载,轴语义与普通 Load 不同),Broadcast 后移跨过此边界会破坏轴对应关系,导致数据读取错误。GE 侧不存在此算子类型。RemovePad:该算子涉及 padding 的逆操作,Broadcast 后移跨过此边界会导致 padding 区域的广播语义不一致。GE 侧不存在此算子类型。这两个屏障是 graph-autofusion 特有算子语义决定的,不是通用优化,无法回传 GE 侧。
3.
MoveBroadcastAfterMerge支持多后继at(0)单个后继GetPeerInDataAnchors()全部后继PR 侧支持 merge_node 输出连接到多个下游节点,GE 侧只处理单后继。
原因分析:
graph-autofusion 的图结构允许 merge 节点(如 Add)输出后继续 fork 到多个下游分支(fork-join-fork 模式),这在融合图中较常见。GE 侧假设 merge 后只有单后继,这在 GE 的图结构中成立(GE 的融合图通常在 merge 后接 Store)。如果 merge 有多个后继,GE 侧只重连第一个后继会导致其他后继丢失数据边,造成图改写不等价。
根因定位:GE 侧
MoveBroadcastAfterMerge的 BugPR 侧支持多后继的直接原因是修复了从 GE 移植时发现的
MoveBroadcastAfterMergeBug。以下是根因分析:图结构:
触发路径:
CollectCandidateMultiRefNodes(L778)选中 Broadcast_0(单输出锚点、2 个消费者、类型是 Broadcast)。CheckAllBranchesCanBackward(L859)从 Abs0 和 Cast1 分别 trace,两条分支都在 Add0 处汇聚(Add0 是多输入→merge),first_merge_node == current_merge_node == Add_0→ 返回 true。BackwardMultiRefBroadcast(L949)调用MoveBroadcastAfterMerge(L932)。Bug:GE 侧
MoveBroadcastAfterMerge只处理了 merge 的第一个消费者:// GE broadcast_backward_pass.cpp:1094 auto merge_next_in_anchor = merge_out_anchor->GetPeerInDataAnchors().at(0); // 只取 Sigmoid_0 GE_ASSERT_GRAPH_SUCCESS(GraphUtils::RemoveEdge(merge_out_anchor, merge_next_in_anchor)); // Add_0 → Broadcast_0 → Sigmoid_0 ✓ // Add_0 → Sqrt_0(仍然直连) ✗ 没有被处理Add_0 有两个消费者 Sigmoid_0 和 Sqrt_0,但
GetPeerInDataAnchors().at(0)只取了第一个(Sigmoid_0)。Sqrt0 仍然直接连在 Add0 上。UpdateComputeNodesAscTensorAttr(L994)把 Add_0 的输出 attr 更新为 Load0 的 attr(z1 维 repeat=1,广播前)。Sqrt0 的输入继承 Add0 的新输出 attr(repeat=1),但 Sqrt0 的输出 attr 仍保持原来的 repeat=18 → codegen 校验失败:修复方向:
MoveBroadcastAfterMerge(L932)需要处理 merge 节点的所有下游消费者,对每个消费者都插入 Broadcast(或克隆 Broadcast 链),而不是只处理.at(0)。PR 侧已修复:遍历GetPeerInDataAnchors()全部后继,逐个断开并重连到 Broadcast 链输出。三、算法层差异(核心逻辑不同)
4.
GetBroAxisFromNode— 轴匹配方式重写这是最核心的算法差异。
GE 侧(L93-127):按索引对齐匹配,要求
bro_attr_axis.size() == pre_bro_axis.size()GE_ASSERT_TRUE(bro_attr_axis.size() == pre_bro_axis.size()); for (auto index = 0U; index < bro_attr_axis.size(); index++) { if (IsEqOne(pre_bro_repeats[index]) && IsEqZero(pre_bro_strides[index])) { if (!IsEqOne(bro_repeats[index])) { bro_axis = bro_attr_axis[index]; return SUCCESS; } } }PR 侧(L202-242):按 axis 值查找匹配,允许大小不对齐
GE_ASSERT_TRUE(pre_bro_repeats.size() == pre_bro_strides.size()); GE_ASSERT_TRUE(bro_repeats.size() == bro_attr_axis.size()); GE_ASSERT_TRUE(bro_strides.size() == bro_attr_axis.size()); GE_ASSERT_TRUE(pre_bro_repeats.size() == pre_bro_axis.size()); for (size_t index = 0U; index < bro_attr_axis.size(); index++) { if (IsEqOne(bro_repeats[index])) { continue; } const auto pre_axis_iter = std::find(pre_bro_axis.begin(), pre_bro_axis.end(), bro_attr_axis[index]); if (pre_axis_iter == pre_bro_axis.end()) { // A missing input axis is a scalar/implicit broadcast dimension. bro_axis = bro_attr_axis[index]; return SUCCESS; } const size_t pre_index = static_cast<size_t>(std::distance(pre_bro_axis.begin(), pre_axis_iter)); if (IsEqOne(pre_bro_repeats[pre_index]) && IsEqZero(pre_bro_strides[pre_index])) { bro_axis = bro_attr_axis[index]; return SUCCESS; } }差异要点:
bro和pre_bro的 axis 列表等长且按序对齐,通过索引直接访问。bro的 axis 在pre_bro中不存在(标量/隐式广播维度),直接返回该轴。std::find按 axis 值匹配,支持两侧 axis 顺序不一致。原因分析:
graph-autofusion 的图结构允许 Broadcast 的输入 axis 与前驱节点的 axis 不对齐,典型场景包括:
GE 侧假设两侧 axis 等长且按序对齐,这在 GE 的图结构中成立(GE 的前端已保证 axis 对齐)。graph-autofusion 的图结构更灵活,需要按值查找匹配并处理 axis 不存在的情况。PR 侧先跳过
bro_repeats==1的轴是因为这些轴不产生实际广播,无需查找前驱对应轴。5. Topo ID 管理 — PR 侧完全移除
GE 侧有完整的 topo id 维护机制(因 GE 要求 pass 内拓扑序正确):
CollectBroNodestopo_listCollectCmpNodestopo_listUpdateTopoIdUpdateTopoIdsForMultiRefBackwardCreateAndUpdateBroadcastNodeUpdateNodeTopoInfoInsertBroadcastNodeasc_adapt::UpdateTopoIdBackwardMultiRefBroadcastUpdateTopoIdsForMultiRefBackwardBroadcastBackwardReallyUpdateTopoIdPR 侧策略:pass 内不管 topo id,在
RunPass入口和BroadcastBackward结尾统一调用ScheduleUtils::TopologicalSorting(graph)重排。GE 侧在BroadcastBackward结尾调用asc_adapt::TopologicalSorting。这也是 PR 侧行数更少的主要原因(1515 → 1324,减少约 190 行 topo 管理代码)。
原因分析:
GE 侧的 pass 框架要求 pass 执行后图内拓扑序仍然正确,因为 GE 后续 pass(如 codegen)直接依赖节点 topo id 顺序遍历,如果 topo id 错乱会导致代码生成顺序错误。因此 GE 侧在每次改写(调序、插入、删除节点)后都需即时更新 topo id。
graph-autofusion 的 pass 框架在 pass 结束后统一调用
TopologicalSorting重排,后续 pass 基于拓扑排序结果而非节点 id 遍历,不依赖 pass 内的 topo id 正确性。因此 PR 侧可以移除所有 topo id 维护代码,简化实现。代价是 pass 内如果需要按拓扑序遍历,必须等待 pass 结束后的统一排序。6.
CreateAndUpdateBroadcastNode— 节点创建方式不同GE 侧:通过适配层创建
const auto b_node = asc_adapt::CreateBroadcastNode(asc_graph, node, broadcast_info, index); asc_adapt::UpdateBroadcastNodeSchedInfo(b_node, tensor_info.sched_axis); asc_adapt::FromDtypeToOtherDtype(b_node, DT_FLOAT, tensor_info.dtype); // 从 FLOAT 验证PR 侧:内联创建,手动设置属性
Broadcast brc_op(brc_name.c_str()); auto b_node = asc_graph.AddNode(brc_op); brc_op.attr.sched = node->attr.sched; brc_op.attr.api.compute_type = af::ComputeType::kComputeBroadcast; brc_op.attr.api.type = af::ApiType::kAPITypeCompute; UpdateBroadcastNodeSchedInfo(b_node, ref_node); // 直接复制整个 sched FromDtypeToOtherDtype(b_node, tensor_info.dtype, tensor_info.dtype); // 相同 dtype 验证dtype 验证语义不同:GE 从
DT_FLOAT验证到目标 dtype,PR 用目标 dtype 自验证。原因分析:
asc_adapt::CreateBroadcastNode封装,因为 GE 的节点创建需要适配层处理 OpDesc、anchor 注册等 GE 框架细节;graph-autofusion 的AscGraph::AddNode接口更直接,传入 op 对象即可完成创建。tensor_info.sched_axis单独传递调度轴信息,因为适配层需要按 GE 的方式设置;PR 侧直接复制整个ref_node->attr.sched,因为 graph-autofusion 的 sched 结构可以直接整体复制。DT_FLOAT验证到目标 dtype,确认 Broadcast 算子支持从 float 转换到实际 dtype(GE 的 Cast 节点可能从 float 中间态转换);PR 侧用目标 dtype 自验证(from_dtype == to_dtype),因为 graph-autofusion 中 Broadcast 节点的输入输出 dtype 一致,只需确认该 dtype 被 Broadcast 支持,无需从 float 转换。四、Scalar 判定差异
GetNodeScalarInputListaf::IsScalarInputBySerializedExprs(SerializeExprs(repeats))— 序列化后判断ascgen_utils::IsScalarInput(repeats)— 直接判断ProcessSingleInputBranchascgen_utils::IsScalarInput(repeats)GE 侧有
SerializeExprs辅助函数,通过序列化表达式字符串判断标量。PR 侧直接调用ascgen_utils::IsScalarInput。原因分析:
GE 侧缺少直接判断标量输入的工具接口,只能通过序列化表达式字符串后用
IsScalarInputBySerializedExprs比较字符串来判断(间接方式,性能较差)。graph-autofusion 的ascgen_utils工具库提供了IsScalarInput接口,直接对Expression对象做符号化判断,无需序列化中间步骤,更高效且语义更清晰。相同的逻辑(移植保持一致)
以下函数两侧算法逻辑完全相同,仅适配层调用方式不同:
GetSingleNextNodeGetPeerOutNodeSafeIsNextViewOpReverseCollectBrcNodesGetBroAxisesGetBroAxisesIndexIsSameBroNodesRemoveAndRelinkNodeEdgeRemoveBroadcastOneByOneRemoveBroadcasts(单链)FindSubSetCollectSameBrcAxisCheckBackwardCommon(除新增屏障)CanBackwardSimplifiedIsScalarInputCheckNodeSupportsScalarInputCheckScalarInputSupportIsMulInputsCanBackwardCanBackwardReorderBroadcastsUpdateComputeNodesAscTensorAttrUpdateBroadcastNodesDataTypeCollectBranchBroadcastNodesHasCommonBroadcastAxisGetPreBroadcastNodeProcessSingleInputBranchProcessOtherInputBranchesJudgeNextCompOpSupportsScalarInputContainsBroadcastNodeCollectCandidateMultiRefNodesExtractBroadcastChainFromNodeTraceBranchToMergeNodeCheckBranchNodesSupportBackwardCheckAllBranchesCanBackwardCheckAllBranchesSupportBackwardDisconnectBranchesFromBroadcastBackwardMultiRefBroadcast(除 topo)ProcessMultiRefBroadcastBackwardCollectBackwardStartNodesCollectBackwardSatisfyStartNodesRemoveBroadcasts(多链)GetBackwardBrcNodesUpdateOutputTensorJudgePartBackwardProcessOriginalBackwardLogicBroadcastBackward主循环主循环结构对比
两侧
BroadcastBackward主循环结构完全相同:GE 侧无步骤 2(shared split 预处理),其余流程一致。
共有遗留问题
以下问题在 GE 侧和 PR 侧都存在,是从 GE 移植时继承的:
view_op_type全局变量未声明为const static— 两侧相同FindSubSet原地修改输入参数 — 两侧相同BroadcastBackwardwhile 循环无迭代上限 — 两侧相同JudgePartBackward递归调用可能栈溢出 — 两侧相同GetPeerOutNodeSafe冗余工作 — 两侧相同(GE 通过asc_adapt::GetPeerOutNode二次获取,PR 通过GetPeerOutNode二次获取)原因分析:
这些问题在 GE 侧原始代码中就存在,PR 侧移植时为保持算法一致性未做修复。具体原因:
view_op_type未const static:GE 侧原始代码遗留,每次进入匿名命名空间时重新构造 vector,存在轻微性能开销和可变性风险。移植时保持一致未修改。FindSubSet原地修改参数:GE 侧原始设计为节省内存,参数按非 const 引用传入。当前调用方传入局部变量无 bug,但接口契约易错。移植时保持一致。JudgePartBackward通过递归处理级联 mul_input 链,深层图可能导致栈溢出。移植时保持一致。GetPeerOutNodeSafe冗余工作:GE 侧原始代码在验证 anchor 后又调用GetPeerOutNode重复获取,是代码冗余。PR 侧移植时保持了相同结构。建议:这些问题应在后续维护中统一修复,且修复时 GE 侧和 graph-autofusion 侧应同步,避免两侧进一步分叉。
总结
GetBroAxisFromNode轴匹配重写、topo id 管理移除PR 侧本质上是 GE 侧的移植+增强:去掉了
asc_adapt适配层和 topo id 管理,新增了 shared split 预处理和更宽松的轴匹配逻辑。核心后移算法(单链后移、多输入后移、多引用后移、部分公共轴后移)的流程完全一致。差异的根本原因在于两套框架的图模型和 pass 框架设计不同:GE 侧 AscGraph 是 GE ComputeGraph 的子图,受 GE 框架约束(适配层、topo id、单后继假设);graph-autofusion 中 AscGraph 是原生图结构,更灵活(axis 不对齐、多后继、共享链),且 pass 框架简化了拓扑序管理。