CANN ops-nn aclnnMseLossOut 算子详解NPU 上均方误差损失的两段式接口开发指南【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn本篇技术指南以 CANN 开源神经网络算子库 ops-nn 中loss/mse_loss模块的官方算子文档 aclnnMseLossOut.md 为核心系统讲解 aclnnMseLossOut 算子的数学定义、产品支持情况、两段式 API 原型、全部参数约束、返回码语义与调用示例并结合仓库内 op_api、op_host、op_kernel 源码与 ATK 测试用例深入剖析其在 NPU 上的真实实现与调用链。读完本文你将掌握在 Ascend 平台上通过 aclnn 接口完成 MSE Loss 计算含 none/mean/sum 三种缩减模式的完整开发与验证方法。产品支持情况aclnnMseLossOut 在 CANN 当前版本中的产品支持矩阵如下来源于算子文档产品支持情况Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品支持Atlas 训练系列产品支持从 算子定义文件 可以看到MseLoss算子为ascend950配置了 AiCore 算子配置DynamicCompileStaticFlag、DynamicRankSupportFlag、DynamicShapeSupportFlag均开启与文档中Ascend 950 支持的说明一致同时支持动态 rank 与动态 shape。功能说明与计算公式接口功能aclnnMseLossOut 用于计算输入xself与目标ytarget逐元素之间的均方误差Mean Squared Error。reduction参数指定应用到输出的缩减方式支持none、mean、sum三种none不做缩减逐元素输出平方误差mean将输出沿0 轴求均值sum将输出沿0 轴求和。计算公式当reduction为none时$$ \ell(x, y) L {l_1,\dots,l_N}^\top, \quad l_n \left( x_n - y_n \right)^2, $$其中 $x$ 是 self$y$ 是 target$N$ 是 batch 的大小。当reduction不是none时$$ \ell(x, y) \begin{cases} \operatorname{mean}(L), \text{if reduction} \text{mean;}\ \operatorname{sum}(L), \text{if reduction} \text{sum.} \end{cases} $$需要特别说明的是与torch.nn.MSELoss的全元素标量归约不同aclnnMseLossOut 在 mean/sum 模式下仅对 0 轴做缩减。这一点在 ATK 测试实现 的注释中有明确说明aclnnMseLossOut固定以reduction0none调用底层 l0 接口计算逐元素 MSE再根据用户传入的 reduction 对结果沿 dim 0 执行sum(dim0)或mean(dim0)而aclnnMseLoss则直接将 reduction 透传给底层算子。因此当 reduction 为 1 或 2 时aclnnMseLoss输出 0 维标量而aclnnMseLossOut仅压缩掉 0 轴。函数原型两段式接口aclnnMseLossOut 属于 CANN 算子的两段式接口必须先调用aclnnMseLossOutGetWorkspaceSize获取计算所需 workspace 大小以及包含算子计算流程的执行器executor再调用aclnnMseLossOut执行计算。aclnnStatus aclnnMseLossOutGetWorkspaceSize( const aclTensor* self, const aclTensor* target, int64_t reduction, aclTensor* out, uint64_t* workspaceSize, aclOpExecutor** executor)aclnnStatus aclnnMseLossOut( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)两个接口的原型声明可在 aclnn_mse_loss_out.h 中查看其 domain 标记为aclnn_ops_train即用于训练场景的 aclnn 算子接口。aclnnMseLossOutGetWorkspaceSize 参数说明第一段接口完成入参校验与计算流程构建各参数含义如下参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续 TensorselfaclTensor*输入公式中的输入 xshape 需要与 target 满足 broadcast 关系数据类型与 target 满足数据类型推导规则参见互转换关系FLOAT、FLOAT16、BFLOAT16ND1-8√targetaclTensor*输入公式中的输入 yshape 需要与 self 满足 broadcast 关系数据类型与 self 满足数据类型推导规则参见互转换关系FLOAT、FLOAT16、BFLOAT16ND1-8√reductionint64_t输入指定要应用到输出的缩减公式中的参数 reduction支持 0(none)|1(mean)|2(sum)。none 表示不应用缩减mean 表示输出的总和将除以输出中的元素数sum 表示输出将被求和INT64--√outaclTensor*输出公式中的输出 ℓ(x,y)数据类型需要是 self 与 target 推导之后可转换的数据类型参见互转换关系当 reduction 的值为 0 时out 的 shape 与 self 和 target broadcast 后的 shape 一致当 reduction 的值不为 0 时out 与 self、target 做 broadcast 后再 reduce 0 轴得到的 tensor 的 shape 一致FLOAT、FLOAT16、BFLOAT16-1-8√workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程-----其中reduction的三种取值在 aclnn_mse_loss_out.cpp 中被映射为常量REDUCTION_NONE_NUM 0、REDUCTION_MEAN_NUM 1、REDUCTION_SUM_NUM 2并进一步映射为字符串none/mean/sum供底层算子消费。返回值与错误码第一段接口返回aclnnStatus状态码具体参见 aclnn 返回码。第一段接口完成入参校验出现以下场景时返回对应错误返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 self、target 或 out 是空指针时ACLNN_ERR_PARAM_INVALID161002self、target 或 out 的数据类型不在支持的范围之内ACLNN_ERR_PARAM_INVALID161002self 和 target 的 shape 无法做 broadcastACLNN_ERR_PARAM_INVALID161002self 或 target 的 shape 超过 8 维ACLNN_ERR_PARAM_INVALID161002reduction 值不在 0~2 范围之内ACLNN_ERR_PARAM_INVALID161002当 reduction 的值为 0 时self 和 target 做 broadcast 后的 shape 与 out 的 shape 不一致ACLNN_ERR_PARAM_INVALID161002当 reduction 的值为 1 或 2 时out 的维度大于 0上述校验逻辑可以在 aclnn_mse_loss_out.cpp 的 CheckParams 函数 中逐一对应空指针检查、数据类型范围检查CheckDtypeValidMseLoss、reduction 范围检查CheckReductionMseLoss与输出 shape 检查CheckShape。其中CheckShape会先对 self、target 做 broadcast 推导第 57-73 行reduction 为 0 时要求 out 与 broadcast 后 shape 完全一致否则要求 out 与broadcast 结果去掉 0 轴后的 shape 一致且维度必须为 0 维即被压缩为空。数据类型支持的分平台差异从源码看第一段接口按芯片平台区分数据类型支持范围aclnn_mse_loss_out.cpp 第 50-55 行Ascend 910ASCEND910_DTYPE_SUPPORT_LISTFLOAT、FLOAT16Ascend 910B 等ASCEND910B_DTYPE_SUPPORT_LISTFLOAT、FLOAT16、BF16。这与文档数据类型 FLOAT、FLOAT16、BFLOAT16的总体描述对应也说明 BF16 支持与具体产品型号相关。此外第一段接口会对 self、target 做隐式类型提升op::PromoteType将两个输入提升为同一精度后再计算这正是测试中 (fp16, fp32)、(bf16, fp32) 等混合精度组合得以支持的底层机制。aclnnMseLossOut 参数说明第二段接口执行实际计算参数如下参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnMseLossOutGetWorkspaceSize 获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream返回值同样为aclnnStatus具体参见 aclnn 返回码。其实现非常简洁aclnn_mse_loss_out.cpp 第 194-199 行直接调用框架的CommonOpExecutorRun完成计算下发。约束说明确定性计算aclnnMseLossOut 默认确定性实现即相同输入在同一硬件上多次运行结果一致。关于 CANN 确定性计算的更多背景可参见 determinism_compute.md。源码实现解析算子内部调用链在第一段接口内部aclnnMseLossOut 的实际计算流程aclnn_mse_loss_out.cpp 第 119-191 行依次为空 tensor 短路若 self 或 target 为空 tensor直接返回 workspaceSize 0 的成功结果连续化与类型提升分别对 self、target 调用l0op::Contiguous处理非连续 tensor与l0op::Cast提升到 promoteTypeBroadcastTo当两个输入 shape 不同时通过BroadcastInferShape推导广播后 shape并对两个输入分别执行l0op::BroadcastToMseLoss 核心计算固定以reductionnone调用l0op::MseLoss得到逐元素平方误差按需缩减reduction 为 1 时对 dim 0 执行l0op::ReduceMean为 2 时对 dim 0 执行l0op::ReduceSumOp否则直接透传类型转换与写回l0op::Cast将结果转为 out 的数据类型最后l0op::ViewCopy将结果拷贝到 outout 可能为非连续 tensor通过GetWorkspaceSize汇总整个计算图所需 workspace并将 executor 移交给调用方。在底层算子分派上mse_loss.cppMseLoss会根据当前 SoC 版本与输入条件自动选择实现当运行在 ASCEND910B / ASCEND910_93 / ASCEND310P / ASCEND950 且两个输入 shape、formatND 或 NCL相同时走MSELossV2Launcher 为MSELossV2AiCore否则走通用路径MseLossV1。对应的 NPU kernel 实现在 op_kernel/mse_loss.cpp通过模板参数Reduction区分三种模式Reduction 0走逐元素计算ElementwiseSchReduction 1走ReduceSch求和Reduction 2走ReduceSch求均值以NAN初始化累加器。而算子图侧的 shape 推导在 mse_loss_infershape.cpp 中定义reduction none时输出 shape 等于输入 shape否则输出为 0 维标量这与 aclnn 层仅 reduce 0 轴的 API 语义通过外层 l0 组合算子共同实现。调用示例以下完整示例来源于 算子文档仓库中另有可直接编译的样例 test_aclnn_mse_loss_out.cpp。示例计算self [[0,1],[2,3]]与target [[1,1],[1,1]]的 MSE Lossreduction 1mean输出 shape 为{2}对 0 轴求均值。具体编译与执行过程请参考编译与运行样例。#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_mse_loss_out.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::vectorint64_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::vectorT hostData, const std::vectorint64_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::vectorint64_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::vectorint64_t selfShape {2, 2}; std::vectorint64_t targetShape {2, 2}; std::vectorint64_t outShape {2}; void* selfDeviceAddr nullptr; void* targetDeviceAddr nullptr; void* outDeviceAddr nullptr; aclTensor* self nullptr; aclTensor* target nullptr; aclTensor* out nullptr; std::vectorfloat selfHostData {0, 1, 2, 3}; std::vectorfloat targetHostData {1, 1, 1, 1}; std::vectorfloat outHostData(2, 0); // 创建self aclTensor ret CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_FLOAT, self); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建target aclTensor ret CreateAclTensor(targetHostData, targetShape, targetDeviceAddr, aclDataType::ACL_FLOAT, target); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建out aclTensor ret CreateAclTensor(outHostData, outShape, outDeviceAddr, aclDataType::ACL_FLOAT, out); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建reduction int64_t reduction 1; // 3. 调用CANN算子库API需要修改为具体的API名称 uint64_t workspaceSize 0; aclOpExecutor* executor; // 调用aclnnMseLossOut第一段接口 ret aclnnMseLossOutGetWorkspaceSize(self, target, reduction, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnMseLossOutGetWorkspaceSize 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); } // 调用aclnnMseLossOut第二段接口 ret aclnnMseLossOut(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnMseLossOut 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侧需要根据具体API的接口定义修改 auto size GetShapeSize(outShape); std::vectorfloat resultData(size, 0); ret aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), outDeviceAddr, size * sizeof(resultData[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, resultData[i]); } // 6. 释放aclTensor和aclScalar需要根据具体API的接口定义修改 aclDestroyTensor(self); aclDestroyTensor(target); aclDestroyTensor(out); // 7. 释放device资源需要根据具体API的接口定义修改 aclrtFree(selfDeviceAddr); aclrtFree(targetDeviceAddr); aclrtFree(outDeviceAddr); if(workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例结果推演以上述数据为例逐元素平方误差为(0-1)²1、(1-1)²0、(2-1)²1、(3-1)²4即[[1,0],[1,4]]reduction1mean对 0 轴求均值后得到[1, 2]与示例中outShape {2}的设计吻合。示例关键点说明两段式调用顺序不可颠倒必须先调GetWorkspaceSize拿到workspaceSize与executor再申请 device 侧 workspace 内存并调aclnnMseLossOutworkspace 按需申请当workspaceSize 0时无需申请 workspace 内存对应空 tensor 等场景资源释放成对出现aclrtMalloc与aclrtFree、aclCreateTensor与aclDestroyTensor、aclrtCreateStream与aclrtDestroyStream必须一一对应并在aclFinalize前完成非连续 tensor 支持接口允许 self/target/out 为非连续 tensor内部通过Contiguous与ViewCopy自动处理见上文源码解析。测试与验证与 torch 对齐仓库为 aclnnMseLossOut 提供了完善的 ATKAscend Test Kit算子级测试黄金数据生成executor_aclnnMseLossOut.py使用torch.nn.MSELoss(reductionnone)计算逐元素损失再根据 reduction 分别执行sum(dim0)/mean(dim0)与 aclnn 接口仅 reduce 0 轴的语义完全对齐该测试同样覆盖了混合精度fp16/fp32/bf16 任意组合提升到 fp32的行为一致性。用例矩阵atk_aclnnMseLossOut.json包含 201 条测试用例覆盖 fp16/fp32/bf16 三种 dtype 及其组合、reduction 取 0/1/2 三种模式、1~8 维 shape、broadcast 场景如[17,7]与[1,1]、以及131073等大元素规模的分片场景精度标准为vector_fused。UT 用例test_aclnn_mse_loss_out.cpp在 op_api 单测层验证接口调用流程与返回码。开发者可以在本地编写类似上述示例的 C 程序或复用 examples/arch35 目录下的样例在支持的产品上编译运行并将输出与 torch 原生torch.nn.MSELoss的结果对比以验证算子行为。总结aclnnMseLossOut 是 CANN ops-nn 中实现均方误差损失的标准 aclnn 接口采用两段式调用模型支持 none/mean/sum 三种缩减方式、broadcast 输入、混合精度提升与确定性计算并覆盖 Ascend 950、A3、A2、Atlas 推理/训练系列等多款产品。通过本文对参数语义、错误码、源码调用链与测试用例的系统梳理你可以直接在 NPU 上接入该算子完成训练或推理中的 MSE Loss 计算。同目录下的 aclnnMseLoss 文档提供了与之对应的透传 reduction版本接口可按需对照使用。【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考