CANN/pyasc昇腾SoftMax算子API文档

asc.language.adv.softmax

【免费下载链接】pyasc本项目为Python用户提供算子编程接口，支持在昇腾AI处理器上加速计算，接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc

asc.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, 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

对应的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 LocalTensor<T>& dstTensor, const LocalTensor<T>& sumTensor, const LocalTensor<T>& maxTensor, const LocalTensor<T>& 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 LocalTensor<half>& dstTensor, const LocalTensor<float>& sumTensor, const LocalTensor<float>& maxTensor, const LocalTensor<half>& 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 LocalTensor<T>& dstTensor, const LocalTensor<T>& 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 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 = {})

  • LocalTensor的数据类型不同

    template <typename T, bool isReuseSource = false, bool isBasicBlock = false, bool isDataFormatNZ = false, const SoftmaxConfig& config = SOFTMAX_DEFAULT_CFG> __aicore__ inline void SoftMax(const LocalTensor<half>& dstTensor, const LocalTensor<float>& sumTensor, const LocalTensor<float>& maxTensor, const LocalTensor<half>& srcTensor, const LocalTensor<uint8_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 LocalTensor<T>& dstTensor, const LocalTensor<T>& srcTensor, const LocalTensor<uint8_t>& sharedTmpBuffer, const SoftMaxTiling& tiling, const SoftMaxShapeInfo& softmaxShapeInfo = {})

参数说明

• dst：目的操作数。
• sum：目的操作数。
• max：目的操作数。
• src：源操作数。
• tiling：SoftMax计算所需Tiling信息。
• tmp_buffer：临时空间。
• reuse_source：该参数预留，传入默认值false即可。
• basic_block：src和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