UnsortedSegmentSum 算子 API 描述【免费下载链接】cann-bench评测AI在处理CANN领域代码任务的能力涵盖算子生成、算子优化等领域支撑模型选型、训练效果评估统一量化评估标准识别Agent能力短板构建CANN领域评测平台推动AI能力在CANN领域的持续演进。项目地址: https://gitcode.com/cann/cann-bench1. 算子简介沿 segment_ids 指定的段对数据进行求和。主要应用场景图神经网络中的节点特征聚合按邻居分段求和点云处理中的体素化聚合稀疏特征的按组求和与池化嵌入表梯度的按 ID 累加算子特征难度等级L2ScatterUpdate双输入单输出根据 segment_ids 将 data 中的元素按段分组求和2. 算子定义数学公式$$ y[i] \sum_{j: \text{segment_ids}[j] i} \text{data}[j] $$对于每个段 $i \in [0, \text{num_segments})$将所有 segment_ids 等于 $i$ 的 data 元素在第 0 维上求和。若某段没有对应的元素则输出为零。3. 接口规范算子原型cann_bench.unsorted_segment_sum(Tensor data, Tensor segment_ids, int num_segments) - Tensor y输入参数说明参数类型默认值描述dataTensor必选输入数据张量segment_idsTensor必选段 ID 张量值在 [0, num_segments) 范围内num_segmentsint必选段数量输出参数Shapedtype描述y(num_segments, *data.shape[1:])与 data 相同输出张量段求和结果数据类型data dtypesegment_ids dtype输出 dtypefloat16int32 / int64float16bfloat16int32 / int64bfloat16float32int32 / int64float32int32int32 / int64int32int64int32 / int64int64规则与约束segment_ids 必须为 1D 张量长度等于 data 的第 0 维大小即segment_ids.shape (data.shape[0],)segment_ids 中的值必须在[0, num_segments)范围内segment_ids 的 dtype 必须为 int32 或 int64输出的第 0 维大小为 num_segments其余维度与 data 的尾部维度一致y.shape (num_segments, *data.shape[1:])输出 dtype 与 data 一致若某个段 ID 在 segment_ids 中未出现对应输出段为全零num_segments 必须为正整数支持范围输入 tensor 各维度与参数的支持范围维度 / 参数范围备注data.rank1 ~ 8cases.csv 实测 1 ~ 5data.shape[0]段轴长度N1 ~ 4194304cases.csv 实测 2 ~ 2097152约束segment_ids.shape[0] data.shape[0]data.shape[1:]每个尾部维度1 ~ 16384cases.csv 实测 7 ~ 8193segment_ids.rank1固定为 1D长度等于data.shape[0]segment_ids.shape[0]1 ~ 4194304cases.csv 实测 2 ~ 2097152与data.shape[0]相等num_segmentsattr1 ~ 32768cases.csv 实测 1 ~ 16384约束segment_ids ∈ [0, num_segments)约束shape 关系segment_ids.shape (data.shape[0],)必须为 1D 且长度等于 data 第 0 维y.shape (num_segments, *data.shape[1:])。dtype 关系y.dtype data.dtype输出沿 data dtypesegment_ids.dtype ∈ {int32, int64}与 data dtype 解耦。值域约束segment_ids中每个元素v必须满足0 ≤ v num_segments越界元素的行为未定义。稀疏段未在segment_ids中出现的段 ID 对应输出行为全零zero-fill。attr 约束num_segments必须为正整数。数值精度浮点 dtype 下内部使用 fp32 累加以减小累积误差参见 §5 golden 代码NaN/Inf 在求和后可能传播为输出 NaN/Inf。4. 精度要求采用生态算子精度标准进行验证。误差指标平均相对误差MERE采样点中相对误差平均值$$ \text{MERE} \text{avg}(\frac{\text{abs}(actual - golden)}{\text{abs}(golden)\text{1e-7}}) $$最大相对误差MARE采样点中相对误差最大值$$ \text{MARE} \max(\frac{\text{abs}(actual - golden)}{\text{abs}(golden)\text{1e-7}}) $$通过标准数据类型FLOAT16BFLOAT16FLOAT32HiFLOAT32FLOAT8 E4M3FLOAT8 E5M2通过阈值(Threshold)2^-102^-72^-132^-112^-32^-2当平均相对误差 MERE Threshold最大相对误差 MARE 10 * Threshold 时判定为通过。5. 标准 Golden 代码import torch UnsortedSegmentSum算子Torch Golden参考实现 沿segment_ids指定的段对数据进行求和 公式: y[i] sum(data[j]) where segment_ids[j] i def unsorted_segment_sum( data: torch.Tensor, segment_ids: torch.Tensor, num_segments: int ) - torch.Tensor: 沿segment_ids指定的段对数据进行求和 公式: y[i] sum(data[j]) where segment_ids[j] i 对于 FP16/BF16 输入使用 FP32 进行内部累加以保证精度 其他类型保持原样 Args: data: 输入数据张量 segment_ids: 段ID张量 num_segments: 段数量 Returns: 输出张量段求和结果 output_shape (num_segments,) data.shape[1:] # FP16/BF16 输入升精度到 FP32 进行累加以保证精度 if data.dtype in (torch.float16, torch.bfloat16): y_fp32 torch.zeros(output_shape, dtypetorch.float32, devicedata.device) data_fp32 data.to(torch.float32) y_fp32.index_add_(0, segment_ids, data_fp32) y y_fp32.to(data.dtype) else: y torch.zeros(output_shape, dtypedata.dtype, devicedata.device) y.index_add_(0, segment_ids, data) return y6. 额外信息算子调用示例import torch import cann_bench data torch.randn(1048576, dtypetorch.float16, devicenpu) segment_ids torch.randint(0, 1024, (1048576,), dtypetorch.int32, devicenpu) y cann_bench.unsorted_segment_sum(data, segment_ids, num_segments1024) # 2D 数据按段求和 data torch.randn(1024, 1024, dtypetorch.float32, devicenpu) segment_ids torch.randint(0, 256, (1024,), dtypetorch.int32, devicenpu) y cann_bench.unsorted_segment_sum(data, segment_ids, num_segments256) # int32 数据类型 data torch.randint(-1000, 1000, (2048, 512), dtypetorch.int32, devicenpu) segment_ids torch.randint(0, 512, (2048,), dtypetorch.int32, devicenpu) y cann_bench.unsorted_segment_sum(data, segment_ids, num_segments512)【免费下载链接】cann-bench评测AI在处理CANN领域代码任务的能力涵盖算子生成、算子优化等领域支撑模型选型、训练效果评估统一量化评估标准识别Agent能力短板构建CANN领域评测平台推动AI能力在CANN领域的持续演进。项目地址: https://gitcode.com/cann/cann-bench创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
