ReduceAll 算子设计文档【免费下载链接】elec-ops-predictionelec-ops-prediction 是 CANN 社区 Electrical Engineering SIG电力行业兴趣小组旗下的电力负荷预测算子库 聚焦于电力系统运行、调度、规划与市场交易中的预测核心需求面向华为昇腾Ascend硬件平台进行深度优化。项目地址: https://gitcode.com/cann/elec-ops-prediction1. 目标将 TBE 实现的 ReduceAll 算子用 Ascend C 重新实现并提交到算子开源仓。支持 int8 (bool)性能不低于 TBE 算子的 95%TBE kernel 实现/usr/local/Ascend/ascend-toolkit/latest/opp/built-in/op_impl/ai_core/tbe/impl/dynamic/TBE 算子原型定义/usr/local/Ascend/ascend-toolkit/latest/opp/built-in/op_proto/inc/TBE 算子信息库910B 配置/usr/local/Ascend/ascend-toolkit/latest/opp/built-in/op_impl/ai_core/tbe/config/ascend910b2. 算子概述ReduceAll 算子沿指定轴对布尔张量执行逻辑与AND归约操作。当指定轴上所有元素均为 True 时输出 True否则输出 False。设输入张量为 $self$归约轴为 $axes$则输出 $out$ 的计算公式为$$out \bigwedge_{axes} self$$其中 $\bigwedge$ 表示逻辑与操作。3. TBE 的实现根据算子信息库TBE 算子支持的数据类型和数据格式如下参数名称数据类型数据格式参数类型input0xboolNDrequiredinput1axesint32, int64NDrequiredoutput0yboolNDrequiredattrkeep_dimsbool-optional (默认值: false)TBE 的 reduce_all 的计算流程参数预处理处理 keep_dims 默认值检查输入数据类型bool 转为 int8检查 axes 数据类型int32 或 int64设置关系位置标记记录原始 shape规约后的 shape需要规约的轴调用 classify 进行输入分类获取不同的输入组合对每个输入组合创建占位符调用 reduce_all_compute 计算生成调度并构建在 reduce_all_compute 计算中先判断输入张量如果存在零维度zero_tensor_flg 0 in x.shape走快速路径否则采用等价的数值计算方式开始计算将 int8 (bool) 转换为 float16调用 tbe.cast_to(x, float16)调用 tbe.vabs(data_fp16)确保 True1 和 False0 都为非负值沿指定轴调用 tbe.reduce_min(data_abs, axisaxes, keepdimskeepdims)求该轴的最小值调用 tbe.cast_to(res_any, dtype)将结果转回 int84. Ascend C 算子设计4.1 算子原型Ascend C 实现的算子原型设计如下与 TBE 算子原型保持一致参数类型名称数据类型格式说明输入selfint8 (bool)ND布尔张量输入axesint32, int64ND归约轴输出outint8 (bool)ND归约结果属性keepdimbool-是否保留归约维度默认 false4.2 算法设计4.2.1 计算路径当输入张量 shape 中包含 0 时直接输出 True对空集合的与操作返回 True否则走数值模拟计算路径核心计算步骤为CopyIn根据坐标计算将数据从 Global Memory 搬运到 Local MemoryCast将 int8 数据转换为 float16Abs对 float16 数据取绝对值确保 True 的任意非零表示都转为正数ReduceMin求最小值Cast将 float16 结果转换回 int8CopyOut将结果从 Local Memory 写回 Global Memory4.2.2 核心算法聚合操作将原来的多个元素归约为一个元素例如 100 个元素每 10 个求和后得到 10 个元素shape 和轴的数量都变了。核心要素是新旧坐标的转换使用坐标和步长 strides 来计算偏移量。由于聚合操作需要将元素的步长分为两段base_offset和relative_offset。因为聚合意味着一个输出元素对应一组输入元素组长就是total_reduce_elements// 计算给定形状的连续内存步长 std::vectorsize_t compute_strides(const std::vectorint shape) { if (shape.empty()) return {}; std::vectorsize_t strides(shape.size()); size_t stride 1; for (int i static_castint(shape.size()) - 1; i 0; --i) { strides[i] stride; stride * shape[i]; } return strides; } // 多维坐标转一维偏移根据多维坐标和步长计算一维内存偏移量 size_t coords_to_offset(const std::vectorint coords, const std::vectorsize_t strides) { size_t offset 0; for (size_t i 0; i coords.size(); i) { offset coords[i] * strides[i]; } return offset; } // 一维索引转多维坐标将线性索引转换为多维坐标 std::vectorint index_to_coords(size_t index, const std::vectorint shape) { int ndim static_castint(shape.size()); std::vectorint coords(ndim); for (int d ndim - 1; d 0; --d) { coords[d] index % shape[d]; index / shape[d]; } return coords; }准备阶段分离保留轴和聚合轴的 shape 和 strides// 计算原始数据的完整步长 std::vectorsize_t in_strides compute_strides(in_shape); // vector 转 set 方便后面 in 判断 std::setint reduce_set(reduce_axes.begin(), reduce_axes.end()); std::vectorint kept_shape; // 保留轴的 shape std::vectorsize_t kept_strides; // 保留轴的 strides std::vectorint reduced_shape; // 聚合轴的 shape std::vectorsize_t reduced_strides; // 聚合轴的 strides // 分离保留轴和聚合轴 for (int i 0; i ndim; i) { if (reduce_set.count(i)) { reduced_shape.push_back(in_shape[i]); reduced_strides.push_back(in_strides[i]); } else { kept_shape.push_back(in_shape[i]); kept_strides.push_back(in_strides[i]); } } // 计算所有输出元素的个数 size_t total_output_elements 1; for (int s : kept_shape) total_output_elements * s; // 计算每次聚合的元素的个数 size_t total_reduce_elements 1; for (int s : reduced_shape) total_reduce_elements * s;核心逻辑// 外层循环逐个遍历每个输出的元素这就是一次聚合 for (size_t i 0; i total_output_elements; i) { // 计算第 i 个输出元素的保留轴坐标 (index_to_coords) std::vectorint kept_coords index_to_coords(i, kept_shape); // 计算 base_offset (coords_to_offset) // 聚合出第 i 个输出元素的那些原始元素是一个不连续组该组的起始地址记为 base_offset size_t base_offset coords_to_offset(kept_coords, kept_strides); float min_val FLT_MAX; // 内层循环遍历本次聚合的每个元素 for (size_t j 0; j total_reduce_elements; j) { // 计算聚合轴坐标 (index_to_coords) std::vectorint reduced_coords index_to_coords(j, reduced_shape); // 计算 relative_offset (coords_to_offset) size_t relative_offset coords_to_offset(reduced_coords, reduced_strides); // 使用 base_offset relative_offset 访问原始数据 float val abs(input_data[base_offset relative_offset]); if (val min_val) min_val val; } output_data[i] (int8_t)min_val; }4.3 Tiling 策略4.3.1 多核划分按照输出元素数量进行多核划分total_output_elements 所有保留轴维度的乘积均匀分配:coreOutputNumtotal_output_elements/coreNum尾部处理: 余数分配给前面几个核心4.3.2 Tile 大小计算UB 空间分配输入 buffer:tileSize× sizeof(int8) × 2 (Ping-Pong)计算 buffer:tileSize× sizeof(float16) × 2 (数据 reduce 工作区)输出 buffer:tileBatchNum× sizeof(int8) × 2 (Ping-Pong)计算 tileCapacitybytesPerElement sizeof(int8) × 2 sizeof(float16) × 2 6 bytes tileCapacity availableUBSize / bytesPerElement4.3.3 统一计算模式采用统一的 Batch 计算逻辑正常情况 (total_reduce_elementstileCapacity)tileBatchNumtileCapacity/total_reduce_elementstileSizetileBatchNum×total_reduce_elements一个 tile 处理多个完整聚合大 reduce 情况 (total_reduce_elementstileCapacity)tileBatchNum 1tileSizetileCapacity多个 tile 处理一个聚合维护跨 tile 的partialMin状态4.4 UB 布局设计输入 buffer使用 TQue 开启 double buffer每块存放一个 tile 的 int8 输入数据输出 buffer使用 TQue 开启 double buffer每块存放一个 tile 的 int8 输出数据计算 buffer不参与流水用于向量计算与归约操作5. TilingData 结构struct ReduceAllTilingData { // 基础形状参数 int64_t totalOutputElements; // 输出元素总数 int64_t reduceElements; // 每次聚合的元素数 // 多核划分参数 int64_t coreOutputNum; // 每核处理的输出数 int64_t tailCoreNum; // 余数核数 // Tile 分块参数 int64_t tileSize; // 每个 tile 的元素数 int64_t tileNum; // tile 数量 int64_t tailTileSize; // 尾 tile 大小 int64_t tileBatchNum; // 每个 tile 处理的完整聚合数 // 保留轴参数用于计算 base_offset int32_t keptAxisNum; // 保留轴数量 int64_t keptShape[MAX_DIM_NUM]; // 保留轴形状 int64_t keptStrides[MAX_DIM_NUM]; // 保留轴步长 // 聚合轴参数用于计算 relative_offset int32_t reducedAxisNum; // 聚合轴数量 int64_t reducedShape[MAX_DIM_NUM]; // 聚合轴形状 int64_t reducedStrides[MAX_DIM_NUM]; // 聚合轴步长 };6. 算子约束限制数据类型约束输入输出仅支持 int8 (bool) 类型axes 类型约束归约轴索引支持 int32 和 int64 类型维度约束支持最大 8D 的输入张量归约轴约束支持任意轴组合的归约包括不连续轴空张量约束支持包含零维度的张量输入返回 True【免费下载链接】elec-ops-predictionelec-ops-prediction 是 CANN 社区 Electrical Engineering SIG电力行业兴趣小组旗下的电力负荷预测算子库 聚焦于电力系统运行、调度、规划与市场交易中的预测核心需求面向华为昇腾Ascend硬件平台进行深度优化。项目地址: https://gitcode.com/cann/elec-ops-prediction创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考