AdjustSoftMaxRes
功能说明
本接口用于调整SoftMax的计算结果为指定的值。主要用于对SoftMax相关计算结果做后处理。当输入的max中存在指定的值的时候,会调整对应的softmaxres中的结果为输入的自定义的值。以上调整方式为按行进行,即当max某一行的值为某个值时,调整当前softmaxres对应一行的值都为输入的值。
为方便理解,通过python脚本实现的方式,表达其计算公式如下,其中res是输入也是输出,max\from\to\res_shape都为输入。
def adjust_softmax_res(res, max, from, to, res_shape): for i in res_shape[0]: if max[i] == from: for j in res_shape[1]: res[i][j] = to return
函数原型
1 2 |
template <typename T1, typename T2, bool isDataFormatNZ = false> __aicore__ inline bool AdjustSoftMaxRes(const LocalTensor<T1>& softMaxRes, const LocalTensor<T2>& maxTensor, const uint32_t from, const T1 to, const SoftMaxShapeInfo& softmaxShapeInfo) |
参数说明
参数名 |
描述 |
---|---|
T1 |
softMaxRes的数据类型。 |
T2 |
maxTensor的数据类型。 |
isDataFormatNZ |
当前输入输出的数据格式是否为NZ格式,默认数据格式为ND,即默认取值为false。 |
参数名 |
输入/输出 |
描述 |
---|---|---|
softMaxRes |
输入/输出 |
既是源操作数也是目的操作数 类型为LocalTensor,支持的TPosition为VECIN/VECCALC/VECOUT。 LocalTensor数据结构的定义请参考LocalTensor Atlas A2训练系列产品/Atlas 800I A2推理产品,支持的数据类型为:half/float Atlas推理系列产品(Ascend 310P处理器)AI Core,支持的数据类型为:half/float last轴长度需要32Byte对齐 一般为softmax计算的输出结果 |
maxTensor |
输入 |
源操作数。 类型为LocalTensor,支持的TPosition为VECIN/VECCALC/VECOUT。 Atlas A2训练系列产品/Atlas 800I A2推理产品,支持的数据类型为:half/float Atlas推理系列产品(Ascend 310P处理器)AI Core,支持的数据类型为:half/float softmax计算过程中reducemax的结果。
|
from |
输入 |
源操作数,类型为uint32_t 需要判断的maxTensor中的值。需要注意的是,由于maxTensor中的值均为浮点数类型,因此此处需要填入的值为浮点数类型对应十六进制的值。比如当需要判断maxTensor是否有1.0这个值时,from值需要填入1.0对应的十六进制值0x3f800000。 |
to |
输入 |
源操作数,类型和softMaxRes的数据类型保持一致。 需要往softMaxRes中填充的值。 |
softmaxShapeInfo |
输入 |
softMaxRes的shape信息,结构定义如下: struct SoftMaxShapeInfo { uint32_t srcM; // 非尾轴乘积长度 uint32_t srcK; // 尾轴长度,必须32Byte对齐 uint32_t oriSrcM; // 原始非尾轴乘积长度 uint32_t oriSrcK; // 原始尾轴长度 }; 需要注意,目前仅支持ND输入。 |
返回值
bool类型,当返回true时,表示maxTensor中存在需要判断的值,若返回false,则表示maxTensor中不存在需要判断的值。
支持的型号
Atlas A2训练系列产品/Atlas 800I A2推理产品
Atlas推理系列产品(Ascend 310P处理器)AI Core
注意事项
- 操作数地址偏移对齐要求请参见通用约束。
调用示例
本样例中需要对SoftMax计算结果做后处理,判断maxTensor中是否存在0xFF7FFFFF,如果存在刷新对应结果为0。其中,输入softMaxRes的shape大小为[320,64],中间计算结果maxTensor的shape大小为[320,8],数据类型均为float,需要对SoftMax计算结果做后处理,判断maxTensor中是否存在0xFF7FFFFF,如果存在刷新对应结果为0。
#include "kernel_operator.h" 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() { AscendC::LocalTensor<T> srcLocal = inQueueSrc.AllocTensor<T>(); AscendC::DataCopy(srcLocal, src1Global, height * width); inQueueSrc.EnQue(srcLocal); } __aicore__ inline void Compute() { AscendC::LocalTensor<T> srcLocal = inQueueSrc.DeQue<T>(); AscendC::LocalTensor<T> sumTempLocal = sumQueue.AllocTensor<T>(); AscendC::LocalTensor<T> maxTempLocal = maxQueue.AllocTensor<T>(); AscendC::LocalTensor<T> dstLocal = outQueueDst.AllocTensor<T>(); AscendC::SoftMaxShapeInfo srcShape = {height, width, height, width}; AscendC::SoftMax<T>(dstLocal, sumTempLocal, maxTempLocal, srcLocal, tiling, srcShape); AscendC::AdjustSoftMaxRes<T, T>(dstLocal, maxTempLocal, 0xFF7FFFFF, 0.0, srcShape); outQueueDst.EnQue<T>(dstLocal); maxQueue.FreeTensor(maxTempLocal); sumQueue.FreeTensor(sumTempLocal); inQueueSrc.FreeTensor(srcLocal); } __aicore__ inline void CopyOut() { AscendC::LocalTensor<T> dstLocal = outQueueDst.DeQue<T>(); AscendC::DataCopy(dstGlobal, dstLocal, height * width); outQueueDst.FreeTensor(dstLocal); } private: AscendC::TPipe pipe; AscendC::TQue<AscendC::QuePosition::VECIN, 1> inQueueSrc; AscendC::TQue<AscendC::QuePosition::VECIN, 1> maxQueue; AscendC::TQue<AscendC::QuePosition::VECIN, 1> sumQueue; AscendC::TQue<AscendC::QuePosition::VECOUT, 1> outQueueDst; AscendC::GlobalTensor<T> src1Global, dstGlobal; uint32_t elementNumPerBlk = 0; uint32_t width = 64; uint32_t height = 320; SoftMaxTiling tiling; }; extern "C" __global__ __aicore__ void softmax_kernel_float( __gm__ uint8_t *srcGm, __gm__ uint8_t *dstGm, __gm__ uint8_t *tiling) { GET_TILING_DATA(tilingData, tiling); KernelSoftmax<float> op; op.Init(srcGm, dstGm, tilingData.softmaxTilingData); op.Process(); }