cann/asc-devkit 元素对归约

ReducePairElem

产品支持情况

功能说明

头文件路径为："basic_api/kernel_operator_vec_reduce_intf.h"。

ReducePairElem接口根据模板参数reduceType，对相邻两个（奇偶）元素进行归约操作，结果按顺序写入目标地址。当前仅支持求和操作。

例如，[a1, a2, a3, a4, a5, a6, ...]的求和结果为[a1+a2, a3+a4, a5+a6, ...]。计算过程如下图所示，灰色部分代表无效元素。

图 1ReducePairElem<ReduceType::SUM>示意图

![ReducePairElem<ReduceType::SUM>示意图](https://raw.gitcode.com/cann/asc-devkit/raw/1ea24e5536a6095446f2c42868e8b9f2d6dafdcc/docs/api/figures/ReducePairElem.png "ReducePairElemReduceType::SUM示意图"?utm_source=gitcode_repo_files)

函数原型

• mask逐bit模式：

  template <ReduceType reduceType, typename T, typename U, bool isSetMask = true> __aicore__ inline void ReducePairElem(const LocalTensor<T>& dst, const LocalTensor<U>& src, const uint64_t mask[],const int32_t repeatTime, const int32_t dstRepStride, const int32_t srcBlkStride, const int32_t srcRepStride)

• mask连续模式：

  template <ReduceType reduceType, typename T, typename U, bool isSetMask = true> __aicore__ inline void ReducePairElem(const LocalTensor<T>& dst, const LocalTensor<U>& src, const int32_t mask, const int32_t repeatTime, const int32_t dstRepStride, const int32_t srcBlkStride, const int32_t srcRepStride)

参数说明

表 1模板参数说明

表 2参数说明

注：以上高维切分相关参数mask，repeatTime，dstRepStride，srcBlkStride，srcRepStride请参考高维切分中的介绍。

数据类型

支持的数据类型如下：

• Ascend 950PR/Ascend 950DT，支持half、float。
• Atlas A3 训练系列产品/Atlas A3 推理系列产品，支持half、float。
• Atlas A2 训练系列产品/Atlas A2 推理系列产品，支持half、float。
• Atlas 200I/500 A2 推理产品，支持half、float。
• Atlas 推理系列产品 AI Core，支持half、float。
• Atlas 训练系列产品，支持half。

目的操作数与源操作数的数据类型需要保持一致。

返回值说明

无

约束说明

• 操作数的起始地址对齐约束请参考通用地址对齐约束，需要32字节对齐。
• 操作数地址重叠约束请参考通用地址重叠约束。
• dstRepStride、srcBlkStride、srcRepStride取值范围为[0, $2^{16}-1$]，需要结合UB的实际大小避免出现越界。

• 针对如下型号，当mask=0或repeatTime=0时，不会执行归约操作，不会对目的操作数进行写入，该接口将被视为NOP（空操作）。
  • Atlas A3 训练系列产品/Atlas A3 推理系列产品
  • Atlas A2 训练系列产品/Atlas A2 推理系列产品

• 针对如下型号，当一对相邻元素的掩码均为0，对应的目的元素将跳过写入，保持不变。
  • Atlas A3 训练系列产品/Atlas A3 推理系列产品
  • Atlas A2 训练系列产品/Atlas A2 推理系列产品
  • Atlas 200I/500 A2 推理产品
  • Atlas 推理系列产品 AI Core
  • Atlas 训练系列产品

• 对于Atlas 200I/500 A2 推理产品，当一对相邻元素的掩码均为0，对应的目的操作数中的值会置为0。

调用示例

本样例中只展示Compute流程中的部分代码。更多样例可参考ReducePairElem样例。

• ReducePairElemtensor高维切分计算样例-mask连续模式：

  // 设定mask为最多的128个全部元素参与计算 int32_t mask = 256/sizeof(half); // 每个repeat128个元素，一共128个元素。 int repeat = 1; // dstLocal: 目的操作数tensor // srcLocal: 源操作数tensor // srcBlkStride = 1，在一个repeat中，block间没有空隙。 // dstRepStride = 1，srcRepStride = 8，repeat间没有空隙。 AscendC::ReducePairElem<AscendC::ReduceType::SUM, half>(dstLocal, srcLocal, mask, repeat, 1, 1, 8);

• ReducePairElemtensor高维切分计算样例-mask逐bit模式：

  // 设定mask为最多的128个全部元素参与计算 uint64_t mask[2] = { UINT64_MAX, UINT64_MAX }; // 每个repeat128个元素，一共128个元素。 int repeat = 1; // dstLocal: 目的操作数tensor // srcLocal: 源操作数tensor // srcBlkStride = 1，在一个repeat中，block间没有空隙。 // dstRepStride = 1，srcRepStride = 8，repeat间没有空隙。 AscendC::ReducePairElem<AscendC::ReduceType::SUM, half>(dstLocal, srcLocal, mask, repeat, 1, 1, 8);

• 示例结果

  输入数据src_gm：

  [1, 1, 1, -1, 2, 2, -1, 2, 3, 3, 3, -1, 4, 4, -2, 4, .... > > ]

  输出数据dst_gm：

  [2, 0, 4, 1, 6, 2, 8, 2, .... > > ]