MrgSort

函数功能

将已经排好序的最多4条队列,合并排列成1条队列,结果按照score域由大到小排序,排布方式如下:

Atlas A2训练系列产品/Atlas 800I A2推理产品采用方式一

Atlas推理系列产品AI Core采用方式二

函数原型

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

参数说明

表1 模板参数说明

接口

功能

T

操作数的数据类型。

isExhaustedSuspension

某条队列耗尽(即该队列已经全部排序到目的操作数)后,是否需要停止合并。类型为bool,参数取值如下:

  • false:直到所有队列耗尽完才停止合并。
  • true:某条队列耗尽后,停止合并。

默认值为false。

表2 接口参数说明

参数名称

输入/输出

含义

dstLocal

输出

目的操作数,存储经过排序后的数据。

类型为LocalTensor,支持的TPosition为VECIN/VECCALC/VECOUT。

Atlas A2训练系列产品/Atlas 800I A2推理产品,支持的数据类型为:half/float

Atlas推理系列产品AI Core,支持的数据类型为:half/float

sortList

输入

源操作数,支持2-4个队列,并且每个队列都已经排好序,类型为MrgSortSrcList结构体,具体请参考表3。MrgSortSrcList中传入要合并的队列。

template <typename T>
struct MrgSortSrcList {
    LocalTensor<T> src1;
    LocalTensor<T> src2;
    LocalTensor<T> src3; // 当要合并的队列个数小于3,可以为空tensor
    LocalTensor<T> src4; // 当要合并的队列个数小于4,可以为空tensor
};

Atlas A2训练系列产品/Atlas 800I A2推理产品,支持的数据类型为:half/float

Atlas推理系列产品AI Core,支持的数据类型为:half/float

elementCountList

输入

四个源队列的长度(8Bytes结构的数目),类型为长度为4的uint16_t数据类型的数组,理论上每个元素取值范围[0, 4095],但不能超出UB的存储空间。

sortedNum

输出

耗尽模式下(即isExhaustedSuspension为true时),停止合并时每个队列已排序的元素个数。

validBit

输入

有效队列个数,取值如下:
  • 0b11:前两条队列有效
  • 0b111:前三条队列有效
  • 0b1111:四条队列全部有效

repeatTimes

输入

迭代次数,每一次源操作数和目的操作数跳过四个队列总长度。取值范围:repeatTimes∈[1,255]。

repeatTimes参数生效是有条件的,需要同时满足以下四个条件:
  • srcLocal包含四条队列并且validBit=15
  • 四个源队列的长度一致
  • 四个源队列连续存储
  • isExhaustedSuspension为false
表3 MrgSortSrcList参数说明

参数名称

输入/输出

含义

src1

输入

源操作数,第一个已经排好序的队列。

类型为LocalTensor,支持的TPosition为VECIN/VECCALC/VECOUT。

数据类型与目的操作数保持一致。

Atlas A2训练系列产品/Atlas 800I A2推理产品,支持的数据类型为:half/float

Atlas推理系列产品AI Core,支持的数据类型为:half/float

src2

输入

源操作数,第二个已经排好序的队列。

类型为LocalTensor,支持的TPosition为VECIN/VECCALC/VECOUT。

数据类型与目的操作数保持一致。

Atlas A2训练系列产品/Atlas 800I A2推理产品,支持的数据类型为:half/float

Atlas推理系列产品AI Core,支持的数据类型为:half/float

src3

输入

源操作数,第三个已经排好序的队列。

类型为LocalTensor,支持的TPosition为VECIN/VECCALC/VECOUT。

数据类型与目的操作数保持一致。

Atlas A2训练系列产品/Atlas 800I A2推理产品,支持的数据类型为:half/float

Atlas推理系列产品AI Core,支持的数据类型为:half/float

src4

输入

源操作数,第四个已经排好序的队列。

类型为LocalTensor,支持的TPosition为VECIN/VECCALC/VECOUT。

数据类型与目的操作数保持一致。

Atlas A2训练系列产品/Atlas 800I A2推理产品,支持的数据类型为:half/float

Atlas推理系列产品AI Core,支持的数据类型为:half/float

返回值

支持的型号

Atlas A2训练系列产品/Atlas 800I A2推理产品

Atlas推理系列产品AI Core

约束说明

调用示例