aclnnMaskedSoftmaxWithRelPosBias

支持的产品型号

• Atlas 推理系列产品。
• Atlas A2训练系列产品/Atlas 800I A2推理产品。

接口原型

每个算子分为undefined，必须先调用“aclnnMaskedSoftmaxWithRelPosBiasGetWorkspaceSize”接口获取计算所需workspace大小以及包含了算子计算流程的执行器，再调用“aclnnMaskedSoftmaxWithRelPosBias”接口执行计算。

• aclnnstatus aclnnMaskedSoftmaxWithRelPosBiasGetWorkspaceSize(const aclTensor *x, const aclTensor *attenMaskOptional, const aclTensor *relativePosBias, double scaleValueOptional, int64_t innerPrecisionModeOptional, aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor)
• aclnnstatus aclnnMaskedSoftmaxWithRelPosBias(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

功能描述

• 算子功能：替换在swinTransformer中使用window attention计算softmax的部分。

• 计算公式：

$$out = softmax(scaleValueOptional * x + attenMaskOptional + relativePosBias)$$

aclnnMaskedSoftmaxWithRelPosBiasGetWorkspaceSize

• 参数说明：

  • x(const aclTensor*, 计算输入)：Device侧的aclTensor，数据类型支持FLOAT、FLOAT16、BFLOAT16，暂不支持undefined，undefined支持ND。其中BFLOAT16仅Atlas A2训练系列产品/Atlas 800I A2推理产品支持。支持shape为4维(B*W, N, S1, S2)和5维(B, W, N, S1, S2)。
  • attenMaskOptional(const aclTensor*, 计算输入)：Device侧的aclTensor，数据类型支持FLOAT、FLOAT16、BFLOAT16，暂不支持undefined，undefined支持ND。其中BFLOAT16仅Atlas A2训练系列产品/Atlas 800I A2推理产品支持。支持shape为3维(W, S1, S2)、4维(W, 1, S1, S2)和5维(1, W, 1, S1, S2)，可选输入，当不需要时为空指针, 数据类型需要保持和x一致。。
  • relativePosBias(const aclTensor*, 计算输入)：Device侧的aclTensor，数据类型支持FLOAT、FLOAT16、BFLOAT16，暂不支持undefined，undefined支持ND。其中BFLOAT16仅Atlas A2训练系列产品/Atlas 800I A2推理产品支持。支持shape为3维(N, S1, S2)、4维(1, N, S1, S2)和5维(1, 1, N, S1, S2), 数据类型需要保持和x一致。
  • scaleValueOptional(double, 计算输入)：Host侧的整型，数据类型支持DOUBLE。
  • innerPrecisionModeOptional(int64_t, 预留参数)：Host侧的整型，数据类型支持INT64。
  • out(aclTensor*, 计算输出)：Device侧的aclTensor，数据类型支持FLOAT、FLOAT16、BFLOAT16，暂不支持undefined，undefined支持ND。其中BFLOAT16仅Atlas A2训练系列产品/Atlas 800I A2推理产品支持。shape与数据类型需要保持和入参x一致。
  • workspaceSize(uint64_t*, 出参)：返回需要在Device侧申请的workspace大小。
  • executor(aclOpExecutor**, 出参)：返回op执行器，包含了算子计算流程。

• 返回值：

  aclnnStatus：返回状态码，具体参见undefined。

  [object Object]

aclnnMaskedSoftmaxWithRelPosBias

• 参数说明：

  • workspace(void *, 入参)：在Device侧申请的workspace内存地址。
  • workspaceSize(uint64_t, 入参)：在Device侧申请的workspace大小，由第一段接口aclnnMaskedSoftmaxWithRelPosBiasGetWorkspaceSize获取。
  • executor(aclOpExecutor *, 入参)：op执行器，包含了算子计算流程。
  • stream(aclrtStream, 入参)：指定执行任务的 AscendCL Stream流。

• 返回值：

  aclnnStatus：返回状态码，具体参见undefined。

约束与限制

• Atlas 推理系列产品不支持入参x的最后一个维度S2非32Byte对齐的场景。

• 需要保证传递给算子的shape所需要的ub空间小于AI处理器版本总ub的大小，该算子所需要的ub空间的总大小minComputeSize如下，其中s2AlignedSize 表示S2对齐32Byte后的结果：

  • 对于attenMaskOptional存在的情况：

    [object Object]

  • 对于attenMaskOptional不存在的情况：

    [object Object]

调用示例

示例代码如下，仅供参考，具体编译和执行过程请参考undefined。