CANN ops-nn 算子解析:aclnnNLLLossBackward 负对数似然损失反向传播两段式接口实战指南
CANN ops-nn 算子解析aclnnNLLLossBackward 负对数似然损失反向传播两段式接口实战指南【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn本篇技术指南围绕 CANN ops-nn 算子库中loss/nll_loss_grad模块对外暴露的aclnnNLLLossBackward接口展开系统讲解负对数似然损失Negative Log Likelihood LossNLLLoss反向传播的数学原理、两段式调用流程、全部参数约束与错误码语义并结合仓库源码剖析其参数校验、算子注册、shape 推导、SIMT tiling 与 AIV 内核实现。读完本文你将能够独立完成该接口的 workspace 申请、执行器获取、设备侧内存管理与梯度回传计算并理解其在 NPU 上的底层执行链路。一、功能说明与数学原理aclnnNLLLossBackward接口实现的是负对数似然损失函数的反向传播即根据正向损失输出的梯度gradOutput计算输入 logitsself对应的梯度out。该算子与 PyTorch 框架中的NLLLossBackward算子语义对齐见 op_graph/nll_loss_grad_proto.h 中 Compatible with pytorch NLLLossGrad operator 的注释。反向传播的核心逻辑是只有真实标签target(t)对应的位置才回传梯度其余位置梯度为 0回传的梯度值由gradOutput、类别权重weight与totalWeight共同决定且依据reduction模式存在三种计算形式。reductionmean均值模式$$ x_grad_{target(t)} \begin{cases} (-gradOutput \cdot w_{(target(t))}) / totalweight , target(t)1 \ 0 , target(t)0 \end{cases} $$reductionsum求和模式$$ x_grad_{target(t)} \begin{cases} -gradOutput \cdot w_{(target(t))} , target(t)1 \ 0 , target(t)0 \end{cases} $$reductionnone不缩减模式$$ x_grad_{target(t)} \begin{cases} -gradOutput_t \cdot w_{(target(t))} , target(t)1 \ 0 , target(t)0 \end{cases} $$从 op_host/nll_loss_grad_def.cpp 的算子定义看底层 GE 算子NLLLossGrad的输入为x、y_grad、target、weight、total_weight输出为x_grad其中reduction为字符串属性默认meanignore_index为整型属性默认-100。aclnnNLLLossBackward正是对底层NLLLossGrad算子的 aclnn 上层封装负责将 int64 形式的reduction枚举值映射为字符串并补齐数据转换、连续性处理等前处理逻辑。二、产品支持情况根据关联文档及 README.mdaclnnNLLLossBackward的软硬件支持矩阵如下产品是否支持Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品不支持Atlas 训练系列产品支持从源码层面印证nll_loss_grad_def.cpp中通过AICore().AddConfig(ascend950, ...)与AICore().AddConfig(ascend350, ...)为两款 AI Core 平台注册了算子配置并开启了DynamicCompileStaticFlag、DynamicRankSupportFlag、DynamicShapeSupportFlag等动态编译与动态 shape 支持能力。三、两段式接口调用模型与 CANN 大部分 aclnn 算子一致aclnnNLLLossBackward采用两段式接口设计必须先调用第一段接口aclnnNLLLossBackwardGetWorkspaceSize完成参数校验、计算流程编排并获取 workspace 大小与算子执行器再调用第二段接口aclnnNLLLossBackward真正下发计算任务。两段式接口的通用说明可参考 两段式接口。3.1 第一段接口原型aclnnStatus aclnnNLLLossBackwardGetWorkspaceSize( const aclTensor *gradOutput, const aclTensor *self, const aclTensor *target, const aclTensor *weight, int64_t reduction, int64_t ignoreIndex, const aclTensor *totalWeight, aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor)3.2 第二段接口原型aclnnStatus aclnnNLLLossBackward( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)第二段接口本质上是框架的通用执行入口从 op_api/aclnn_nll_loss_backward.cpp 可见其实现为return CommonOpExecutorRun(workspace, workspaceSize, executor, stream);即复用 CANN 的统一执行器运行机制完成计算。四、aclnnNLLLossBackwardGetWorkspaceSize 参数详解第一段接口共 10 个参数下表完整列出各参数的输入/输出属性、使用约束、数据类型、格式、shape 与非连续 Tensor 支持情况参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorgradOutputaclTensor*输入输入 aclTensorshape 为 (N)、一维且元素个数为 1或 ()。当 reduction 为 0 且 self 的 shape 为 (N,C) 时shape 支持 (N)当 reduction 为 0 且 self 的 shape 为 (C) 时shape 支持一维且元素个数为 1或 ()当 reduction 不为 0 时shape 支持一维且元素个数为 1或 ()FLOAT、FLOAT16、BFLOAT16ND-√selfaclTensor*输入输入 aclTensor数据类型与出参 out 一致shape 与 out 一致。shape 为 (N,C) 或 (C)N 为 batch sizeC 为类别数当 target 的 shape 为 (N) 时self 需为 (N,C)当 target 的 shape 为 () 时self 需为 (C)FLOAT、FLOAT16、BFLOAT16ND-√targetaclTensor*输入真实标签shape 为 (N) 或 ()每个元素取值范围为 [0, C - 1]INT64、UINT8、INT32ND-√weightaclTensor*输入每个类别的缩放权重公式中的 wshape 为 (C,)与 self 保持一致ND(C)√reductionint64_t输入指定损失函数的计算方式支持 0(none) / 1(mean) / 2(sum)none 不应用缩减mean 输出总和除以元素数sum 输出被求和INT64---ignoreIndexint64_t输入指定一个被忽略且不影响输入梯度的目标值-INT64---totalWeightaclTensor*输入公式中的 totalWeight仅当 reduction 为 mean 时生效通过 target 取相应位置的 weight去除 ignoreIndex 对应的 weight 后求和reduction 为其他值时该参数默认不处理与 weight 相同ND(1,)-outaclTensor*输出公式中的 outshape 与 self 相同与 self 一致ND-√workspaceSizeuint64_t*输出需在 Device 侧申请的 workspace 大小-----executoraclOpExecutor**输出op 执行器包含算子计算流程-----4.1 关键参数语义补充gradOutput 与 reduction 的联动约束从 op_api/aclnn_nll_loss_backward.cpp 的CheckShape实现可见当reduction 0且self为二维 (N,C) 时gradOutput必须是一维且长度等于 N其他场景下reduction 非 0或 self 为一维gradOutput必须是单元素 tensorshape 为 (1,) 或 ()。这一点与文档中的描述完全一致。totalWeight 的语义它仅在 mean 模式下参与除法运算计算公式中x_grad -gradOutput * w / totalWeight。源码校验其 shape size 必须为 1单元素 tensor见CheckShape中OP_CHECK(totalWeight-GetViewShape().GetShapeSize() 1, ...)。ignoreIndex 默认值文档调用示例中使用-100这正是底层算子属性ignore_index的默认值见 op_host/nll_loss_grad_def.cpp 的DEFAULT_IGNORE_IDX -100与 PyTorch 的F.nll_loss默认ignore_index-100保持一致。五、返回值与错误码说明第一段接口返回aclnnStatus状态码具体返回码定义参见 aclnn返回码。第一段接口会完成入参校验以下场景会报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 gradOutput、self、target、weight、out、reduction、totalWeight 为空指针ACLNN_ERR_PARAM_INVALID161002gradOutput、self、target、weight、totalWeight 的数据类型不在支持范围内ACLNN_ERR_PARAM_INVALID161002gradOutput、self、weight、totalWeight 的数据类型不一致ACLNN_ERR_PARAM_INVALID161002gradOutput、self、weight、out、totalWeight 的 shape 不正确ACLNN_ERR_PARAM_INVALID161002reduction 值不在 0~2 范围之内上述校验逻辑在源码中均有对应实现op_api/aclnn_nll_loss_backward.cpp 的CheckParams依次执行空指针检查CheckNotNull、数据类型检查CheckDtypeValid、reduction 范围检查CheckReduction要求0 reduction 2与 shape 检查CheckShape任一环节失败即返回对应的ACLNN_ERR_PARAM_NULLPTR或ACLNN_ERR_PARAM_INVALID。5.1 数据类型支持的平台差异从 op_api/aclnn_nll_loss_backward.cpp 可以看到 dtype 支持列表是按 NPU 架构区分的ASCEND910_DTYPE_SUPPORT_LISTFLOAT、FLOAT16对应 DAV_1001 等老架构ASCEND910B_DTYPE_SUPPORT_LISTFLOAT、FLOAT16、BF16对应 DAV_2201、DAV_3510 等新架构即 Atlas A2/A3 系列与 950 系列。target 的 dtype 统一支持INT64、UINT8、INT32。此外非 regbase 环境下接口内部还会将 FLOAT16 提升为 FLOAT 参与内部计算并将 target 统一转换为 INT32见 op_api/aclnn_nll_loss_backward.cpp。六、aclnnNLLLossBackward 第二段接口参数参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnNLLLossBackwardGetWorkspaceSize 获取executor输入op 执行器包含算子计算流程stream输入指定执行任务的 Stream七、约束说明确定性计算aclnnNLLLossBackward默认采用确定性实现即相同输入在多次执行下结果可复现。八、完整调用示例以下示例代码展示了从资源初始化、tensor 构造、两段式接口调用到结果回拷与资源释放的完整流程。具体编译与执行过程请参考 编译与运行样例完整可运行示例亦可参考 examples/test_aclnn_nll_loss_grad.cpp。#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_nll_loss_backward.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 gradOutputShape {2}; std::vectorint64_t selfShape {2, 3}; std::vectorint64_t targetShape {2}; std::vectorint64_t weightShape {3}; std::vectorint64_t totalWeightShape {1}; std::vectorint64_t outShape {2, 3}; void* gradOutputDeviceAddr nullptr; void* selfDeviceAddr nullptr; void* targetDeviceAddr nullptr; void* weightDeviceAddr nullptr; void* totalWeightDeviceAddr nullptr; void* outDeviceAddr nullptr; aclTensor* gradOutput nullptr; aclTensor* self nullptr; aclTensor* target nullptr; aclTensor* weight nullptr; aclTensor* totalWeight nullptr; aclTensor* out nullptr; std::vectorfloat gradOutputHostData {3.1, 6.5}; std::vectorfloat selfHostData {0, 1, 2, 3, 4, 5}; std::vectorint32_t targetHostData {0, 2}; std::vectorfloat weightHostData {1.1, 1.2, 1.3}; std::vectorfloat totalWeightHostData {0}; std::vectorfloat outHostData(6, 0); int64_t reduction 0; int64_t ignoreIndex -100; // 创建gradOutput aclTensor ret CreateAclTensor(gradOutputHostData, gradOutputShape, gradOutputDeviceAddr, aclDataType::ACL_FLOAT, gradOutput); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建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_INT32, target); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建weight aclTensor ret CreateAclTensor(weightHostData, weightShape, weightDeviceAddr, aclDataType::ACL_FLOAT, weight); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建totalWeight aclTensor ret CreateAclTensor(totalWeightHostData, totalWeightShape, totalWeightDeviceAddr, aclDataType::ACL_FLOAT, totalWeight); 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); // 3. 调用CANN算子库API需要修改为具体的API名称 uint64_t workspaceSize 0; aclOpExecutor* executor; // 调用aclnnNLLLossBackward第一段接口 ret aclnnNLLLossBackwardGetWorkspaceSize(gradOutput, self, target, weight, reduction, ignoreIndex, totalWeight, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnNLLLossBackwardGetWorkspaceSize 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); } // 调用aclnnNLLLossBackward第二段接口 ret aclnnNLLLossBackward(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnNLLLossBackward 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(gradOutput); aclDestroyTensor(self); aclDestroyTensor(target); aclDestroyTensor(weight); aclDestroyTensor(totalWeight); aclDestroyTensor(out); // 7. 释放device资源 aclrtFree(gradOutputDeviceAddr); aclrtFree(selfDeviceAddr); aclrtFree(targetDeviceAddr); aclrtFree(weightDeviceAddr); aclrtFree(totalWeightDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }8.1 示例数据推演以示例数据为例N2, C3reduction0ignoreIndex-100gradOutput {3.1, 6.5}target {0, 2}weight {1.1, 1.2, 1.3}按 none 模式公式x_grad[target(t)] -gradOutput_t * w[target(t)]样本 0 的标签为 0x_grad[0][0] -3.1 * 1.1 -3.41其余位置为 0样本 1 的标签为 2x_grad[1][2] -6.5 * 1.3 -8.45其余位置为 0。由于target元素均在[0, C-1]范围内且不等于ignoreIndex所有位置均正常回传梯度。九、源码级实现剖析9.1 接口内部的计算图编排aclnnNLLLossBackwardGetWorkspaceSize并非直接调用单个算子而是在执行器内动态编排了一组底层 l0 算子见 op_api/aclnn_nll_loss_backward.cpp依次对gradOutput、self、weight、target、totalWeight执行l0op::Contiguous转为连续 Tensor支持非连续输入按平台提升类型执行l0op::CastFLOAT16 提升为 FLOATtarget 转为 INT32非 regbase 场景self为一维 (C) 时通过l0op::UnsqueezeNd扩展为二维 (1,C)使内核统一以 (N,C) 视角处理最后再以l0op::SqueezeNd还原调用l0op::NLLLossGrad完成核心梯度计算reduction通过GetReductionStr映射为none/mean/sum字符串结果经l0op::Cast转回out的目标 dtype再通过l0op::ViewCopy写回可能非连续的输出out通过uniqueExecutor-GetWorkspaceSize()汇总全流程所需 workspace 大小并释放执行器。这也解释了文档中所有输入/输出均支持非连续 Tensor的设计来源——连续性处理统一由接口内部完成。9.2 shape 与 dtype 推导op_host/nll_loss_grad_infershape.cpp 中InferShapeForNLLLossGrad展示了底层算子的 shape 推导规则selfx为一维 (C) 时输出 shape 为 (C)self为二维 (N,C) 时输出 shape 为 (N,C)同时支持四维 (N,C,H,W) 的扩展场景对应 NLLLoss2d 的维度底层算子复用了同一实现dtype 推导InferDtypeForNLLLossGrad仅允许 FLOAT、FLOAT16、BF16输出与输入x保持同 dtype。9.3 内核与 tiling 实现内核入口op_kernel/nll_loss_grad.cpp 注册了nll_loss_grad的 AI Core 内核入口采用KERNEL_TYPE_MIX_AIV_1_0混合任务类型通过SetSysWorkspace配置系统 workspace 后调用KernelNLLLossGrad模板类执行模板参数为输入 dtype 与 target dtype 的组合。Tiling 策略op_host/arch35/nll_loss_grad_tiling_arch35.cpp 实现 SIMT tiling按输入 rank1D/2D/4D解析batchNum、classNum、height、width将reduction字符串映射为NONE_MODE(0)/MEAN_MODE(1)/SUM_MODE(2)枚举写入 tilingData按输出总元素数与 AIV 核数计算blockPerCore/blockTailCore做核间任务切分并为 workspace 预留 16MBASCENDC_TOOLS_WORKSPACE。算子属性注册op_host/nll_loss_grad_def.cpp 中以REQUIRED声明 5 个输入与 1 个输出reduction属性默认meanignore_index属性默认-100与 aclnn 接口的默认语义一一对应。9.4 测试与验证仓库为aclnnNLLLossBackward提供了完整的测试支撑单测覆盖 op_host 层tests/ut/op_host/test_nll_loss_grad_infershape.cpp 验证 shape/dtype 推导tests/ut/op_host/arch35/test_nll_loss_grad_tiling.cpp 验证 tiling 计算ST 用例tests/st/aclnnNLLLossBackward/executor_aclnnNLLLossBackward.py 与配套的atk_aclnnNLLLossBackward.json覆盖接口级端到端行为kernel 级 ST 场景表tests/st/arch35/ttk_kernel_nll_loss_grad_st.csv 列出了内核在各 dtype/shape 组合下的执行场景。十、常见问题与使用建议workspace 申请第一段接口返回的workspaceSize可能为 0例如self为空 Tensor 时见源码中if (self-IsEmpty())分支直接返回workspaceSize 0调用第二段接口前需按if (workspaceSize 0)判断后再申请避免无效分配。reduction 取值范围reduction必须落在 0~2对应none/mean/sum越界会返回ACLNN_ERR_PARAM_INVALID(161002)。totalWeight 的构造mean 模式下它必须是单元素 Tensor且数据类型与weight一致其数值应等于正向NLLLoss中按 target 取权并剔除 ignoreIndex 后求和的结果否则反向梯度会被错误缩放。梯度语义只有target命中位置回传梯度ignoreIndex指定的目标值及其对应位置不参与梯度计算这与 PyTorchNLLLoss的反向行为一致可用于标签平滑、mask 类训练场景。【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考