CANN/pyasc昇腾SoftMax算子API文档

asc.language.adv.softmax【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyascasc.language.adv.softmax(dst: LocalTensor, sum: LocalTensor, max: LocalTensor, src: LocalTensor, tiling: SoftmaxTiling, temp_buffer: LocalTensor | None None, reuse_source: bool False, basic_block: bool False, data_format_nz: bool False) → None将输入tensor[m0, m1, …mt, n]t大于等于0的非尾轴长度相乘的结果看作m则输入tensor的shape看作[m, n]。 为方便理解通过Python脚本实现的方式表达其计算公式以输入为ND格式为例如下其中src是源操作数输入dst、sum、max为目的操作数输出。def softmax(src): # 基于last轴进行rowmax按行取最大值处理 max np.max(src, axis-1, keepdimsTrue) sub src - max exp np.exp(sub) # 基于last轴进行rowsum按行求和处理 sum np.sum(exp, axis-1, keepdimsTrue) dst exp / sum return dst, max, sum对应的Ascend C函数原型接口框架申请临时空间LocalTensor的数据类型相同template typename T, bool isReuseSource false, bool isBasicBlock false, bool isDataFormatNZ false, const SoftmaxConfig config SOFTMAX_DEFAULT_CFG __aicore__ inline void SoftMax(const LocalTensorT dstTensor, const LocalTensorT sumTensor, const LocalTensorT maxTensor, const LocalTensorT srcTensor, const SoftMaxTiling tiling, const SoftMaxShapeInfo softmaxShapeInfo {})LocalTensor的数据类型不同template typename T, bool isReuseSource false, bool isBasicBlock false, bool isDataFormatNZ false, const SoftmaxConfig config SOFTMAX_DEFAULT_CFG __aicore__ inline void SoftMax(const LocalTensorhalf dstTensor, const LocalTensorfloat sumTensor, const LocalTensorfloat maxTensor, const LocalTensorhalf srcTensor, const SoftMaxTiling tiling, const SoftMaxShapeInfo softmaxShapeInfo {})不带sumTensor和maxTensor参数template typename T, bool isReuseSource false, bool isBasicBlock false, const SoftmaxConfig config SOFTMAX_DEFAULT_CFG __aicore__ inline void SoftMax(const LocalTensorT dstTensor, const LocalTensorT srcTensor, const SoftMaxTiling tiling, const SoftMaxShapeInfo softmaxShapeInfo {})通过sharedTmpBuffer入参传入临时空间LocalTensor的数据类型相同template typename T, bool isReuseSource false, bool isBasicBlock false, bool isDataFormatNZ false, const SoftmaxConfig config SOFTMAX_DEFAULT_CFG __aicore__ inline void SoftMax(const LocalTensorT dstTensor, const LocalTensorT sumTensor, const LocalTensorT maxTensor, const LocalTensorT srcTensor, const LocalTensoruint8_t sharedTmpBuffer, const SoftMaxTiling tiling, const SoftMaxShapeInfo softmaxShapeInfo {})LocalTensor的数据类型不同template typename T, bool isReuseSource false, bool isBasicBlock false, bool isDataFormatNZ false, const SoftmaxConfig config SOFTMAX_DEFAULT_CFG __aicore__ inline void SoftMax(const LocalTensorhalf dstTensor, const LocalTensorfloat sumTensor, const LocalTensorfloat maxTensor, const LocalTensorhalf srcTensor, const LocalTensoruint8_t sharedTmpBuffer, const SoftMaxTiling tiling, const SoftMaxShapeInfo softmaxShapeInfo {})不带sumTensor和maxTensor参数template typename T, bool isReuseSource false, bool isBasicBlock false, const SoftmaxConfig config SOFTMAX_DEFAULT_CFG __aicore__ inline void SoftMax(const LocalTensorT dstTensor, const LocalTensorT srcTensor, const LocalTensoruint8_t sharedTmpBuffer, const SoftMaxTiling tiling, const SoftMaxShapeInfo softmaxShapeInfo {})参数说明dst目的操作数。sum目的操作数。max目的操作数。src源操作数。tilingSoftMax计算所需Tiling信息。tmp_buffer临时空间。reuse_source该参数预留传入默认值false即可。basic_blocksrc和dst的shape信息和Tiling切分策略满足基本块要求的情况下可以使能该参数用于提升性能默认不使能。data_format_nz当前输入输出的数据格式是否为NZ格式默认数据格式为ND即默认取值为false。约束说明src和dst的Tensor空间可以复用。sum和max为输出并且last轴长度必须固定32Byte非last轴大小需要和src以及dst保持一致。sum和max的数据类型需要保持一致。操作数地址对齐要求请参见 《Ascend C算子开发接口》 中的“通用说明和约束-通用地址对齐约束”。不支持tmp_buffer与源操作数和目的操作数地址重叠。开发者需要对GM上的原始输入(ori_src_M, ori_src_K)在M或K方向补齐数据到(src_M, src_K)补齐的数据会参与部分运算 在输入输出复用的场景下API的计算结果会覆盖src中补齐的原始数据在输入输出不复用的场景下 API的计算结果会覆盖dst中对应src补齐位置的数据。调用示例src_local in_queue_src.deque(T) sum_temp_local sum_queue.alloc_tensor(T) max_temp_local max_queue.alloc_tensor(T) dst_local out_queue_dst.alloc_tensor(T) src_shape asc.SoftMaxShapeInfo(height, width, height, width); asc.adv.softmax(dst_local, sum_temp_local, max_temp_local, srcLocal, tiling, src_shape); out_queue_dst.EnQue(dstLocal) max_queue.free_tensor(max_temp_local) sum_queue.free_tensor(sum_temp_local) in_queue_src.free_tensor(src_local)【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考