Gather

Applicability

Function Usage

• Collects elements in Unified Buffer.

  Collects the source operand into the result register tensor by element based on the index positions after the base address and index of the source operand in Unified Buffer are given. The following figure shows the collection process.

  

• Collects elements in RegTensor.

  Collects srcReg into dstReg by element based on the index positions in indexReg. The following figure shows the collection process.

  

Prototype

• Collecting elements in Unified Buffer

  template <typename T0 = DefaultType, typename T1, typename T2 = DefaultType, typename T3, typename T4>
  __simd_callee__ inline void Gather(T3& dstReg, __ubuf__ T1* baseAddr, T4& index, MaskReg& mask)

• Collecting elements in RegTensor

  1 2	template <typename T = DefaultType, typename U = DefaultType, typename S, typename V> __simd_callee__ inline void Gather(S& dstReg, S& srcReg, V& indexReg)

Parameters

• Collecting elements in Unified Buffer

  Table 1 Template parameters

  Parameter	Description
  T0	Data type of the destination operand.
  T1	Data type of the source operand.
  T2	Data type of the index value.
  T3	RegTensor type of the destination operand, for example, RegTensor<half>. It is automatically inferred by the compiler and does not need to be specified.
  T4	RegTensor type of the index value, for example, RegTensor<uint16_t>. It is automatically inferred by the compiler and does not need to be specified.

  Table 2 Function parameters

  Parameter	Input/Output	Description
  dstReg	Output	Destination operand. The type is RegTensor. The T0 data type must be used together with the T1 and T2 data types. For details about the type mapping table, see the constraints. For details about the supported data types, see Table 5.
  baseAddr	Input	Base address of the source operand in Unified Buffer. The type is Unified Buffer pointer. The T1 data type must be used together with the T0 and T2 data types. For details about the type mapping table, see the constraints. For details about the supported data types, see Table 5.
  index	Input	Index of each element in dstReg relative to baseAddr in Unified Buffer. The address offset must be greater than or equal to 0. The type is RegTensor. The T2 data type must be used together with the T0 and T1 data types. For details about the type mapping table, see the constraints. For details about the supported data types, see Table 5.
  mask	Input	Validity indicator for the src element. For details, see MaskReg.

• Collecting elements in RegTensor

  Table 3 Template parameters

  Parameter	Description
  T	Data type of the destination and source operands. For the Atlas 350 Accelerator Card, the supported data types are b8, b16, and b32.
  U	Data type of the index value. For the Atlas 350 Accelerator Card, the supported data types are uint8_t, uint16_t, and uint32_t.
  S	RegTensor type of the destination operand, for example, RegTensor<half>. It is automatically inferred by the compiler and does not need to be specified.
  V	RegTensor type of the index value, for example, RegTensor<uint16_t>. It is automatically inferred by the compiler and does not need to be specified.

  Table 4 Function parameters

  Parameter	Input/Output	Description
  dstReg	Output	Destination operand. The type is RegTensor.
  srcReg	Input	Source operand. The type is RegTensor. The data type must be the same as that of the destination operand.
  indexReg	Input	Data index. The type is RegTensor. The bit width of the data type must be the same as that of the destination operand. srcReg is of the RegTensor type. The bit width is fixed to VL, and the number of stored elements is fixed. If the index value in indexReg exceeds the maximum number of data elements that can be stored in the current RegTensor, the following processing is performed: Assume that the maximum number of data elements that can be stored in the current RegTensor is vlLength, and the index value in indexReg is i. The index value is updated to i % vlLength.

Constraints

• Collecting elements in Unified Buffer
  • The T0, T1, and T2 data types must be used together. The following table lists the mapping relationships.

    Table 5 Data types of the Gather operand

    T0	T1	T2
    int16_t	int8_t	uint16_t
    uint16_t	uint8_t
    int16_t	int16_t
    uint16_t	uint16_t
    half	half
    bfloat16_t	bfloat16_t
    int32_t	int32_t	uint32_t
    uint32_t	uint32_t
    float	float
    uint64_t	uint64_t	uint32_t
    int64_t	int64_t
    uint64_t	uint64_t	uint64_t
    int64_t	int64_t

  • When T1 is of the B8 data type and T0 is of the B16 data type, the lower 8 bits of the destination operand are the same as those of the source operand, and the upper 8 bits are automatically padded with 0s. For example, if T1 is of the int8 data type:

    40=0b00101000 -> 0b0000000000101000: The value is 40 after being extended to 16 bits.

    -40=0b11011000 -> 0b0000000011011000: The value is 216 after being extended to 16 bits.

  • When T1 is of the B64 data type, the data types of T0, T1, T2, T4, and T3 can only be one of the following combinations:

    T0	T1	T2	T4	T3	Remarks
    B64	B64	uint32_t	RegTensor<uint32_t>	RegTensor<uint64_t, RegTraitNumOne>	The first 32 index values are valid.
    				RegTensor<int64_t, RegTraitNumOne>
    				RegTensor<uint64_t, RegTraitNumTwo>	-
    				RegTensor<int64_t, RegTraitNumTwo>
    B64	B64	uint64_t	RegTensor<uint64_t, RegTraitNumOne>	RegTensor<uint64_t, RegTraitNumOne>	-
    				RegTensor<int64_t, RegTraitNumOne>
    			RegTensor<uint64_t, RegTraitNumTwo>	RegTensor<uint64_t, RegTraitNumOne>	The first 32 index values are valid.
    				RegTensor<int64_t, RegTraitNumOne>
    				RegTensor<uint64_t, RegTraitNumTwo>	-
    				RegTensor<int64_t, RegTraitNumTwo>

• Collecting elements in RegTensor

  No constraints

Examples

• Collecting elements in Unified Buffer

  template <typename T, typename U>
  __simd_vf__ inline void GatherVF(__ubuf__ T* dstAddr, __ubuf__ T* src0Addr, __ubuf__ U* src1Addr, uint32_t count, uint32_t oneRepeatSize, uint16_t repeatTimes)
  {
      AscendC::Reg::RegTensor<U> srcReg;
      AscendC::Reg::RegTensor<T> dstReg;
      AscendC::Reg::MaskReg mask;    
      for (uint16_t i = 0; i < repeatTimes; i++) {
          mask = AscendC::Reg::UpdateMask<T>(count);
          AscendC::Reg::LoadAlign(srcReg, src1Addr + i * oneRepeatSize);
          AscendC::Reg::Gather(dstReg, src0Addr, srcReg, mask);
          AscendC::Reg::StoreAlign(dstAddr + i * oneRepeatSize, dstReg, mask);
      }
  }

• Collecting elements in RegTensor

  template <typename T, typename U>
  __simd_vf__ inline void GatherVF(__ubuf__ T* dstAddr, __ubuf__ T* src0Addr, __ubuf__ U* src1Addr, uint32_t count, uint32_t oneRepeatSize, uint16_t repeatTimes)
  {
      AscendC::Reg::RegTensor<T> srcReg0, dstReg;
      AscendC::Reg::RegTensor<U> srcReg1;
      AscendC::Reg::MaskReg mask;
      AscendC::Reg::LoadAlign(srcReg1, src1Addr);
      for (uint16_t i = 0; i < repeatTimes; i++) {
          mask = AscendC::Reg::UpdateMask<T>(count);
          AscendC::Reg::LoadAlign(srcReg0, src0Addr + i * oneRepeatSize);
          AscendC::Reg::Gather(dstReg, srcReg0, srcReg1);
          AscendC::Reg::StoreAlign(dstAddr + i * oneRepeatSize, dstReg, mask);
      }
  }