Softmax样例

概述

本样例在激活函数场景下，基于SoftMax高阶API对输入Tensor按行做SoftMax计算，并使用AdjustSoftMaxRes对计算结果做后处理。SoftMax将输入tensor[m, n]按行执行reducemax、sub、exp、reducesum、div等步骤，得到归一化概率分布，常用于注意力机制和分类任务的输出层。

接口提示： 除本样例使用的 SoftMax 接口外，Ascend C 还提供了 SimpleSoftMax 接口，使用预计算的sum和max数据对输入Tensor做SoftMax计算，内部无reduce过程。将 softmax.asc 中 AscendC::SoftMax 替换为 AscendC::SimpleSoftMax 即可切换。

AdjustSoftMaxRes用于对SoftMax计算结果做后处理，当输入的max中存在指定值时，按行调整对应softmaxres结果为自定义的值。本样例使用float数据类型，输入Tensor shape为[32, 32]，完成SoftMax计算及AdjustSoftMaxRes后处理。

本样例支持的产品及CANN软件版本

产品	CANN软件版本
Ascend 950PR/Ascend 950DT	>= CANN 9.1.0
Atlas A3 训练系列产品/Atlas A3 推理系列产品	>= CANN 9.0.0
Atlas A2 训练系列产品/Atlas A2 推理系列产品	>= CANN 9.0.0

目录结构介绍

├── softmax
│   ├── scripts
│   │   └── gen_data.py         // 输入数据和真值数据生成脚本
│   ├── CMakeLists.txt          // 编译工程文件
│   ├── data_utils.h            // 数据读入写出函数
│   ├── README.md               // 样例说明文档
│   └── softmax.asc             // Ascend C样例实现 & 调用样例

样例描述

• 样例功能：

  本样例对输入 Tensor 按行做 SoftMax 计算，并使用 AdjustSoftMaxRes 对计算结果做后处理。当输入的 max 中存在指定的值（0xFF7FFFFF，即 float 类型的最大有限值）时，调整输出中对应位置的数据为自定义的值（0.0，即浮点数零）。该机制常用于注意力掩码（attention mask）场景，将无效位置的 softmax 输出置零。

• 样例规格：

• 样例实现：
  本样例中实现的是固定shape为输入x[32, 32]，输出y[32, 32]、max[32, 8]、sum[32, 8]的softmax样例。

  • Kernel实现
    核心计算步骤：将输入数据搬入后，调用 AscendC::SoftMax 完成 SoftMax 计算，再调用 AscendC::AdjustSoftMaxRes 对结果做后处理，最后将结果搬出。

  • Tiling实现
    softmax样例的tiling实现流程如下：首先对shape按照行数进行分核，使用平均分配法先按照核数向上对齐分配，确定主核的计算行数，再确定尾核计算行数，对主核计算的shape调用GetSoftMaxMinTmpSize获取API所需临时缓冲区大小，再调用SoftMax高阶API的tiling函数获取API所需tiling参数，尾核计算所需的高阶API的tiling由kernel侧自行计算。

  • 调用实现
    使用内核调用符<<<>>>调用核函数。

编译运行

在本样例根目录下执行如下步骤，编译并执行样例。

• 配置环境变量 请根据当前环境上CANN开发套件包的安装方式，配置环境变量。

  source ${install_path}/cann/set_env.sh
  

  说明： ${install_path} 为CANN包安装目录，未指定安装目录时默认安装至 /usr/local/Ascend 下。

• 样例执行

  在本样例目录下执行如下命令。

  mkdir -p build && cd build;
  cmake -DCMAKE_ASC_ARCHITECTURES=dav-2201 ..;make -j; # 默认npu模式
  python3 ../scripts/gen_data.py
  ./demo
  

  使用 CPU调试 或 NPU仿真 模式时，添加 -DCMAKE_ASC_RUN_MODE=cpu 或 -DCMAKE_ASC_RUN_MODE=sim 参数即可。

  示例如下：

  cmake -DCMAKE_ASC_RUN_MODE=cpu -DCMAKE_ASC_ARCHITECTURES=dav-2201 ..;make -j; # cpu调试模式
  cmake -DCMAKE_ASC_RUN_MODE=sim -DCMAKE_ASC_ARCHITECTURES=dav-2201 ..;make -j; # NPU仿真模式
  

  注意： 切换编译模式前需清理 cmake 缓存，可在 build 目录下执行 rm CMakeCache.txt 后重新 cmake。

• 编译选项说明

  选项	可选值	说明
  CMAKE_ASC_RUN_MODE	npu（默认）、cpu、sim	运行模式：NPU 运行、CPU调试、NPU仿真
  CMAKE_ASC_ARCHITECTURES	dav-2201（默认）、dav-3510	NPU 架构：dav-2201 对应 Atlas A2 训练系列产品/Atlas A2 推理系列产品和 Atlas A3 训练系列产品/Atlas A3 推理系列产品，dav-3510 对应 Ascend 950PR/Ascend 950DT

• 执行结果

  执行结果如下，说明精度对比成功。

  test pass!