ReduceMeanV2 算子深度解析：CANN ops-math 均值规约实现与 aclnn 调用指南

ReduceMeanV2 算子深度解析：CANN ops-math 均值规约实现与 aclnn 调用指南

【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math

导读

本文以 CANN ops-math 开源仓库中的 experimental/math/reduce_mean_v2/README.md 为主体，系统讲解 ReduceMeanV2 算子（对输入张量沿指定维度求平均）的功能语义、参数规格、产品支持范围与 aclnn 接口调用方式，并结合仓库内算子定义、形状推导、tiling 切分与 AscendC kernel 实现源码，深入剖析该算子在 NPU 上的多核归约实现原理。读完本文，你将掌握在 Atlas 训练/推理产品上通过aclnnReduceMeanV2接口完成 2 维张量均值规约的完整流程，并理解其当前的能力边界与演进方向。

功能说明

ReduceMeanV2 的算子功能为：对输入张量 x 沿指定维度计算给定张量的求平均（reduce mean）。它属于规约（reduce）类算子，输出张量的形状由被规约的轴决定：沿某轴求平均后，该轴要么被压缩为长度 1（keepdims=1），要么从输出形状中移除（keepdims=0）。

在项目中，该算子位于 experimental/math/reduce_mean_v2 目录，按 CANN 算子开发的标准五层结构组织：

• op_host/：算子属性定义、形状推导与 tiling 计算（Host 侧）；
• op_kernel/：AscendC 编写的 NPU 侧 kernel 实现；
• examples/：aclnn 接口调用样例；
• tests/：单元测试框架；
• 顶层CMakeLists.txt与各子目录CMakeLists.txt负责构建接入。

从算子注册代码 op_host/reduce_mean_v2_def.cpp 可以看到，该算子通过OP_ADD(ReduceMeanV2)注册为名为ReduceMeanV2的算子，并为ascend910b平台配置了 AICore 执行信息，支持动态 Rank、动态 Shape 与精度规约（PrecisionReduceFlag(true)）。

产品支持情况

与算子注册中this->AICore().AddConfig("ascend910b", aicoreConfig)（见 reduce_mean_v2_def.cpp）相对应，该算子当前面向 Ascend 910B 系列硬件（即 A2 系列芯片）提供支持，其他平台的算子二进制配置尚未在仓库中出现。

参数说明

README 中给出的参数规格如下：

结合源码，可对该表做更精确的说明：

• 输入x：在算子定义中实际命名为x，ParamType(REQUIRED)必选输入，支持DT_FLOAT（float32）与DT_FLOAT16（half）两种数据类型，格式为FORMAT_ND，并设置了AutoContiguous()（见 reduce_mean_v2_def.cpp）。
• 属性axes：规约轴索引，AttrType(OPTIONAL)可选属性、int 类型。形状推导与 tiling 中均以-1作为默认值，表示对所有轴求平均（全量规约，输出为标量）；取0表示沿第 0 维（行方向）规约，取1表示沿第 1 维（列方向）规约。当前实现仅支持单根轴。
• 属性keepdims：是否保留被规约轴，AttrType(OPTIONAL)可选属性、int 类型，默认值为1（保留，规约后该轴长度变为 1）。支持取值0/1。
• 输出y：ParamType(REQUIRED)必选输出，数据类型与输入一致（float32/float16），格式为 ND。

需要说明的是：README 参数表中的 "self" 与 "axis" 分别是 Python 侧接口/文档中的习惯叫法，仓库注册的实际属性名为axes（见 reduce_mean_v2_def.cpp），在使用 aclnn 接口与算子二进制配置文件时应以axes为准。

输出形状推导规则

从 op_host/reduce_mean_v2_infershape.cpp 的实现可以确认输出形状的具体推导逻辑：

• 当keepdims == 1时，输出与输入同维数：
  • axes == -1（默认全规约）：所有维度均置为 1；
  • axes == 0：第 0 维置 1，第 1 维保持输入的列数；
  • 其他情形：axes所指维度置 1，其余维度保持输入尺寸。
• 当keepdims == 0时，输出为一维：
  • axes == 0：输出长度为输入列数；
  • axes == 1：输出长度为输入行数；
  • 其他情形：输出长度为 1。

最终会调用公共的Ops::Base::InferShape4Reduce(context)完成规约类算子的通用形状推导收尾。

约束说明

README 明确标注：约束说明：无，即算子层面未声明额外的输入格式或取值约束。但需要注意，下文"待更新"一节中列出的当前能力限制属于实际的实现约束，使用时应予以遵守。

当前能力限制（待更新）

README 的"待更新"一节明确指出当前版本的实现边界，这也是该算子处于experimental（实验性）目录下的原因：

目前支持 float32、float16 数据类型，仅支持 2 维，axis 仅支持 1 根轴 reduce，keepdims 支持 0/1，后续会持续改进。

即：

1. 数据类型：仅支持 float32 与 float16；
2. 维度：输入张量仅支持 2 维（shape 形如[rows, cols]），这在 kernel 实现中以rows、cols两个字段直接读取存储形状的第 0、1 维可得到印证（见 reduce_mean_v2_tiling.cpp）；
3. 规约轴：axes仅支持单根轴（0、1或默认的-1全量规约），不支持多轴同时规约；
4. keepdims：仅支持0/1两个取值。

尽管算子定义中已开启DynamicRankSupportFlag(true)与DynamicShapeSupportFlag(true)（见 reduce_mean_v2_def.cpp），为后续动态形态扩展预留了能力，但当前 kernel 与 tiling 仍按 2 维实现，使用 3 维及以上输入需等待后续版本。

调用说明：aclnn 接口

README 中给出了该算子的调用方式：通过aclnnReduceMeanV2接口调用 reduce_mean 算子，对应样例为 examples/test_aclnn_reduce_mean_v2.cpp。

aclnn（Ascend Computing Language NN）是 CANN 提供的上层算子调用接口，遵循"先查询 workspace 大小、再申请 workspace、最后执行算子"的三段式异步调用模式。该样例完整展示了这一流程，可概括为以下几个步骤。

步骤一：ACL 环境初始化

auto ret = aclInit(nullptr); ret = aclrtSetDevice(deviceId); ret = aclrtCreateStream(stream);

依次调用aclInit完成 ACL 初始化、aclrtSetDevice绑定设备（样例中deviceId = 0）、aclrtCreateStream创建执行流。所有返回值均需用CHECK_RET宏检查是否为ACL_SUCCESS。

步骤二：构造输入输出 Tensor

std::vector<int64_t> selfXShape = {4, 8}; // 输入形状 4 行 8 列 std::vector<float> selfXHostData(32, 2); // 32 个元素，全部初始化为 2 ret = CreateAclTensor(selfXHostData, selfXShape, &selfXDeviceAddr, ACL_FLOAT, &selfX); std::vector<int64_t> outShape = {1, 8}; // keepdims=1 且 axes=0 时的输出形状 std::vector<float> outHostData(8, 1); ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, ACL_FLOAT, &out);

样例通过模板函数CreateAclTensor完成"申请显存（aclrtMalloc）→ 拷贝 Host 数据到设备（aclrtMemcpy，ACL_MEMCPY_HOST_TO_DEVICE）→ 依据形状计算连续 strides →aclCreateTensor创建aclTensor"的完整流程，创建时明确指定了ACL_FORMAT_ND格式。本例输入为 4×8 的 float32 张量（所有元素为 2），输出预申请为 1×8，对应axes=0、keepdims=1的规约结果形状。

步骤三：查询 workspace 并执行

uint64_t workspaceSize = 0; aclOpExecutor* executor; int32_t axes = 0; int32_t keepdims = 1; ret = aclnnReduceMeanV2GetWorkspaceSize(selfX, axes, keepdims, out, &workspaceSize, &executor); if (workspaceSize > 0) { ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); } ret = aclnnReduceMeanV2(workspaceAddr, workspaceSize, executor, stream); ret = aclrtSynchronizeStream(stream);

aclnnReduceMeanV2GetWorkspaceSize在传入输入selfX、属性axes/keepdims与输出out后，返回所需的 workspace 大小与算子执行器executor；随后按需申请 workspace 显存，调用aclnnReduceMeanV2异步下发算子，最后通过aclrtSynchronizeStream等待执行完成。需要注意：即便workspaceSize为 0，也必须将workspaceAddr传入执行接口（样例中传入的是空指针）。

步骤四：结果回拷与资源释放

PrintOutResult(outShape, &outDeviceAddr); // aclrtMemcpy 将设备结果拷回 Host 并逐元素打印 aclDestroyTensor(selfX); aclDestroyTensor(out); aclrtFree(selfXDeviceAddr); aclrtFree(outDeviceAddr); aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize();

PrintOutResult通过ACL_MEMCPY_DEVICE_TO_HOST将输出拷回 Host 内存后打印每个均值结果。以本样例数据（4×8 全 2 矩阵、axes=0）为例，沿第 0 维求平均后每列均值仍为 2，输出 8 个元素均为 2.0。样例最后依次销毁aclTensor、释放设备内存与流、复位设备并调用aclFinalize收尾。

源码级实现原理

算子注册与属性声明

op_host/reduce_mean_v2_def.cpp 定义了算子的输入/输出/属性契约：输入x与输出y均支持DT_FLOAT/DT_FLOAT16与 ND 格式；属性axes、keepdims均为可选 int。同时通过OpAICoreConfig声明了DynamicCompileStaticFlag(true)、DynamicRankSupportFlag(true)、DynamicShapeSupportFlag(true)等能力标记，为后续动态化演进预留接口。

二进制配置

op_host/config/ascend910b/reduce_mean_v2_binary.json 定义了 float32 与 float16 两套算子二进制条目，输入输出均为 ND 格式、format_match_mode为FormatAgnostic，属性值固定为axes=0、keepdims=1；reduce_mean_v2_simplified_key.ini 则为 opc 工具编译二进制 kernel 时配置--simplified_key_mode（ReduceMeanV2默认 mode 为 0）。

Tiling：多核切分策略

op_host/reduce_mean_v2_tiling.cpp 是 Host 侧的核心性能逻辑，其策略可概括为：

1. 平台信息获取：通过platform_ascendc::PlatformAscendC获取 UB（Unified Buffer）大小与可用核数coreNum（见 GetPlatformInfo）；
2. 输入数据归一：读取输入元素总数、数据类型字节长度，并按 UB 容量与双缓冲（BUFFER_NUM = 2）计算单次搬运的tileDataNum，同时对输入长度做 32B 对齐（inputLengthAlign32，见 GetShapeAttrsInfo）；
3. 核数自适应：若单核 UB 即可容纳全部数据则只用 1 核，否则按"每核至少 32B 数据"的原则把数据均摊到coreNum个核上（见 ReduceMeanV2TilingFunc）；
4. 负载均衡：CalculateCoreBlockNums计算每核的大/小块数据量bigCoreDataNum/smallCoreDataNum、大/小核的 tile 数与尾块数tailBlockNum，tiling 结果写入ReduceMeanV2TilingData（结构定义见 reduce_mean_v2_tiling_data.h）；
5. Workspace 申请：固定申请 16MB 用户 workspace 外加框架系统 workspace（GetLibApiWorkSpaceSize），用于核间原子累加的结果汇总（见 GetWorkspaceSize）。

Kernel：三类规约路径

op_kernel/reduce_mean_v2.cpp 是 kernel 入口，模板实例化NsReduceMeanV2::ReduceMeanV2<DTYPE_X>后调用Init与Process。核心实现位于 op_kernel/reduce_mean_v2.h，根据axes取值分流到三条计算路径：

• ReduceMeanV2Axes0（沿第 0 维）：每个核将本核分到的数据按"行索引 = 全局索引 / 列数"累加进colSum局部缓冲区，随后用SetAtomicAdd+DataCopy把列和写入 workspace 做核间原子累加；SyncAll后由 0 号核读出总列和，乘以1.0f / rows得到均值，float16 场景经Cast(CAST_ROUND)舍入后通过DataCopyPad写回输出（见 reduce_mean_v2.h）；
• ReduceMeanV2Axes1（沿第 1 维）：逻辑对称，按"列索引 = 全局索引 % 列数"累加进rowSum，最终乘以1.0f / cols得到每行均值（见 reduce_mean_v2.h）；
• ReduceMeanV2AxesAll（全量规约，axes 为默认 -1）：各核先累加本核局部和并原子累加到 workspace，0 号核读出全局和除以totalElements后写回单元素输出（见 reduce_mean_v2.h）。

三个分支共同体现了该算子的设计要点：先在核内用标量累加/临时缓冲区完成局部归约，再借助 workspace + 原子加（SetAtomicAdd）与SyncAll完成跨核汇总，最后由 0 号核统一做除法求均值与类型转换，从而在减少核间通信的同时保证多核结果一致。float16 路径始终先将数据Cast到 float 累加，避免半精度累加误差。

总结

ReduceMeanV2 是 CANN ops-math 中面向 Atlas A2 系列产品提供的均值规约算子，当前支持 float32/float16 的 2 维 ND 张量、单轴（axes取 0/1/-1）与keepdims0/1 两种模式。通过 aclnn 接口（aclnnReduceMeanV2GetWorkspaceSize+aclnnReduceMeanV2）即可完成端到端调用；其底层采用"Host 侧 tiling 多核切分 + AscendC kernel 核内归约 + workspace 原子汇总"的实现架构，为理解 CANN 规约类算子的标准实现范式提供了完整参考。该算子仍处于实验阶段，动态 Rank、更多维度与多轴规约能力将随版本持续演进。

【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math