Gather (Supporting Registers as Source Operands)

Applicability

Product	Supported
Atlas 350 Accelerator Card	√
Atlas A3 training product/Atlas A3 inference product	x
Atlas A2 training product/Atlas A2 inference product	x
Atlas 200I/500 A2 inference product	x
Atlas inference product AI Core	x
Atlas inference product Vector Core	x
Atlas training product	x

Function Usage

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

Figure 1 Gather function usage

Prototype

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

Parameters

**Table 1** 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 2** 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

None

Example

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);
    }
}

Parent topic: Gather and Scatter