Gather
产 品 支 持 情 况
功 能 说 明
头 文 件 路 径 为："basic_api/kernel_operator_vec_gather_intf.h"。
Gather接 口 接 受 输 入 张 量（src）、地 址 偏 移 张 量（srcOffset）和 基 地 址（srcBaseAddr），根 据 基 地 址 和 地 址 偏 移 量 确 定 输 入 张 量 的 索 引，将 输 入 张 量 中 对 应 元 素 收 集 到 结 果 张 量（dst）中。
接 口 支 持 使 用 前n个 数 据 计 算 和 高 维 切 分 计 算 的 方 式，前n个 数 据 计 算 原 理 和 参 考 伪 代 码 如 下：
Python
import numpy as np
def Gather(src, dst, count, srcOffset, srcBaseAddr):
inputType = np.dtype("uint16")
for i in range(count):
dst[i] = src[(srcBaseAddr + srcOffset[i]) // inputType.itemsize]
完 整 示 例 请 参 考：Gather真 值 计 算。
函 数 原 型
tensor前n个 数 据 计 算。
C++template <typename T> __aicore__ inline void Gather(const LocalTensor<T>& dst, const LocalTensor<T>& src, const LocalTensor<uint32_t>& srcOffset, const uint32_t srcBaseAddr, const uint32_t count)tensor高 维 切 分 计 算。
mask逐bit模 式。
C++template <typename T> __aicore__ inline void Gather(const LocalTensor<T>& dst, const LocalTensor<T>& src,const LocalTensor<uint32_t>& srcOffset, const uint32_t srcBaseAddr, const uint64_t mask[], const uint8_t repeatTime, const uint16_t dstRepStride)mask连 续 模 式。
C++template <typename T> __aicore__ inline void Gather(const LocalTensor<T>& dst, const LocalTensor<T>& src, const LocalTensor<uint32_t>& srcOffset, const uint32_t srcBaseAddr, const uint64_t mask, const uint8_t repeatTime, const uint16_t dstRepStride)
参 数 说 明
表 1 模 板 参 数 说 明
| 参 数 名 | 描 述 |
|---|---|
| T | 操 作 数 数 据 类 型。 |
表 2 参 数 说 明
| 参 数 名 称 | 输 入/输 出 | 含 义 |
|---|---|---|
| dst | 输 出 | 目 的 操 作 数，类 型 为LocalTensor，支 持 的TPosition为VECIN/VECCALC/VECOUT（存 储 位 置 为Unified Buffer）。 LocalTensor的 起 始 地 址 需 要 按 照32字 节 对 齐。 |
| src | 输 入 | 源 操 作 数，类 型 为LocalTensor，支 持 的TPosition为VECIN/VECCALC/VECOUT（存 储 位 置 为Unified Buffer）。 LocalTensor的 起 始 地 址 需 要 按 照32字 节 对 齐。 数 据 类 型 和dst保 持 一 致。 |
| srcOffset | 输 入 | 每 个 元 素 在src中 对 应 的 地 址 偏 移，类 型 为LocalTensor，支 持 的TPosition为VECIN/VECCALC/VECOUT（存 储 位 置 为Unified Buffer）。 LocalTensor的 起 始 地 址 需 要32字 节 对 齐。 该 偏 移 量 相 对 于src的 起 始 基 地 址 而 言。单 位 为 字 节。取 值 要 求 见约 束 说 明。 |
| srcBaseAddr | 输 入 | 用 于 指 定Gather操 作 中 源 操 作 数 的 起 始 位 置，单 位 为 字 节。取 值 应 保 证src元 素 类 型 位 宽 对 齐，否 则 会 导 致 非 预 期 行 为。 |
| count | 输 入 | 参 与 计 算 的 元 素 个 数。 参 数 取 值 范 围 和 操 作 数 的 数 据 类 型 有 关，数 据 类 型 不 同，能 够 处 理 的 元 素 个 数 最 大 值 不 同，最 大 处 理 的 数 据 量 不 能 超 过UB大 小 限 制。 注：count = 0表 示 不 会 执 行 计 算 操 作，不 会 对 目 的 操 作 数 进 行 写 入，该 接 口 将 被 视 为NOP（空 操 作）。 |
| mask/mask[] | 输 入 | mask用 于 控 制 每 次 迭 代 内 参 与 计 算 的 元 素。 设 置 详 见掩 码 操 作。 |
| repeatTime | 输 入 | 指 令 迭 代 次 数。针 对 不 同 的 型 号，每 个 迭 代 处 理 的DataBlock可 能 存 在 差 异，详 见约 束 说 明。 注：repeatTime = 0表 示 不 会 执 行 计 算 操 作，不 会 对 目 的 操 作 数 进 行 写 入，该 接 口 将 被 视 为NOP（空 操 作）。 |
| dstRepStride | 输 入 | 相 邻 迭 代 间 的 地 址 步 长，单 位 是DataBlock（32Bytes）。 |
数 据 类 型
Ascend 950PR/Ascend 950DT，支 持 的 数 据 类 型 为：int8_t、uint8_t、int16_t、uint16_t、half、bfloat16_t、int32_t、uint32_t、float、int64_t、uint64_t。
Atlas A3 训 练 系 列 产 品/Atlas A3 推 理 系 列 产 品，支 持 的 数 据 类 型 为：int16_t、uint16_t、half、bfloat16_t、int32_t、uint32_t、float。
Atlas A2 训 练 系 列 产 品/Atlas A2 推 理 系 列 产 品，支 持 的 数 据 类 型 为：int16_t、uint16_t、half、bfloat16_t、int32_t、uint32_t、float。
Atlas 200I/500 A2 推 理 产 品，支 持 的 数 据 类 型 为：int8_t、uint8_t、int16_t、uint16_t、half、int32_t、uint32_t、float。
Atlas 推 理 系 列 产 品AI Core，支 持 的 数 据 类 型 为：int16_t、uint16_t、half、int32_t、uint32_t、float。
约 束 说 明
- 操 作 数 地 址 对 齐 要 求 请 参 见Unified Buffer地 址 对 齐 约 束。
- 不 支 持 源 操 作 数 与 目 的 操 作 数 使 用 同 一 块 内 存 地 址。
- srcOffset的 取 值 要 求 如 下：
- 取 值 应 保 证src元 素 类 型 位 宽 对 齐。
- 偏 移 地 址 后 不 能 超 出UB大 小 数 据 的 范 围。
- 针 对 以 下 型 号，地 址 偏 移 的 取 值 范 围：当 操 作 数 为8位 时，取 值 范 围 为[0, 2^16-1]；当 操 作 数 为16位 时，取 值 范 围 为[0, 2^17-1]，当 操 作 数 为32位 或 者64位 时，不 能 超 出uint32_t的 范 围。 -
Ascend 950PR/Ascend 950DT
- Atlas 200I/500 A2 推 理 产 品
- 针 对 以 下 型 号，地 址 偏 移 的 取 值 范 围：不 能 超 出uint32_t的 范 围。 -
- repeatTime取 值 范 围：repeatTime∈[0,255]。 Atlas A3 训 练 系 列 产 品/Atlas A3 推 理 系 列 产 品
- Atlas A2 训 练 系 列 产 品/Atlas A2 推 理 系 列 产 品
- Atlas 推 理 系 列 产 品AI Core
- 针 对 以 下 型 号，当 操 作 数 为**8位**时，每 次 迭 代 完 成**4个DataBlock**的 数 据 收 集；其 他 情 况 下，每 次 迭 代 完 成8个DataBlock的 数 据 收 集。 -
Ascend 950PR/Ascend 950DT
- Atlas 200I/500 A2 推 理 产 品
- 针 对 以 下 型 号，每 次 迭 代 完 成8个DataBlock的 数 据 收 集。 -
- Atlas A3 训 练 系 列 产 品/Atlas A3 推 理 系 列 产 品
- Atlas A2 训 练 系 列 产 品/Atlas A2 推 理 系 列 产 品
- Atlas 推 理 系 列 产 品AI Core
针 对Ascend 950PR/Ascend 950DT，int8\_t、uint8\_t数 据 类 型 仅 支 持tensor前n个 数 据 计 算 接 口。
调 用 示 例
本 样 例 中 只 展 示Compute流 程 中 的 部 分 代 码。如 果 您 需 要 运 行 样 例 代 码，请 将 该 代 码 段 拷 贝 并 替 换Gather类 样 例场 景 三 中Compute函 数 的 部 分 代 码 即 可。
tensor高 维 切 分 计 算 样 例-mask连 续 模 式。
C++// repeatTime = 4, mask = 128, 128 elements one repeat, 512 elements total // srcLocal数 据 类 型 为half，srcOffsetLocal数 据 类 型 为uint32_t，dstLocal数 据 类 型 为half // srcBaseAddr = 0, srcLocal的 起 始 基 地 址 为0 // dstRepStride = 8, no gap between repeats AscendC::Gather(dstLocal, srcLocal, srcOffsetLocal, (uint32_t)0, 128, 4, 8);tensor高 维 切 分 计 算 样 例-mask逐bit模 式。
C++uint64_t mask[2] = { 0xFFFFFFFFFFFFFFFF, 0xFFFFFFFFFFFFFFFF }; // repeatTime = 4, 128 elements one repeat, 512 elements total // srcLocal数 据 类 型 为half，srcOffsetLocal数 据 类 型 为uint32_t，dstLocal数 据 类 型 为half // srcBaseAddr = 0, srcLocal的 起 始 基 地 址 为0 // dstRepStride = 8, no gap between repeats AscendC::Gather(dstLocal, srcLocal, srcOffsetLocal, (uint32_t)0, mask, 4, 8);tensor前n个 数 据 计 算 样 例。
C++uint32_t count = 512; // 参 与 计 算 的 元 素 个 数 // srcLocal数 据 类 型 为half，srcOffsetLocal数 据 类 型 为uint32_t，dstLocal数 据 类 型 为half // srcBaseAddr = 0, srcLocal的 起 始 基 地 址 为0 AscendC::Gather(dstLocal, srcLocal, srcOffsetLocal, (uint32_t)0, count);
结 果 示 例 如 下：
输 入 数 据srcOffsetLocal:
[254 252 250 ... 4 2 0]
输 入 数 据srcLocal（128个half类 型 数 据）:
[0 1 2 ... 125 126 127]
输 出 数 据(dstLocal)初 始 值:
[0. 0. 0. 0. 0. 0. ... 0.]
进 行Gather计 算 后，输 出 数 据(dstLocal):
[127 126 125 ... 2 1 0]