CANN社区SoftmaxCrossEntropyWithLogits算子设计

SoftmaxCrossEntropyWithLogits算子设计文档

【免费下载链接】cann-ops-competitions本仓库用于 CANN 开源社区各类竞赛、开源课题、社区任务等课题发布、开发者作品提交和展示。项目地址: https://gitcode.com/cann/cann-ops-competitions

一、需求背景

1.1 需求来源

社区任务03-8要求参考内置TBE SoftmaxCrossEntropyWithLogits，在昇腾 NPU 上完成SoftmaxCrossEntropyWithLogits算子的 Ascend C 设计、开发与测试交付。

1.2 背景介绍

1.2.1 算子目标

SoftmaxCrossEntropyWithLogits用于分类训练场景。输入features表示 logits，输入labels表示标签分布，算子输出每个样本的交叉熵损失loss，同时输出对 logits 的梯度backprop。

本次设计目标如下：

1. 功能语义与任务环境中的内置TBE SoftmaxCrossEntropyWithLogits对齐。
2. 交付 Ascend C 原生算子工程，包含 Host 定义、InferShape、Tiling、AiCore Kernel、aclnn两段式接口、样例与测试框架。
3. 面向 Atlas A2 训练系列产品，支持FLOAT、FLOAT16、BFLOAT16三类浮点输入输出。
4. 按任务环境内置 TBE 基线实现ND二维输入和NCHW四维输入支持，features与labels同 shape、同 dtype，backprop与输入同 shape，loss在类别维归约后按 TBE 输出 format 写回。
5. 对 TBE 配置中的 dtype、format、rank 和 shape 推导规则进行设计对齐，不把任务书明确排除的数据类型和输入广播纳入设计承诺。

1.2.2 TBE基线来源说明

内置 TBE 算子是本任务的功能、精度与性能对齐基线。基线信息需要在任务 CANN 环境中从以下位置确认：

本工作区已在 CANN 8.5.2、910B3 容器中完成基线探测。探测结论为：任务环境内置 TBE 的op_proto/inc/nn_norm_ops.h中保留ND or NHWC、2D/4D等原型描述；op_impl/ai_core/tbe/config/ascend910b/aic-ascend910b-ops-info-legacy.json中features/labels/backprop的 format 为NCHW,ND，loss的 format 为NHWC,ND；op_impl/ai_core/tbe/kernel/config/ascend910b/ops_legacy/softmax_cross_entropy_with_logits.json的动态 kernel 配置列出FLOAT、FLOAT16、BFLOAT16的ND格式，aclnnop/aclnn_softmax_cross_entropy_with_logits.h声明 aclnn 入参与出参数据格式支持ND。本设计按 TBE 口径支持FLOAT、FLOAT16、BFLOAT16，输入支持二维ND和四维NCHW同形输入，loss输出支持ND/NHWC，backprop输出支持ND/NCHW；任务书明确排除int64、double与 broadcast，因此这三类能力不纳入本次设计。

1.2.3 TBE算子现状分析

1.2.3.1 支持的数据类型、数据格式与形状

本次 Ascend C 设计范围如下：

其中N表示 batch 数量，C表示类别数量，H/W表示四维布局下的空间维。features与labels必须同 shape、同 dtype、同 format；loss和backprop的 dtype 必须与features一致。

格式边界说明如下：

关于审核反馈中的NCHW：CANN 8.5.2 的ascend910bops-info 中features/labels输入为NCHW/ND，backprop输出为NCHW/ND，loss输出为NHWC/ND；原型说明中四维输入对应loss为 3D。因此本设计按 TBE 口径处理四维 layout：输入不声明NHWC支持；ND使用最后一维为类别维，NCHW使用第 1 维为类别维；backprop与输入同 shape；loss去除类别维后保留其余维度，NCHW [N,C,H,W]对应loss [N,H,W]，输出 format 与 TBE 的NHWC/ND元数据对齐。Kernel 侧统一折叠为逻辑[A, R]计算域，其中R=C为类别归约维，A为其余维度乘积。

本任务不支持整数输入、不支持双精度输入、不支持features与labels间的 shape 扩展或广播。

1.2.3.2 算子数学语义

设输入features和labels按 format 折叠为逻辑二维计算域[A, R]，其中R=C为类别归约维，A为 batch 与非类别维乘积。对第a个逻辑样本，先求类别维最大值：

$$ m_a = \max_{0 \le r < R} x_{a,r}. $$

使用稳定 softmax 形式计算指数和：

$$ s_a = \sum_{r=0}^{R-1}\exp(x_{a,r} - m_a). $$

softmax 概率为：

$$ p_{a,r} = \frac{\exp(x_{a,r} - m_a)}{s_a}. $$

loss输出为：

$$ \text{loss}a = -\sum{r=0}^{R-1} y_{a,r}\log(p_{a,r}). $$

Kernel 中使用等价形式避免直接计算log(softmax)：

$$ \text{loss}a = \sum{r=0}^{R-1} y_{a,r}\left(\log(s_a) - (x_{a,r} - m_a)\right). $$

backprop输出为：

$$ \text{backprop}{a,r} = p{a,r} - y_{a,r}. $$

1.2.3.3 TBE算子计算语义分析

TBE / 参考实现的核心调度抽象是把输出backprop的有效计算域折叠为[A, R]，其中R是类别归约维，A是其余维度乘积。二维ND [N, C]场景下即A=N、R=C；四维NCHW [N, C, H, W]场景下即A=N*H*W、R=C。本设计沿用该[A, R]抽象，Host 侧负责根据 format 做逻辑维度折叠和输出 shape 推导。

该算子的计算可以拆为三类操作：

1. 沿类别维C做ReduceMax，得到稳定 softmax 所需的行最大值。
2. 沿类别维C做ReduceSum，得到指数和以及 loss 归约结果。
3. 对逻辑[A, R]元素做逐元素指数、除法、乘法和减法，按原 layout 写回backprop，并对每个逻辑样本写回loss。

由于不同逻辑样本之间没有数据依赖，分核主轴选取A维。类别维R决定 UB 组织方式：R可完整进入 UB 时使用 FullLoad；R较大时使用 SplitR 分段重扫。

TBE 基线流程图如下：

二、需求分析

2.1 外部组件依赖

不引入新的第三方组件，复用开源算子仓已有能力：

2.2 内部适配模块

代码交付采用社区任务独立算子工程组织，设计文档、接口文档、Host、Kernel、样例和测试脚本在同一提交目录内保持可追溯关系。当前工程目录职责如下：

2.3 需求模块设计

2.3.1 Ascend C算子原型

原生算子输入输出契约：

aclnn两段式接口如下：

aclnnStatus aclnnSoftmaxCrossEntropyWithLogitsGetWorkspaceSize( const aclTensor* features, aclTensor* labels, aclTensor* loss, aclTensor* backprop, uint64_t* workspaceSize, aclOpExecutor** executor); aclnnStatus aclnnSoftmaxCrossEntropyWithLogits( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, const aclrtStream stream);

2.3.2 设计范围与约束

三、需求详细设计

3.1 使能方式

本设计面向 Ascend C 原生算子工程与aclnn两段式接口调用场景，硬件范围为 Atlas A2 训练系列产品。

3.2 需求总体设计

整体链路分为四层：

1. aclnn接口层：完成非空校验、dtype 校验、shape 校验、空 tensor 快速返回、输入连续化与输出ViewCopy。
2. L0 原生调用层：根据featuresshape 分配loss和backprop临时 tensor，并把 AICore 算子加入执行器。
3. Host Tiling 层：读取 dtype、format、shape 与平台信息，根据 layout 确定逻辑A和R=C，完成按逻辑样本分核、UB 规划和tilingKey选择。
4. AiCore Kernel 层：按tilingKey进入 FullLoad、SplitR 或 Empty 分支，完成稳定 softmax cross entropy 计算并写回。

Host-Tiling-Kernel 策略图如下：

3.2.1 Host侧设计

3.2.1.1 接口校验与规整策略

aclnn接口层校验顺序如下：

1. features、labels、loss、backprop非空。
2. 四个 tensor 的 dtype 均属于支持列表。
3. labels、loss、backprop的 dtype 与features一致。
4. features与labelsformat 为ND或NCHW。
5. features与labelsshape 完全一致。
6. ND输入要求 rank 为 2；NCHW输入要求 rank 为 4。
7. lossshape 满足 TBE 类别维归约规则：二维ND为[N]，四维NCHW为[N,H,W]。
8. backpropshape 与features完全一致。

合法输入进入计算前，接口层执行：

1. 若features或labels为空，返回空执行器。
2. 对features和labels执行Contiguous，保证原生 Kernel 面向连续存储。
3. L0 层分配连续临时输出。
4. 原生算子完成计算后，通过ViewCopy写回用户传入的loss和backprop。

3.2.1.2 Shape推导策略

设输入 format 为ND时：

$$ \text{features.shape} = \text{labels.shape} = [N, C]. $$

则输出 shape 为：

$$ \text{loss.shape} = [N], $$

$$ \text{backprop.shape} = [N, C]. $$

设输入 format 为NCHW时：

$$ \text{features.shape} = \text{labels.shape} = [N, C, H, W]. $$

则backprop与输入同 shape，loss按类别维归约规则去除C维并保留N/H/W：

$$ \text{backprop.shape} = [N, C, H, W],\quad \text{loss.shape} = [N, H, W]. $$

对于四维 layout，loss的输出 format 与 shape 推导同时对齐 TBE 元数据和原型契约：输入为NCHW，loss输出 format 为NHWC，逻辑上保留 batch 与非类别维，去除类别维；backprop始终与输入同 shape。以NCHW [2,3,4,5]为例，按 TBE 四维loss为 3D 的约定，归约类别维后得到逻辑 loss 形状[2,4,5]。Kernel 内部可用[N,1,H,W]形式保存行归约结果以便广播，但对外输出 shape 按[N,H,W]校验与写回。

Shape 推导要求features与labels同 shape，不做 shape 扩展或 broadcast。

3.2.1.3 分核策略

Host tiling 将输入按 format 折叠为逻辑二维计算域：

$$ \text{ND}: A=N,\quad R=C. $$

$$ \text{NCHW}: A=N\cdot H\cdot W,\quad R=C. $$

其中A是逻辑样本数，是多核并行主轴；R是类别维，是每个逻辑样本的归约维。由于逻辑样本之间互不依赖，分核只按A维进行。

设可用 AIV 核数为coreNum：

$$ \text{useCoreNum} = \min(A, \text{coreNum}). $$

每核基础行数：

$$ \text{blockFactor} = \left\lceil \frac{A}{\text{useCoreNum}} \right\rceil. $$

实际使用核数：

$$ \text{realCoreNum} = \left\lceil \frac{A}{\text{blockFactor}} \right\rceil. $$

尾核行数：

$$ \text{tailBlockFactor} = A - \text{blockFactor}\cdot(\text{realCoreNum}-1). $$

Kernel 侧根据blockIdx计算本核起始行，普通核处理blockFactor行，尾核处理tailBlockFactor行。

3.2.1.4 UB分块策略

设输入类型字节数为：

$$ b = \begin{cases} 4, & \text{FLOAT}, \ 2, & \text{FLOAT16 或 BFLOAT16}. \end{cases} $$

32B对齐下单个 block 能容纳的元素数为：

$$ \text{perBlock} = \frac{32}{b}. $$

类别维对齐长度为：

$$ R_{\text{align}} = \left\lceil \frac{R}{\text{perBlock}} \right\rceil \cdot \text{perBlock}. $$

FullLoad 路径要求单行完整Ralign能在 UB 中同时容纳输入队列、输出队列和 float32 中间缓冲。估算可容纳行数：

$$ A_{\text{ub}} = \left\lfloor \frac{UB}{M_{AR}+M_A} \right\rfloor. $$

其中MAR表示按[A, R]组织的输入、输出和中间缓冲开销，MA表示行级归约结果开销。当Aub至少能覆盖一个32B对齐粒度时，选择 FullLoad。

若 FullLoad 不满足 UB 容量要求，则进入 SplitR。SplitR 固定按较小的Atile 处理行，剩余 UB 用于分段容纳类别维：

$$ \text{rLoopTime} = \left\lceil \frac{R}{\text{rUbNumFactor}} \right\rceil. $$

尾段有效长度为：

$$ \text{rLoopTile} = \begin{cases} \text{rUbNumFactor}, & R \bmod \text{rUbNumFactor} = 0,\ R \bmod \text{rUbNumFactor}, & \text{otherwise}. \end{cases} $$

3.2.1.5 tilingKey规划策略

tilingKey由模板参数组成：

schId, featuresBrc, labelsBrc, db

本设计不支持features与labels间的 shape 扩展或 broadcast，因此featuresBrc和labelsBrc固定使用非扩展路径；db暂保留模板位。format 信息由 Host 侧转换为类别维位置、逻辑A/R、GM stride 和输出 shape 参数写入 tilingData。

Host 写入SoftmaxCrossEntropyWithLogitsTilingData后执行：

SetBlockDim(realCoreNum)

Kernel 侧根据blockIdx处理本核负责的行块，行间互不依赖，loss与backprop均由负责对应行的核直接写回。

3.2.2 Kernel侧设计

3.2.2.1 Kernel入口

Kernel 入口读取SoftmaxCrossEntropyWithLogitsTilingData，根据schId分发：

Kernel 计算流程图如下：

3.2.2.2 FullLoad路径

FullLoad 每次处理若干行和完整R维。单行处理流程如下：

1. 从 GM 读取当前行features与labels。
2. 将输入转为 float32 中间值。
3. 执行ReduceMax得到m_n。
4. 计算z = features - m_n。
5. 计算exp(z)。
6. 执行ReduceSum得到s_n。
7. 计算backprop = exp(z) / s_n - labels并写回。
8. 计算lossPart = labels * (log(s_n) - z)。
9. 对lossPart做ReduceSum得到loss[n]并写回。

该路径每行只读取一次features和一次labels，中间数据保留在 UB 内，是当前主要性能路径。

针对常见小类别数场景，FullLoad 内部进一步细分：

1. C==1时 softmax 恒为1，直接计算backprop = 1 - labels，loss = 0，避免指数、对数和归约。
2. 大N小C场景按行块处理，并对C=2/3/4/5/8等形状使用列紧凑和行级标量复用子路径，减少通用Reduce和行级结果展开的固定开销。
3. C非32B对齐时 Host 记录有效列数，Kernel 右侧 padding 只参与 UB 对齐，不写回无效列，保证backprop按原 shape 输出不越界。

3.2.2.3 SplitR路径

SplitR 面向R维较大、完整行无法同时容纳输入和多个 float32 中间缓冲的场景。该路径使用三阶段分段重扫：

1. 第一轮按R分段读取features，对每段做局部ReduceMax，归并为全行最大值m_n。
2. 第二轮按R分段读取features，计算exp(features - m_n)并累加全行指数和s_n。
3. 第三轮按R分段读取features与labels，计算并写回backprop，同时累加各段 loss 贡献，最后写回loss[n]。

该路径面向大类别维输入的泛化执行能力设计。相比 FullLoad，它增加了 GM 读取次数，性能需要在 runtime 对比中单独评估。

3.2.2.4 精度策略

精度设计如下：

1. FLOAT输入以 float32 参与中间计算。
2. FLOAT16和BFLOAT16输入读取后转换为 float32，完成max、exp、sum、log、乘法和减法。
3. 输出 dtype 与输入 dtype 一致，写回前转换回对应类型。
4. 使用log(sum) - (features - max)计算 loss，避免直接计算log(softmax)。
5. FullLoad 与 SplitR 使用相同数学公式，SplitR 的分段累加误差按 AscendOpTest 默认阈值验证。

3.2.2.5 Ascend C路径与TBE路径的差异点及原因

四、特性交叉分析

五、可维可测分析

5.1 验收标准与验证口径

5.2 验证矩阵

当前交付目录保留 Host UT 与 runtime custom/TBE 对比入口。Host UT 覆盖 InferShape、Tiling 和基础错误路径；runtime 用例覆盖 dtype、输入 format、常规 shape、小 shape、四维 shape、非对齐C、多核A、大C、空 tensor、非连续输入和非法输入等场景，用于验证loss与backprop的正确性、精度和性能表现。

5.3 兼容性分析

本设计的兼容性结论如下：

1. SoftmaxCrossEntropyWithLogits的正式语义与任务环境内置 TBE 保持一致，输入features与labels同 shape、同 dtype、同 format，backprop与输入同 shape，loss为类别维归约后的 shape。
2. 支持FLOAT、FLOAT16、BFLOAT16三类任务范围内 dtype；输入支持ND二维[N,C]与NCHW四维 layout，loss输出 format 按 TBE 对齐ND/NHWC。
3. aclnn接口层兼容非连续输入的逻辑 view 语义，进入原生 Kernel 前完成连续化，输出通过 ViewCopy 写回。
4. 兼容空 tensor、小 shape、非对齐C、多核A、四维 layout 和大C输入，分别由接口快速返回、FullLoad、小C专用路径、layout stride 和 SplitR 路径覆盖。
5. 任务范围外 dtype、rank 和 shape 关系由接口与 Host 侧校验拒绝，不进入 Kernel 主执行路径。

【免费下载链接】cann-ops-competitions本仓库用于 CANN 开源社区各类竞赛、开源课题、社区任务等课题发布、开发者作品提交和展示。项目地址: https://gitcode.com/cann/cann-ops-competitions