aclnnBatchMatMul

支持的产品型号

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

接口原型

每个算子分为undefined，必须先调用“aclnnBatchMatMulGetWorkspaceSize”接口获取入参并根据流程计算所需workspace大小，再调用“aclnnBatchMatMul”接口执行计算。

• aclnnStatus aclnnBatchMatMulGetWorkspaceSize(const aclTensor *self, const aclTensor *mat2, aclTensor *out, int8_t cubeMathType, uint64_t *workspaceSize, aclOpExecutor **executor)
• aclnnStatus aclnnBatchMatMul(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

功能描述

• 算子功能： 完成两个Tensor的矩阵乘。 仅支持三维的Tensor传入。第一维是batch维度，仅支持一维batch做broadcast（例如左矩阵batch为1，右矩阵batch为n），最后两个维度做矩阵乘法。

• 计算公式：

$$out = self@mat2$$

• 示例：

  self的shape是[A, M, K]，mat2的shape是[A, K, N]，计算输出out的shape是[A, M, N]。第一维相等，后两维做矩阵乘运算。

aclnnBatchMatMulGetWorkSpaceSize

• 参数说明：

  • self(const aclTensor *, 计算输入): 做矩阵乘运算的第一个输入Tensor，Device侧aclTensor。数据类型支持FLOAT16、FLOAT、BFLOAT16（仅Atlas A2训练系列产品/Atlas 800I A2推理产品支持），且数据类型应该与mat2一致，支持undefined，undefined支持ND。
  • mat2(const aclTensor *, 计算输入): 做矩阵乘运算的第二个输入Tensor，Device侧aclTensor。数据类型支持FLOAT16、FLOAT、BFLOAT16（仅Atlas A2训练系列产品/Atlas 800I A2推理产品支持），且数据类型应该与self一致，支持undefined，undefined支持ND。
  • out(aclTensor *, 计算输出): 输出Tensor，是Device侧aclTensor。数据类型支持FLOAT16、FLOAT、BFLOAT16（仅Atlas A2训练系列产品/Atlas 800I A2推理产品支持），数据类型应该与self和mat2保持一致。undefined支持ND，与self保持一致。
  • cubeMathType(int8_t *, 计算输入): 用于判断Cube单元使用哪种计算逻辑进行运算，数据类型支持INT8，支持的枚举值如下：
    • 0:KEEP_DTYPE，保持输入的数据类型进行计算。当输入是FLOAT，Atlas 训练系列产品和Atlas 推理系列产品暂不支持，取0时会报错。
    • 1:ALLOW_FP32_DOWN_PRECISION，允许将输入数据降精度计算。当输入是FLOAT，Atlas 训练系列产品和Atlas 推理系列产品转换为FLOAT16计算，Atlas A2训练系列产品/Atlas 800I A2推理产品转换为HFLOAT32计算。
    • 2:USE_FP16，允许转换为数据类型FLOAT16进行计算。当输入数据类型是FLOAT，转换为FLOAT16计算。
    • 3:USE_HF32，允许转换为数据类型HFLOAT32计算。当输入是FLOAT，Atlas 训练系列产品和Atlas 推理系列产品暂不支持，取3时会报错，Atlas A2训练系列产品/Atlas 800I A2推理产品转换为HFLOAT32计算。
  • workspaceSize(uint64_t *, 出参): 返回需要在Device侧申请的workspace大小。
  • executor(aclOpExecutor **, 出参): 返回op执行器，包含了算子计算流程。

• 返回值：

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

  [object Object]

aclnnBatchMatMul

• 参数说明：

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

• 返回值：

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

约束与限制

对于Atlas 训练系列产品和Atlas 推理系列产品，Cube单元不支持FLOAT32计算。当输入为FLOAT32，可通过设置cubeMathType=1（Allow_FP32_DOWN_PRECISION）来允许接口内部cast到FLOAT16进行计算。

调用示例

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