CANN/pyasc Gather算子接口文档

asc.language.basic.gather

【免费下载链接】pyasc本项目为Python用户提供算子编程接口，支持在昇腾AI处理器上加速计算，接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc

asc.language.basic.gather(dst: LocalTensor, src: LocalTensor, src_offset: LocalTensor, src_base: int, mask: int, repeat_times: int, dst_rep_stride: int) → None

asc.language.basic.gather(dst: LocalTensor, src: LocalTensor, src_offset: LocalTensor, src_base: int, mask: List[int], repeat_times: int, dst_rep_stride: int) → None

asc.language.basic.gather(dst: LocalTensor, src: LocalTensor, src_offset: LocalTensor, src_base: int, count: int) → None

给定输入的张量和一个地址偏移张量，本接口根据偏移地址将输入张量按元素收集到结果张量中。

对应的Ascend C函数原型

• tensor前n个数据计算

  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模式

    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连续模式

    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)

参数说明

• dst：目的操作数。类型为LocalTensor，支持的TPosition为VECIN/VECCALC/VECOUT。LocalTensor的起始地址需要32字节对齐。
• src: 源操作数。类型为LocalTensor，支持的TPosition为VECIN/VECCALC/VECOUT。LocalTensor的起始地址需要32字节对齐。
• src_offset：每个元素在src中对应的地址偏移。类型为LocalTensor，支持的TPosition为VECIN/VECCALC/VECOUT。LocalTensor的起始地址需要32字节对齐。 该偏移量相对于src的起始基地址而言。单位为Bytes。取值要求如下：
  • 取值应保证src元素类型位宽对齐。
  • 偏移地址后不能超出UB大小数据的范围。
  • 地址偏移的取值范围：不能超出uint32_t的范围。
• src_base：src的起始基地址，用于指定Gather操作中源操作数的起始位置，单位为Bytes。取值应保证src元素类型位宽对齐，否则会导致非预期行为。
• count：执行处理的数据个数。
• mask： 控制每次迭代内参与计算的元素。
  • 逐bit模式：mask为数组形式。数组长度和数组元素的取值范围和操作数的数据类型有关。可以按位控制哪些元素参与计算，bit位的值为1表示参与计算，0表示不参与。
    • 操作数 16 位：数组长度 2，mask[0], mask[1] ∈ [0, 2⁶⁴-1]，且不能同时为 0
    • 操作数 32 位：数组长度 1，mask[0] ∈ (0, 2⁶⁴-1]
    • 操作数 64 位：数组长度 1，mask[0] ∈ (0, 2³²-1]
    • 例如：mask = [8, 0]，表示仅第 4 个元素参与计算
  • 连续模式：mask为整数形式。表示前面连续多少个元素参与计算。取值范围和操作数的数据类型有关，数据类型不同，每次迭代内能够处理的元素个数最大值不同。
    • 操作数 16 位：mask ∈ [1, 128]
    • 操作数 32 位：mask ∈ [1, 64]
    • 操作数 64 位：mask ∈ [1, 32]
• repeat_time：指令迭代次数，每次迭代完成8个datablock（32Bytes）的数据收集，数据范围：repeat_time∈[0,255]。
• dst_rep_stride：相邻迭代间的地址步长，单位是datablock（32Bytes）。

约束说明

• 操作数地址对齐要求请参见 《Ascend C算子开发接口》 中的“通用说明和约束-通用地址对齐约束”。
• 操作数地址重叠约束请参考 《Ascend C算子开发接口》 中的“通用说明和约束-通用地址重叠约束”。

调用示例

z_local = asc.LocalTensor(dtype=asc.float16, pos=asc.TPosition.VECOUT, addr=0, tile_size=512) src_offset = asc.LocalTensor(dtype=asc.uint32, pos=asc.TPosition.VECIN, addr=0, tile_size=512) x_local = asc.LocalTensor(dtype=asc.float16, pos=asc.TPosition.VECIN, addr=0, tile_size=512) asc.gather(z_local, x_local, src_offset, src_base=0, count=512) asc.gather(z_local, x_local, src_offset, src_base=0, mask=512, repeat_times=1, dst_rep_stride=8) uint64_max = 2**64 - 1 mask_bits = [uint64_max, uint64_max] asc.gather(z_local, x_local, src_offset, src_base=0, mask=mask_bits, repeat_times=1, dst_rep_stride=8)

【免费下载链接】pyasc本项目为Python用户提供算子编程接口，支持在昇腾AI处理器上加速计算，接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc