aclnnMaskedSoftmaxWithRelPosBias

产品支持情况

产品	是否支持
[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]	√
[object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]	√
[object Object]Atlas 200I/500 A2 推理产品[object Object]	×
[object Object]Atlas 推理系列产品 [object Object]	√
[object Object]Atlas 训练系列产品[object Object]	×

功能说明

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

• 计算公式：

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

函数原型

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

• aclnnStatus aclnnMaskedSoftmaxWithRelPosBiasGetWorkspaceSize(const aclTensor *x, const aclTensor *attenMaskOptional, const aclTensor *relativePosBias, double scaleValue, int64_t innerPrecisionMode, const aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor)
• aclnnStatus aclnnMaskedSoftmaxWithRelPosBias(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

aclnnMaskedSoftmaxWithRelPosBiasGetWorkspaceSize

• 参数说明：

  • x（aclTensor*，计算输入）：Device侧的aclTensor，暂不支持undefined，undefined支持ND。支持shape为4维(B*W, N, S1, S2)和5维(B, W, N, S1, S2)。
    • [object Object]Atlas 推理系列产品[object Object]：数据类型支持FLOAT、FLOAT16。
    • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：数据类型支持FLOAT、FLOAT16、BFLOAT16。
  • attenMaskOptional（aclTensor*，计算输入）：Device侧的aclTensor，暂不支持undefined，undefined支持ND。支持shape为3维(W, S1, S2)、4维(W, 1, S1, S2)和5维(1, W, 1, S1, S2)，可选输入，当不需要时为空指针，数据类型需要保持和x一致。
    • [object Object]Atlas 推理系列产品[object Object]：数据类型支持FLOAT、FLOAT16。
    • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：数据类型支持FLOAT、FLOAT16、BFLOAT16。
  • relativePosBias（aclTensor*，计算输入）：Device侧的aclTensor，暂不支持undefined，undefined支持ND。支持shape为3维(N, S1, S2)、4维(1, N, S1, S2)和5维(1, 1, N, S1, S2)，数据类型需要保持和x一致。
    • [object Object]Atlas 推理系列产品[object Object]：数据类型支持FLOAT、FLOAT16。
    • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：数据类型支持FLOAT、FLOAT16、BFLOAT16。
  • scaleValue（double，计算输入）：Host侧的整型，数据类型支持DOUBLE。
  • innerPrecisionMode（int64_t，预留参数）：Host侧的整型，数据类型支持INT64。
  • out（aclTensor*，计算输出）：Device侧的aclTensor，暂不支持undefined，undefined支持ND。shape与数据类型需要保持和入参x一致。
    • [object Object]Atlas 推理系列产品[object Object]：数据类型支持FLOAT、FLOAT16。
    • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：数据类型支持FLOAT、FLOAT16、BFLOAT16。
  • 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，入参）：指定执行任务的Stream。

• 返回值：

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

约束说明

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

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

  • 对于attenMaskOptional存在的情况：
    [object Object]
  • 对于attenMaskOptional不存在的情况：
    [object Object]
  • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：如果为BFLOAT16类型，其与FLOAT16类型的公式保持一致。

调用示例

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