MrgSort

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	x
Atlas inference product AI Core	√
Atlas inference product Vector Core	x
Atlas training product	x

Function Usage

Merges at most four sorted lists into one. The results are sorted in descending order of the score fields. The layout modes are as follows:

For the Atlas 350 Accelerator Card, mode 1 is used.

For the Atlas A3 training product/Atlas A3 inference product, mode 1 is used.

For the Atlas A2 training product/Atlas A2 inference product, mode 1 is used.

For the Atlas inference product AI Core, mode 2 is used.

Layout mode 1:
Generally, the data to be processed by MrgSort is the data processed by Sort, or the output of Sort. The list structure is as follows:
- When the data type is float, each structure occupies 8 bytes.
- When the data type is half, each structure occupies 8 bytes, but 2 bytes in the middle are reserved.

Layout mode 2: region proposal
The input and output data are region proposals. For details, see mode 2 in Sort.

Prototype

template <typename T, bool isExhaustedSuspension = false>
__aicore__ inline void MrgSort(const LocalTensor<T> &dst, const MrgSortSrcList<T> &sortList, const uint16_t elementCountList[4], uint32_t sortedNum[4], uint16_t validBit, const int32_t repeatTime)

Parameters

**Table 1** Template parameters
Parameter	Description
T	Operand data type. For the Atlas 350 Accelerator Card, the supported data types are half and float. For the Atlas A3 training product/Atlas A3 inference product, the supported data types are half and float. For the Atlas A2 training product/Atlas A2 inference product, the supported data types are half and float. For the Atlas inference product AI Core, the supported data types are half and float.
isExhaustedSuspension	Whether to stop merging after a list is exhausted (that is, all operations in the list have been sorted to the destination operand list). The type is bool. The options are as follows: false: The merging stops only when all lists are exhausted. true: The merging stops after a list is exhausted. The default value is false.

Table 2 API parameters

Parameter

Input/Output

Description

dst

Output

Destination operand, which stores sorted data.

The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT.

sortList

Input

Source operand of the MrgSortSrcList struct type, which can contain two to four sorted lists. For details, see Table 3. The lists to be merged are passed to MrgSortSrcList.

template <typename T>
struct MrgSortSrcList {
    LocalTensor<T> src1;
    LocalTensor<T> src2;
    LocalTensor<T> src3; // If the number of lists to be merged is less than 3, the tensor can be empty.
    LocalTensor<T> src4; // If the number of lists to be merged is less than 4, the tensor can be empty.
};

elementCountList

Input

Lengths of the four source lists (sorting mode 1: number of 8-byte structures; sorting mode 2: number of 16 × sizeof(T)Bytes structures). The type is an array of the uint16_t data type with a length of 4. Theoretically, the value range of each element is [0, 4095], but the value cannot exceed the storage space of the UB.

sortedNum

Output

Number of sorted elements in each list when merging is stopped in exhaustion mode (that is, when isExhaustedSuspension is true).

validBit

Input

Number of valid lists. The values are as follows:

0b11: The first two lists are valid.
0b111: The first three lists are valid.
0b1111: All the four lists are valid.

repeatTime

Input

Number of iteration repeats. The total length of the four lists is skipped for the source and destination operands in each iteration. Value range: repeatTime ∈ [1, 255]

The repeatTime parameter takes effect only when the following conditions are met:

srcLocal contains four lists and validBit is 15.
The lengths of the four source lists are the same.
The four source lists are stored consecutively.
isExhaustedSuspension is false.

**Table 3** **MrgSortSrcList** parameters
Parameter	Input/Output	Description
src1	Input	Source operand, which stores the first sorted list. The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT. The source operand must have the same data type as the destination operand. For the Atlas 350 Accelerator Card, the supported data types are half and float. For the Atlas A3 training product/Atlas A3 inference product, the supported data types are half and float. For the Atlas A2 training product/Atlas A2 inference product, the supported data types are half and float. For the Atlas inference product AI Core, the supported data types are half and float.
src2	Input	Source operand, which stores the second sorted list. The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT. The source operand must have the same data type as the destination operand. For the Atlas 350 Accelerator Card, the supported data types are half and float. For the Atlas A3 training product/Atlas A3 inference product, the supported data types are half and float. For the Atlas A2 training product/Atlas A2 inference product, the supported data types are half and float. For the Atlas inference product AI Core, the supported data types are half and float.
src3	Input	Source operand, which stores the third sorted list. The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT. The source operand must have the same data type as the destination operand. For the Atlas 350 Accelerator Card, the supported data types are half and float. For the Atlas A3 training product/Atlas A3 inference product, the supported data types are half and float. For the Atlas A2 training product/Atlas A2 inference product, the supported data types are half and float. For the Atlas inference product AI Core, the supported data types are half and float.
src4	Input	Source operand, which stores the fourth sorted list. The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT. The source operand must have the same data type as the destination operand. For the Atlas 350 Accelerator Card, the supported data types are half and float. For the Atlas A3 training product/Atlas A3 inference product, the supported data types are half and float. For the Atlas A2 training product/Atlas A2 inference product, the supported data types are half and float. For the Atlas inference product AI Core, the supported data types are half and float.

Returns

None

Constraints

When score[i] is the same as score[j], if i > j, score[j] is selected first. That is, the index sequence is the same as the input sequence.
Data within each iteration is sorted, but data among different iterations is not sorted.
For details about the operand address alignment requirements, see General Address Alignment Restrictions.

Examples

Processing 128 pieces of half-type data

This example applies to:

Atlas 350 Accelerator Card

Atlas A2 training product/Atlas A2 inference product

Atlas A3 training product/Atlas A3 inference product

For details about a complete example, see MrgSort operator sample.

uint32_t elementCount = 128; // Number of elements.
uint32_t calcBufferSize = elementCount * 8; // Each element occupies 8 bytes.
uint32_t tmpBufferSize = elementCount * 8;
uint32_t sortedLocalSize = elementCount * 4; 
uint32_t sortRepeatTimes = elementCount / 32;
uint32_t extractRepeatTimes = elementCount / 32;
uint32_t sortTmpLocalSize = elementCount * 4;

uint32_t singleMergeTmpElementCount = elementCount / 4;
uint32_t baseOffset = AscendC::GetSortOffset<half>(singleMergeTmpElementCount);
AscendC::MrgSortSrcList sortList = AscendC::MrgSortSrcList(sortedLocal[0], sortedLocal[baseOffset], sortedLocal[2 * baseOffset], sortedLocal[3 * baseOffset]); // sortList: list of ordered queues to be merged
uint16_t singleDataSize = elementCount / 4; // Queue length
const uint16_t elementCountList[4] = {singleDataSize, singleDataSize, singleDataSize, singleDataSize}; // Length of 4 queues
uint32_t sortedNum[4];
// Merge four queues in sortList.
AscendC::MrgSort<half, false>(sortTmpLocal, sortList, elementCountList, sortedNum, 0b1111, 1);

Result example:
Input data (srcValueGm): 128 pieces of half-type data
[31 30 29 ... 2 1 0
 63 62 61 ... 34 33 32
 95 94 93 ... 66 65 64
 127 126 125 ... 98 97 96]
Input (srcIndexGm):
[31 30 29 ... 2 1 0
 63 62 61 ... 34 33 32
 95 94 93 ... 66 65 64
 127 126 125 ... 98 97 96]
Output (dstValueGm):
[127 126 125 ... 2 1 0]
Output (dstIndexGm):
[127 126 125 ... 2 1 0]

Processing 64 pieces of half-type data

This example applies to:

Atlas inference product AI Core

uint32_t elementCount = 64; // Number of elements
// Number of elements in a single queue
uint32_t singleMergeTmpElementCount = elementCount / 4;
uint32_t baseOffset = AscendC::GetSortOffset<half>(singleMergeTmpElementCount);
AscendC::MrgSortSrcList sortList = AscendC::MrgSortSrcList(sortedLocal[0], sortedLocal[baseOffset], sortedLocal[2 * baseOffset], sortedLocal[3 * baseOffset]); // sortList: list of ordered queues to be merged
uint16_t singleDataSize = elementCount / 4; // Queue length
const uint16_t elementCountList[4] = {singleDataSize, singleDataSize, singleDataSize, singleDataSize}; // Length of 4 queues
uint32_t sortedNum[4];
// Merge four queues in sortList.
AscendC::MrgSort<half, false>(sortTmpLocal, sortList, elementCountList, sortedNum, 0b1111, 1);

Result example:
Input data (srcValueGm): 64 pieces of half-type data
[15 14 13 ... 2 1 0
 31 30 29 ... 18 17 16
 47 46 45 ... 34 33 32
 63 62 61 ... 50 49 48]
Input (srcIndexGm):
[15 14 13 ... 2 1 0
 31 30 29 ... 18 17 16
 47 46 45 ... 34 33 32
 63 62 61 ... 50 49 48]
Output (dstValueGm):
[63 62 61 ... 2 1 0]
Output (dstIndexGm):
[63 62 61 ... 2 1 0]

Parent topic: Sorting Operations