SoftMax

功能说明

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

def softmax(src):
    #基于last轴进行rowmax(按行取最大值)处理
    max = np.max(src, axis=-1, keepdims=True)
    sub = src - max
    exp = np.exp(sub)
    #基于last轴进行rowsum(按行求和)处理
    sum = np.sum(exp, axis=-1, keepdims=True)
    dst = exp / sum
    return dst, max, sum

当输入的数据排布格式不同时，内部的reduce过程会有所不同：当输入为ND格式时，内部的reduce过程按last轴进行；当输入为NZ格式时，内部的reduce过程按照last轴和first轴进行，reduce过程如下图所示：

图1 ND格式的reduce过程

图2 NZ格式的reduce过程

函数原型

• 接口框架申请临时空间

  template <typename T, bool isReuseSource = false, bool isBasicBlock = false, bool isDataFormatNZ = false>

  __aicore__ inline void SoftMax(const LocalTensor<T>& dstTensor, const LocalTensor<T>& sumTensor,

  const LocalTensor<T>& maxTensor, const LocalTensor<T>& srcTensor, const SoftMaxTiling& tiling,

  const SoftMaxShapeInfo& softmaxShapeInfo = {})

  template <typename T, bool isReuseSource = false, bool isBasicBlock = false, bool isDataFormatNZ = false>

  __aicore__ inline void SoftMax(const LocalTensor<half>& dst, const LocalTensor<float>& sumTensor,

  const LocalTensor<float>& maxTensor, const LocalTensor<half>& src, const SoftMaxTiling& tiling,

  const SoftMaxShapeInfo& softmaxShapeInfo = {})

  template <typename T, bool isReuseSource = false, bool isBasicBlock = false>

  __aicore__ inline void SoftMax(const LocalTensor<T>& dstTensor, const LocalTensor<T>& srcTensor, const SoftMaxTiling& tiling,

  const SoftMaxShapeInfo& softmaxShapeInfo = {})

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

  template <typename T, bool isReuseSource = false, bool isBasicBlock = false, bool isDataFormatNZ = false>

  __aicore__ inline void SoftMax(const LocalTensor<T>& dstTensor, const LocalTensor<T>& sumTensor,

  const LocalTensor<T>& maxTensor, const LocalTensor<T>& srcTensor, const LocalTensor<uint8_t>& sharedTmpBuffer,

  const SoftMaxTiling& tiling, const SoftMaxShapeInfo& softmaxShapeInfo = {})

  template <typename T, bool isReuseSource = false, bool isBasicBlock = false, bool isDataFormatNZ = false>

  __aicore__ inline void SoftMax(const LocalTensor<half>& dst, const LocalTensor<float>& sumTensor,

  const LocalTensor<float>& maxTensor, const LocalTensor<half>& src, const LocalTensor<uint8_t>& sharedTmpBuffer,

  const SoftMaxTiling& tiling, const SoftMaxShapeInfo& softmaxShapeInfo = {})

  template <typename T, bool isReuseSource = false, bool isBasicBlock = false>

  __aicore__ inline void SoftMax(const LocalTensor<T>& dstTensor, const LocalTensor<T>& 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推理系列产品AI Core

Atlas 200I/500 A2推理产品

注意事项

• srcTensor和dstTensor的Tensor空间可以复用。
• sumTensor和maxTensor为输出，并且last轴长度必须固定32Byte，非last轴大小需要和srcTensor以及dstTensor保持一致。
• sumTensor和maxTensor的数据类型需要保持一致。

• 操作数地址偏移对齐要求请参见通用约束。

调用示例

本样例中输入srcTensor和输出dstTensor的shape大小为[320,64]，中间计算结果sumTensor和maxTensor的shape大小为[320,16]，数据类型均为half，输入输出的数据排布格式为ND，srcTensor和dstTensor空间不复用，不使能基本块。

#include "kernel_operator.h"

namespace AscendC {
template <typename T> class KernelSoftmax {
public:
    __aicore__ inline KernelSoftmax() {}
    __aicore__ inline void Init(__gm__ uint8_t* srcGm, __gm__ uint8_t* dstGm, const SoftMaxTiling& tilingData)
    {
        elementNumPerBlk = 32 / sizeof(T);
        src1Global.SetGlobalBuffer((__gm__ T*)srcGm);
        dstGlobal.SetGlobalBuffer((__gm__ T*)dstGm);
        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>();
        DataCopy(srcLocal, src1Global, height * width);
        inQueueSrc.EnQue(srcLocal);
    }
    __aicore__ inline void Compute()
    {
        LocalTensor<T> srcLocal = inQueueSrc.DeQue<T>();
        LocalTensor<T> sumTempLocal = sumQueue.AllocTensor<T>();
        LocalTensor<T> maxTempLocal = maxQueue.AllocTensor<T>();
        LocalTensor<T> dstLocal = outQueueDst.AllocTensor<T>();

        SoftMaxShapeInfo srcShape = { height, width, height, width };
        SoftMax<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> src1Global, dstGlobal;
    uint32_t elementNumPerBlk = 0;
    uint32_t width = 64;
    uint32_t height = 320;
    SoftMaxTiling tiling;
};
} // namespace AscendC

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