昇腾CANN/GE Concat No Task特性分析

Concat No Task 特性分析【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge1. 特性背景1.1 问题场景在深度学习模型中ConcatD/ConcatV2D算子用于将多个输入 Tensor 沿指定维度拼接为一个输出 Tensor。传统的 Concat 算子执行流程为输入A ──┐ 输入B ──┼──► Concat Task ──► 输出(拼接结果) 输入C ──┘其中Concat Task需要在昇腾 AI Core 上启动一个计算任务通过 DataMove 指令将各输入数据搬运到输出地址的连续内存区域。1.2 优化思路当 Concat 算子的输入 Tensor 在内存中天然连续排列时实际上不需要执行任何数据搬运操作——输出地址直接复用第一个输入的地址即可。Concat No Task 特性通过在编译期识别这种场景将 Concat 算子标记为虚拟算子Virtual Op从而不生成硬件 Task跳过 AI Core 任务下发不搬运内存输出直接复用输入内存地址消除冗余计算避免无意义的数据搬移2. 用户使用场景2.1 典型场景分布式训练中的 AllGather Concat在分布式训练中多个卡通过 AllGather 收集各自的数据后往往需要拼接为完整批次Card0: Data ──► AllGather ──┐ ├──► ConcatD ──► 完整Batch Card1: Data ──► AllGather ──┘AllGather 输出的数据在内存中已经按卡号顺序连续排列Concat 只是逻辑上的拼接物理上不需要搬运数据。2.2 典型场景分块计算后的结果合并将大 Tensor 拆分到多个算子并行计算后合并结果Input ──► Split ──┬──► Relu ──┐ ├──► Relu ──┼──► ConcatD ──► Output └──► Relu ──┘当 Split 沿 batch 维度拆分、各分支计算不改变内存布局时Concat 的输入天然连续。2.3 适用条件Concat No Task 优化需要同时满足以下严格条件条件类别具体要求算子类型仅限ConcatD和ConcatV2DShape 约束concat_dim 轴之前的所有维度值必须为 1对齐约束concat_dim 轴原始尺寸必须是 align_shape 对应维度的整数倍无 padding输入约束不能有 Scalar 输入所有输入 Tensor 大小必须是 32 字节对齐来源约束不能有多个输入来源于同一个输出锚点前驱节点不能是 DATA、REFDATA、VARIABLE、CONSTANTOP、CONSTANT 等节点类型前驱节点不能包含子图subgraph前驱属性不能有 continuous_input、continuous_output、reference 等属性后继节点后继节点不能已经是虚拟算子已有 _no_task 属性输出约束输入不能同时是模型输出NetOutput内存类型所有输入的内存类型必须一致LxFusion不能涉及 LxFusionL1/L2/UB 地址、lxslice 算子Shape 模式不支持 Unknown Shape动态 Shape场景图模式不支持 Single Op 场景和内存非连续分配场景3. 对外接口3.1 编译期属性标记Concat No Task 通过以下 Graph 属性与系统其他模块交互属性名类型设置对象说明_no_taskboolConcat 算子 OpDesc标记该算子不生成硬件 Task_nopadding_continuous_inputboolConcat 算子 OpDesc标记输入为无 padding 的连续内存_output_reuse_inputboolConcat 算子 OpDesc标记输出复用输入内存_reuse_input_on_dim_indexint64Concat 算子 OpDesc指定复用输入内存的维度索引固定为 0can_reused_for_concat_optimizebool前驱节点的输出 TensorDesc标记该输出已被 Concat 优化占用不可再被其他 Concat 复用3.2 Pass 注册ConcatNotaskPass 注册为 O3 优化级别的 GraphPass REG_PASS_OPTION(ConcatNotaskPass).LEVELS(OoLevel::kO3);在编译流程中该 Pass 运行于OptimizeStage2的最后阶段在子图合并完成后、内存冲突处理之前执行graph_manager.cc: OptimizeAfterMergeSubGraph() ├── ... (前期优化) ├── BufferPoolMemoryPass ├── ParallelGroupPass └── ConcatNotaskPass ← 图稳定后执行3.3 运行时行为在运行时RT1 和 RT2标记了_no_task的算子会被特殊处理RT1 (Davinci)TbeKernelHandle::NeedInit检测到_no_task属性后返回 false跳过 Kernel 初始化RT2 (V2)IsVirtualOp函数检测_no_task属性在 AICore Node Converter 中跳过 Task 生成4. 具体实现4.1 整体架构┌─────────────────────────────────────────────────────────────┐ │ 编译期 (Compiler) │ │ │ │ ┌──────────────────┐ ┌──────────────────────────────┐ │ │ │ ConcatNotaskPass │───►│ 属性标记 │ │ │ │ │ │ _no_task │ │ │ │ 校验链: │ │ _nopadding_continuous_input │ │ │ │ ├─ InputCheck │ │ _output_reuse_input │ │ │ │ ├─ CheckConcatDim│ │ _reuse_input_on_dim_index │ │ │ │ ├─ OutputCheck │ └──────────────────────────────┘ │ │ │ └─ LxFusionCheck│ │ │ └──────────────────┘ │ └─────────────────────┬───────────────────────────────────────┘ │ 属性传递 ▼ ┌─────────────────────────────────────────────────────────────┐ │ 内存分配 (Memory Assigner) │ │ │ │ ┌─────────────────────┐ ┌──────────────────────────┐ │ │ │ BlockMemAssigner │ │ GraphMemAssigner │ │ │ │ │ │ │ │ │ │ 检测 NOPADDING_ │ │ 计算 continuous_type │ │ │ │ CONTINUOUS_INPUT │ │ kTypeInputNoPadding │ │ │ │ │ │ │ │ │ │ no_assign_memtrue │ │ 按 dim_index 计算 │ │ │ │ (不分配独立内存) │ │ nopadding_size │ │ │ └─────────────────────┘ └──────────────────────────┘ │ └─────────────────────┬───────────────────────────────────────┘ │ ▼ ┌─────────────────────────────────────────────────────────────┐ │ Task 生成 (Task Generator) │ │ │ │ 检测到 _no_task 属性 → 跳过该节点的 Task 生成 │ │ MarkFirstAndLastOps 中跳过 notask 节点 │ └─────────────────────┬───────────────────────────────────────┘ │ ▼ ┌─────────────────────────────────────────────────────────────┐ │ 运行时 (Runtime) │ │ │ │ RT1: TbeKernelHandle 跳过初始化 │ │ RT2: AICoreNodeConverter 跳过转换 │ │ 输出地址直接复用第一个输入地址 │ └─────────────────────────────────────────────────────────────┘4.2 ConcatNotaskPass 核心校验链ConcatNotaskPass::Run对图中每个 ConcatD/ConcatV2D 节点执行以下校验链全部通过后才设置属性4.2.1 场景排除Single Op 场景单算子模式无需优化内存非连续分配设置了ATTR_NAME_MEMORY_DISCONTIGUOUS_ALLOCATION的图跳过Unknown Shape输入或输出包含动态 Shape 的算子跳过Dynamic Shape 图所属图标记了动态 Shape 分区的跳过4.2.2 InputCheck输入校验对每个输入锚点依次检查IsScalarInput排除维度数为 0 的 Scalar 输入CheckTensorAlign多输入场景下每个 Tensor 大小必须是 32 字节对齐HasSameSourceAnchor通过GetFirstOutAnchorNotInRefNode追溯 RefNode 链确保没有多个输入追溯到同一个输出锚点IsPreNodeTypeValid通过GetFirstNotRefNode找到实际生产节点排除 DATA/REFDATA/VARIABLE/CONSTANTOP/CONSTANTIsPreNodeWithSubgraph前驱节点不能包含子图实例IsPreOutAnchorCanReuseForConcatOptimize检查前驱输出 TensorDesc 的can_reused_for_concat_optimize属性确保未被其他 Concat 占用IsPreOutAnchorValidMultiRef前驱输出如果同时连接到 NetOutput则不能优化IsPreNodeAttrValid前驱节点不能有 continuous_input、continuous_output、reference、_no_task、_output_reuse_input、_nopadding_continuous_input 等属性也不能有 atomic outputIsSameInputMemType所有输入的内存类型必须一致通过ATTR_NAME_OUTPUT_MEM_TYPE_LIST检查4.2.3 CheckConcatDimConcat 维度校验这是最核心的校验逻辑确保 concat_dim 轴前面的维度全为 1原始格式 (如 NCHW) ──► 运行时格式 (如 NC1HWC0) │ │ │ GetTransferDims() │ │ (调用 FE 接口) │ ▼ ▼ src_to_dst_transfer_dims dst_to_src_transfer_dims {0},{1,4},{2},{3} {0},{1},{2},{3},{1}GetAlignedShape调用transformer::TransferShapeUtils::GetAlignedShape获取对齐后的 shapeGetTransferDims调用transformer::TransferShapeUtils::TransferDims获取原始格式到运行时格式的轴映射关系CheckRealConcatDim在运行时格式下找到实际的 concat_dim 轴验证该轴之前的所有维度值都为 1CheckConcatDimAlignment验证 concat_dim 轴原始尺寸是 align_shape 对应维度的整数倍确保无 paddingCheckRealConcatDim 的关键逻辑通过src_to_dst_transfer_dims[concat_dim][0]找到运行时格式中的 real_concat_dim如果 real_concat_dim 是由合轴产生的dst_to_src_transfer_dims中对应多个源轴需要额外验证real_concat_dim 轴之前的所有轴值都为 1concat_dim 所在的合轴中、concat_dim 之前的值都为 1如果 real_concat_dim 不是合轴产生的只需验证之前的轴值都为 14.2.4 OutputCheck输出校验遍历 Concat 节点的所有后继节点如果后继是 Reshape 且有输出节点则继续检查 Reshape 的输出节点后继节点不能已有_no_task、_output_reuse_input、_nopadding_continuous_input属性4.2.5 LxFusionCheckLxFusion 校验IsLxFusionMem检查输入/输出内存类型是否为 L1/L2/UB这些是 LxFusion 使用的片上内存IsLxFusionOp检查算子名称是否包含 lxslice4.3 属性设置所有校验通过后SetAttrForConcatNotask执行以下操作// 设置 Concat 算子自身属性 AttrUtils::SetBool(op_desc, ATTR_NAME_NOTASK, true); AttrUtils::SetBool(op_desc, ATTR_NAME_NOPADDING_CONTINUOUS_INPUT, true); AttrUtils::SetBool(op_desc, ATTR_NAME_OUTPUT_REUSE_INPUT, true); AttrUtils::SetInt(op_desc, ATTR_NAME_REUSE_INPUT_ON_DIM_INDEX, 0); // 标记前驱节点的输出 TensorDesc 不可再被复用 for each input: AttrUtils::SetBool(output_tensor_desc, can_reused_for_concat_optimize, false);4.4 内存分配联动4.4.1 BlockMemAssigner在AnalyzeSymbolMemReuseInfo中当检测到节点具有ATTR_NAME_NOPADDING_CONTINUOUS_INPUT属性时if (is_input_continuous) { symbol_mem_reuse_info_[symbol].no_assign_mem_ true; }这意味着该符号内存块不会被独立分配内存而是复用上游的内存地址。GetContinuousNodeLifeTimeBegin方法处理级联的 continuous input 场景a b c (第一层) | | | d e f (第二层) |___|___| | g h i (第三层, h 是 continuous input) |___|___| | j (第四层, j 是 continuous input)g 不能与 a/b/c 复用内存因为 d/e/f 的内存会被 g 的地址替换级联 continuous input。因此 g 的生命期起点要追溯到 d/e/f 中的最早者。4.4.2 GraphMemAssigner在GetContinuousType中识别 continuous 类型kTypeInputNoPadding _nopadding_continuous_input _output_reuse_input kTypeOutputNoPadding _nopadding_continuous_output _output_reuse_input在GetMemorySize中对于 nopadding 类型通过_reuse_input_on_dim_index获取维度索引计算nopadding_size实际数据大小和tensor_size完整带 padding 大小内存分配器据此分配精确大小的内存块4.4.3 SetInputOutputOffsetPass对于 no_task 的 Concat 节点如果具有ATTR_NAME_CONTINUOUS_INPUT或满足 BufferFusion 条件会调用SetOutputOffsetForConcat设置输出偏移确保输出地址正确指向输入数据的起始位置。4.5 Task 生成跳过在TaskGenerator::MarkFirstAndLastOps中bool attr_notask false; if (ge::AttrUtils::GetBool(op_desc, ATTR_NAME_NOTASK, attr_notask) attr_notask) { continue; // 跳过 notask 节点不纳入连续节点列表 }这确保 no_task 节点不会被视为流中连续计算节点的一部分不会影响首尾算子的标记。4.6 运行时处理4.6.1 RT1 (Davinci Model)在TbeKernelHandle::NeedInit中bool attr_no_task false; const bool get_attr_no_task_flag AttrUtils::GetBool(op_desc, ATTR_NAME_NOTASK, attr_no_task); if (get_attr_no_task_flag attr_no_task) { GELOGI(Node[name:%s, type:%s] does not generate task, skip initialization., ...); return false; // 跳过 Kernel 初始化 }4.6.2 RT2 (V2 Engine)在aicore_node_converter.cc的IsVirtualOp中bool attr_no_task false; (void)ge::AttrUtils::GetBool(op_desc, ge::ATTR_NAME_NOTASK, attr_no_task); return attr_no_task; // 返回 true 表示虚拟算子跳过转换5. 与其他优化的协同5.1 与 Split No Task 的配合Concat No Task 经常与 Split No Task 配合使用形成切分-计算-合并的零拷贝流水线Input ──► Split(no_task) ──► 计算 ──► Concat(no_task) ──► OutputSplit 沿 concat_dim 的反方向切分输出地址直接指向输入地址的不同偏移Concat 则将这些偏移地址视为连续内存直接复用。5.2 与内存复用的关系Concat No Task 节点在内存复用检查中被特殊对待ReuseChecker中具有_no_task属性的节点被视为 buffer_pool 类型mem_reuse_strategy.cc中nopadding continuous input 的节点参与特殊的内存复用策略5.3 与地址刷新的关系在MemLayoutConflictUtil中Concat No Task 场景下需要考虑地址刷新data | identity (插入 identity 算子支持地址刷新) | split(no_task, no_padding_continuous_output) / \ op1 op2当涉及用户输入且需要地址刷新时会在 Data 和 Split 之间插入 Identity 算子。6. 测试验证单元测试文件位于tests/ge/ut/ge/graph/passes/concat_notask_pass_unittest.cc覆盖了以下场景测试用例验证内容allgather_connect_to_concatAllGather 输出连接到 Concat验证属性设置正确allgather_connect_to_concat_reshapeAllGather → Reshape → Concat 链路多组 RefData 测试验证 RefNode 链追溯逻辑不同 Format 测试ND、NCHW、NC1HWC0 等格式下的 dim 校验7. 设计思考7.1 为什么选择属性标记而非图改写Concat No Task 采用属性标记而非删除节点的方式原因是保留调试信息Dump 功能需要保留算子的 OpDesc 和地址信息见InitNoTaskAndDumpNeededNode保持图结构完整删除节点会破坏图的拓扑关系影响其他 Pass 的执行支持动态场景属性标记可以在不同编译阶段灵活处理7.2 为什么校验条件如此严格Concat No Task 的校验条件多达十余项这是因为内存安全如果输入不连续却标记为 no_task会导致读取到错误数据对齐约束昇腾硬件对内存对齐有严格要求不满足 32B 对齐可能导致硬件异常级联影响一个 no_task 节点会影响下游的内存分配策略错误标记会传播7.3 为什么放在 Stage2 最后执行注释明确说明图稳定了再做 ConcatNotaskPass在 Stage2 之前图结构可能还在变化子图合并、算子融合等前驱节点的属性可能在后续 Pass 中被修改放在最后执行可以确保基于最终稳定的图结构做判断【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考