企业官网建设流程全解析

• 人工智能
• 算子库
• 深度学习
• CANN
• Ascend

LogSoftmaxGrad 是 CANN ops-nn 算子库中负责 LogSoftmax 反向传播的 AscendC 算子，它接收正向输出的梯度 gradOutput 与正向输出 output，还原出正向输入 x 的梯度 gradInput。本文以仓库中 LogSoftmaxGrad 算子 README 与 aclnnLogSoftmaxBackward 接口文档 为主线，结合算子定义、tiling、kernel 源码与测试用例，系统讲解该算子的数学原理、两段式 aclnn 调用方式、参数与约束，以及它在 NPU 上"合轴分块、分核并行、float32 统一计算"的工程实现思路，帮助读者在 Atlas A2 系列产品上正确调用并深入理解其底层机制。

一、算子定位与产品支持情况

LogSoftmaxGrad 属于激活函数（activation）模块的梯度类算子，位于仓库 experimental/activation/log_softmax_grad 目录，在算子注册层面对应的内部算子名为LogSoftmaxGrad，对外暴露的 aclnn 接口为aclnnLogSoftmaxBackward。

根据 README 产品支持情况 与 aclnnLogSoftmaxBackward 文档 的记录，其支持情况如下：

其中 Atlas A2 训练系列产品与 Atlas 800I A2 推理产品对应 NPU 架构DAV_2201，这也是 bfloat16 支持与否的分水岭（详见下文"数据检测"一节）。从算子定义文件 log_softmax_grad_def.cpp 可以看到，该算子通过this->AICore().AddConfig("ascend910b")完成 AICore 侧的算子配置注册。

二、功能与数学原理

2.1 接口功能

该算子完成 LogSoftmax 的反向传播，即给定：

• gradOutput：LogSoftmax 正向输出的上游梯度；
• output：LogSoftmax 正向输出，即 $\text{log_softmax}(x)$；

计算出 LogSoftmax 正向输入 x 的梯度gradInput。

2.2 计算公式

README 与接口文档给出的核心公式为：

$$ gradInput = gradOutput - \exp(output) \cdot \text{reduce_sum}(gradOutput, dim) $$

由于 $\exp(output)$ 恰好就是 softmax 概率分布，因此公式等价于：

$$ gradInput = gradOutput - softmax(x) \cdot \text{reduce_sum}(gradOutput, dim) $$

其中reduce_sum(gradOutput, dim)表示沿dim指定的轴对上游梯度求和，softmax(x)为 LogSoftmax 正向输入的 softmax 概率。两种写法在数值上是等价的，第二种写法更直观地揭示了"LogSoftmax 反向 = 上游梯度减去概率加权的梯度总和"这一语义。

从 kernel 源码 no_need_reduce.h 可以印证该公式的实现：当 reduce 维度长度为 1 时（即不需要真正做 reduce 求和），公式可以化简为：

$$ output_z = (input_dy - 1) \cdot \exp(input_x) $$

对应的 kernel 计算序列为：Exp(x, x, count)计算 $\exp(x)$ →Mul(x, dy, x, count)计算 $dy \cdot \exp(x)$ →Sub(z, dy, x, count)得到 $z = dy - dy \cdot \exp(x) = (dy-1) \cdot \exp(x)$。

2.3 支持的数据类型

算子支持 FLOAT（float32）、FLOAT16（float16）、BFLOAT16（bfloat16）三种数据类型，且要求gradOutput、output、out三者数据类型与 shape 完全一致，不支持广播。该约束在算子定义 log_softmax_grad_def.cpp 与 aclnn 入口的参数校验中均有体现。

三、两段式 aclnn 接口与完整调用示例

3.1 接口原型

本算子提供两段式 aclnn 接口：必须先调用aclnnLogSoftmaxBackwardGetWorkspaceSize获取 workspace 大小与执行器，再调用aclnnLogSoftmaxBackward执行计算。

aclnnStatus aclnnLogSoftmaxBackwardGetWorkspaceSize( const aclTensor* gradOutput, const aclTensor* output, int64_t dim, aclTensor* out, uint64_t* workspaceSize, aclOpExecutor** executor)

aclnnStatus aclnnLogSoftmaxBackward( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)

两段式接口的具体机制可参考仓库 docs/zh/context/two_phase_api.md 的说明。

3.2 第一段接口参数说明

值得注意的是，接口文档标注三个 Tensor 均支持"非连续 Tensor"（√），即允许传入 stride 非默认的视图张量，aclnn 入口会通过aclnn_kernels/contiguous.h等基础设施做兼容处理。维度范围 0-8 与 aclnn 入口处AXIS_LIMIT = 8（见 aclnn_logsoftmax_backward.cpp）一致，即底层算子不支持超过 8 维的输入。

3.3 第二段接口参数说明

3.4 返回值与常见报错

两段接口均返回aclnnStatus状态码，具体取值可参考 docs/zh/context/aclnn_return_code.md。第一段接口会在入参校验阶段直接报错，常见错误码如下：

3.5 完整可编译调用示例

仓库提供了完整可编译的调用示例 test_aclnn_log_softmax_grad.cpp，其核心流程如下：

#include "acl/acl.h" #include "aclnn_logsoftmax_backward.h" // 1. 初始化：aclInit、aclrtSetDevice、aclrtCreateStream int Init(int32_t deviceId, aclrtStream* stream); // 2. 准备输入/输出 Tensor（host 数据拷贝到 device，再 aclCreateTensor 包装） // gradOutput: shape {5}, 数据 {-0.5, 0.7, 0.1, -0.8, 0.3}, ACL_FLOAT // output: shape {5}, 数据 {-1.5, -3.2, -2.3, -4.4, -5.5}, ACL_FLOAT // out: shape {5}, ACL_FLOAT int64_t axis = -1; // 指定 reduce 维度为最后一维 // 3. 两段式调用 aclOpExecutor* executor; uint64_t workspaceSize = 0; ret = aclnnLogSoftmaxBackwardGetWorkspaceSize(input_dyTensor, input_xTensor, axis, output_zTensor, &workspaceSize, &executor); void* workspaceAddr = nullptr; if (workspaceSize > 0) { ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); } ret = aclnnLogSoftmaxBackward(workspaceAddr, workspaceSize, executor, stream); ret = aclrtSynchronizeStream(stream); // 4. 结果回拷 host，销毁 Tensor 与资源，aclFinalize

示例中的关键点：

• 示例默认deviceId = 0，axis = -1（最后一维做 reduce），输入为 1 维、shape{5}的 float32 张量；
• workspace 申请使用ACL_MEM_MALLOC_HUGE_FIRST策略，且仅当workspaceSize > 0时才申请，避免无谓的内存分配；
• 输出通过aclrtMemcpy从 device 回拷到 host 后再做后续校验。

该示例的编译与执行步骤可参考仓库 docs/zh/context/compile_and_run_sample.md 的通用流程。UT 目录下还有对应的 aclnn 级测试 test_aclnn_logsoftmax_backward.cpp，该测试以 torch golden 结果做逐元素对比，覆盖三种数据类型与多种 shape。

四、约束说明

综合 README 与接口文档，算子使用约束如下：

1. 输入gradOutput、output与输出out的 shape 必须一致，不支持广播；
2. 指定维度dim需在 gradOutput 的维度范围内，即 $[-dimNum, dimNum-1]$；
3. 支持的数据类型为 FLOAT、FLOAT16、BFLOAT16；当硬件不支持 bfloat16 向 float32 的 Cast 时，使用 BFLOAT16 会报错；
4. 输入 tensor 的维度数不能超过 8；
5. 当dim为多个轴时，reduce 轴必须连续（如 [2,3,4] 合法，[2,4] 不合法），这是算子语义对多轴 reduce 的限制，在 tiling 参数校验中强制检查（见 log_softmax_grad_tiling.cpp）。

五、Host 侧详细设计

5.1 算子定义与 shape 推导

算子定义文件 log_softmax_grad_def.cpp 完成了三个要素的声明：

• 两个输入input_dy（对应 gradOutput）、input_x（对应 output）与一个输出output_z（对应 out），均为REQUIRED，数据类型限定DT_FLOAT / DT_FLOAT16 / DT_BF16，格式FORMAT_ND；
• 属性axis为OPTIONAL的ListInt，默认值{-1}；
• 通过OP_ADD(LogSoftmaxGrad)完成算子注册。

shape 推导实现位于 log_softmax_grad_infershape.cpp，逻辑非常直接：*zShape = *dyShape，即输出 shape 直接继承 gradOutput 的 shape，这也与"反向梯度与正向输入同 shape"的语义一致。

5.2 tiling 策略：合轴与三维化

tiling 函数入口为 LogSoftmaxGradTilingFunc，其执行流程为 Init（含参数校验）→ CheckParams → DoTiling。

参数校验（CheckParams）确保：dy、x、z 三者 shape 相同，三者 dtype 相同，且 dtype 只能是 float32、float16、bfloat16 之一。

reduce 轴解析（GetReduceAxes）会把负数轴归一化为正轴（dimVal = axesData[i] >= 0 ? axesData[i] : dimNum_ + axesData[i]），并强制校验多 reduce 轴的连续性。

合轴操作：根据 reduce 轴的起始与结束位置，将任意 shape 合并为三个轴：

$$ [mergedDim0,\ mergedDim1,\ mergedDim2] $$

• mergedDim0：reduce 轴之前的全部维度乘积；
• mergedDim1：全部 reduce 轴的维度乘积；
• mergedDim2：reduce 轴之后的全部维度乘积。

特殊情况下：若 reduce 的最小轴是 0，则mergedDim0 = 1；若 reduce 的最大轴是最后一个轴，则mergedDim2 = 1。

平台动态适配：tiling 通过平台信息动态获取 UbSize、CoreNum 与 UB block 大小（GetUbBlockSize），不再硬编码 BLOCK_SIZE，以保证不同芯片平台的兼容性。这一点在 Init 函数 与 GetUbSizeAndCoreNum 中可见：GetCoreMemSize(CoreMemType::UB, ubSize_)获取 UB 大小，GetCoreNum()获取核数，blockSize_来自Ops::Base::GetUbBlockSize(context_)。

5.3 分核策略

根据 README 的说明，分核遵循以下原则：

1. 根据mergedDim0分核，优先满核使用；核间若能均分则无大小核区分，数据块一致；不能均分时余出的数据块分配到前几个核；若mergedDim0 < CoreNum，实际使用核数取mergedDim0；
2. 数据分块与内存优化：充分利用 UB 空间。算子共有两个输入一个输出，input_dy需要做累加和 reduce 操作，为其开启 double buffer 分配 2 块 UB 空间，同时为 reduce 操作预留 1 块空间，input_x与output_z各 1 块，合计 5 块；将 UB 均分为 5 份并对齐 BLOCK_SIZE，计算出每个 UB 块可容纳的数据量。对于 bfloat16 和 half 类型，计算前需要做 cast，搬入的 UB 空间与 Cast 结果的 UB 空间复用：搬入时数据放在 UB 块后半部分，再 cast 到整个 UB 块；
3. tilingkey 规划：若mergedDim1 == 1（无需 reduce），tilingkey 为 0（NO_NEED_REDUCE）；否则 tilingkey 为 1，具体再由 DoTiling 细分出REDUCE_TAIL与REDUCE_MID两种模式。

上述缓冲区数量策略在 log_softmax_grad_tiling.cpp 中有对应常量定义：

constexpr uint64_t BUFFER_NUM = 2; // 是否开启 Double buffer constexpr uint64_t NO_REDUCE_BUF_NUM = 3 * BUFFER_NUM; // 6 constexpr uint64_t REDUCE_BUF_NUM = 3 * BUFFER_NUM + 1; // 7 constexpr uint64_t REDUCE_MID_NODB_BUF_NUM = 3; constexpr uint64_t REDUCE_TAIL_NODB_BUF_NUM = 4;

5.4 三种调度模式

DoTiling（log_softmax_grad_tiling.cpp）根据合轴结果选择调度模式：

• mergedDim1 == 1→NO_NEED_REDUCE（ProcNoNeedReduce）：无需 reduce，直接按totalElems = mergedDim0 * mergedDim2线性切块；
• mergedDim2 == 1→REDUCE_TAIL（ProcReduceTail）：reduce 轴位于尾部，把mergedDim2与mergedDim1交换后按行 reduce；
• 其他情况 →REDUCE_MID（ProcReduceMid）：reduce 轴在中间，需要对dim1、dim2两个方向分别切 tile 与循环。

每种模式都会计算dimXTile / dimXLoopTime / dimXRemained（X=0,1,2）等分块参数，写入 LogSoftmaxGradTilingData 结构体，最终通过context_->SetTilingKey(...)与context_->SetBlockDim(usedCoreNum_)下发到 kernel。tilingkey 由三个模板参数组合而成：调度模式SCH_MOD（3 选 1）、IS_SMALL（小 shape 标志）、IS_CONTIGUOUS（连续搬移标志），声明见 log_softmax_grad_tiling_key.h。

5.5 数据检测：bfloat16 平台校验

README 强调：对不支持AscendC::Cast()bfloat16 向 float32 转换的硬件（即非 Atlas A2 系列），在参数校验阶段直接失败并打印报错提示。对应实现位于 aclnn 入口的 CheckNpuArchIsSupportBf16：

static inline bool CheckNpuArchIsSupportBf16(void) { // 禁用 GetSocVersion() 判断芯片平台，改用 NPU 架构判断；910B 系列对应 DAV_2201 return GetCurrentPlatformInfo().GetCurNpuArch() == NpuArch::DAV_2201; }

即仅当当前 NPU 架构为DAV_2201（对应 Atlas A2/910B 系列）时才允许 bfloat16，否则报ACLNN_ERR_PARAM_INVALID。

六、Kernel 侧详细设计

6.1 核函数框架与模板实例化

Kernel 入口为 log_softmax_grad.cpp，是一个带三个模板参数的__global__ __aicore__核函数：

template <uint32_t SCH_MOD, bool IS_SMALL, bool IS_CONTIGUOUS> __global__ __aicore__ void log_softmax_grad(GM_ADDR x, GM_ADDR y, GM_ADDR z, GM_ADDR workspace, GM_ADDR tiling) { REGISTER_TILING_DEFAULT(LogSoftmaxGradTilingData); GET_TILING_DATA_WITH_STRUCT(LogSoftmaxGradTilingData, tilingData, tiling); TPipe pipe; constexpr int BUF_NUM = IS_SMALL ? 1 : 2; if constexpr (SCH_MOD == REDUCE_MID) { ReduceMid<DTYPE_INPUT_DY, IS_SMALL, IS_CONTIGUOUS, BUF_NUM> op; op.Init(tilingData, &pipe); op.Process(x, y, z); } else if constexpr (SCH_MOD == REDUCE_TAIL) { ReduceTail<...> op; op.Init(tilingData, &pipe); op.Process(x, y, z); } else if constexpr (SCH_MOD == NO_NEED_REDUCE) { NoNeedReduce<DTYPE_INPUT_DY> op; op.Init(tilingData, &pipe); op.Process(x, y, z); } }

即根据 tilingkey 在编译期实例化不同的核函数对象，运行期执行 Init 与 Process 两个阶段，其中 Process 内部包含数据搬入（CopyIn）、计算（Compute）、搬出（CopyOut）三个阶段。小 shape 场景（IS_SMALL=true）buffer 数降为 1，节省 UB 开销。

6.2 基类 LogSoftmaxGradBase：统一的搬入/搬出与 float32 化

基类 log_softmax_grad_base.h 提供了通用的数据通路：

• CopyInDy / CopyInX：根据 golden 实现，所有数据都需要 cast 为 float32 计算。对于 float16/bfloat16，先把数据搬入 UB 块后半部分（ReinterpretCast<T>()[singleBufElems_]），再通过Cast转成 float32 铺满整个 UB 块，实现"搬入空间与 Cast 结果空间复用"；对于 float32 则直接搬入。连续数据用DataCopy（按 block 对齐），非连续数据用DataCopyPad处理；
• CopyOutZ：计算结果（float32）在搬出前按需Cast回原类型（CAST_RINT舍入），连续场景先搬整 block 再处理剩余尾部，非连续场景走DataCopyPad；
• 通过SetFlag / WaitFlag与PipeBarrier同步 MTE2（搬入）、V（向量计算）、MTE3（搬出）流水阶段。

6.3 NO_NEED_REDUCE：免 reduce 通路

no_need_reduce.h 实现mergedDim1 == 1的场景。Process 以jumpStride = blockNum * singleBufElems步长做核间轮转，先处理完整块再处理剩余尾部。Compute 如上文所述三步完成：

Exp(x, x, count); // x = exp(x) Mul(x, dy, x, count); // x = dy * exp(x) Sub(z, dy, x, count); // z = dy - dy * exp(x) = (dy - 1) * exp(x)

6.4 REDUCE_TAIL：reduce 轴在尾部

reduce_tail.h 处理mergedDim2 == 1（reduce 轴在尾部）的场景，按dim2Tile是否超过单 buffer 容量分成小 shape（ProcSmallDim2）与大 shape（ProcLargeDim2）两条路径：

• 小 shape：整行搬入后，用BinaryAddReduceDy对每行 dy 做累加 reduce，BroadcastSum把行和广播回整行，RectMul完成"行和 × exp(x)"的矩形乘，最后Sub得到结果；当dim2Tile不是 32 字节 block 的整数倍时，走ComputeSmallUseTrans，借助TransDataTo5HD转置 +BinaryAddReduceAndCopy完成 reduce 与广播；
• 大 shape：先在t1_临时区把整列 dy 累加得到标量行和（Duplicate清零 + 逐块Add），随后逐块搬入 dy/x，Brcb广播行和，RectMul乘 exp(x)，Sub输出。

过程中用到的MAX_REPEAT_TIME = 255、REPEAT_SIZE = 256（即单次向量指令的 repeat 上限与 repeat 长度）等常量，体现了向量指令的"repeat 切分"编程约束（见 log_softmax_grad_base.h）。

6.5 REDUCE_MID：reduce 轴在中间

reduce_mid.h 处理最一般的情形。Process 依据dim0Tile是否非零分为两条路径：

• TileDim0（dim0Tile != 0）：先对mergedDim0方向分 tile，每个核处理若干 dim0 切片，内部按ProcSmallDim1完成"每行 dy reduce → 广播 → 乘 exp(x) → 减"的流水；
• LoopDim0（dim0Tile == 0）：把任务平铺成mergedDim0 × dim2LoopTime个任务，核间轮转调度，每个任务处理一块dim1 × dim2Tile的数据；dim2Remained剩余部分在 ProcessDim2Remained 单独处理。

非连续场景下，inParams_ / outParams_会按dim2Tile、mergedDim2计算 srcStride/dstStride 等DataCopyExtParams，保证跨行搬移的地址正确性。

6.6 核函数级 UT 验证

Kernel 级测试 test_log_softmax_grad.cpp 采用单核（block_dim = 1）冒烟方式，通过FillTiling手工构造LogSoftmaxGradTilingData，让整个张量在单个 tile 内被完整处理，从而覆盖NO_NEED_REDUCE / REDUCE_TAIL / REDUCE_MID三条 kernel 分支的真实计算路径；其枚举定义与 log_softmax_grad_tiling_key.h 严格对应。更全面的正确性（多 dtype、多 shape、与 torch golden 逐元素对比）由 op_api 级 UT test_aclnn_logsoftmax_backward.cpp 承担。此外 host 侧还有 tiling 与 infershape 的单元测试（test_log_softmax_grad_tiling.cpp、test_log_softmax_grad_infershape.cpp），以及 kernel 数据生成/比对脚本 gen_data.py 与 compare_data.py，共同构成从定义、推导、tiling 到 kernel 的完整验证闭环。

七、目录结构与二次开发指引

LogSoftmaxGrad 算子的完整代码布局如下，可供二次开发参考：

experimental/activation/log_softmax_grad/ ├── README.md # 算子主文档（本文主体） ├── docs/ │ └── aclnnLogSoftmaxBackward.md # aclnn 接口文档 ├── examples/ │ └── test_aclnn_log_softmax_grad.cpp # 完整可编译调用示例 ├── op_host/ │ ├── log_softmax_grad_def.cpp # 算子定义与注册（AICore config: ascend910b） │ ├── log_softmax_grad_infershape.cpp # shape 推导（输出继承 dy shape） │ ├── log_softmax_grad_tiling.cpp # tiling：合轴、分核、缓冲区规划 │ └── op_api/ # aclnn 两段式接口实现与参数校验 ├── op_kernel/ │ ├── log_softmax_grad.cpp # 核函数入口（按 tilingkey 实例化） │ ├── log_softmax_grad_base.h # 基类：CopyIn/Cast/CopyOut 数据通路 │ ├── no_need_reduce.h # 免 reduce 分支 │ ├── reduce_tail.h # reduce 轴在尾部 │ ├── reduce_mid.h # reduce 轴在中间 │ ├── log_softmax_grad_tiling_data.h # tiling 数据结构 │ └── log_softmax_grad_tiling_key.h # tilingkey 模板参数声明 └── tests/ut/ # op_api / op_host / op_kernel 三级 UT

从源码结构可以推断：该算子的设计体现了 CANN AscendC 算子开发的通用范式——"算子定义（def）→ shape 推导（infershape）→ tiling（host 侧分块调度）→ kernel（device 侧流水计算）→ 两级 UT（golden 对比 + 分支冒烟）"五件套，并且通过模板参数（SCH_MOD / IS_SMALL / IS_CONTIGUOUS）在编译期展开不同调度路径，兼顾代码复用与运行期性能。若要在其他产品形态上适配该算子，需要重点核对 bfloat16 的 Cast 平台支持（DAV_2201架构判断）以及 UB block 大小、核数等平台动态参数的获取逻辑。

• 人工智能
• 算子库
• 深度学习
• CANN
• Ascend