SimpleSoftMax

功能说明

与SoftMax接口不同，该接口内部没有reduce过程计算sum和max数据，而是使用计算好的sum和max数据对输入tensor做Softmax计算。计算公式如下：

为方便理解，通过python脚本实现的方式，表达其计算公式如下，其中src、max、sum是源操作数（输入），dst为目的操作数（输出）。

def simple_softmax(src, max, sum):
    dst = np.exp(src - max)/sum
return dst

实现原理

以float类型，ND格式，shape为[m, k]的输入Tensor为例，描述SimpleSoftMax高阶API内部算法框图，如下图所示。

图1 SimpleSoftMax算法框图

计算过程分为如下几步，均在Vector上进行：

1.sub步骤：对输入x的所有数据按行减去输入的max；

2.exp步骤：对sub之后的所有数据求exp；

3.div步骤：对exp结果的所有数据按行除以输入的sum，得到结果；

函数原型

• 接口框架申请临时空间

  1 2

  template <typename T, bool isReuseSource = false, bool isBasicBlock = false, bool isDataFormatNZ = false> __aicore__ inline void SimpleSoftMax(const LocalTensor<T>& dstTensor, const LocalTensor<T>& inSumTensor, const LocalTensor<T>& inMaxTensor, const LocalTensor<T>& srcTensor, const SoftMaxTiling& tiling, const SoftMaxShapeInfo& softmaxShapeInfo = {})

  1 2

  template <typename T, bool isReuseSource = false, bool isBasicBlock = false, bool isDataFormatNZ = false> __aicore__ inline void SimpleSoftMax(const LocalTensor<half>& dstTensor, const LocalTensor<float>& inSumTensor, const LocalTensor<float>& inMaxTensor, const LocalTensor<half>& srcTensor, const SoftMaxTiling& tiling, const SoftMaxShapeInfo& softmaxShapeInfo = {})

• 通过sharedTmpBuffer入参传入临时空间

  1 2

  template <typename T, bool isReuseSource = false, bool isBasicBlock = false, bool isDataFormatNZ = false> __aicore__ inline void SimpleSoftMax(const LocalTensor<T>& dstTensor, const LocalTensor<T>& inSumTensor, const LocalTensor<T>& inMaxTensor, const LocalTensor<T>& srcTensor, const LocalTensor<uint8_t>& sharedTmpBuffer, const SoftMaxTiling& tiling, const SoftMaxShapeInfo& softmaxShapeInfo = {})

  1 2

  template <typename T, bool isReuseSource = false, bool isBasicBlock = false, bool isDataFormatNZ = false> __aicore__ inline void SimpleSoftMax(const LocalTensor<half>& dstTensor, const LocalTensor<float>& inSumTensor, const LocalTensor<float>& inMaxTensor, const LocalTensor<half>& srcTensor, const LocalTensor<uint8_t>& sharedTmpBuffer, const SoftMaxTiling& tiling, const SoftMaxShapeInfo& softmaxShapeInfo = {})

由于该接口的内部实现中涉及复杂的计算，需要额外的临时空间来存储计算过程中的中间变量。临时空间支持接口框架申请和开发者通过sharedTmpBuffer入参传入两种方式。

• 接口框架申请临时空间，开发者无需申请，但是需要预留临时空间的大小。

• 通过sharedTmpBuffer入参传入，使用该tensor作为临时空间进行处理，接口框架不再申请。该方式开发者可以自行管理sharedTmpBuffer内存空间，并在接口调用完成后，复用该部分内存，内存不会反复申请释放，灵活性较高，内存利用率也较高。

接口框架申请的方式，开发者需要预留临时空间；通过sharedTmpBuffer传入的情况，开发者需要为tensor申请空间。临时空间大小BufferSize的获取方式如下：通过SoftMax/SimpleSoftMax Tiling中提供的GetSoftMaxMaxTmpSize/GetSoftMaxMinTmpSize接口获取所需最大和最小临时空间大小，最小空间可以保证功能正确，最大空间用于提升性能。

参数说明

struct SoftMaxShapeInfo {
uint32_t srcM; // 非尾轴长度的乘积
uint32_t srcK; // 尾轴长度，必须32Byte对齐
uint32_t oriSrcM; // 原始非尾轴长度的乘积
uint32_t oriSrcK;  // 原始尾轴长度
};

返回值

无

支持的型号

Atlas A2训练系列产品/Atlas 800I A2推理产品

Atlas推理系列产品（Ascend 310P处理器）AI Core

Atlas 200/500 A2推理产品

注意事项

• srcTensor和dstTensor的Tensor空间可以复用。
• inSumTensor和inMaxTensor为输入，并且last轴长度必须固定32Byte。
• inSumTensor和inMaxTensor的数据类型需要保持一致。
• 操作数地址偏移对齐要求请参见通用约束。

调用示例

#include "kernel_operator.h"

namespace AscendC {
template <typename T> class KernelSimpleSoftmax {
public:
    __aicore__ inline KernelSimpleSoftmax() {}
    __aicore__ inline void Init(__gm__ uint8_t* srcGm, __gm__ uint8_t* inMaxGm, __gm__ uint8_t* inSumGm,
        __gm__ uint8_t* dstGm, const SoftMaxTiling& tilingData)
    {
        elementNumPerBlk = 32 / sizeof(T);
        srcGlobal.SetGlobalBuffer((__gm__ T*)srcGm);
        dstGlobal.SetGlobalBuffer((__gm__ T*)dstGm);
        maxGlobal.SetGlobalBuffer((__gm__ T*)inMaxGm);
        sumGlobal.SetGlobalBuffer((__gm__ T*)inSumGm);

        pipe.InitBuffer(inQueueSrc, 1, height * width * sizeof(T));
        pipe.InitBuffer(maxQueue, 1, height * elementNumPerBlk * sizeof(T));
        pipe.InitBuffer(sumQueue, 1, height * elementNumPerBlk * sizeof(T));
        pipe.InitBuffer(outQueueDst, 1, height * width * sizeof(T));
        tiling = tilingData;
    }
    __aicore__ inline void Process()
    {
        CopyIn();
        Compute();
        CopyOut();
    }

private:
    __aicore__ inline void CopyIn()
    {
        LocalTensor<T> srcLocal = inQueueSrc.AllocTensor<T>();
        LocalTensor<T> sumTempLocal = sumQueue.AllocTensor<T>();
        LocalTensor<T> maxTempLocal = maxQueue.AllocTensor<T>();
        DataCopy(srcLocal, srcGlobal, height * width);
        DataCopy(sumTempLocal, sumGlobal, height * elementNumPerBlk);
        DataCopy(maxTempLocal, maxGlobal, height * elementNumPerBlk);
        inQueueSrc.EnQue(srcLocal);
        sumQueue.EnQue(sumTempLocal);
        maxQueue.EnQue(maxTempLocal);
    }
    __aicore__ inline void Compute()
    {
        LocalTensor<T> srcLocal = inQueueSrc.DeQue<T>();
        LocalTensor<T> sumTempLocal = sumQueue.DeQue<T>();
        LocalTensor<T> maxTempLocal = maxQueue.DeQue<T>();
        LocalTensor<T> dstLocal = outQueueDst.AllocTensor<T>();

        SoftMaxShapeInfo srcShape = { height, width, height, width };
        SimpleSoftMax<T>(dstLocal, sumTempLocal, maxTempLocal, srcLocal, tiling, srcShape);

        outQueueDst.EnQue<T>(dstLocal);
        maxQueue.FreeTensor(maxTempLocal);
        sumQueue.FreeTensor(sumTempLocal);
        inQueueSrc.FreeTensor(srcLocal);
    }
    __aicore__ inline void CopyOut()
    {
        LocalTensor<T> dstLocal = outQueueDst.DeQue<T>();
        DataCopy(dstGlobal, dstLocal, height * width);
        outQueueDst.FreeTensor(dstLocal);
    }

private:
    TPipe pipe;
    TQue<QuePosition::VECIN, 1> inQueueSrc;
    TQue<QuePosition::VECIN, 1> maxQueue;
    TQue<QuePosition::VECIN, 1> sumQueue;
    TQue<QuePosition::VECOUT, 1> outQueueDst;
    GlobalTensor<T> srcGlobal, dstGlobal;
    GlobalTensor<T> maxGlobal, sumGlobal;
    uint32_t elementNumPerBlk = 0;
    uint32_t width = 64;
    uint32_t height = 320;
    SoftMaxTiling tiling;
};
} // namespace AscendC

extern "C" __global__ __aicore__ void simple_softmax_kernel_half(__gm__ uint8_t* srcGm, __gm__ uint8_t* inMaxGm, __gm__ uint8_t* inSumGm, __gm__ uint8_t* dstGm, __gm__ uint8_t* tiling)
{
    GET_TILING_DATA(tilingData, tiling);
    AscendC::KernelSimpleSoftmax<half> op;
    op.Init(srcGm, inMaxGm, inSumGm, dstGm, tilingData.softmaxTilingData);
    op.Process();
}