CANN ops-nn 算子解读：SigmoidCrossEntropyWithLogits 的数学原理、接口约束与 NPU 实现-平芜编程栈

CANN ops-nn 算子解读：SigmoidCrossEntropyWithLogits 的数学原理、接口约束与 NPU 实现

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

SigmoidCrossEntropyWithLogits 是 CANN ops-nn 神经网络算子库（loss/sigmoid_cross_entropy_with_logits）中用于二分类场景的逐元素损失算子，它直接在 NPU 上计算 logits 与标签之间的 Sigmoid Cross Entropy 损失，与 TensorFlow 同名算子语义兼容。本文以该算子官方文档为主体，结合仓库中的算子原型、Shape 推导、Tiling 与 Kernel 源码，系统讲解其数学公式、参数与约束、调用方式，以及从图构图到 AIV 向量核执行的全链路实现原理，帮助读者在 CANN 环境下正确使用与二次开发该算子。

一、算子功能与数学原理

1.1 功能定位

SigmoidCrossEntropyWithLogits 算子接收两个输入：预测值predict（logits）与标签值target，输出逐元素计算的 Sigmoid Cross Entropy 损失loss。它常用于多标签二分类、多任务分类等场景：对每个 logit 独立应用 Sigmoid 激活，再计算与对应标签的交叉熵，最后对整批结果做归约即可得到训练损失。

从算子原型注册文件 op_graph/sigmoid_cross_entropy_with_logits_proto.h 的注释可以确认其框架兼容性：

Compatible with TensorFlow operator SigmoidCrossEntropyWithLogits.

也就是说，在 CANN 图中可直接承接 TensorFlow 模型导出图中的同名算子，实现无损迁移。

1.2 计算公式

官方文档给出的计算公式为：

$$ \text{loss} = \max(\text{predict}, 0) - \text{predict} \times \text{target} + \log(1 + \exp(-|\text{predict}|)) $$

其中：

predict为输入的 logits 值；
target为标签值；
loss为计算得到的损失值。

该公式在数学上等价于先对 logits 求 Sigmoid 概率 $p = \frac{1}{1+e^{-x}}$，再计算二元交叉熵 $-(t \cdot \log p + (1-t) \cdot \log(1-p))$，但采用了数值稳定的写法：用 $\max(x,0)$ 与 $\log(1+e^{-|x|})$ 组合，避免 $e^{-x}$ 在 $x$ 为较大负数时下溢、在 $x$ 为较大正数时上溢的问题，这也是 TensorFlow 官方实现的同款稳定化处理。

1.3 Kernel 层的指令级实现

数值稳定的公式并不是只在文档层面描述，在 NPU Kernel 中得到了逐条指令的落实。核心计算位于 op_kernel/arch35/sigmoid_cross_entropy_with_logits_dag.h 的CalcSigmoidCrossEntropyWithLogits结构中，其向量计算序列为：

Reg::Maxs(vregMaxPredict, vregPredict, 0.0f)—— 计算 $\max(\text{predict}, 0)$；
Reg::Abs(vregAbsPredict, vregPredict)→Reg::Neg(vregNegAbsPredict, ...)→Reg::Exp(vregExpNegAbs, ...)—— 计算 $e^{-|\text{predict}|}$；
Reg::Duplicate(vregOneAddExp, 1.0f)+Reg::Add(...)—— 计算 $1 + e^{-|\text{predict}|}$；
Reg::Log(vregLog, vregOneAddExp)—— 计算 $\log(1 + e^{-|\text{predict}|})$；
Reg::Mul(vregMulPredictTarget, vregPredict, vregTarget)—— 计算 $\text{predict} \times \text{target}$；
Reg::Sub(vregOutput, vregMaxPredict, vregMulPredictTarget)+Reg::Add(vregOutput, vregOutput, vregLog)—— 组合出最终损失。

可见文档公式与 Kernel 指令序列一一对应，逐元素、逐向量地完成稳定化损失计算。

二、产品支持情况

官方文档中的产品支持矩阵如下：

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

从源码结构看，该算子的 Tiling 与 Kernel 均位于arch35目录（op_host/arch35 与 op_kernel/arch35），算子定义中通过OpAICoreConfig为ascend950平台添加了 AICore 配置（见 op_host/sigmoid_cross_entropy_with_logits_def.cpp），与文档中 Ascend 950 系列产品支持情况一致。文档明确标注 Atlas 200I/500 A2 推理产品不支持，部署前请务必核对目标产品型号。

三、参数说明

官方文档参数表如下：

参数名	输入/输出/属性	描述	数据类型	数据格式
predict	输入	预测值 logits。	FLOAT16、FLOAT、BFLOAT16	ND
target	输入	标签值。	FLOAT16、FLOAT、BFLOAT16	ND
loss	输出	损失值。shape 和输入 predict 一致。	FLOAT16、FLOAT、BFLOAT16	ND

3.1 原型层的类型与格式约束

上述参数表在算子原型中有完全对应的硬性约束，op_graph/sigmoid_cross_entropy_with_logits_proto.h 中的注册信息为：

REG_OP(SigmoidCrossEntropyWithLogits) .INPUT(predict, TensorType({DT_FLOAT16, DT_FLOAT, DT_BF16})) .INPUT(target, TensorType({DT_FLOAT16, DT_FLOAT, DT_BF16})) .OUTPUT(loss, TensorType({DT_FLOAT16, DT_FLOAT, DT_BF16})) .OP_END_FACTORY_REG(SigmoidCrossEntropyWithLogits)

在算子定义文件 op_host/sigmoid_cross_entropy_with_logits_def.cpp 中，两个输入与一个输出均被声明为REQUIRED（必选），数据类型限定为{DT_FLOAT16, DT_FLOAT, DT_BF16}三种，数据格式限定为{FORMAT_ND, FORMAT_ND, FORMAT_ND}，即仅支持 ND 格式。

3.2 动态能力配置

同一文件中的 AICore 配置还揭示了算子的动态特性（sigmoid_cross_entropy_with_logits_def.cpp）：

OpAICoreConfig aicoreConfig; aicoreConfig.DynamicCompileStaticFlag(true) .DynamicRankSupportFlag(true) .DynamicShapeSupportFlag(true) .PrecisionReduceFlag(false); this->AICore().AddConfig("ascend950", aicoreConfig);

DynamicCompileStaticFlag(true)：支持编译期静态确定部分信息后的动态编译；
DynamicRankSupportFlag(true)：支持动态 Rank（维度数量在运行期确定）；
DynamicShapeSupportFlag(true)：支持动态 Shape；
PrecisionReduceFlag(false)：不做精度降级处理，FLOAT32 输入按原生精度计算。

这解释了算子为何能以shape: [-2]的形式出现在二进制配置中（见下文），并支持运行期任意维度的张量。

3.3 输出与输入的一致性推导

输出loss的 Shape 和数据类型并非独立指定，而是由 Shape 推导（InferShape）逻辑从输入推导而来。op_host/sigmoid_cross_entropy_with_logits_infershape.cpp 中的实现为：

static ge::graphStatus InferShapeForSigmoidCrossEntropyWithLogits(gert::InferShapeContext* context) { ge::graphStatus ret = Ops::Base::InferShape4Elewise(context); ... } static graphStatus InferDataTypeForSigmoidCrossEntropyWithLogits(gert::InferDataTypeContext* context) { context->SetOutputDataType(LOSS_INDEX, context->GetInputDataType(PREDICT_INDEX)); return GRAPH_SUCCESS; }

Shape 推导复用InferShape4Elewise（逐元素通用推导工具），保证输出与输入同 Shape；
数据类型推导直接将输出类型设置为predict输入的类型，保证loss与predict数据类型一致。

这一点也被单元测试覆盖，tests/ut/op_host/test_sigmoid_cross_entropy_with_logits_infershape.cpp 中分别以二维形状{96, 256}和一维形状{1024}构造输入输出并断言推导成功。

四、约束说明

官方文档明确了两条约束，且这些约束在 Tiling 源码中有对应的强制校验逻辑（op_host/arch35/sigmoid_cross_entropy_with_logits_tiling.cpp）：

predict 和 target 必须具有相同的数据类型和形状。
- CalcInputDtype()校验输入类型必须是 FLOAT16/BF16/FLOAT 三者之一，且target的 dtype 必须与predict相同；
- CheckShape()校验target的存储 Shape 与predict一致，且输出loss的存储 Shape 也与predict一致；同时拒绝空张量（shape size 为 0 时报错）。
支持 FLOAT16、FLOAT、BFLOAT16 数据类型。
- 除 dtype 校验外，CalcOutputDtype()还会校验输出loss的 dtype 与输入predict相同。

从实现看，标量输入（Shape 为 0 维）会被EnsureNotScalar统一视为{1}形状参与一致性比较，属于对边界情况的兼容处理。

五、调用说明

官方文档给出的调用方式为图模式：通过算子 IR 构图方式调用，即引用 op_graph/sigmoid_cross_entropy_with_logits_proto.h 中的算子声明完成构图。仓库提供了完整可编译的构图示例 examples/test_geir_sigmoid_cross_entropy_with_logits.cpp，其核心构图流程如下：

// 1. 创建算子实例 auto sigmoidCrossEntropyWithLogits1 = op::SigmoidCrossEntropyWithLogits("sigmoidCrossEntropyWithLogits1"); // 2. 构造 predict 输入（占位符 + 形状 {4, 2}，ND 格式，DT_FLOAT） std::vector<int64_t> predictShape = {4, 2}; auto placeholder1 = op::Data("placeholder1").set_attr_index(0); TensorDesc placeholder1_desc = TensorDesc(ge::Shape(predictShape), FORMAT_ND, inDtype); sigmoidCrossEntropyWithLogits1.set_input_predict(placeholder1); // 3. 构造 target 输入（形状同样为 {4, 2}） auto placeholder2 = op::Data("placeholder2").set_attr_index(1); sigmoidCrossEntropyWithLogits1.set_input_target(placeholder2); // 4. 声明输出并建图运行 TensorDesc loss_desc = TensorDesc(ge::Shape(predictShape), FORMAT_ND, inDtype); sigmoidCrossEntropyWithLogits1.update_output_desc_loss(loss_desc); graph.SetInputs(inputs).SetOutputs(outputs); session->AddGraph(graph_id, graph, graph_options); session->RunGraph(graph_id, input, output);

构图要点可归纳为：

使用op::Data创建输入占位符，并通过set_attr_index指定输入索引；
通过set_input_predict/set_input_target将两个占位符挂到算子输入；
通过update_output_desc_loss预先声明输出张量描述（Shape、格式、类型）；
将算子加入Graph，设置图输入输出后，经Session::AddGraph与Session::RunGraph完成 NPU 上的执行；
示例中ge.exec.deviceId指定运行设备，ge.graphRunMode指定图运行模式，运行结果会 dump 为tc_ge_irrun_test_0008_npu_input_*.bin/_output_*.bin文件并打印每个输出元素。

示例中使用的输入形状为{4, 2}、数据类型DT_FLOAT，与算子支持的数据类型完全匹配；由于算子支持动态 Shape，实际业务中可将{4, 2}替换为任意形状的 predict/target 对。

六、NPU 执行链路：从 Tiling 到 Kernel

6.1 Tiling 准备与数据切分

在图编译阶段，Tiling 负责根据输入形状、平台资源决定 Kernel 的切分参数。op_host/arch35/sigmoid_cross_entropy_with_logits_tiling.cpp 的执行链路为：

TilingPrepareForSigmoidCrossEntropyWithLogits：通过platform_ascendc::PlatformAscendC获取 AIV 核数（GetCoreNumAiv）与 UB 内存大小（GetCoreMemSize(CoreMemType::UB, ...)），写入编译信息；
Tiling4SigmoidCrossEntropyWithLogits→RunTiling()：依次执行CalcInputDtype（输入类型校验）→CalcOutputDtype（输出类型校验）→CheckShape（形状一致性校验）→DoElewiseTiling；
DoElewiseTiling：按 dtype 选择对应的模板参数（FP16/BF16/FP32），调用通用逐元素 Tiling 工具ElewiseBaseTiling的DoTiling生成EleBaseTilingData16B类型的 Tiling 数据，并设置tilingKey与blockDim（核数）；
申请固定大小 workspace：const size_t ASCEND_WORKSPACE = 16777216;，即 16 MB 的工作内存。

6.2 Kernel 分派与模板实例化

Kernel 入口 op_kernel/sigmoid_cross_entropy_with_logits.cpp 通过编译期常量dType做三路分派：

template <uint64_t dType> __global__ __aicore__ void sigmoid_cross_entropy_with_logits(GM_ADDR predict, GM_ADDR target, GM_ADDR loss, GM_ADDR workspace, GM_ADDR tiling) { REGISTER_TILING_DEFAULT(EleBaseTilingData16B); GET_TILING_DATA_PTR_WITH_STRUCT(EleBaseTilingData16B, tilingData, tiling); KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY); if constexpr (dType == TPL_FP16) { ElementwiseSch16B<1, SigmoidCrossEntropyWithLogitsDagWithCast<half>::OpDag> sch(tilingData); sch.Init(predict, target, loss); sch.Process(); } else if constexpr (dType == TPL_BF16) { ElementwiseSch16B<1, SigmoidCrossEntropyWithLogitsDagWithCast<bfloat16_t>::OpDag> sch(tilingData); ... } else if constexpr (dType == TPL_FP32) { ElementwiseSch16B<1, SigmoidCrossEntropyWithLogitsDagNoCast<float>::OpDag> sch(tilingData); ... } }

任务类型为KERNEL_TYPE_AIV_ONLY，即由 AIV 向量核执行逐元素计算；
三种 dtype 的模板参数在 op_kernel/arch35/sigmoid_cross_entropy_with_logits_struct.h 中定义：TPL_FP16=1、TPL_BF16=2、TPL_FP32=3；
调度器统一使用ElementwiseSch16B<1, OpDag>（16 字节对齐的逐元素调度器），实际计算逻辑封装在 DAG 描述中。

6.3 DAG 计算图：FP16/BF16 的精度提升策略

op_kernel/arch35/sigmoid_cross_entropy_with_logits_dag.h 中定义了两套计算 DAG：

SigmoidCrossEntropyWithLogitsDagNoCast<T>：FP32 专用，输入输出直接以 float 参与计算，不做类型转换；
SigmoidCrossEntropyWithLogitsDagWithCast<U, T = float>：FP16/BF16 专用，先将 16 位输入Cast提升为 float 计算（CAST_MODE_NONE），完成损失计算后再以CAST_MODE_RINT（四舍五入）模式将结果 Cast 回 16 位输出。

这一设计意味着：FP16/BF16 输入在 NPU 上实际以float32 中间精度完成 $\max$、$\exp$、$\log$ 等运算，只在边界处做一次 Cast，可有效降低 16 位浮点逐元素累积带来的精度损失；FP32 输入则直接原生计算，无精度降级（与PrecisionReduceFlag(false)配置相互印证）。两份 DAG 的MemCfg均为MemOptCfg<MemLevel::LEVEL_2>，即允许中间数据驻留二级缓存优化访存。

6.4 二进制配置与运行期分派

平台侧以 JSON 形式维护算子的二进制库分派配置，op_host/config/ascend950/sigmoid_cross_entropy_with_logits_binary.json 中为三种数据类型分别登记了二进制文件：

二进制文件	predict dtype	target dtype	loss dtype	格式
SigmoidCrossEntropyWithLogits_a1b2c3d4e5f67890	float16	float16	float16	ND
SigmoidCrossEntropyWithLogits_b2c3d4e5f6789012	float32	float32	float32	ND
SigmoidCrossEntropyWithLogits_c3d4e5f678901234	bfloat16	bfloat16	bfloat16	ND

配置中所有输入输出均声明为shape: [-2]（动态 Rank 通配），paramType: required，运行期根据实际输入 dtype 选择对应二进制，与算子支持的动态 Shape 能力相配套。

七、测试与验证

仓库为算子提供了两级验证手段：

InferShape 单元测试：tests/ut/op_host/test_sigmoid_cross_entropy_with_logits_infershape.cpp 使用gert::InferShapeContextFaker构造{96, 256}、{1024}等形状，验证算子实现已注册、InferShape 函数非空且推导返回GRAPH_SUCCESS；
Tiling 单元测试：tests/ut/op_host/arch35/test_sigmoid_cross_entropy_with_logits_tiling.cpp 覆盖 Tiling 数据生成逻辑；
端到端 GEIR 示例：examples/test_geir_sigmoid_cross_entropy_with_logits.cpp 完整演示从 GE 初始化、构图、建图到RunGraph并在 NPU 上执行、导出输入输出 bin 文件的全流程。

对于开发者而言，验证算子的标准步骤为：使用示例构图代码构造 predict/target 输入 → 运行图得到 loss 输出 → 与公式 $\max(x,0) - x \cdot t + \log(1+e^{-|x|})$ 的逐元素计算结果比对，由于 Kernel 指令序列与公式严格对应，逐元素误差应在所选数据类型的精度范围内。

八、小结

SigmoidCrossEntropyWithLogits 是 ops-nn 中一个典型的"小而完整"的逐元素损失算子：文档层面给出了清晰的功能、公式、参数与约束；源码层面则完整覆盖了算子原型注册、Shape 推导、动态 Tiling、AIV 向量 Kernel、二进制分派配置与多级测试。其核心设计要点可归纳为：

数值稳定公式：以 $\max(x,0) + \log(1+e^{-|x|}) - x \cdot t$ 形式规避 Sigmoid 交叉熵在极端 logits 下的溢出；
严格输入约束：predict/target/loss 三者 dtype 与 shape 必须一致，仅支持 FLOAT16/FLOAT/BFLOAT16 与 ND 格式；
动态能力完备：支持动态 Shape、动态 Rank、动态编译，可承接 TensorFlow 导出的同名算子；
16 位精度增强：FP16/BF16 输入以 float32 中间精度计算，FP32 原生计算不降级；
完整工程配套：图模式示例、InferShape/Tiling 单测与二进制分派配置一应俱全，可直接作为 CANN 算子开发与集成的参考模板。

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

创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考