GatherB

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 the source operand into the result register tensor by data block based on the index positions after the base address and index of the source operand in Unified Buffer are given. The length of each data block is 32 bytes. The following figure shows the collection process.

Prototype

template <typename T = DefaultType, typename U, typename S>
__simd_callee__ inline void GatherB(U& dstReg, __ubuf__ T* baseAddr, S& index, MaskReg& mask)

Parameters

**Table 1** Template parameters
Parameter	Description
T	Data type of the destination and source operands.
U	RegTensor type of the destination operand, for example, RegTensor<half>. It is automatically inferred by the compiler and does not need to be specified.
S	RegTensor type of the index value, for example, RegTensor<uint32_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. For the Atlas 350 Accelerator Card, the supported data types are uint8_t, int8_t, uint16_t, int16_t, uint32_t, int32_t, half, float, bfloat16_t, uint64_t, and int64_t.
baseAddr	Input	Base address of the source operand in Unified Buffer. 32-byte alignment is required. The type is Unified Buffer pointer. For the Atlas 350 Accelerator Card, the supported data types are uint8_t, int8_t, uint16_t, int16_t, uint32_t, int32_t, half, float, bfloat16_t, uint64_t, and int64_t.
index	Input	Index position of each data block in dstReg relative to baseAddr in Unified Buffer. The index position must be greater than or equal to 0 and 32-byte aligned. The type is RegTensor. For the Atlas 350 Accelerator Card, the supported data type is uint32_t.
mask	Input	Validity indicator for the src element. For details, see MaskReg.

Constraints

The destination operand has the same data type as the source operand.
The base address of the source operand in Unified Buffer must be 32-byte aligned.
The index position must be greater than or equal to 0 and 32-byte aligned.
The index register can contain the same value. This means the data of the same data block in the source operand can be read multiple times.

Example

template<typename T, typename U>
__simd_vf__ inline void GatherBVF(__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::GatherB(dstReg, src0Addr, srcReg, mask);
        AscendC::Reg::StoreAlign(dstAddr + i * oneRepeatSize, dstReg, mask);
    }
}

Parent topic: Gather and Scatter