CANN ops-nn 分段求和算子 SegmentSum 详解：原理、约束与 GE 图模式调用实战

人工智能算子库深度学习CANNAscend【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-nn点击查看免费下载SegmentSum分段求和是 CANN ops-nn 算子库中的索引类计算算子按分段索引对输入 Tensor 的若干行进行分组求和。本文以 index/segment_sum/README.md 为核心骨架结合仓库内算子定义、Tiling、Kernel、配置与测试源码系统讲解 SegmentSum 的功能语义、参数与约束、底层实现机制并通过完整示例演示如何在 GE 图模式下调用该算子。产品支持情况SegmentSum 在当前仓库中的产品适配情况如下表所示产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品×Atlas A2 训练系列产品/Atlas A2 推理系列产品×Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×当前仅有 Ascend 950 系列950PR/950DT支持 SegmentSum其余 Atlas 训练/推理产品暂不支持。这一结论同样可以从算子注册代码中得到印证op_host/segment_sum_def.cpp 中仅对ascend950与ascend350两个 SoC 配置注册了 AICore 配置且 tiling 与 kernel 实现均位于arch35目录下arch35对应 Ascend 950 的架构代号。同时 op_host/config/ 下只提供了ascend950与ascend350两套算子二进制 JSON 配置与 README 的支持矩阵保持一致。功能说明SegmentSum 的功能是对输入 Tensorx按分段索引segment_ids进行求和。其数学定义如下$$ y[i] \sum_{\substack{j\ \text{segment_ids}[j] i}} x[j] $$即遍历所有满足segment_ids[j] i的索引j将对应的x[j]逐元素累加到y[i]若某个段i没有任何元素属于它则y[i] 0。该算子与 TensorFlow 的tf.math.segment_sum语义兼容这一点在 op_graph/segment_sum_proto.h 的注释中有明确说明测试脚本 tests/assets/golden.py 也直接以tf.math.segment_sum的结果作为 Golden 基准进行对拍验证。计算示例README 中给出的标准用例为输入 Tensor$$ x \begin{bmatrix} [1 2] \ [3 4] \ [5 6] \ [7 8] \end{bmatrix} $$分段索引 Tensor$$ segment_ids [0, 0, 1, 2] $$输出 Tensor$$ y \begin{bmatrix} [4 6] \ [5 6] \ [7 8] \end{bmatrix} $$推导过程segment_ids中索引 0、1 对应段 0因此y[0] x[0] x[1] [13, 24] [4, 6]索引 2 对应段 1y[1] x[2] [5, 6]索引 3 对应段 2y[2] x[3] [7, 8]。由于每个元素按行第 0 维分组SegmentSum 的求和粒度是“整行累加”即逐列分别相加。该示例在 GE 图模式的调用样例 examples/test_geir_segment_sum.cpp 中被完整复现输入x为 shape{4, 2}、值全为 2.0 的 FLOAT32 Tensorsegment_ids为 shape{4}、值{0, 0, 1, 2}的 INT64 Const 节点输出y的 shape 为{3, 2}与 README 中的示例一一对应。关键语义要点segment_ids必须按升序排序但允许存在重复值重复值表示多个行归入同一个段segment_ids指示当前分段的值归属于哪个段segment_ids的值必须 0且各段编号可以不连续例如[0, 0, 2]也是合法输入此时y[1]将全为 0输出 shape 为[max(segment_ids) 1, x.shape[1:]]即输出第 0 维由最大段编号决定其余维度与x保持一致。参数说明参数名输入/输出/属性描述数据类型数据格式x输入输入数据即公式中的xFLOAT32、FLOAT16、BFLOAT16、INT32、INT64、UINT32、UINT64NDsegment_ids输入分段索引即公式中的segment_idsINT32、INT64NDy输出输出值信息即公式中的yFLOAT32、FLOAT16、BFLOAT16、INT32、INT64、UINT32、UINT64ND从算子定义源码看op_host/segment_sum_def.cpp 通过valueDataTypeXY声明了x/y支持的 7 种数值类型FLOAT16、FLOAT、INT32、INT64、UINT32、UINT64、BF16并通过valueDataTypeIds将segment_ids限定为 INT32、INT64 两种索引类型输入输出均要求 ND 格式并开启了动态 Shape、动态 Rank 支持。此外segment_ids被标记为ValueDepend(OPTIONAL)表示索引张量的值在编译期可参与 tiling 决策。约束说明x维度至少为 1rank 1从算子定义看x支持 1D~8D见 op_graph/segment_sum_proto.h 注释ND 格式下逐行累加。segment_ids必须是 INT32 或 INT64 类型必须为 1D Tensor且segment_ids.shape[0] x.shape[0]即每一行都有一个段归属值必须按升序排序且segment_ids.value 0。y类型必须与x相同维度与x相同shape 为[max(segment_ids) 1, x.shape[1:]]。这些约束在 op_graph/segment_sum_proto.h 的attention注释中逐条列出同时单测与 ST 用例也对非法输入如乱序索引进行了覆盖。调用说明SegmentSum 当前支持的调用方式为GE 图模式通过算子 IR 构图方式调用。调用方式调用样例说明GE 图模式test_geir_segment_sum.cpp通过算子 IRsegment_sum_proto.h构图方式调用 SegmentSum 算子算子 IR 定义op_graph/segment_sum_proto.h 使用REG_OP宏注册算子原型REG_OP(SegmentSum) .INPUT(x, TensorType::NumberType()) .INPUT(segment_ids, TensorType::IndexNumberType()) .OUTPUT(y, TensorType::NumberType()) .OP_END_FACTORY_REG(SegmentSum)x与y使用NumberType数值类型segment_ids使用IndexNumberType索引类型与参数表中的类型范围一致。GE 图模式调用示例详解examples/test_geir_segment_sum.cpp 展示了完整的调用流程核心步骤如下初始化 GE通过ge::GEInitialize传入全局配置包括ge.exec.deviceId0与ge.graphRunMode1std::mapAscendString, AscendString global_options {{ge.exec.deviceId, 0}, {ge.graphRunMode, 1}}; Status ret ge::GEInitialize(global_options);构造算子节点创建op::SegmentSum(segmentSum)算子实例并连接数据节点与常量节点auto segmentSum op::SegmentSum(segmentSum); // input x: Data 节点shape {4, 2}填充值 2.0 vectorint64_t xShape {4, 2}; auto xData op::Data(placeholder1).set_attr_index(0); TensorDesc xDesc(ge::Shape(xShape), FORMAT_ND, inDtype); xData.update_input_desc_x(xDesc); xData.update_output_desc_y(xDesc); graph.AddOp(xData); segmentSum.set_input_x(xData); // input segment_ids: Const 节点shape {4}值 {0, 0, 1, 2} vectorint64_t segmentIdsShape {4}; auto segmentIdsConst op::Const(placeholder2); TensorDesc segmentIdsDesc(ge::Shape(segmentIdsShape), FORMAT_ND, DT_INT64); segmentIdsConst.SetAttr(value, segmentIdsTensor); graph.AddOp(segmentIdsConst); segmentSum.set_input_segment_ids(segmentIdsConst);声明输出为segmentSum设置输出y的 TensorDescshape{3, 2}并加入输出算子列表构建并运行图将图添加至ge::Session后调用RunGraph执行随后读取输出 Tensor 数据、打印结果并落盘为.bin文件ret session-AddGraph(graph_id, graph, graph_options); std::vectorge::Tensor output; ret session-RunGraph(graph_id, input, output);该样例还给出了数据生成辅助函数如GenFloat32Data、GenInt64Data与GetDataTypeSize字节宽度映射方便读者替换为任意受支持的数据类型组合进行验证。源码实现机制算子定义与注册OpHost 侧op_host/segment_sum_def.cpp 中除了声明输入输出类型外还通过OpAICoreConfig配置了算子行为OpAICoreConfig aicoreConfig; aicoreConfig.DynamicCompileStaticFlag(true) .DynamicFormatFlag(true) .DynamicRankSupportFlag(true) .DynamicShapeSupportFlag(true) .NeedCheckSupportFlag(false) .ExtendCfgInfo(opFile.value, segment_sum_apt); this-AICore().AddConfig(ascend950, aicoreConfig); this-AICore().AddConfig(ascend350, aicoreConfig);其中DynamicShapeSupportFlag(true)与DynamicRankSupportFlag(true)表示算子支持动态 Shape 与动态 RankExtendCfgInfo(opFile.value, segment_sum_apt)将 Kernel 实现指向 op_kernel/segment_sum_apt.cpp。Tiling 策略Host 侧Tiling 的入口在 op_host/arch35/segment_sum_tiling.cppTilingPrepare4SegmentSum在编译期通过PlatformAscendC获取 AIV 核数core_num与 UB 内存大小ub_size存入SegmentSumCompileInfoTiling4SegmentSum调用TilingRegistry::GetInstance().DoTilingImpl(context)完成实际 tiling 计算。op_host/arch35/segment_sum_tiling_base.h 中的SegmentSumBaseTiling定义了核心 tiling 参数包括外层维度outerDim_即分段行数、内层维度innerDim_每行元素数、段数量segmentNum_、x数据类型字节数valueTypeBytes_、索引类型字节数idTypeBytes_等为后续核内计算划分提供依据。Kernel 实现NPU 侧op_kernel/segment_sum_apt.cpp 是 SegmentSum 的 AICore 内核入口通过 Tiling Key 分发到三种计算路径#define TEMPLATE_SIMT_TILING_KEY 1000 #define SIMD_ATOMIC_SUPPORT_TILING_KEY 2000 #define SIMD_DETERM_TILING_KEY 2002TILING_KEY 1000SIMT 路径调用SegmentSumSimtDTYPE_X, DTYPE_SEGMENT_IDS适用于无原子操作的 SIMT 计算TILING_KEY 2000SIMD 原子操作路径先通过AllClear清空输出再调用SegmentSumSimd执行求和。从代码看该路径不适用于 UINT32、UINT64、INT64 三种类型通过constexpr编译期分支排除推测与原子加法类型支持范围有关TILING_KEY 2002SIMD 确定性路径依次执行AllClear清空输出、SegmentSumSimdDeterm确定性求和、SegmentSumMultiCoreAdd多核结果累加用于保证多核场景下求和结果的确定性与可复现性。对应 Kernel 实现在 op_kernel/arch35/ 下segment_sum_simt.h、segment_sum_simd.h、segment_sum_simd_determ.h、segment_sum_simd_mult_core_add.h、clear_output.h与segment_sum_struct.htiling 结构体定义。算子二进制配置op_host/config/ascend950/segment_sum_binary.json 列出了各 (x 类型, segment_ids 类型) 组合对应的二进制文件映射共 14 组覆盖 FLOAT32/FLOAT16/BFLOAT16/INT32/INT64/UINT32/UINT64 × INT32/INT64。所有条目的 shape 均为[-2]-2 表示动态 Shapeformat 均为 NDparamType 均为 required说明该算子在 Ascend 950 上支持动态 Shape 的运行时编译。ascend350下另有同构的 segment_sum_binary.json。测试与验证SegmentSum 在仓库内提供了完整的测试资产与用例输入生成脚本tests/assets/input.py随机生成x与segment_ids。生成的索引值范围受输出第 0 维约束随后调用np.sort保证segment_ids升序符合算子约束要求Golden 脚本tests/assets/golden.py直接调用 TensorFlow 的tf.math.segment_sum计算期望输出bfloat16 输入先转为 float32 计算再转回用于与 NPU 实际输出进行精度对拍ST 用例tests/st/arch35/ttk_kernel_segment_sum_st.csv覆盖 FLOAT32×INT32/INT64、BFLOAT16×INT32、INT64×INT64 等多种类型组合Shape 从 1D如(5070,)到 8D如(4, 3, 1, 7, 2, 5, 5, 5)不等精度容忍度统一为 1e-8同时验证了输出 shape 恒等于[max(segment_ids)1, x.shape[1:]]UT 用例tests/ut/op_host/arch35/test_segment_sum_tiling.cpp针对 Tiling 逻辑做宿主侧单元测试。小结SegmentSum 是 ops-nn 中一个语义简洁但实现路径丰富的分段聚合算子Host 侧通过动态 Tiling 感知 AIV 核数与 UB 容量Kernel 侧依据数据规模在 SIMT 与多种 SIMD 路径间选择并通过确定性累加保证多核结果可复现。使用时需严格遵守segment_ids升序、非负、1D 且长度等于x.shape[0]的约束输出维度则由最大段编号决定。读者可基于 examples/test_geir_segment_sum.cpp 直接构造 GE 图进行功能验证并借助仓库内的 tiling 单测与 ST 对拍用例深入理解其行为。赞分享人工智能算子库深度学习CANNAscend【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-nn点击查看免费下载相关推荐CANN ops-nn 稀疏分段求和算子 SparseSegmentSum 详解原理、参数约束与 GE IR 图模式调用CANN ops nn 稀疏分段求和算子 SparseSegmentSum 详解原理、参数约束与 GE IR 图模式调用 导读 SparseSegmentSu人工智能算子库深度学习CANNAscendCANN ops-nn 算子解析UnsortedSegmentMax 分段最大值算子原理、约束与图模式调用指南CANN ops nn 算子解析UnsortedSegmentMax 分段最大值算子原理、约束与图模式调用指南 UnsortedSegmentMax 是 CA人工智能算子库深度学习CANNAscendCANN ops-nn UnsortedSegmentSum 算子全解析分段求和原理、图模式调用与 NPU 实现CANN ops nn UnsortedSegmentSum 算子全解析分段求和原理、图模式调用与 NPU 实现 导读 本文围绕 CANN ops nn 神经人工智能算子库深度学习CANNAscend创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考