reduce-all-ascend【免费下载链接】elec-ops-predictionelec-ops-prediction 是 CANN 社区 Electrical Engineering SIG电力行业兴趣小组旗下的电力负荷预测算子库 聚焦于电力系统运行、调度、规划与市场交易中的预测核心需求面向华为昇腾Ascend硬件平台进行深度优化。项目地址: https://gitcode.com/cann/elec-ops-prediction项目简介本项目为基于华为 CANN 计算框架开发的布尔归约算子实现与torch.all功能完全对齐的 ReduceAll 算子。算子支持对任意维度的布尔张量bool/int8沿指定轴执行逻辑与归约∧支持多轴归约与全归约两种模式并通过轴融合、多模式分发、AtomicMin 多核写合并等技术在昇腾 910B NPU 上实现高效计算。主要功能接口对齐完整实现torch.all(input, dim, keepdim)的全部参数语义支持单轴、多轴、全归约三种调用形式可直接替换 PyTorch 原生调用。轴融合优化将输入维度中连续的 Reduce 轴R与 Keep 轴K合并消除冗余的下标计算并根据融合后的维度模式自动选择专用 Kernel 分支。六种归约模式针对融合后常见的轴排列模式FullReduce / KR / RK / RKR / KRK / General提供独立优化实现各分支针对对应的内存访问模式做专项加速。AtomicMin 多核安全写回各核将归约结果写入独立的 WorkGM 缓冲区最终通过SetAtomicMin原子操作安全合并到输出 GM避免多核写冲突。Double Buffer 流水线利用 AscendC 双缓冲机制在计算当前 Tile 的同时后台预取下一 Tile 数据隐藏 HBM 访问延迟提升硬件利用率。数据类型支持支持DT_BOOL与DT_INT8两种输入类型内部统一转换为half参与向量计算。应用场景应用领域典型场景说明模型推理掩码有效性校验批量检查 attention mask 中各序列的 padding 掩码是否全部有效数据预处理批量条件过滤对大规模布尔条件矩阵按行/列归约快速筛选满足所有条件的样本训练监控梯度健康检查检测梯度张量中是否存在 NaN/Inf转布尔后全归约图神经网络边存在性判断对邻接矩阵按节点维度归约判断节点间连通性强化学习终止状态判断对 batch done 标记向量做全归约判断当前 episode 是否全部结束昇腾原生 × 电力行业赋能场景电力负荷预测elec-ops-prediction业务背景 区域电网实际业务中针对特定台区或园区的超短期分钟级负荷预测直接关系到换流站的功率动态调节与储能调度。这类预测需融合传感器实时采集的温湿度、光照、电流等高维特征矩阵。然而边缘采集设备极易受电磁干扰导致时序数据常出现断点或跳变前置的“数据有效性校验”往往会卡住整个 NPU 算力流水线导致推理延迟超标。算子价值提供与torch.all语义对齐的硬件级加速将传统 HBM 访存次数从 6 次降低至 2 次消除时序矩阵在 NPU 上的逻辑归约瓶颈为模型推理提供更高效的数据吞吐支持。对应模型LSTM / GRU 时序网络、Informer / Autoformer 长时序预测大模型、GNN-Transformer 电网空间拓扑模型。应用落地传感器异常的极速熔断应用关口表计分钟级功率动态预测在分钟级预测流水线中将当前时间窗口内多维特征的“是否缺失/异常”状态映射为布尔矩阵利用ReduceAll执行多轴并行归约。一旦判定核心特征全为False立即在 NPU 侧触发短路退出Short-circuit阻止脏数据喂入 LSTM 导致功率预测值剧烈震荡并瞬间调起上层应用的历史均值平滑插值策略。大模型高维特征的掩码过滤应用跨区域气象-负荷联合中长期预测在使用 Informer 等大模型的注意力机制处理长达数周的时序特征时需对大量零值或无效时间步进行有效性校验Padding Mask。利用本算子的轴融合优化Fuse Axes能极速归约剔除无效的冷空气/降雨影响因子序列大幅提升 Transformer 结构对高维时序特征的注意力计算效率。优化说明1. 内存搬运优化传统分步torch.all实现在 NPU 上通常被拆解为多个独立 Kernel 依次调用典型流程如下读取输入布尔张量 → 类型转换为浮点 → 写回中间张量2 次搬运读取中间张量 → 执行 ReduceMin/ReduceSum → 写回归约结果2 次搬运读取归约结果 → 与阈值比较 → 写回最终布尔输出2 次搬运总计约6 次HBM 搬运且存在多个中间张量占据 HBM。本算子融合后的优化路径一次从 HBM 读取输入数据到片上 UB在 UB 内完成 Cast → Abs → ReduceMin 全流程直接将结果写回 WorkGM最后一次 AtomicMin 合并到输出。总计 2 次 HBM 搬运中间结果全部驻留片上。优化点传统方案融合算子HBM 搬运次数6 次2 次中间张量数量2 个浮点中间、归约结果0 个全部片上完成多核写冲突处理需要额外同步 BarrierAtomicMin 原子写无需额外同步2. 计算优化优化技术说明轴融合将连续同类型轴R-R 或 K-K合并为单轴从 N 维退化为最多 3 维大幅降低下标计算开销模式专用分支根据融合后的轴排列KR/RK/RKR/KRK选择专用计算路径避免 General 路径的逐元素坐标映射开销快速跳过优化每个 Tile 先做 ReduceMin 预检若整块 Tile 全为非零全 True则直接跳过无需逐元素扫描短路退出FullReduce 模式下一旦检测到第一个 False 元素立即终止遍历无需扫描剩余数据Tile 自适应分块根据 UB 容量自动计算 tile 面积公式为tileSize AlignDown((ubSize - FIXED_EXPENSES) / UB_BUFFER_FACTOR, BLOCK_SIZE)充分利用片上缓存Double BufferBUFFER_NUM2双缓冲流水DMA 搬运与向量计算并发执行隐藏内存访问延迟核间负载均衡的分配策略如下设总输入元素数为 $N$使用核数为 $C$则$$\text{elementsPerCore} \lfloor N / C \rfloor, \quad \text{largeCoreCount} N \bmod C$$前largeCoreCount个核各处理elementsPerCore 1个元素其余核各处理elementsPerCore个元素实现严格均衡。3. 精度保证特性说明初始化保护WorkGM 与 OutputGM 均初始化为全 1TrueAtomicMin 只能将其改为 0False确保未被覆盖的位置保持正确初始值AtomicMin 原子性多核并发写同一输出位置时通过硬件原子操作保证结果正确无竞争条件SyncAll 同步屏障初始化 OutputGM 后显式调用SyncAll()确保所有核在开始计算前看到一致的初始值Cache 刷新WorkGM 写回前调用DataCacheCleanAndInvalid强制刷出缓存行防止 SetValue 写入遗留在 L2 Cache 中未落地到 GM零维张量处理显式处理含零维输入hasZeroDim的退化情况输出符合 PyTorch 语义参数说明参数名输入/输出描述数据类型数据格式self输入待归约的布尔张量任意维度DT_BOOL / DT_INT8NDdim属性可选指定归约的轴列表为空时执行全归约ListIntint64—keepdim属性可选是否保留归约后的维度默认 falseBool—out输出归约结果张量DT_BOOL / DT_INT8ND约束说明输入张量数据类型必须为DT_BOOL或DT_INT8输出类型与输入保持一致。dim属性为可选不传或传空列表时默认对所有维度执行全归约。dim中的轴值支持负数索引负值会自动归一化为axis rank越界时返回错误。融合后的维度数fusedDimCount不得超过MAX_DIMS编译期常量通常为 8。当前仅支持 Ascend 910B AI Core 配置其他硬件需补充AddConfig。输入张量元素总数为 0含零维时走特殊路径直接输出空张量或标量 True不启动向量计算。架构设计aclnn 层C 两段式接口 │ aclnnReduceAllGetWorkspaceSize() │ aclnnReduceAll() ▼ Host 侧Tilingreduce_all_tiling.cpp │ GetReduceAxesFromAttr() —— 读取 dim 属性归一化轴列表 │ ParseAndValidate() —— 计算 totalInputSize检测零维 │ FuseAxes() —— 连续同类轴融合得到 fusedDims/fusedIsReduce │ ComputeFusedInputStrides() —— 计算融合维的 inputStrides │ ComputeOutputMetaByFused() —— 计算 totalOutputSize 和 outputStrides │ DetermineReduceMode() —— 根据轴模式选择 ReduceMode6 种 │ ComputeMaxTileSize() —— 根据 UB 容量计算最大 Tile 字节数 │ ComputeCoreAllocation() —— 均匀分核计算 elementsPerCore / largeCoreCount ▼ Kernel 侧AscendCreduce_all.h │ Init() │ ├── 绑定 inputGM / outputGM / workGM 全局内存 │ ├── 初始化 outputGM 和 workGM 为全 1TrueSyncAll 同步 │ └── InitBufferinputQueue(×2) / workBuf / tmpBuf / minBuf │ └── Process() ├── MODE_FULL_REDUCE → ProcessAllReduce() 全归约短路退出 ├── MODE_KR → ProcessKR() [Keep][Reduce] 模式 ├── MODE_RK → ProcessRK() [Reduce][Keep] 模式 ├── MODE_RKR → ProcessRKR() [R][K][R] 三段模式 ├── MODE_KRK → ProcessKRK() [K][R][K] 三段模式 └── MODE_GENERAL → ProcessGeneral() 通用坐标映射模式 │ └── CopyWorkGmToOutput() DataCacheCleanAndInvalid → DataCopyPad(workGM→UB) → DataCopy(UB→UB) → SetAtomicMin DataCopyPad(UB→outputGM)层级职责aclnn两段式 C 接口workspace 计算算子注册Host/Tiling轴解析、轴融合、模式选择、核间均匀分配、Workspace 计算Kernel/InitGM 绑定、WorkGM 初始化、UB Buffer 分配Kernel/Process六模式分发、Tile 分块遍历、ReduceMin 快速跳过、AtomicMin 写回六种归约模式详解ReduceAll 算子在 Host 侧完成轴融合后根据融合结果中 R/K 轴的排列自动选择以下六种专用模式之一。MODE_FULL_REDUCE全归约是输入所有维度均为归约轴的情况输出单个标量。该模式采用短路策略遍历过程中一旦发现第一个 False 元素即立刻中止无需扫描剩余数据时间复杂度最优情况为 $O(1)$最差情况为 $O(N)$。MODE_KR[Keep][Reduce]对应融合后的二维布局[outerK][innerR]即先保留维再归约维。每个 K 位置对应innerR个连续元素归约结果写入workGm[k]。该模式对每个 Tile 先做全局 ReduceMin 预检若整块均非零则跳过否则按 K 分段独立归约。MODE_RK[Reduce][Keep]对应融合后的二维布局[outerR][innerK]即先归约维再保留维。每个 K 位置的元素在内存中不连续步长为innerK因此按 K 列逐步扫描对每个 K 列内的 Tile 元素做下标映射后写回workGm[k]。MODE_RKR[R1][K][R2]对应融合后的三维布局[outerR1][middleK][innerR2]输出维度仅为middleK。该模式通过公式k (globalIdx % (K×R2)) / R2直接从全局下标反推 K 坐标同时维护缓存变量lastK避免对同一输出位置的重复 GM 读写。MODE_KRK[K1][R][K2]对应融合后的三维布局[outerK1][middleR][innerK2]输出维度为K1 × K2输出下标为outputIdx k1 × innerK2 k2。通过双缓存变量(lastK1, lastK2)去除对同一输出位置的重复操作。MODE_GENERAL通用模式处理不符合上述任何模式的任意维度组合通过GlobalIdxToOutputIdx()函数按融合维逐级做除余映射将输入全局下标转换为输出下标适用性最广但性能相对较低。该模式还针对末尾为 R 轴的情况做了额外优化检测到 False 后跳过当前 R 段剩余元素减少无效坐标计算。调用说明// C aclnn 接口调用示例 #include aclnn_reduce_all.h // 第一段获取 workspace 大小 uint64_t workspaceSize 0; aclOpExecutor* executor nullptr; aclnnReduceAllGetWorkspaceSize( selfTensor, // 输入布尔张量 dimArray, // 归约轴列表 (aclIntArray*) keepdim, // bool outTensor, // 输出张量 workspaceSize, executor ); // 第二段分配 workspace 并执行 void* workspace nullptr; aclrtMalloc(workspace, workspaceSize, ACL_MEM_MALLOC_NORMAL_ONLY); aclnnReduceAll(workspace, workspaceSize, executor, stream);# Python 等价调用PyTorch 语义对齐 import torch x torch.tensor([[True, True], [True, False]], dtypetorch.bool) # 全归约 result torch.all(x) # tensor(False) # 按列归约 result torch.all(x, dim0) # tensor([True, False]) # 按行归约保留维度 result torch.all(x, dim1, keepdimTrue) # tensor([[True], [False]])调用方式样例入口说明aclnn 接口test_aclnn_reduce_all通过aclnnReduceAll直接调用底层 KernelPyTorch 算子注册torch.ops.npu.reduce_all通过 NPU 算子注册后在 PyTorch 中透明调用测试数据测试用例输入形状归约轴输出形状CPU 耗时 (ms)NPU 耗时 (ms)加速比结果2D 全归约[4, 3000000][0, 1][1]2588.8890.5814452.47x✓2D 全归约[4, 3000000][0, 1][1]2584.7910.16215927.97x✓2D 全归约[4, 3000000][0, 1][1]2585.5970.13618969.90x✓2D 全归约[4, 3000000][0, 1][1]2588.5170.15117157.27x✓2D 单轴归约[4, 3000000][1][4]2433.5050.16115113.06x✓3D 中间轴归约[4, 30000, 5][1][4, 5]121.4442.09058.11x✓3D 末尾轴归约小张量[4, 3, 5][2][4, 3]0.0220.0820.27x✓4D 单轴归约[2, 4, 500, 200][3][2, 4, 500]162.9320.597273.12x✓4D 多轴归约[2, 4, 500, 200][2, 3][2, 5]173.0100.0772253.32x✓4D 单轴归约[2, 4, 5, 2000][1][2, 5, 2000]21.3111.79311.88x✓4D 单轴归约[2, 400, 5, 2][2][2, 400, 2]2.0820.4155.01x✓4D 多轴归约[2, 4, 5, 200][1, 2][2, 200]1.8270.3814.80x✓4D 多轴归约小张量[2, 4, 5, 2][0, 2][4, 2]0.0270.0580.46x✓4D 超大规模多轴归约[2, 4, 8, 1600000][2, 3][2, 4]22064.4070.23792984.98x✓4D 大规模单轴归约[200, 4, 8, 1600][1][200, 8, 1600]2708.9729.559283.40x✓4D 单轴归约小张量[2, 4, 8, 8][1][2, 8, 8]0.1430.1610.89x✓4D 单轴归约小张量[2, 8, 8, 8][1][2, 8, 8]0.2500.1032.43x✓4D 大规模单轴归约[2, 4, 800000, 16][2][2, 4, 16]21137.06087.086242.71x✓4D 大规模多轴归约[2, 4, 8, 160000][1, 2][2, 160000]2330.2859.906235.24x✓4D 多轴归约[2, 4, 800, 16][0, 2][4, 16]22.0422.8437.75x✓8D 多轴归约[2, 4, 8, 3, 8, 28, 18, 26][0, 2, 4, 6][4, 3, 28, 26]4989.25741.511120.19x✓【免费下载链接】elec-ops-predictionelec-ops-prediction 是 CANN 社区 Electrical Engineering SIG电力行业兴趣小组旗下的电力负荷预测算子库 聚焦于电力系统运行、调度、规划与市场交易中的预测核心需求面向华为昇腾Ascend硬件平台进行深度优化。项目地址: https://gitcode.com/cann/elec-ops-prediction创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
CANN电力负荷预测ReduceAll算子
reduce-all-ascend【免费下载链接】elec-ops-predictionelec-ops-prediction 是 CANN 社区 Electrical Engineering SIG电力行业兴趣小组旗下的电力负荷预测算子库 聚焦于电力系统运行、调度、规划与市场交易中的预测核心需求面向华为昇腾Ascend硬件平台进行深度优化。项目地址: https://gitcode.com/cann/elec-ops-prediction项目简介本项目为基于华为 CANN 计算框架开发的布尔归约算子实现与torch.all功能完全对齐的 ReduceAll 算子。算子支持对任意维度的布尔张量bool/int8沿指定轴执行逻辑与归约∧支持多轴归约与全归约两种模式并通过轴融合、多模式分发、AtomicMin 多核写合并等技术在昇腾 910B NPU 上实现高效计算。主要功能接口对齐完整实现torch.all(input, dim, keepdim)的全部参数语义支持单轴、多轴、全归约三种调用形式可直接替换 PyTorch 原生调用。轴融合优化将输入维度中连续的 Reduce 轴R与 Keep 轴K合并消除冗余的下标计算并根据融合后的维度模式自动选择专用 Kernel 分支。六种归约模式针对融合后常见的轴排列模式FullReduce / KR / RK / RKR / KRK / General提供独立优化实现各分支针对对应的内存访问模式做专项加速。AtomicMin 多核安全写回各核将归约结果写入独立的 WorkGM 缓冲区最终通过SetAtomicMin原子操作安全合并到输出 GM避免多核写冲突。Double Buffer 流水线利用 AscendC 双缓冲机制在计算当前 Tile 的同时后台预取下一 Tile 数据隐藏 HBM 访问延迟提升硬件利用率。数据类型支持支持DT_BOOL与DT_INT8两种输入类型内部统一转换为half参与向量计算。应用场景应用领域典型场景说明模型推理掩码有效性校验批量检查 attention mask 中各序列的 padding 掩码是否全部有效数据预处理批量条件过滤对大规模布尔条件矩阵按行/列归约快速筛选满足所有条件的样本训练监控梯度健康检查检测梯度张量中是否存在 NaN/Inf转布尔后全归约图神经网络边存在性判断对邻接矩阵按节点维度归约判断节点间连通性强化学习终止状态判断对 batch done 标记向量做全归约判断当前 episode 是否全部结束昇腾原生 × 电力行业赋能场景电力负荷预测elec-ops-prediction业务背景 区域电网实际业务中针对特定台区或园区的超短期分钟级负荷预测直接关系到换流站的功率动态调节与储能调度。这类预测需融合传感器实时采集的温湿度、光照、电流等高维特征矩阵。然而边缘采集设备极易受电磁干扰导致时序数据常出现断点或跳变前置的“数据有效性校验”往往会卡住整个 NPU 算力流水线导致推理延迟超标。算子价值提供与torch.all语义对齐的硬件级加速将传统 HBM 访存次数从 6 次降低至 2 次消除时序矩阵在 NPU 上的逻辑归约瓶颈为模型推理提供更高效的数据吞吐支持。对应模型LSTM / GRU 时序网络、Informer / Autoformer 长时序预测大模型、GNN-Transformer 电网空间拓扑模型。应用落地传感器异常的极速熔断应用关口表计分钟级功率动态预测在分钟级预测流水线中将当前时间窗口内多维特征的“是否缺失/异常”状态映射为布尔矩阵利用ReduceAll执行多轴并行归约。一旦判定核心特征全为False立即在 NPU 侧触发短路退出Short-circuit阻止脏数据喂入 LSTM 导致功率预测值剧烈震荡并瞬间调起上层应用的历史均值平滑插值策略。大模型高维特征的掩码过滤应用跨区域气象-负荷联合中长期预测在使用 Informer 等大模型的注意力机制处理长达数周的时序特征时需对大量零值或无效时间步进行有效性校验Padding Mask。利用本算子的轴融合优化Fuse Axes能极速归约剔除无效的冷空气/降雨影响因子序列大幅提升 Transformer 结构对高维时序特征的注意力计算效率。优化说明1. 内存搬运优化传统分步torch.all实现在 NPU 上通常被拆解为多个独立 Kernel 依次调用典型流程如下读取输入布尔张量 → 类型转换为浮点 → 写回中间张量2 次搬运读取中间张量 → 执行 ReduceMin/ReduceSum → 写回归约结果2 次搬运读取归约结果 → 与阈值比较 → 写回最终布尔输出2 次搬运总计约6 次HBM 搬运且存在多个中间张量占据 HBM。本算子融合后的优化路径一次从 HBM 读取输入数据到片上 UB在 UB 内完成 Cast → Abs → ReduceMin 全流程直接将结果写回 WorkGM最后一次 AtomicMin 合并到输出。总计 2 次 HBM 搬运中间结果全部驻留片上。优化点传统方案融合算子HBM 搬运次数6 次2 次中间张量数量2 个浮点中间、归约结果0 个全部片上完成多核写冲突处理需要额外同步 BarrierAtomicMin 原子写无需额外同步2. 计算优化优化技术说明轴融合将连续同类型轴R-R 或 K-K合并为单轴从 N 维退化为最多 3 维大幅降低下标计算开销模式专用分支根据融合后的轴排列KR/RK/RKR/KRK选择专用计算路径避免 General 路径的逐元素坐标映射开销快速跳过优化每个 Tile 先做 ReduceMin 预检若整块 Tile 全为非零全 True则直接跳过无需逐元素扫描短路退出FullReduce 模式下一旦检测到第一个 False 元素立即终止遍历无需扫描剩余数据Tile 自适应分块根据 UB 容量自动计算 tile 面积公式为tileSize AlignDown((ubSize - FIXED_EXPENSES) / UB_BUFFER_FACTOR, BLOCK_SIZE)充分利用片上缓存Double BufferBUFFER_NUM2双缓冲流水DMA 搬运与向量计算并发执行隐藏内存访问延迟核间负载均衡的分配策略如下设总输入元素数为 $N$使用核数为 $C$则$$\text{elementsPerCore} \lfloor N / C \rfloor, \quad \text{largeCoreCount} N \bmod C$$前largeCoreCount个核各处理elementsPerCore 1个元素其余核各处理elementsPerCore个元素实现严格均衡。3. 精度保证特性说明初始化保护WorkGM 与 OutputGM 均初始化为全 1TrueAtomicMin 只能将其改为 0False确保未被覆盖的位置保持正确初始值AtomicMin 原子性多核并发写同一输出位置时通过硬件原子操作保证结果正确无竞争条件SyncAll 同步屏障初始化 OutputGM 后显式调用SyncAll()确保所有核在开始计算前看到一致的初始值Cache 刷新WorkGM 写回前调用DataCacheCleanAndInvalid强制刷出缓存行防止 SetValue 写入遗留在 L2 Cache 中未落地到 GM零维张量处理显式处理含零维输入hasZeroDim的退化情况输出符合 PyTorch 语义参数说明参数名输入/输出描述数据类型数据格式self输入待归约的布尔张量任意维度DT_BOOL / DT_INT8NDdim属性可选指定归约的轴列表为空时执行全归约ListIntint64—keepdim属性可选是否保留归约后的维度默认 falseBool—out输出归约结果张量DT_BOOL / DT_INT8ND约束说明输入张量数据类型必须为DT_BOOL或DT_INT8输出类型与输入保持一致。dim属性为可选不传或传空列表时默认对所有维度执行全归约。dim中的轴值支持负数索引负值会自动归一化为axis rank越界时返回错误。融合后的维度数fusedDimCount不得超过MAX_DIMS编译期常量通常为 8。当前仅支持 Ascend 910B AI Core 配置其他硬件需补充AddConfig。输入张量元素总数为 0含零维时走特殊路径直接输出空张量或标量 True不启动向量计算。架构设计aclnn 层C 两段式接口 │ aclnnReduceAllGetWorkspaceSize() │ aclnnReduceAll() ▼ Host 侧Tilingreduce_all_tiling.cpp │ GetReduceAxesFromAttr() —— 读取 dim 属性归一化轴列表 │ ParseAndValidate() —— 计算 totalInputSize检测零维 │ FuseAxes() —— 连续同类轴融合得到 fusedDims/fusedIsReduce │ ComputeFusedInputStrides() —— 计算融合维的 inputStrides │ ComputeOutputMetaByFused() —— 计算 totalOutputSize 和 outputStrides │ DetermineReduceMode() —— 根据轴模式选择 ReduceMode6 种 │ ComputeMaxTileSize() —— 根据 UB 容量计算最大 Tile 字节数 │ ComputeCoreAllocation() —— 均匀分核计算 elementsPerCore / largeCoreCount ▼ Kernel 侧AscendCreduce_all.h │ Init() │ ├── 绑定 inputGM / outputGM / workGM 全局内存 │ ├── 初始化 outputGM 和 workGM 为全 1TrueSyncAll 同步 │ └── InitBufferinputQueue(×2) / workBuf / tmpBuf / minBuf │ └── Process() ├── MODE_FULL_REDUCE → ProcessAllReduce() 全归约短路退出 ├── MODE_KR → ProcessKR() [Keep][Reduce] 模式 ├── MODE_RK → ProcessRK() [Reduce][Keep] 模式 ├── MODE_RKR → ProcessRKR() [R][K][R] 三段模式 ├── MODE_KRK → ProcessKRK() [K][R][K] 三段模式 └── MODE_GENERAL → ProcessGeneral() 通用坐标映射模式 │ └── CopyWorkGmToOutput() DataCacheCleanAndInvalid → DataCopyPad(workGM→UB) → DataCopy(UB→UB) → SetAtomicMin DataCopyPad(UB→outputGM)层级职责aclnn两段式 C 接口workspace 计算算子注册Host/Tiling轴解析、轴融合、模式选择、核间均匀分配、Workspace 计算Kernel/InitGM 绑定、WorkGM 初始化、UB Buffer 分配Kernel/Process六模式分发、Tile 分块遍历、ReduceMin 快速跳过、AtomicMin 写回六种归约模式详解ReduceAll 算子在 Host 侧完成轴融合后根据融合结果中 R/K 轴的排列自动选择以下六种专用模式之一。MODE_FULL_REDUCE全归约是输入所有维度均为归约轴的情况输出单个标量。该模式采用短路策略遍历过程中一旦发现第一个 False 元素即立刻中止无需扫描剩余数据时间复杂度最优情况为 $O(1)$最差情况为 $O(N)$。MODE_KR[Keep][Reduce]对应融合后的二维布局[outerK][innerR]即先保留维再归约维。每个 K 位置对应innerR个连续元素归约结果写入workGm[k]。该模式对每个 Tile 先做全局 ReduceMin 预检若整块均非零则跳过否则按 K 分段独立归约。MODE_RK[Reduce][Keep]对应融合后的二维布局[outerR][innerK]即先归约维再保留维。每个 K 位置的元素在内存中不连续步长为innerK因此按 K 列逐步扫描对每个 K 列内的 Tile 元素做下标映射后写回workGm[k]。MODE_RKR[R1][K][R2]对应融合后的三维布局[outerR1][middleK][innerR2]输出维度仅为middleK。该模式通过公式k (globalIdx % (K×R2)) / R2直接从全局下标反推 K 坐标同时维护缓存变量lastK避免对同一输出位置的重复 GM 读写。MODE_KRK[K1][R][K2]对应融合后的三维布局[outerK1][middleR][innerK2]输出维度为K1 × K2输出下标为outputIdx k1 × innerK2 k2。通过双缓存变量(lastK1, lastK2)去除对同一输出位置的重复操作。MODE_GENERAL通用模式处理不符合上述任何模式的任意维度组合通过GlobalIdxToOutputIdx()函数按融合维逐级做除余映射将输入全局下标转换为输出下标适用性最广但性能相对较低。该模式还针对末尾为 R 轴的情况做了额外优化检测到 False 后跳过当前 R 段剩余元素减少无效坐标计算。调用说明// C aclnn 接口调用示例 #include aclnn_reduce_all.h // 第一段获取 workspace 大小 uint64_t workspaceSize 0; aclOpExecutor* executor nullptr; aclnnReduceAllGetWorkspaceSize( selfTensor, // 输入布尔张量 dimArray, // 归约轴列表 (aclIntArray*) keepdim, // bool outTensor, // 输出张量 workspaceSize, executor ); // 第二段分配 workspace 并执行 void* workspace nullptr; aclrtMalloc(workspace, workspaceSize, ACL_MEM_MALLOC_NORMAL_ONLY); aclnnReduceAll(workspace, workspaceSize, executor, stream);# Python 等价调用PyTorch 语义对齐 import torch x torch.tensor([[True, True], [True, False]], dtypetorch.bool) # 全归约 result torch.all(x) # tensor(False) # 按列归约 result torch.all(x, dim0) # tensor([True, False]) # 按行归约保留维度 result torch.all(x, dim1, keepdimTrue) # tensor([[True], [False]])调用方式样例入口说明aclnn 接口test_aclnn_reduce_all通过aclnnReduceAll直接调用底层 KernelPyTorch 算子注册torch.ops.npu.reduce_all通过 NPU 算子注册后在 PyTorch 中透明调用测试数据测试用例输入形状归约轴输出形状CPU 耗时 (ms)NPU 耗时 (ms)加速比结果2D 全归约[4, 3000000][0, 1][1]2588.8890.5814452.47x✓2D 全归约[4, 3000000][0, 1][1]2584.7910.16215927.97x✓2D 全归约[4, 3000000][0, 1][1]2585.5970.13618969.90x✓2D 全归约[4, 3000000][0, 1][1]2588.5170.15117157.27x✓2D 单轴归约[4, 3000000][1][4]2433.5050.16115113.06x✓3D 中间轴归约[4, 30000, 5][1][4, 5]121.4442.09058.11x✓3D 末尾轴归约小张量[4, 3, 5][2][4, 3]0.0220.0820.27x✓4D 单轴归约[2, 4, 500, 200][3][2, 4, 500]162.9320.597273.12x✓4D 多轴归约[2, 4, 500, 200][2, 3][2, 5]173.0100.0772253.32x✓4D 单轴归约[2, 4, 5, 2000][1][2, 5, 2000]21.3111.79311.88x✓4D 单轴归约[2, 400, 5, 2][2][2, 400, 2]2.0820.4155.01x✓4D 多轴归约[2, 4, 5, 200][1, 2][2, 200]1.8270.3814.80x✓4D 多轴归约小张量[2, 4, 5, 2][0, 2][4, 2]0.0270.0580.46x✓4D 超大规模多轴归约[2, 4, 8, 1600000][2, 3][2, 4]22064.4070.23792984.98x✓4D 大规模单轴归约[200, 4, 8, 1600][1][200, 8, 1600]2708.9729.559283.40x✓4D 单轴归约小张量[2, 4, 8, 8][1][2, 8, 8]0.1430.1610.89x✓4D 单轴归约小张量[2, 8, 8, 8][1][2, 8, 8]0.2500.1032.43x✓4D 大规模单轴归约[2, 4, 800000, 16][2][2, 4, 16]21137.06087.086242.71x✓4D 大规模多轴归约[2, 4, 8, 160000][1, 2][2, 160000]2330.2859.906235.24x✓4D 多轴归约[2, 4, 800, 16][0, 2][4, 16]22.0422.8437.75x✓8D 多轴归约[2, 4, 8, 3, 8, 28, 18, 26][0, 2, 4, 6][4, 3, 28, 26]4989.25741.511120.19x✓【免费下载链接】elec-ops-predictionelec-ops-prediction 是 CANN 社区 Electrical Engineering SIG电力行业兴趣小组旗下的电力负荷预测算子库 聚焦于电力系统运行、调度、规划与市场交易中的预测核心需求面向华为昇腾Ascend硬件平台进行深度优化。项目地址: https://gitcode.com/cann/elec-ops-prediction创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
