CANN/asc-devkit SoftmaxFlashV2激活函数API

SoftmaxFlashV2【免费下载链接】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 训练系列产品不支持Kirin X90支持Kirin 9030支持功能说明SoftmaxFlash增强版本对应FlashAttention-2算法。将输入tensor[m0, m1, ...mt, n]t大于等于0的非尾轴长度相乘的结果看作m则输入tensor的shape看作[m, n]。对输入tensor[m,n]按行做如下计算不同的update值对应不同的计算公式其中x、inmax和insum为输入M、S、E均为输出。update为falseupdate为true当输入shape为ND格式时内部的reduce过程按last轴进行当输入shape为NZ格式时内部的reduce过程按照last轴和first轴进行reduce过程可以参考SoftMax中的图示说明。为方便理解通过Python脚本实现的方式表达其计算公式如下其中src、inmax、 insum、update为输入dst、x_sum、x_max、exp_max为输出。def softmax_flash_2(src, inmaxNone, insumNone, updateNone): if update is None: x_max np.max(src, axis-1, keepdimsTrue) x_sub src - x_max dst np.exp(x_sub) x_sum np.sum(dst, axis-1, keepdimsTrue) exp_max None return dst, x_max, x_sum, exp_max else: x_max np.max(np.concatenate((inmax, src), axis-1), axis-1, keepdimsTrue) dst np.exp(src - x_max) exp_max np.exp(inmax - x_max) x_sum np.sum(dst, axis-1, keepdimsTrue) x_sum exp_max * insum x_sum return dst, x_max, x_sum, exp_max实现原理以float类型ND格式shape为[m, k]的输入Tensor为例描述SoftmaxFlashV2高阶API内部算法框图如下图所示。图1SoftmaxFlashV2算法框图计算过程根据isUpdate是否开启分为两个分支处理均在Vector上进行。当isUpdate为False时分为如下几步reducemax步骤对输入x的每一行数据求最大值得到[m, 1]计算结果会保存到一个临时空间temp中broadcast步骤对temp中的数据[m, 1]做一个按datablock为单位的填充比如float类型下把[m, 1]扩展成[m, 8]同时输出maxsub步骤对输入x的所有数据按行减去maxexp步骤对sub之后的所有数据求exp并且输出yreducesum步骤对exp结果的每一行数据求和得到[m, 1]计算结果会保存到临时空间temp中broadcast步骤对temp[m, 1]做一个按datablock为单位的填充比如float类型下把[m, 1]扩展成[m, 8]同时输出sum。当isUpdate为True时分为如下几步reducemax步骤对输入x的每一行数据求最大值得到[m, 1]计算结果会保存到一个临时空间temp中broadcast步骤对temp中的数据[m, 1]做一个按datablock为单位的填充比如float类型下把[m, 1]扩展成[m, 8]保存为max;max步骤对输入inmax和上一步计算的max做max操作得到新的max并输出sub步骤将输入inmax和新的max相减然后做exp计算得到expmax并输出sub步骤将输入x和新的max按行相减exp步骤对sub之后的所有数据求exp并且输出yreducesum步骤对exp结果的每一行数据求和得到[m, 1]计算结果会保存到临时空间temp中broadcast步骤对temp数据[m, 1]做一个按datablock为单位的填充比如float类型下把[m, 1]扩展成[m, 8]保存到sum中mul步骤将输入insum和expmax结果相乘add步骤将相乘结果和sum相加保存到sum并输出。函数原型接口框架申请临时空间LocalTensor的数据类型相同不输出ReduceMaxtemplate typename T, bool isUpdate false, bool isReuseSource false, bool isBasicBlock false, bool isDataFormatNZ false, const SoftmaxConfig config SOFTMAX_DEFAULT_CFG __aicore__ inline void SoftmaxFlashV2(const LocalTensorT dstTensor, const LocalTensorT expSumTensor, const LocalTensorT maxTensor, const LocalTensorT srcTensor, const LocalTensorT expMaxTensor, const LocalTensorT inExpSumTensor, const LocalTensorT inMaxTensor, const SoftMaxTiling tiling, const SoftMaxShapeInfo softmaxShapeInfo {})LocalTensor的数据类型相同且输出ReduceMaxtemplate typename T, bool isUpdate false, bool isReuseSource false, bool isBasicBlock false, bool isDataFormatNZ false, const SoftmaxConfig config SOFTMAX_DEFAULT_CFG __aicore__ inline void SoftmaxFlashV2(const LocalTensorT dstTensor, const LocalTensorT outReduceMax, const LocalTensorT outExpSum, const LocalTensorT outMax, const LocalTensorT srcTensor, const LocalTensorT outExpMax, const LocalTensorT inExpSum, const LocalTensorT inMax, const SoftMaxTiling tiling, const SoftMaxShapeInfo softmaxShapeInfo {})Atlas 200I/500 A2 推理产品不支持该接口。Atlas 推理系列产品AI Core不支持该接口。LocalTensor的数据类型不同不输出ReduceMaxtemplate typename T, bool isUpdate false, bool isReuseSource false, bool isBasicBlock false, bool isDataFormatNZ false, const SoftmaxConfig config SOFTMAX_DEFAULT_CFG __aicore__ inline void SoftmaxFlashV2(const LocalTensorhalf dstTensor, const LocalTensorfloat expSumTensor, const LocalTensorfloat maxTensor, const LocalTensorhalf srcTensor, const LocalTensorhalf expMaxTensor, const LocalTensorfloat inExpSumTensor, const LocalTensorfloat inMaxTensor, const SoftMaxTiling tiling, const SoftMaxShapeInfo softmaxShapeInfo {})通过sharedTmpBuffer入参传入临时空间LocalTensor的数据类型相同不输出ReduceMaxtemplate typename T, bool isUpdate false, bool isReuseSource false, bool isBasicBlock false, bool isDataFormatNZ false, const SoftmaxConfig config SOFTMAX_DEFAULT_CFG __aicore__ inline void SoftmaxFlashV2(const LocalTensorT dstTensor, const LocalTensorT outExpSum, const LocalTensorT outMax, const LocalTensorT srcTensor, const LocalTensorT outExpMax, const LocalTensorT inExpSum, const LocalTensorT inMax, const LocalTensoruint8_t sharedTmpBuffer, const SoftMaxTiling tiling, const SoftMaxShapeInfo softmaxShapeInfo {})LocalTensor的数据类型相同且输出ReduceMaxtemplate typename T, bool isUpdate false, bool isReuseSource false, bool isBasicBlock false, bool isDataFormatNZ false, const SoftmaxConfig config SOFTMAX_DEFAULT_CFG __aicore__ inline void SoftmaxFlashV2(const LocalTensorT dstTensor, const LocalTensorT outReduceMax, const LocalTensorT expSumTensor, const LocalTensorT maxTensor, const LocalTensorT srcTensor, const LocalTensorT expMaxTensor, const LocalTensorT inExpSumTensor, const LocalTensorT inMaxTensor, const LocalTensoruint8_t sharedTmpBuffer, const SoftMaxTiling tiling, const SoftMaxShapeInfo softmaxShapeInfo {})Atlas 200I/500 A2 推理产品不支持该接口。Atlas 推理系列产品AI Core不支持该接口。LocalTensor的数据类型不同不输出ReduceMaxtemplate typename T, bool isUpdate false, bool isReuseSource false, bool isBasicBlock false, bool isDataFormatNZ false, const SoftmaxConfig config SOFTMAX_DEFAULT_CFG __aicore__ inline void SoftmaxFlashV2(const LocalTensorhalf dstTensor, const LocalTensorfloat expSumTensor, const LocalTensorfloat maxTensor, const LocalTensorhalf srcTensor, const LocalTensorhalf expMaxTensor, const LocalTensorfloat inExpSumTensor, const LocalTensorfloat inMaxTensor, const LocalTensoruint8_t sharedTmpBuffer, const SoftMaxTiling tiling, const SoftMaxShapeInfo softmaxShapeInfo {})由于该接口的内部实现中涉及复杂的计算需要额外的临时空间来存储计算过程中的中间变量。临时空间支持接口框架申请和开发者通过sharedTmpBuffer入参传入两种方式。接口框架申请临时空间开发者无需申请但是需要预留临时空间的大小。通过sharedTmpBuffer入参传入使用该tensor作为临时空间进行处理接口框架不再申请。该方式开发者可以自行管理sharedTmpBuffer内存空间并在接口调用完成后复用该部分内存内存不会反复申请释放灵活性较高内存利用率也较高。接口框架申请的方式开发者需要预留临时空间通过sharedTmpBuffer传入的情况开发者需要为tensor申请空间。临时空间大小BufferSize的获取方式如下通过SoftmaxFlashV2 Tiling接口中提供的GetSoftMaxFlashV2MinTmpSize/GetSoftMaxFlashV2MaxTmpSize接口获取所需最小和最大临时空间大小最小空间可以保证功能正确最大空间用于提升性能。另外提供了一个kernel侧tiling计算的接口当kernel侧的输入shape与通过host侧TilingData传入的shape不一致时可使用该接口在kernel侧重新计算tiling。该接口的参数含义请参考SoftmaxFlashV2 Tiling接口。kernel侧tiling计算接口__aicore__ inline constexpr SoftMaxTiling SoftMaxFlashV2TilingFunc(const SoftMaxShapeInfo shapeInfo, const uint32_t dataTypeSize1, const uint32_t dataTypeSize2, const uint32_t localWorkSpaceSize, const bool isUpdate false, const bool isBasicBlock false, const bool isDataFormatNZ false, const bool isFlashOutputBrc false)参数说明表1模板参数说明参数名描述T操作数的数据类型。支持的数据类型为half、float。isUpdate是否开启update部分中的计算。isReuseSource该参数预留传入默认值false即可。isBasicBlocksrcTensor和dstTensor的shape信息和Tiling切分策略满足基本块要求的情况下可以设置为true开启该参数用于提升性能默认为false表示不开启。是否满足基本块的要求可以采用如下两种方式之一判断srcTensor和dstTensor的shape信息[m,n]需要满足如下条件尾轴长度n小于2048并且大于等于256/sizeof(T)即half场景下n最小为128float场景下n最小为64同时n是64的倍数非尾轴长度的乘积m为8的倍数。在Tiling实现中通过调用IsBasicBlockInSoftMax判断Tiling切分策略是否满足基本块的切分要求。针对Atlas 200I/500 A2 推理产品该参数为预留参数暂未启用为后续的功能扩展做保留保持默认值即可。isDataFormatNZ当前输入输出的数据格式是否为NZ格式默认数据格式为ND即默认取值为false。针对Atlas 200I/500 A2 推理产品不支持配置为NZ格式。config结构体模板参数此参数可选配SoftmaxConfig类型具体定义如下方代码所示其中参数的含义为isCheckTiling是否需要检查shape和tiling的一致性若不一致API内会根据shape重新计算所需tiling。默认取值trueAPI内部会检查一致性。oriSrcM原始非尾轴长度的乘积。设置该参数后将shape常量化编译过程中使用常量化的shape。oriSrcK原始尾轴长度。设置该参数后将shape常量化编译过程中使用常量化的shape。mode输出shape的处理模式。当输入输出的数据格式为NZ格式时不支持配置mode参数。SoftmaxMode类型取值如下SOFTMAX_NORMAL 默认值常规模式对输出数据做Broadcast使得输出shape由(m, 1)拓展成(m, 8)输出为float数据类型或者(m, 16)输出为half数据类型。SOFTMAX_OUTPUT_WITHOUT_BRC 非拓展模式不对输出数据做Broadcast输出shape均为(m, 1)相应的输入参数例如inExpSumTensor、inMaxTensorshape也均为(m, 1)。此参数一般用于配合kernel侧tiling计算的接口使用。注意设置了oriSrcM与oriSrcK后模板参数isBasicBlock不生效计算数据是否为基本块由API内部判断并处理。针对Ascend 950PR/Ascend 950DT支持该参数。Atlas A3 训练系列产品/Atlas A3 推理系列产品支持该参数。Atlas A2 训练系列产品/Atlas A2 推理系列产品支持该参数。针对Atlas 200I/500 A2 推理产品该参数为预留参数暂未启用保持默认值即可。Atlas 推理系列产品AI Core支持该参数不支持配置mode。针对Kirin X90支持该参数。针对Kirin 9030支持该参数。struct SoftmaxConfig { bool isCheckTiling true; uint32_t oriSrcM 0; uint32_t oriSrcK 0; SoftmaxMode mode SoftmaxMode::SOFTMAX_NORMAL; };配置示例如下。constexpr SoftmaxConfig SOFTMAX_DEFAULT_CFG {true, 0, 0, SoftmaxMode::SOFTMAX_NORMAL};表2接口参数说明参数名输入/输出描述dstTensor输出目的操作数。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。dstTensor的shape和源操作数srcTensor一致。outReduceMax输出目的操作数。用于保存softmax计算过程中reducemax第一次计算的结果。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。outReduceMax的shape与目的操作数maxTensor一致。对于输出该结果的接口模板参数isUpdate为false时不输出该结果。仅支持输入输出的数据格式为ND模板参数isDataFormatNZ为预留参数传入默认值false即可。模板参数config.isCheckTiling为预留参数传入默认值false即可。模板参数config.mode仅支持配置为非拓展模式SoftmaxMode::SOFTMAX_OUTPUT_WITHOUT_BRC。expSumTensor、outExpSum输出目的操作数。用于保存softmax计算过程中reducesum的结果。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。除模板参数config配置为非拓展模式SoftmaxMode::SOFTMAX_OUTPUT_WITHOUT_BRC的场景外expSumTensor的last轴长度固定为32Byte即一个datablock长度。该datablock中的所有数据为同一个值比如half数据类型下该datablock中的16个数均为相同的reducesum的值。非last轴的长度与dstTensor保持一致。maxTensor、outMax输出目的操作数。用于保存softmax计算过程中reducemax的结果。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。除模板参数config配置为非拓展模式SoftmaxMode::SOFTMAX_OUTPUT_WITHOUT_BRC的场景外maxTensor的last轴长度固定为32Byte即一个datablock长度。该datablock中的所有数据为同一个值。比如half数据类型下该datablock中的16个数均为相同的reducemax的值。非last轴的长度与dstTensor保持一致。srcTensor输入源操作数。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。last轴长度需要32Byte对齐。expMaxTensor、outExpMax输出目的操作数。用于保存inmax与reducemax差值的e的指数幂的结果。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。除模板参数config配置为非拓展模式SoftmaxMode::SOFTMAX_OUTPUT_WITHOUT_BRC的场景外expMaxTensor的last轴长度固定为32Byte即一个datablock长度。该datablock中的所有数据为同一个值比如half数据类型下该datablock中的16个数均为相同的值。非last轴的长度需要与dstTensor保持一致。inExpSumTensor、inExpSum输入源操作数。softmax计算所需要的sum值。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。除模板参数config配置为非拓展模式SoftmaxMode::SOFTMAX_OUTPUT_WITHOUT_BRC的场景外inExpSumTensor的last轴长度固定为32Byte即一个datablock长度。该datablock中的所有数据为同一个值比如half数据类型下该datablock中的16个数均为相同的值。非last轴的长度需要与dstTensor保持一致。inMaxTensor、inMax输入源操作数。softmax计算所需要的max值。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。除模板参数config配置为非拓展模式SoftmaxMode::SOFTMAX_OUTPUT_WITHOUT_BRC的场景外inMaxTensor的last轴长度固定为32Byte即一个datablock长度。该datablock中的所有数据为同一个值比如half数据类型下该datablock里的16个数均为相同的值。非last轴的长度需要与dstTensor保持一致。sharedTmpBuffer输入临时空间。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。该操作数的数据类型固定uint8_t。接口内部复杂计算时用于存储中间变量由开发者提供。临时空间大小BufferSize的获取方式请参考SoftmaxFlashV2 Tiling接口。tiling输入softmaxflashv2接口计算所需tiling信息Tiling信息的获取请参考SoftmaxFlashV2 Tiling接口。softmaxShapeInfo输入srcTensor的shape信息。SoftMaxShapeInfo类型具体定义如下方代码所示其中参数的含义为srcM非尾轴长度的乘积。srcK尾轴长度必须32Byte对齐。oriSrcM原始非尾轴长度的乘积。oriSrcK原始尾轴长度。需要注意当输入输出的数据格式为NZ格式时尾轴长度为reduce轴长度即图2中的W0*W1非尾轴为H0*H1。struct SoftMaxShapeInfo { uint32_t srcM; uint32_t srcK; uint32_t oriSrcM; uint32_t oriSrcK; };返回值说明无约束说明srcTensor和dstTensor的Tensor的空间可以复用maxTensor和inMaxTensor的空间可以复用expSumTensor和inExpSumTensor的空间可以复用。除模板参数config配置为非拓展模式SoftmaxMode::SOFTMAX_OUTPUT_WITHOUT_BRC的场景外expSumTensor、maxTensor、expMaxTensor、inExpSumTensor、inMaxTensor的Tensor空间last轴长度必须固定32Byte。对于输出ReduceMax的接口模板参数isReuseSource、isDataFormatNZ、config.isCheckTiling均为预留参数config.mode只支持配置为非拓展模式SOFTMAX_OUTPUT_WITHOUT_BRC其配置为SOFTMAX_NORMAL模式时接口功能不执行不保存各输出模板参数isUpdate为false时outReduceMax不输出除outReduceMax外其余每个输出的计算结果与不输出ReduceMax的接口相同。操作数地址对齐要求请参见通用地址对齐约束。不支持sharedTmpBuffer与源操作数和目的操作数地址重叠。当参数softmaxShapeInfo中srcM ! oriSrcM或者srcK ! oriSrcK时开发者需要对GM上的原始输入(oriSrcM, oriSrcK)在M或K方向补齐数据到(srcM, srcK)补齐的数据会参与部分运算在输入输出复用的场景下API的计算结果会覆盖srcTensor中补齐的原始数据在输入输出不复用的场景下API的计算结果会覆盖dstTensor中对应srcTensor补齐位置的数据。调用示例srcK对齐本样例中输入srcTensor和输出dstTensor的shape大小为[320,64]输入inSumTensor、inMaxTensor的shape大小为[320,16]输出expMaxTensor的shape大小为[320,16]数据类型均为half输入输出的数据排布格式为NDsrcTensor和dstTensor空间不复用不开启基本块isUpdate为true。// dstLocal: 存放SoftMax计算结果的Tensor // expSumLocal存放softmax计算过程中reducesum的结果 // maxLocal存放softmax计算过程中reducemax的结果 // srcLocal存放SoftMax计算的输入Tensor // expMaxLocal存放inmax与reducemax差值的e的指数幂的结果 // inExpSumLocal存放softmax计算所需要的sum值 // inMaxLocal存放softmax计算所需要的max值 // sharedTmpBuffer: 存放SoftMax计算过程中临时缓存的Tensor // softmaxTiling存放SoftMax计算所需Tiling信息可通过SoftMaxFlashV2TilingFunc接口获取 AscendC::SoftMaxShapeInfo softmaxInfo( /* 非尾轴长度的乘积 */ 320, /* 尾轴长度必须32Bytes对齐 */ 64, /* 原始非尾轴长度的乘积 */ 320, /* 原始尾轴长度 */ 64); // 通过sharedTmpBuffer入参传入临时空间不输出ReduceMax传入模板参数将shape常量化 AscendC::SoftmaxFlashV2T, true, false, false, false, static_config( dstLocal, expSumLocal, maxLocal, srcLocal, expMaxLocal, inExpSumLocal, inMaxLocal, sharedTmpBuffer, tiling, softmaxInfo); // 通过sharedTmpBuffer入参传入临时空间不输出ReduceMax AscendC::SoftmaxFlashV2T, true( dstLocal, expSumLocal, maxLocal, srcLocal, expMaxLocal, inExpSumLocal, inMaxLocal, sharedTmpBuffer, tiling, softmaxInfo); // 接口框架申请临时空间带sumTensor和maxTensor参数 AscendC::SoftmaxFlashV2T, true( dstLocal, expSumLocal, maxLocal, srcLocal, expMaxLocal, inExpSumLocal, inMaxLocal, tiling, softmaxInfo);结果示例如下输入数据(srcLocal) [[-10. -10. -10. ... -9.94 -9.94 -9.94 ] [ -9.94 -9.94 -9.94 ... -9.875 -9.875 -9.875] [ -9.875 -9.875 -9.875 ... -9.81 -9.81 -9.81 ] ... [ 9.81 9.81 9.81 ... 9.875 9.875 9.875] [ 9.875 9.875 9.875 ... 9.94 9.94 9.94 ] [ 9.94 9.94 9.94 ... 10. 10. 10. ]] 输出数据(expSumLocal) [[62.03 62.03 62.03 ... 62.03 62.03 62.03] [62.03 62.03 62.03 ... 62.03 62.03 62.03] [62.03 62.03 62.03 ... 62.03 62.03 62.03] ... [62.03 62.03 62.03 ... 62.03 62.03 62.03] [62.03 62.03 62.03 ... 62.03 62.03 62.03] [62.03 62.03 62.03 ... 62.03 62.03 62.03]] 输出数据(maxLocal) [[-9.94 -9.94 -9.94 ... -9.94 -9.94 -9.94 ] [-9.875 -9.875 -9.875 ... -9.875 -9.875 -9.875] [-9.81 -9.81 -9.81 ... -9.81 -9.81 -9.81 ] ... [ 9.875 9.875 9.875 ... 9.875 9.875 9.875] [ 9.94 9.94 9.94 ... 9.94 9.94 9.94 ] [10. 10. 10. ... 10. 10. 10. ]] 输出数据(dstLocal) [[0.015144 0.015144 0.015144 ... 0.01611 0.01611 0.01611 ] [0.015144 0.015144 0.015144 ... 0.01611 0.01611 0.01611 ] [0.015144 0.015144 0.015144 ... 0.01611 0.01611 0.01611 ] ... [0.015144 0.015144 0.015144 ... 0.01611 0.01611 0.01611 ] [0.015144 0.015144 0.015144 ... 0.01611 0.01611 0.01611 ] [0.015144 0.015144 0.015144 ... 0.01611 0.01611 0.01611 ]]srcK非对齐本示例中srcTensor和输出dstTensor的shape大小为[320,63]数据类型均为half输入输出的数据排布格式为ND展示非对齐padding补齐的搬入搬出操作和API调用方式。#include kernel_operator.h // init阶段height320, width63 padWidth AlignUp(width * sizeof(T), 32) / sizeof(T); // copyin阶段 AscendC::DataCopyExtParams copyParams{static_castuint16_t(height), static_castuint32_t(width * sizeof(T)), 0, 0, 0}; AscendC::DataCopyPadExtParamsT padParam {true, 0, static_castuint8_t(padWidth - width), 0}; AscendC::DataCopyPad(srcLocal, srcGlobal, copyParams, padParam); // compute阶段 // 由于发生paddingAPI调用时shape和原始shape发生了不一致 AscendC::SoftMaxShapeInfo srcShape {height, padWidth, height, width}; AscendC::SoftmaxFlashV2T, true( dstLocal, expSumLocal, maxLocal, srcLocal, expMaxLocal, inExpSumLocal, inMaxLocal, tiling, srcShape); // copyout阶段 AscendC::DataCopyExtParams copyParams{static_castuint16_t(height), static_castuint32_t(width * sizeof(T)), 0, 0, 0}; AscendC::DataCopyPad(dstGlobal, dstLocal, copyParams);结果示例如下输入数据(srcLocal) [[-10. -10. -10. ... -9.94 -9.94 0. ] [ -9.94 -9.94 -9.94 ... -9.875 -9.875 0. ] [ -9.875 -9.875 -9.875 ... -9.81 -9.81 0. ] ... [ 9.81 9.81 9.81 ... 9.875 9.875 0. ] [ 9.875 9.875 9.875 ... 9.94 9.94 0. ] [ 9.94 9.94 9.94 ... 10. 10. 0. ]] 输出数据(expSumLocal) [[61.03 61.03 61.03 ... 61.03 61.03 61.03] [61.03 61.03 61.03 ... 61.03 61.03 61.03] [61.03 61.03 61.03 ... 61.03 61.03 61.03] ... [61.1 61.1 61.1 ... 61.1 61.1 61.1 ] [61.1 61.1 61.1 ... 61.1 61.1 61.1 ] [61.1 61.1 61.1 ... 61.1 61.1 61.1 ]] 输出数据(maxLocal) [[-9.94 -9.94 -9.94 ... -9.94 -9.94 -9.94 ] [-9.875 -9.875 -9.875 ... -9.875 -9.875 -9.875] [-9.81 -9.81 -9.81 ... -9.81 -9.81 -9.81 ] ... [ 9.875 9.875 9.875 ... 9.875 9.875 9.875] [ 9.94 9.94 9.94 ... 9.94 9.94 9.94 ] [10. 10. 10. ... 10. 10. 10. ]] 输出数据(dstLocal) [[0.015396 0.015396 0.015396 ... 0.01639 0.01639 0.01639 ] [0.015396 0.015396 0.015396 ... 0.01639 0.01639 0.01639 ] [0.015396 0.015396 0.015396 ... 0.01639 0.01639 0.01639 ] ... [0.01538 0.01538 0.01538 ... 0.01637 0.01637 0.01637 ] [0.01538 0.01538 0.01538 ... 0.01637 0.01637 0.01637 ] [0.01538 0.01538 0.01538 ... 0.01637 0.01637 0.01637 ]]【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考