CANN ops-math 算子 aclnnBatchNormStats 接口详解：均值与标准差倒数计算的 L2 API 实战指南-平芜编程栈

算子库
人工智能
CANN

【免费下载链接】ops-math

本项目是CANN提供的数学类基础计算算子库，实现网络在NPU上加速计算。

aclnnBatchNormStats 是 CANN ops-math 仓库中 ReduceStdWithMean 数学算子的 L2 单算子调用接口，用于在单卡上计算输入数据的均值与标准差倒数（1/sqrt(var + eps)），是 BatchNorm 训练/推理统计量计算的关键原语。本文基于 experimental/math/reduce_std_with_mean 模块的接口文档与源码实现，完整讲解两段式接口用法、参数约束、返回码语义、底层 L0 kernel 的 Two-Pass 计算管线以及编译测试方法，读完即可独立编写并运行 aclnnBatchNormStats 的 C++ 调用程序。

功能说明与数学原理

aclnnBatchNormStats 的接口功能为：计算单卡输入数据的均值和标准差的倒数。它是 BatchNorm 前向统计量计算的标准数学表达，输出两个一维张量：

均值（mean）：沿除 Channel 轴（第 2 维）外的所有维度归约求平均：

$$ \bar{x} = \frac{\sum_{i=1}^{n} x_i}{n} $$

标准差倒数（invstd）：方差（E[(x - \bar{x})^2]）加数值稳定常数eps后开根号再取倒数：

$$ \frac{1}{\sigma} = \frac{1}{\sqrt{\frac{1}{n}\sum_{i=1}^{n}(x_i - \bar{x})^2 + eps}} $$

其中eps是加在分母上的数值稳定性常数（接口中由调用方以double类型传入），用于避免方差为 0 时的除零问题。从仓库实现看，aclnnBatchNormStats 与同模块的 aclnnStdMeanCorrection 覆盖了两种不同场景：前者对应 BatchNorm 统计量（均值 + 标准差倒数），后者对应torch.std_mean（标准差 + 均值，支持 Bessel 修正），二者在底层共享同一个 L0 kernel reduce_std_with_mean，仅通过invert参数区分输出是sqrt(var)还是1/sqrt(var + eps)。

产品支持情况

接口文档明确列出了各产品形态的支持情况，当前版本仅 Atlas A2 系列支持：

产品	是否支持
Ascend 950PR/Ascend 950DT	×
Atlas A3 训练系列产品/Atlas A3 推理系列产品	×
Atlas A2 训练系列产品/Atlas A2 推理系列产品	√
Atlas 200I/500 A2 推理产品	×
Atlas 推理系列产品	×
Atlas 训练系列产品	×

需要特别注意的是，在 Atlas A2 系列上，参数input的数据类型不支持 BFLOAT16（见下文参数说明）。而从 L0 算子定义 reduce_std_with_mean_def.cpp 看，底层 kernel 本身注册了 FLOAT16、FLOAT、BF16 三种数据类型与ascend910b计算配置，支持范围取决于调用链组合。

两段式接口与函数原型

aclnnBatchNormStats 采用 CANN 单算子调用标准的两段式接口设计，必须先调用第一段接口aclnnBatchNormStatsGetWorkspaceSize完成入参校验、执行器创建和 workspace 大小计算，再调用第二段接口aclnnBatchNormStats真正执行计算。

函数原型如下（声明见 aclnn_batch_norm_stats_experimental.h）：

aclnnStatus aclnnBatchNormStatsGetWorkspaceSize( const aclTensor* input, double eps, aclTensor* mean, aclTensor* invstd, uint64_t* workspaceSize, aclOpExecutor** executor)

aclnnStatus aclnnBatchNormStats( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, const aclrtStream stream)

典型调用流程为：第一段接口得到workspaceSize与executor→ 用aclrtMalloc在 Device 侧按workspaceSize申请内存 → 调用第二段接口执行 →aclrtSynchronizeStream同步等待完成。

aclnnBatchNormStatsGetWorkspaceSize 参数说明

第一段接口共 6 个参数，其中input、eps为输入，mean、invstd为输出，workspaceSize、executor为返回参数：

参数名	输入/输出	描述	使用说明	数据类型	数据格式	维度(shape)	非连续Tensor
input（aclTensor*）	输入	-	支持的 shape 和数据格式有：2 维（NC）；3 维（NCL）；4 维（NCHW）；5 维（NCDHW）；6-8 维（ND，其中第 2 维固定为 channel 轴）	FLOAT、FLOAT16、BFLOAT16	ND	2-8 维	√
eps（double）	输入	为数值稳定性添加到分母中的值	-	double	-	-	-
mean（aclTensor*）	输出	输出均值	-	FLOAT	ND	1 维	√
invstd（aclTensor*）	输出	输出标准差倒数	-	FLOAT	ND	1 维	√
workspaceSize（uint64_t*）	输出	返回需要在 Device 侧申请的 workspace 大小	-	-	-	-	-
executor（aclOpExecutor**）	输出	返回 op 执行器，包含了算子计算流程	-	-	-	-	-

关于该表格需要强调的几点：

input 的维度语义：维度数限定 2-8 维，且第 2 维（索引 1）固定为 channel 轴，归约发生在除 channel 轴外的所有维度上。从 executor_aclnnBatchNormStats.py 的参考实现可以印证这一语义：dim_array = [x for x in range(length) if x != 1]，即对除第 1 维外的所有轴求均值。
mean 与 invstd 的 shape 约束：两者都必须是 1 维，且长度必须等于 input 的 channel 轴长度。例如输入 shape 为{2, 3}时，mean/invstd 的 shape 为{3}。
Atlas A2 系列约束：input不支持 BFLOAT16 数据类型。
非连续 Tensor：input、mean、invstd 均允许非连续 Tensor 输入（标记为 √），源码中会先经l0op::Contiguous处理。

返回值与错误码

两段接口均返回aclnnStatus状态码。第一段接口会在入参校验阶段报错，具体错误场景如下：

返回值	错误码	描述
ACLNN_ERR_PARAM_NULLPTR	161001	传入的 input、mean 或 invstd 是空指针。
ACLNN_ERR_PARAM_INVALID	161002	input、mean 和 invstd 的数据类型和数据格式不在支持的范围之内。
ACLNN_ERR_PARAM_INVALID	161002	input 维度小于 2 或者大于 8，mean 和 invstd 维度不为 1。
ACLNN_ERR_PARAM_INVALID	161002	mean 或 invstd 的 shape 与 input 的 channel 轴不一致。
ACLNN_ERR_PARAM_INVALID	161002	input 的第二维的值为 0。

这些校验规则与源码 aclnn_batch_norm_stats.cpp 中的CheckParams校验链一一对应：CheckNotNull（空指针 → 161001）、CheckDtypeValid（数据类型合法性 → 161002）、CheckFormat（仅支持 ND/NCL/NCHW/NCDHW 等非私有格式 → 161002）、CheckShape（维度范围、mean/invstd 维度为 1、channel 轴与 shape 一致性、第二维非 0 → 161002）。

aclnnBatchNormStats 执行接口参数说明

第二段接口 4 个参数全部为输入，直接消费第一段接口的产出：

参数名	输入/输出	描述
workspace	输入	在 Device 侧申请的 workspace 内存地址。
workspaceSize	输入	在 Device 侧申请的 workspace 大小，由第一段接口 aclnnBatchNormStatsGetWorkspaceSize 获取。
executor	输入	op 执行器，包含了算子计算流程。
stream	输入	指定执行任务的 Stream。

从实现看，第二段接口内部仅通过CommonOpExecutorRun(workspace, workspaceSize, executor, stream)触发执行器运行，同时被L2_DFX_PHASE_2宏包裹用于 DFX 打点，整体是一个标准的执行封装。

约束说明

确定性计算：aclnnBatchNormStats 默认确定性实现，相同输入恒产生相同输出。从 kernel 代码看，这得益于统一的 Pre-computed Mean Two-Pass 算法路径（见下文），不依赖任何不确定性归约。
仅支持 ND 数据格式。
输入输出数据类型在非 RegBase 平台必须一致（均为 FLOAT16/FLOAT/BF16），mean 与 invstd 输出固定为 FLOAT。

调用示例与逐步解析

接口文档提供了完整可编译的 C++ 调用示例，仓库中的可直接运行版本见 test_aclnn_batch_norm_status.cpp。以下给出完整代码并逐段说明：

#include <iostream> #include <vector> #include "acl/acl.h" #include "aclnnop/aclnn_batch_norm_stats.h" #define CHECK_RET(cond, return_expr) \ do { \ if (!(cond)) { \ return_expr; \ } \ } while (0) #define LOG_PRINT(message, ...) \ do { \ printf(message, ##__VA_ARGS__); \ } while (0) int64_t GetShapeSize(const std::vector<int64_t>& shape) { int64_t shapeSize = 1; for (auto i : shape) { shapeSize *= i; } return shapeSize; } int Init(int32_t deviceId, aclrtStream* stream) { // 固定写法，资源初始化 auto ret = aclInit(nullptr); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret); ret = aclrtSetDevice(deviceId); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret); ret = aclrtCreateStream(stream); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret); return 0; } template <typename T> int CreateAclTensor(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr, aclDataType dataType, aclTensor** tensor) { auto size = GetShapeSize(shape) * sizeof(T); // 调用aclrtMalloc申请Device侧内存 auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret); // 调用aclrtMemcpy将Host侧数据拷贝到Device侧内存上 ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret); // 计算连续tensor的strides std::vector<int64_t> strides(shape.size(), 1); for (int64_t i = shape.size() - 2; i >= 0; i--) { strides[i] = shape[i + 1] * strides[i + 1]; } // 调用aclCreateTensor接口创建aclTensor *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; } int main() { // 1.（固定写法）device/stream初始化，参考acl API手册 // 根据自己的实际device填写deviceId int32_t deviceId = 0; aclrtStream stream; auto ret = Init(deviceId, &stream); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret); // 2.构造输入与输出，需要根据API的接口自定义构造 std::vector<int64_t> inputShape = {2, 3}; std::vector<int64_t> outShape = {3,}; void* inputDeviceAddr = nullptr; void* meanDeviceAddr = nullptr; void* invstdDeviceAddr = nullptr; aclTensor* input = nullptr; aclTensor* mean = nullptr; aclTensor* invstd = nullptr; std::vector<float> inputHostData = {1, 2, 3, 4, 5, 6}; std::vector<float> meanHostData = {0, 0, 0}; std::vector<float> invstdHostData = {0, 0, 0}; double eps = 1e-5; // 创建input aclTensor ret = CreateAclTensor(inputHostData, inputShape, &inputDeviceAddr, aclDataType::ACL_FLOAT, &input); CHECK_RET(ret == ACL_SUCCESS, return ret); // 创建mean aclTensor ret = CreateAclTensor(meanHostData, outShape, &meanDeviceAddr, aclDataType::ACL_FLOAT, &mean); CHECK_RET(ret == ACL_SUCCESS, return ret); // 创建invstd aclTensor ret = CreateAclTensor(invstdHostData, outShape, &invstdDeviceAddr, aclDataType::ACL_FLOAT, &invstd); CHECK_RET(ret == ACL_SUCCESS, return ret); uint64_t workspaceSize = 0; aclOpExecutor* executor; // 3.调用aclnnBatchNormStats第一段接口 ret = aclnnBatchNormStatsGetWorkspaceSize(input, eps, mean, invstd, &workspaceSize, &executor); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnBatchNormStatsGetWorkspaceSize failed. ERROR: %d\n", ret); return ret); // 根据第一段接口计算出的workspaceSize申请device内存 void* workspaceAddr = nullptr; if (workspaceSize > 0) { ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret); } // 调用aclnnBatchNormStats第二段接口 ret = aclnnBatchNormStats(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnBatchNormStats failed. ERROR: %d\n", ret); return ret); // 4.（固定写法）同步等待任务执行结束 ret = aclrtSynchronizeStream(stream); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret); // 5.获取输出的值，将Device侧内存上的结果拷贝至Host侧 auto size = GetShapeSize(outShape); std::vector<float> meanData(size, 0); ret = aclrtMemcpy(meanData.data(), meanData.size() * sizeof(meanData[0]), meanDeviceAddr, size * sizeof(meanData[0]), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return ret); for (int64_t i = 0; i < size; i++) { LOG_PRINT("result[%ld] is: %f\n", i, meanData[i]); } std::vector<float> invstdData(size, 0); ret = aclrtMemcpy(invstdData.data(), invstdData.size() * sizeof(invstdData[0]), invstdDeviceAddr, size * sizeof(invstdData[0]), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return ret); for (int64_t i = 0; i < size; i++) { LOG_PRINT("result[%ld] is: %f\n", i, invstdData[i]); } // 6.释放aclTensor aclDestroyTensor(input); aclDestroyTensor(mean); aclDestroyTensor(invstd); // 7.释放device资源 aclrtFree(inputDeviceAddr); aclrtFree(meanDeviceAddr); aclrtFree(invstdDeviceAddr); if (workspaceSize > 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }

示例程序按 7 个步骤组织，是 aclnn 单算子调用的通用骨架：

资源初始化：aclInit→aclrtSetDevice→aclrtCreateStream，固定写法；
构造输入输出：CreateAclTensor模板函数完成 Device 内存申请（aclrtMalloc）、Host 数据拷贝（aclrtMemcpy）、连续 strides 计算，并通过aclCreateTensor以 ND 格式创建aclTensor。示例输入 shape{2, 3}（NC），输出 shape{3}（channel 长度）；
两段式调用：先aclnnBatchNormStatsGetWorkspaceSize，再按返回的workspaceSize申请 workspace（注意workspaceSize > 0才需要申请），最后aclnnBatchNormStats提交执行；
同步等待：aclrtSynchronizeStream保证任务在 stream 上执行完成；
结果回拷：将 mean 与 invstd 从 Device 拷贝到 Host 并打印。以示例数据{1,2,3,4,5,6}、shape{2,3}计算，channel 0 对应{1,4}，均值(1+4)/2 = 2.5，方差((1-2.5)²+(4-2.5)²)/2 = 2.25，invstd =1/sqrt(2.25+1e-5) ≈ 0.66666，其余 channel 同理； 6-7.资源释放：对称地释放aclTensor、Device 内存、workspace、stream，并aclrtResetDevice、aclFinalize。

源码级原理：L2 API 计算管线与 L0 kernel

L2 API 内部管线

从 aclnn_batch_norm_stats.cpp 的实现看，aclnnBatchNormStats 并非直接调用单一 kernel，而是在执行器中编排了一组 L0 算子构成完整管线：

input ──→ Contiguous ──→ Cast(to FLOAT) ──→ ReFormat(ND) │ ├──→ ReduceMean(除channel外所有轴) ──→ ViewCopy → mean(输出) │ │ │ └──→ UnsqueezeNd ──→ BroadcastTo(还原到input同shape) │ │ └──→ （必要时 Transpose + Contiguous） │ │ │ └──────→ ReduceStdWithMean(self, mean) ──→ ViewCopy → invstd(输出)

关键点：

mean 由内部 ReduceMean 计算：dimIndexNoC收集除索引 1 之外的所有维度作为归约轴，l0op::ReduceMean(reformat, axes, false, ...)计算均值后ViewCopy到meanOut；
mean 需要 broadcast 回原 shape：UnsqueezeNd+BroadcastTo把归约后的均值张量扩展为与 input 同 shape，作为 L0 kernel 的第二个输入；
非连续归约维度的处理：注释明确说明 "Welford kernel requires reduce-dimension elements to be contiguous in memory"，当归约维度不在最内侧时，代码构造 perm 将非归约维移到前面、归约维移到末尾，再Transpose+Contiguous，避免调用方手动处理；
精度保护：输入先Cast到 FLOAT32 再参与后续计算，规避 FLOAT16/BF16 中间累加精度损失；
RegBase 分支：IsRegBase()为真时走aclnnBatchNormStatsImplUnify（直接 ReduceVar + Add eps + Pow(-0.5) 的等价路径），此时数据类型支持列表扩展为包含 BF16（对应ASCEND950_DTYPE_SUPPORT_LIST），这与接口文档"Atlas A2 不支持 BFLOAT16、其他平台受限"的说明互为印证；
空 Tensor 处理：input->IsEmpty()时对 mean 填充 0、对 invstd 填充 NaN（ProcessEmptyTensorWithValue），保证边界场景不崩溃。

L0 kernel 的 Two-Pass 计算

底层 reduce_std_with_mean_kernel.h 实现了统一的 Two-Pass 算法：meanGM提供预计算的 broadcast 均值，kernel 对每个输出元素（coreM_个非归约位置）逐 tile 计算diff = self - mean，再Mul平方、ReduceSum累加，最终得到方差并依据invert输出sqrt(var)或1/sqrt(var + eps)。其内部公式：

$$ \begin{aligned} \text{diff} &= \text{self} - \text{mean} \ \text{sum_sqr} &= \sum (\text{diff})^2 \ \text{var} &= \frac{\text{sum_sqr}}{\max(0,\ N - \text{correction})} \ \text{output} &= \begin{cases} \sqrt{\text{var}}, & \text{invert} = \text{false} \ \dfrac{1}{\sqrt{\text{var} + \text{eps}}}, & \text{invert} = \text{true} \end{cases} \end{aligned} $$

BatchNormStats 场景固定使用invert = true、correction = 0，且 kernel 对非 FLOAT 输入先Cast到 FLOAT32 再执行 Sub/Mul/ReduceSum（树形归约），进一步保证精度。

Tiling 策略

Host 侧 tiling 实现见 reduce_std_with_mean_tiling.cpp，核心策略为：

多核沿非归约维拆分：blockFactor = CeilAlign(CeilDiv(nonReduceNum, coreNum), ubBlockSize)，usedCoreNum = CeilDiv(nonReduceNum, blockFactor)，把非归约元素均匀分到 AIV 核；
UB 沿归约维分块：按self双缓冲 +mean双缓冲 + scratch/work/cast 缓冲区占用估算perElemUB，从 UB 容量扣除保留区后计算ubLength，支持超长归约维的分块累加；
TilingKey 按数据类型分派：FLOAT16/FLOAT/BF16 分别对应REDUCE_STD_SCH_FP16/FP32/BF16三种调度模板。

形状推导

reduce_std_with_mean_infershape.cpp 注册了IMPL_OP_INFERSHAPE：根据dim属性（空列表表示全维度归约）与keepdim属性，调用框架的ReduceDimsWithKeepDims/ReduceDimsWithoutKeepDims推导输出 shape；遇到未知 rank 时直接传播 unknown shape。

编译、测试与运行

模块 README（README.md）给出了完整的编译运行流程，前提是已安装 CANN 环境并source set_env.sh：

# 在仓库根目录构建自定义算子包（--experimental 指定 experimental 目录） source /usr/local/Ascend/cann-8.5.1/set_env.sh bash build.sh --pkg --experimental --soc=ascend910b --ops=reduce_std_with_mean -j16 bash build_out/cann-ops-math-custom_linux-*.run --install-path=/usr/local/Ascend/cann-8.5.1

安装后还需将 op_api/op_tiling/kernel 产物同步到 opp 的 vendors 目录并重新加载环境，随后可执行三类验证：

UT 测试：bash build.sh -u --ophost --opapi --opkernel --ops=reduce_std_with_mean --experimental -j16；
ST 测试（ATK 精度对比）：
```
source /usr/local/Ascend/cann-8.5.1/set_env.sh source /usr/local/Ascend/cann-8.5.1/vendors/custom_math/bin/set_env.bash atk aclnn --task accuracy --devices 0 \ -p experimental/math/reduce_std_with_mean/tests/st/aclnnBatchNormStats/executor_aclnnBatchNormStats.py \ experimental/math/reduce_std_with_mean/tests/st/aclnnBatchNormStats/atk_aclnnBatchNormStats.json
```
其中 ATK 用例 JSON（atk_aclnnBatchNormStats.json）内置了上百组覆盖 2 维到 4 维、channel 从 16 到 1024、batch 从 4 到 1570 的 fp32 用例，参考实现（executor_aclnnBatchNormStats.py）用torch.mean与1.0/torch.sqrt(variance+1e-5)构造 golden 值，直接印证了接口"沿除 channel 外所有维归约、方差加 eps 开根号取倒数"的语义；
示例程序运行：bash build.sh --run_example reduce_std_with_mean eager cust --experimental --vendor_name=custom --soc=ascend910b，即编译运行上文 test_aclnn_batch_norm_status.cpp 示例。

小结

aclnnBatchNormStats 以两段式 aclnn 接口封装了均值与标准差倒数计算：接口层通过 ReduceMean → Broadcast → Transpose → ReduceStdWithMean 的 L0 算子编排完成"预计算均值 + 二次遍历求方差"的 Two-Pass 计算，kernel 层以 FLOAT32 中间精度保证 FLOAT16/BF16 输入的累加精度，tiling 层按非归约维多核并行、归约维 UB 分块。使用者在 Atlas A2 训练/推理产品上，仅需遵守"2-8 维 ND 输入、第 2 维为 channel 轴、mean/invstd 输出为与 channel 等长的一维 FLOAT 张量、Atlas A2 不支持 BFLOAT16"等约束，即可按本文示例骨架快速集成到 BatchNorm 相关训练推理流程中。