CANN/asc-devkit SoftmaxFlashV3接口

SoftmaxFlashV3【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit产品支持情况Ascend 950PR/Ascend 950DT支持Atlas A3 训练系列产品/Atlas A3 推理系列产品支持Atlas A2 训练系列产品/Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品AI Core不支持Atlas 推理系列产品Vector Core不支持Atlas 训练系列产品不支持功能说明SoftmaxFlash增强版本对应Softmax PASA算法。将输入tensor[m0, m1, ..., mt, n]t大于或等于0的非尾轴长度m0, m1, ..., mt相乘的结果看作m则输入tensor的shape看作[m, n]。对输入tensor x的尾轴进行切分分块个数为splitMeanCnt切分后的tensor为x_cnti。按如下公式进行计算其中x、inmax、insum、inmean为输入M、S、E、A均为输出。update为falseupdate为true本接口当前只支持ND格式的输入内部的reduce过程按last轴处理。为方便理解通过Python伪代码实现的方式表达其计算公式如下。其中repeatSize为64elementNumPerBlk/BlkcntPerRepeat为8splitMeanCnt为8src、inmean、inmax、 insum、update为输入dst、x_mean、x_sum、x_max、exp_max为输出。def softmax_flash_3(src, height, width, loopCnt, alpha, baseK, inmaxNone, insumNone, inmeanNone, updateFalse): scalar alpha / (1 - alpha) #(m,n)-(m,64) tmpbuffer0 BlockReduceSum(repeatSize, repeatSize, elementNumPerBlk) remain int(width / repeatSize - BlkcntPerRepeat) tmpbuffer0 Add(tmpbuffer0, src, remain, repeatSize * elementNumPerBlk, width) #(m,64)-(m,8) tmpbuffer0 BlockReduceSum(1, elementNumPerBlk, elementNumPerBlk) #width baseK * splitMeanCnt rowMeanLocal tmpbuffer0 / baseK rowMeanGlobal np.mean(src, axis(-1), keepdimsTrue) rowMeanGlobalTmp (rowMeanGlobal - rowMeanLocal) * scalar src src - rowMeanGlobalTmp if update False: x_mean rowMeanGlobal maxTmp np.max(src, axis-1, keepdimsTrue) shiftCurr (rowMeanGlobal - x_mean) * scalar x_max shiftCurr maxTmp maxTmp x_max - shiftCurr x_sub src - maxTmp dst np.exp(x_sub) x_sum np.sum(dst, axis-1, keepdimsTrue) exp_max None return dst, x_max, x_sum, x_mean, exp_max else: x_mean (rowMeanGlobal inmean * (loopCnt - 1)) / loopCnt maxTmp np.max(src, axis-1, keepdimsTrue) shiftCurr (rowMeanGlobal - x_mean) * scalar shiftPrev (inmean - x_mean) * scalar x_max shiftCurr maxTmp maxTmp shiftPrev inmax x_max np.max(np.concatenate((x_max, maxTmp), axis(-1)), axis(-1), keepdimsTrue) maxTmp x_max - shiftCurr x_sub src - maxTmp dst np.exp(x_sub) exp_max np.exp(inmax - x_max shiftPrev) x_sum np.sum(x_exp, axis-1, keepdimsTrue) x_sum exp_max * insum x_sum return x_exp, x_max, x_sum, x_mean, exp_max函数原型接口框架申请临时空间template typename T, typename U, bool isUpdate false, bool isReuseSource false, bool isBasicBlock false, bool isDataFormatNZ false, const SoftmaxConfig config SOFTMAX_DEFAULT_CFG __aicore__ inline void SoftmaxFlashV3(const LocalTensorT dstTensor, const LocalTensorU meanTensor, const LocalTensorU expSumTensor, const LocalTensorU maxTensor, const LocalTensorT srcTensor, const LocalTensorT expMaxTensor, const LocalTensorU inMeanTensor, const LocalTensorU inExpSumTensor, const LocalTensorU inMaxTensor, const SoftMaxTiling tiling, const SoftMaxParams params)通过sharedTmpBuffer入参传入临时空间template typename T, typename U, bool isUpdate false, bool isReuseSource false, bool isBasicBlock false, bool isDataFormatNZ false, const SoftmaxConfig config SOFTMAX_DEFAULT_CFG __aicore__ inline void SoftmaxFlashV3(const LocalTensorT dstTensor, const LocalTensorU meanTensor,const LocalTensorU expSumTensor, const LocalTensorU maxTensor, const LocalTensorT srcTensor,const LocalTensorT expMaxTensor, const LocalTensorU inMeanTensor, const LocalTensorU inExpSumTensor, const LocalTensorU inMaxTensor, const LocalTensoruint8_t sharedTmpBuffer, const SoftMaxTiling tiling, const SoftMaxParams params)由于该接口的内部实现中涉及复杂的计算需要额外的临时空间来存储计算过程中的中间变量。临时空间支持接口框架申请和开发者通过sharedTmpBuffer入参传入两种方式。接口框架申请临时空间开发者无需申请但是需要预留临时空间的大小。通过sharedTmpBuffer入参传入使用该tensor作为临时空间进行处理接口框架不再申请。该方式开发者可以自行管理sharedTmpBuffer内存空间并在接口调用完成后复用该部分内存内存不会反复申请释放灵活性较高内存利用率也较高。接口框架申请的方式开发者需要预留临时空间通过sharedTmpBuffer传入的情况开发者需要为tensor申请空间。临时空间大小BufferSize的获取方式如下通过SoftmaxFlashV3 Tiling接口中提供的GetSoftMaxFlashV3MaxMinTmpSize接口获取所需最小和最大临时空间大小最小空间可以保证功能正确最大空间用于提升性能。参数说明表1模板参数说明参数名描述T输入srcTensor及输出dstTensor、expMaxTensor操作数的数据类型。支持的数据类型为half。U输入inMeanTensor、inExpSumTensor、inMaxTensor及输出meanTensor、expSumTensor、maxTensor操作数的数据类型。支持的数据类型为float。isUpdate是否开启update为true的计算。isReuseSource该参数预留传入默认值false即可。isBasicBlock该参数预留传入默认值false即可。isDataFormatNZ该参数预留传入默认值false即可。config该参数预留传入默认值SOFTMAX_DEFAULT_CFG即可。表2接口参数说明参数名输入/输出描述dstTensor输出目的操作数。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。dstTensor的shape和源操作数srcTensor一致。meanTensor输出目的操作数。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。用于保存softmax计算过程中平均值的结果。meanTensor的last轴长度固定为32Byte即一个datablock长度。该datablock中的所有数据为同一个值。比如float数据类型下该datablock中的8个数均为相同的reducesum求平均后的值。非last轴的长度与dstTensor保持一致。expSumTensor输出目的操作数。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。用于保存softmax计算过程中reducesum的结果。expSumTensor的last轴长度固定为32Byte即一个datablock长度。该datablock中的所有数据为同一个值。比如float数据类型下该datablock中的8个数均为相同的reducesum的值。非last轴的长度与dstTensor保持一致。maxTensor输出目的操作数。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。用于保存softmax计算过程中reducemax的结果。maxTensor的last轴长度固定为32Byte即一个datablock长度。该datablock中的所有数据为同一个值。比如float数据类型下该datablock中的8个数均为相同的reducemax的值。非last轴的长度与dstTensor保持一致。srcTensor输入源操作数。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。last轴长度需要32Byte对齐。expMaxTensor输出目的操作数。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。expMaxTensor的last轴长度固定为32Byte即一个datablock长度。该datablock中的所有数据为同一个值。比如half数据类型下该datablock中的16个数均为相同的值。非last轴的长度需要与dstTensor保持一致。inMeanTensor输入源操作数。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。softmax计算所需要的mean值。inMeanTensor的last轴长度固定为32Byte即一个datablock长度。该datablock中的所有数据为同一个值。比如float数据类型下该datablock中的8个数均为相同的值。非last轴的长度需要与dstTensor保持一致。inExpSumTensor输入源操作数。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。softmax计算所需要的sum值。inExpSumTensor的last轴长度固定为32Byte即一个datablock长度。该datablock中的所有数据为同一个值。比如float数据类型下该datablock中的8个数均为相同的值。非last轴的长度需要与dstTensor保持一致。inMaxTensor输入源操作数。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。softmax计算所需要的max值。inMaxTensor的last轴长度固定为32Byte即一个datablock长度。该datablock中的所有数据为同一个值。比如float数据类型下该datablock中的8个数均为相同的值。非last轴的长度需要与dstTensor保持一致。sharedTmpBuffer输入临时空间。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。该操作数的数据类型固定uint8_t。接口内部复杂计算时用于存储中间变量由开发者提供。临时空间大小BufferSize的获取方式请参考SoftmaxFlashV3 Tiling接口。tiling输入SoftmaxFlashV3接口计算所需Tiling信息Tiling信息的获取请参考SoftmaxFlashV3 Tiling接口。params输入srcTensor的shape信息和计算相关参数。SoftMaxParams类型具体定义如下方代码所示其中参数的含义为srcM非尾轴长度的乘积。srcK尾轴长度必须32Byte对齐。oriSrcM原始非尾轴长度的乘积。oriSrcK原始尾轴长度。loopCntupdate为true时公式中的循环次数loopCnt该参数大于等于1。splitMeanCnt公式中计算每一行平均值时的分块个数当前该参数仅支持取值为8。alpha公式中的计算参数推荐取值0.9375、0.96889、0.984497。注意当前本接口不支持非对齐场景因此参数srcM与oriSrcM相等参数srcK与oriSrcK相等。struct SoftMaxParams { uint32_t srcM; uint32_t srcK; uint32_t oriSrcM; uint32_t oriSrcK; uint32_t loopCnt; uint32_t splitMeanCnt; float alpha; };返回值说明无约束说明操作数地址对齐要求请参见通用地址对齐约束。对于输入srcTensor需要满足尾轴长度n大于等于512同时n是64的倍数非尾轴长度的乘积m为8的倍数。srcTensor和dstTensor的Tensor的空间可以复用meanTensor和inMeanTensor的空间可以复用maxTensor和inMaxTensor的空间可以复用expSumTensor和inExpSumTensor的空间可以复用。meanTensor、expSumTensor、maxTensor、expMaxTensor、inMeanTensor、inExpSumTensor、inMaxTensor的Tensor空间last轴长度必须是32字节。不支持sharedTmpBuffer与源操作数和目的操作数地址重叠。调用示例本样例中输入srcTensor和输出dstTensor的shape大小为[8, 1024]输入inMeanTensor、inExpSumTensor、inMaxTensor的shape大小为[8, 8]数据类型为float输出expMaxTensor的shape大小为[8, 16]数据类型为half输入和输出的数据排布格式为NDsrcTensor和dstTensor空间不复用模板参数isUpdate为true。// dstLocal: 存放SoftMax计算结果的Tensor // meanLocal存放softmax计算过程中平均值的结果 // expSumLocal存放softmax计算过程中reducesum的结果 // maxLocal存放softmax计算过程中reducemax的结果 // srcLocal存放SoftMax计算的输入Tensor // expMaxLocal存放inmax与reducemax差值的e的指数幂的结果 // inMeanLocal存放softmax计算所需要的mean值 // inExpSumLocal存放softmax计算所需要的sum值 // inMaxLocal存放softmax计算所需要的max值 // sharedTmpBuffer: 存放SoftMax计算过程中临时缓存的Tensor // softmaxTiling存放SoftMax计算所需Tiling信息可通过SoftMaxFlashV3TilingFunc接口获取 AscendC::SoftMaxParams params( /* 非尾轴长度的乘积 */ srcM, /* 尾轴长度必须32Bytes对齐 */ srcK, /* 原始非尾轴长度的乘积 */ oriSrcM, /* 原始尾轴长度 */ oriSrcK, /* 循环次数update为true时大于等于1 */ loopCn, /* 每一行平均值时的分块个数仅支持为8 */ splitMeanCnt, /* 计算参数推荐取值0.9375、0.96889、0.984497 */ alpha); // 通过sharedTmpBuffer入参传入临时空间 AscendC::SoftmaxFlashV3T, U, true( dstLocal, meanLocal, expSumLocal, maxLocal, srcLocal, expMaxLocal, inMeanLocal, inExpSumLocal, inMaxLocal, sharedTmpBuffer, tiling, params); // 接口框架申请临时空间 AscendC::SoftmaxFlashV3T, U, true( dstLocal, meanLocal, expSumLocal, maxLocal, srcLocal, expMaxLocal, inMeanLocal, inExpSumLocal, inMaxLocal, tiling, params);结果示例如下输入数据(srcLocal) [[ 0. 0.001221 0.002441 ... 2.496 2.498 2.498 ] [ 2.5 2.502 2.502 ... 4.996 4.996 5. ] [ 5. 5. 5.004 ... 7.496 7.496 7.5 ] ... [12.5 12.5 12.5 ... 15. 15. 15. ] [15. 15. 15. ... 17.5 17.5 17.5 ] [17.5 17.5 17.5 ... 20. 20. 20. ]] 输入数据(inMeanLocal/inExpSumLocal/inMaxLocal) [[0.5 0.5 0.5 0.5 0.5 0.5 0.5 0.5] [0.5 0.5 0.5 0.5 0.5 0.5 0.5 0.5] ... [0.5 0.5 0.5 0.5 0.5 0.5 0.5 0.5] [0.5 0.5 0.5 0.5 0.5 0.5 0.5 0.5]] 输出数据(dstLocal) [[0.0049 0.004906 0.004913 ... 0.998 1. 1. ] [0.00488 0.00489 0.00489 ... 0.996 0.996 1. ] [0.004868 0.004868 0.004887 ... 0.996 0.996 1. ] ... [0.004894 0.004894 0.004894 ... 1. 1. 1. ] [0.00472 0.00472 0.00472 ... 1. 1. 1. ] [0.004684 0.004684 0.004684 ... 1. 1. 1. ]]【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考