Gather

Applicability

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

Function Usage

Gathers elements from the input tensor into the result tensor based on the provided address offset tensor.

Prototype

Computation of the first n data elements of a tensor

        
             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)

High-dimensional tensor sharding computation

Bitwise mask mode

          
               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)

Contiguous mask mode

          
               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)

Parameters

**Table 1** Parameters in the template
Parameter	Description
T	Operand data type. For the Atlas 350 Accelerator Card, the supported data types are uint8_t, int8_t, uint16_t, int16_t, half, bfloat16_t, uint32_t, int32_t, float, uint64_t, and int64_t. For the Atlas A3 training product / Atlas A3 inference product , the supported data types are int16_t, uint16_t, int32_t, uint32_t, float, half, and bfloat16_t. For the Atlas A2 training product / Atlas A2 inference product , the supported data types are int16_t, uint16_t, int32_t, uint32_t, float, half, and bfloat16_t. For the Atlas 200I/500 A2 inference product , the supported data types are uint8_t, int8_t, uint16_t, int16_t, half, uint32_t, int32_t, and float. For the Atlas inference product AI Core, the supported data types are int16_t, uint16_t, int32_t, uint32_t, float, and half.

**Table 2** Parameters
Parameter	Input/Output	Meaning
dst	Output	Destination operand. The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT. The start address of LocalTensor must be 32-byte aligned.
src	Input	Source operand. The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT. The start address of LocalTensor must be 32-byte aligned. Its data type must match that of dst.
srcOffset	Input	Address offset of each element in src. The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT. The start address of LocalTensor must be 32-byte aligned. The offset is relative to the start base address of src. The unit is byte. The value must meet the following requirements: The value must be aligned to the bit width of the src element type. The final address must not exceed the UB range. For the following models, the address offset must be within the range of uint32_t. Atlas inference product AI Core Atlas A2 training product / Atlas A2 inference product Atlas A3 training product / Atlas A3 inference product For the following models, the address offset must be within the range of [0, 2¹⁶ – 1] when the operand is 8-bit, within the range of [0, 2¹⁷ – 1] when the operand is 16-bit, and within the range of uint32_t when the operand is 32-bit or 64-bit. Atlas 200I/500 A2 inference product Atlas 350 Accelerator Card
srcBaseAddr	Input	Start base address of src, which specifies the start position of the source operand in the gather operation. The unit is byte. Ensure that the bit width of the src element type is aligned. Otherwise, unexpected behavior occurs.
count	Input	Number of data elements to be processed.
mask/mask[]	Input	mask controls the elements that participate in computation in each iteration. Contiguous mode: indicates the number of contiguous elements that participate in computation. The value range is related to the operand data type. The maximum number of elements that can be processed in each repeat varies according to the data type. When the operand is 8-bit or 16-bit, mask ∈ [1, 128]. When the operand is 32-bit, mask ∈ [1, 64]. When the operand is 64-bit, mask ∈ [1, 32]. Bitwise mode: controls the elements that participate in computation by bit. If a bit is set to 1, the corresponding element participates in the computation. If a bit is set to 0, the corresponding element is masked in the computation. The parameter type is a uint64_t array with a length of 2. For example, if mask = [0, 8] and 8 = 0b1000, only the fourth element participates in computation. The parameter value range is related to the operand data type. The maximum number of elements that can be processed in each repeat varies according to the data type. When the operand is 8-bit or 16-bit, mask[0] and mask[1] ∈ [0, 2⁶⁴ – 1] and cannot be 0 at the same time. When the operand is 32-bit, mask[1] is 0 and mask[0] ∈ (0, 2⁶⁴ – 1]. When the operand is 64-bit, mask[1] is 0 and mask[0] ∈ (0, 2³² – 1].
repeatTime	Input	Number of instruction iterations. Data of eight data blocks (32 bytes) is collected in each iteration. Data range: repeatTime ∈ [0,255] For the following models: Atlas 200I/500 A2 inference product Atlas 350 Accelerator Card When the operand is 8-bit, four data blocks (32 bytes) are collected in each repeat.
dstRepStride	Input	Address stride of the operand between adjacent iterations. The unit is data block (32 bytes).

Restrictions

For details about the operand address alignment requirements, see General Address Alignment Restrictions.
For details about the constraints on operand address overlapping, see General Address Overlapping Restrictions.

For the Atlas 350 Accelerator Card, only the APIs that compute the first n data elements of a tensor are supported for the uint8_t and int8_t data types.

Examples

These examples show only part of the code used in the computation.

Example of high-dimensional tensor sharding computation (contiguous mask mode)

        
             // repeatTime = 4, mask = 128, 128 elements one repeat, 512 elements total
// srcLocal and dstLocal values are of the half type, and the srcOffsetLocal value is of the uint32_t type.
// srcBaseAddr = 0, indicating that the start base address of srcLocal is 0.
// dstRepStride = 8, no gap between repeats 
AscendC::Gather(dstLocal, srcLocal, srcOffsetLocal, (uint32_t)0, 128, 4, 8);

Example of high-dimensional tensor sharding computation (bitwise mask mode)

        
             uint64_t mask[2] = { 0xFFFFFFFFFFFFFFFF, 0xFFFFFFFFFFFFFFFF };
// repeatTime = 4, 128 elements one repeat, 512 elements total
// srcLocal and dstLocal values are of the half type, and the srcOffsetLocal value is of the uint32_t type.
// srcBaseAddr = 0, indicating that the start base address of srcLocal is 0.
// dstRepStride = 8, no gap between repeats 
AscendC::Gather(dstLocal, srcLocal, srcOffsetLocal, (uint32_t)0, mask, 4, 8);

Example of computing the first n data elements of a tensor

        
             uint32_t count = 512;    // Number of elements involved in computation
// srcLocal and dstLocal values are of the half type, and the srcOffsetLocal value is of the uint32_t type.
// srcBaseAddr = 0, indicating that the start base address of srcLocal is 0.
AscendC::Gather(dstLocal, srcLocal, srcOffsetLocal, (uint32_t)0, count);

Result example:

Input srcOffsetLocal:
[254 252 250 ... 4 2 0]
Input srcLocal (128 data elements of the half type):
[0 1 2 ... 125 126 127]
Initial value of output data (dstLocal):
[0. 0. 0. 0. 0. 0. ... 0.]
Output (dstLocal) after Gather computation:
[127 126 125 ... 2 1 0]

Parent topic: Data Padding