CANN ops-nn NllLoss 算子全面解析负对数似然损失的计算原理、aclnn 两段式调用与 NPU 实现【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn本篇技术指南以 CANN 神经网络算子库 ops-nn 中的 NllLoss负对数似然损失算子为对象完整讲解其数学定义、输入输出参数与约束、aclnnNLLLoss / aclnnNLLLoss2d 两段式接口的原型与调用示例并结合 op_host、op_kernel、op_api 等目录下的源码剖析该算子在 Atlas A2 / A3 系列 NPU 上的 Shape 推导、Tiling 切分与 AICore/AICPU 双路执行原理。读完本文你将能够在自己的训练工程中正确构造 NllLoss 输入、编写可运行的 aclnn 调用代码并理解该算子在仓库中的实现脉络。一、算子概述NllLoss 是什么NllLossNegative Log Likelihood Loss负对数似然损失是深度学习中多分类任务的经典损失函数在仓库中的典型使用方式是配合LogSoftmax一起构成与CrossEntropyLoss等价的损失计算链。该算子在 CANN ops-nn 仓库中的位置为 experimental/loss/nll_loss其产品支持情况如下产品是否支持Atlas A3 训练系列 / Atlas A3 推理系列√Atlas A2 训练系列 / Atlas A2 推理系列√从 nll_loss_def.cpp 的算子注册信息可以看出算子对 AICore 提供了ascend910b与ascend910_93两套硬件配置分别对应 A2Ascend 910B与 A3 系列芯片同时由于算子还实现了 AICPU 兜底路径见下文“双路执行”一节因此在更广泛的平台上也能运行。二、数学定义与计算公式对每个样本i设其目标类别为target_i。当target_i等于ignore_index时该样本不参与计算否则单个样本的损失为$$ loss_i -weight[target_i] \times x[i, target_i] $$其中x的第 2 维最后一维是类别数Cx[i, c]表示样本i属于类别c的对数得分weight是类别权重向量。当未提供weight时各类别权重按1处理。按照reduction属性的取值损失会以三种方式归约none逐样本输出y_i loss_isum对所有样本求和y \sum_i loss_imean对所有样本的损失求和后除以参与样本的权重之和$$ y \frac{\sum_i loss_i}{\sum_i weight[target_i]} $$同时算子还会输出参与计算样本的权重之和$$ total_weight \sum_i weight[target_i] $$这一公式在 aclnnNLLLoss 文档中以 PyTorch 风格给出了更严格的表述w_c weight[c] · 1{c ! ignoreIndex}即被忽略类别对应的权重项视为 0从而保证mean归约的除数为实际参与样本的权重总和与 PyTorch 的F.nll_loss语义一致。三、参数说明3.1 算子级参数IR 定义下表来自 README.md 的算子参数说明是使用该算子时必须遵守的输入输出契约参数名输入/输出/属性描述数据类型数据格式x输入公式中的输入 x最后一维为类别数 C其余维度展平为样本数 NBFLOAT16、FLOAT16、FLOATNDtarget输入公式中的输入 target每个样本的目标类别索引元素个数为样本数 NINT32、INT64NDweight可选输入公式中的输入 weight各类别的权重长度为 C未传入时权重按 1 处理BFLOAT16、FLOAT16、FLOATNDreduction可选属性指定损失函数的计算方式支持 none | mean | sum。none 表示不应用归约mean 表示损失的加权平均sum 表示损失求和。默认为 meanSTRING-ignore_index可选属性指定被忽略且不参与损失计算的目标类别值。默认为 -100INT64-y输出公式中的输出 y。reduction 为 none 时形状与 target 一致否则为标量BFLOAT16、FLOAT16、FLOATNDtotal_weight输出公式中的输出 total_weight参与计算样本的权重之和BFLOAT16、FLOAT16、FLOATND对应地在 nll_loss_def.cpp 中可以看到算子 IR 的定义细节x、weight、输出y、total_weight支持DT_FLOAT16 / DT_FLOAT / DT_BF16格式全部为FORMAT_NDtarget支持DT_INT32 / DT_INT64weight的ParamType为OPTIONAL可选输入属性reduction为可选字符串属性默认值mean属性ignore_index为可选整型属性默认值-100与 PyTorch 的默认值一致所有输入输出均声明了AutoContiguous()由框架保证参与计算时内存连续。3.2 约束说明使用该算子需满足以下约束见 README.mdtarget的取值需落在[0, C)区间内或等于ignore_indexx的最后一维为类别数Ctarget的元素个数为样本数N若传入weight其长度需与类别数C保持一致。此外从 nll_loss_tiling.cpp 的运行时校验还可以看到 Tiling 阶段的两条硬性检查x的维度数必须 ≥ 1类别数classNumx 最后一维与样本数rowNumtarget 元素个数均不能为 0target的数据类型必须是 INT32 或 INT64。四、aclnn 两段式接口函数原型与调用流程CANN 算子库的每个算子都采用两段式接口设计详见 两段式接口说明先调用xxxGetWorkspaceSize获取计算所需的 workspace 大小及包含算子计算流程的执行器再调用xxx真正执行计算。4.1 aclnnNLLLoss一维/通用输入函数原型见 aclnnNLLLoss.mdaclnnStatus aclnnNLLLossGetWorkspaceSize( const aclTensor *self, // 输入 xshape 为 (N,C) 或 (C) const aclTensor *target, // 真实标签shape 为 (N) 或 () const aclTensor *weight, // 每类权重shape 为 (C)可传空 int64_t reduction, // 0none, 1mean, 2sum int64_t ignoreIndex, // 被忽略的目标值默认 -100 aclTensor *out, // 输出 y aclTensor *totalWeightOut, // 输出 total_weight uint64_t *workspaceSize, aclOpExecutor **executor); aclnnStatus aclnnNLLLoss( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream);接口各参数要点详见 aclnnNLLLoss.mdself数据类型支持 FLOAT、FLOAT16、BFLOAT16shape 为(N, C)或(C)单样本时 target 为标量 shape()支持非连续 Tensor。target支持 INT64、UINT8、INT32self为(N,C)时target为(N)每个元素取值范围[0, C-1]支持非连续 Tensor。weight与self数据类型一致shape 为(C)支持非连续 Tensor。reductionint64_t0(none) | 1(mean) | 2(sum)其中mean表示输出总和除以输出元素数。ignoreIndexint64_t被忽略且不影响输入梯度的目标值。outreduction为 0none且self为 2 维时 shape 为(N,)否则为(1,)。totalWeightOut在reduction非 0非 none时输出值有效shape 为(1,)。workspaceSize / executor分别返回 Device 侧所需 workspace 大小与算子执行器。第一段接口的入参校验与返回码与 aclnn 返回码 对应返回值错误码触发场景ACLNN_ERR_PARAM_NULLPTR161001self、target、weight、out、totalWeightOut 为空指针ACLNN_ERR_PARAM_INVALID161002self、target、weight 数据类型不在支持范围self 与 weight 数据类型不一致各 Tensor 的 shape 不正确reduction 不在 0~2 范围4.2 aclnnNLLLoss2d4 维输入NCDHW 语义对于 shape 为 4 维的输入例如图像/序列场景第 2 维为类别数C仓库提供了aclnnNLLLoss2d接口原型与aclnnNLLLoss一致见 aclnnNLLLoss2d.md区别在于 shape 契约self4 维第 2 维是类别数C支持 FLOAT、FLOAT16、BFLOAT16支持非连续 Tensortarget3 维其第 1/2/3 维分别与self的第 1/3/4 维相等元素取值范围[0, C-1]支持 INT64、UINT8、INT32weightshape(C,)数据类型与self一致outreduction为 0none时 shape 与target相同否则为(1,)totalWeightOut非 none 归约下有效shape(1,)。该接口的入参校验返回码与aclnnNLLLoss基本一致161001 空指针、161002 类型/一致性/shape/reduction 非法只是校验描述针对 2d 场景做了适配见 aclnnNLLLoss2d.md。五、完整调用示例test_aclnn_nll_loss仓库在 examples/test_aclnn_nll_loss.cpp 中提供了可直接运行的 aclnn 调用样例同时在 docs/aclnnNLLLoss.md 中给出了带详细注释的等价版本。以下结合示例代码讲解关键步骤具体编译与执行过程参考 编译与运行样例。5.1 初始化 Device 与 Streamint Init(int32_t deviceId, aclrtStream* stream) { auto ret aclInit(nullptr); // 初始化 ACL 运行环境 ret aclrtSetDevice(deviceId); // 指定 device ret aclrtCreateStream(stream); // 创建执行流 return 0; }5.2 构造 aclTensor示例通过模板函数CreateAclTensor完成申请 Device 内存 → Host 数据拷贝到 Device → 计算连续 strides → 调用aclCreateTensor创建aclTensor的完整流程其中strides按行主序连续排布计算。随后在main中构造一个N2、C3的典型输入std::vectorint64_t selfShape {2, 3}; // x: 2 个样本 × 3 个类别 std::vectorint64_t targetShape {2}; // target: 每个样本的目标类别 std::vectorint64_t weightShape {3}; // weight: 每类权重 std::vectorint64_t outShape {1}; // 归约后输出为标量 std::vectorint64_t totalWeightShape {1}; // 权重之和 std::vectorfloat selfHostData {0.1f, 0.2f, 0.7f, 0.5f, 0.3f, 0.2f}; std::vectorint64_t targetHostData {2, 0}; // 样本0目标类2样本1目标类0 std::vectorfloat weightHostData {1.0f, 1.0f, 1.0f}; int64_t reduction 1; // mean int64_t ignoreIndex -100;注意reduction在 aclnn 接口中是一个int64_t数值0/1/2而非字符串ignoreIndex传-100与算子属性默认值一致。5.3 两段式调用与结果回拷uint64_t workspaceSize 0; aclOpExecutor* executor nullptr; // 第一段获取 workspace 大小与执行器 ret aclnnNLLLossGetWorkspaceSize(self, target, weight, reduction, ignoreIndex, out, totalWeightOut, workspaceSize, executor); // 按需申请 workspace void* workspaceAddr nullptr; if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); } // 第二段执行计算 ret aclnnNLLLoss(workspaceAddr, workspaceSize, executor, stream); // 同步等待执行结束 ret aclrtSynchronizeStream(stream); // 将 out 从 Device 拷回 Host 并打印 ret aclrtMemcpy(resultData.data(), ..., outDeviceAddr, ..., ACL_MEMCPY_DEVICE_TO_HOST);最后按顺序释放aclDestroyTensor销毁各aclTensoraclrtFree释放 Device 内存含 workspace再aclrtDestroyStream、aclrtResetDevice、aclFinalize完成收尾。5.4 2d 版本调用差异aclnnNLLLoss2d.md 中的示例代码结构与 1d 版几乎一致仅输入 shape 不同selfShape {1, 2, 3, 2}4 维第 2 维C2、targetShape {1, 3, 2}、weightShape {2}、outShape {1, 3, 2}头文件改为aclnnop/aclnn_nll_loss2d.h。六、源码级实现剖析从调用到 NPU 计算6.1 op_api 层的双路执行AICore 优先、AICPU 兜底op_api/nll_loss.cpp 是 aclnn 接口的 L0 层实现核心逻辑在NLLLoss函数中通过executor-AllocTensor为输出out与totalWeightOut分配与x同数据类型的 ND 张量调用INFER_SHAPE(NllLoss, ...)完成形状推导依据芯片架构选择执行路径AICore 路径NLLLossAiCore通过ADD_TO_LAUNCHER_LIST_AICORE把算子下发到 Vector/Cube 核执行AICPU 路径NLLLossAiCpu通过ADD_TO_LAUNCHER_LIST_AICPU交由 AICPU 执行属性以{reduction, ignore_index}命名传入。路径选择由IsAiCoreSupport决定而各架构支持的 AICore 数据类型从源码结构看存在差异芯片架构AICore 支持的数据类型Ascend 910DAV_1001FLOAT、FLOAT16Ascend 910BDAV_2201FLOAT、BF16、FLOAT16Ascend 910_93 / 91095DAV_3510FLOAT、BF16、FLOAT16当x的数据类型不在对应架构的 AICore 支持列表例如某些架构上的 BF16时自动回退到 AICPU 路径保证算子可用性。6.2 op_host 层Shape 推导nll_loss_infershape.cpp 实现形状与数据类型推导输出y的形状reduction none时直接复制target的 shape否则设为 1 维且dim(0) 1标量total_weight恒为 1 维、长度为 1输出数据类型与输入x保持一致InferDataType直接透传x的 dtype。这与第 3 节参数表以及接口文档中none 时输出 shape 与 target 一致否则为标量的约定完全对应。6.3 op_host 层Tiling 切分策略nll_loss_tiling.cpp 是 NPU 高性能实现的关键它在 Host 侧把计算任务切分到多个 AIV 核Tiling Key 选择SelectNllLossTilingKey按x的数据类型映射模板调度模式——FLOAT16→NLLLOSS_TPL_SCH_MODE_0FLOAT→MODE_1BF16→MODE_2属性解析ParseNllLossAttrs把字符串reduction翻译为整数none→0、mean→1默认、sum→2同时读取ignore_index默认 -100核数切分ComputeNllLossSplit以每核约 16KB 工作量WORK_PER_CORE 16 * 1024为粒度估算usedCoreNum再向上取整得到rowsPerCore每个核处理的样本行数当单核行数 ≥ 128 时启用向量化useVector1UB 空间分块以 140KB UB 预算为上限扣除weight驻留开销后计算每个 Tile 可容纳的行数tileRows保证 Kernel 按 Tile 循环搬数、计算避免 UB 溢出workspace 规划currentWorkspace[0] 系统库 workspace 32 × usedCoreNum同步 2 × 32 × usedCoreNum归约为多核间归约预留内存最终通过context-SetBlockDim(usedCoreNum)设置核数、context-SetTilingKey(tilingKey)绑定 Kernel 模板。切分结果写入NllLossTilingData结构体见 nll_loss_tiling_data.h其中reduction、ignoreIndex等字段的默认值与算子 IR 定义保持一致reduction1、ignoreIndex-100。6.4 op_kernel 层AICore Kernelnll_loss.cpp 是核函数入口通过REGISTER_TILING_DEFAULT读取 Tiling 数据再按schMode模板参数实例化NsNllLoss::RunhalfMODE_0或NsNllLoss::RunfloatMODE_1执行。核函数实现nll_loss.h中包含半精度/BF16 位级转换工具HalfBitsToFloat、Bf16BitsToFloat等说明 Kernel 内部统一按 FLOAT 精度进行损失计算后再写回对应类型输出索引钳制ClampIdx将越界 target 索引安全映射到合法范围BLK_ELEM8、VEC_ALIGN64、BLK_BYTES32等常量体现与 NPU 向量单元32B 块对齐相关的访存优化。6.5 单元测试佐证仓库为算子提供了完整的 UT 覆盖Host 侧 test_nll_loss_tiling.cpp 验证不同 dtype 下的 Tiling Key 选择FLOAT 对应NLLLOSS_TPL_SCH_MODE_1、FLOAT16 对应NLLLOSS_TPL_SCH_MODE_0并校验 Tiling 数据的正确性test_nll_loss_infershape.cpp 验证 Shape 推导逻辑Kernel 侧 test_nll_loss.cpp 配合 gen_data.py 与 compare_data.py 生成输入数据并与期望结果比对从算子级到 Tiling 级形成了完整验证闭环。七、确定性计算说明根据接口文档的约束说明aclnnNLLLoss.md、aclnnNLLLoss2d.mdaclnnNLLLoss与aclnnNLLLoss2d默认均为非确定性实现多核归约顺序可能影响浮点结果支持通过aclrtCtxSetSysParamOpt开启确定性计算。对于需要严格可复现结果的训练/调试场景建议按需开启。八、总结与使用建议NllLoss 算子是 CANN ops-nn 仓库中一个结构完整、实现路径清晰的典型 loss 算子具备以下特征语义与 PyTorch 对齐默认reductionmean、ignore_index-100支持none/sum/mean三种归约权重缺失时按 1 处理双接口覆盖aclnnNLLLoss面向(N,C)/(C)输入aclnnNLLLoss2d面向 4 维输入两者共享两段式调用模型AICore/AICPU 双路实现依据芯片架构与数据类型自动选择执行路径兼顾性能与兼容性工程化完整Shape 推导、Tiling 切分、workspace 规划与多级 UT 一应俱全。实际使用时建议将x构造为LogSoftmax的输出值域为负的对数概率target使用 INT32/INT64 类别索引并确保取值范围在[0, C)内或等于ignore_index若传入weight需保证长度为C且与x同数据类型需要精确复现结果时开启确定性计算模式。【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考