MrgSort

函数功能

将已经排好序的最多4条队列，合并排列成1条队列，结果按照score域由大到小排序。

• Atlas A2训练系列产品/Atlas 800I A2推理产品：MrgSort处理的数据一般是经过Sort处理后的数据，也就是Sort接口的输出，队列的结构如下所示：
  • 数据类型为float，每个结构占据8Bytes。

    

  • 数据类型为half，每个结构也占据8Bytes，其中中间有2Bytes保留。

    

• Atlas 推理系列产品（Ascend 310P处理器）：输入输出数据均为Region Proposal，具体请参见ProposalConcat中的Region Proposal说明。

函数原型

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)

参数说明

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推理产品

Atlas推理系列产品（Ascend 310P处理器）AI Core

约束说明

• 当存在score[i]与score[j]相同时，如果i>j，则score[j]将首先被选出来，排在前面，即index的顺序与输入顺序一致。
• 每次迭代内的数据会进行排序，不同迭代间的数据不会进行排序。
• 操作数地址偏移对齐要求请参见通用约束。

调用示例

• 处理64个half类型数据。

  该样例适用于：

  Atlas 推理系列产品（Ascend 310P处理器）

  #include "kernel_operator.h"
  
  namespace AscendC {
  template <typename T>
  class FullSort {
  public:
      __aicore__ inline FullSort() {}
      __aicore__ inline void Init(__gm__ uint8_t* srcValueGm, __gm__ uint8_t* srcIndexGm, __gm__ uint8_t* dstValueGm, __gm__ uint8_t* dstIndexGm)
      {
          concatRepeatTimes = m_elementCount / 16;
          inBufferSize = m_elementCount * sizeof(uint32_t);
          outBufferSize = m_elementCount * sizeof(uint32_t);
          calcBufferSize = m_elementCount * 8;
          tmpBufferSize = m_elementCount * 8;
          sortedLocalSize = m_elementCount * 8 * sizeof(T);
          sortRepeatTimes = m_elementCount / 16;
          extractRepeatTimes = m_elementCount / 16;
          sortTmpLocalSize = m_elementCount * 8 * sizeof(T);
          m_valueGlobal.SetGlobalBuffer((__gm__ T*)srcValueGm);
          m_indexGlobal.SetGlobalBuffer((__gm__ uint32_t*)srcIndexGm);
          m_dstValueGlobal.SetGlobalBuffer((__gm__ T*)dstValueGm);
          m_dstIndexGlobal.SetGlobalBuffer((__gm__ uint32_t*)dstIndexGm);
          m_pipe.InitBuffer(m_queIn, 2, inBufferSize);
          m_pipe.InitBuffer(m_queOut, 2, outBufferSize);
          m_pipe.InitBuffer(m_queCalc, 1, calcBufferSize*sizeof(T));
          m_pipe.InitBuffer(m_queTmp, 2, tmpBufferSize*sizeof(T));
      }
      __aicore__ inline void Process()
      {
          CopyIn();
          Compute();
          CopyOut();
      }
  private:
      __aicore__ inline void CopyIn()
      {
          LocalTensor<T> valueLocal = m_queIn.AllocTensor<T>();
          DataCopy(valueLocal, m_valueGlobal, m_elementCount);
          m_queIn.EnQue(valueLocal);
  
          LocalTensor<uint32_t> indexLocal = m_queIn.AllocTensor<uint32_t>();
          DataCopy(indexLocal, m_indexGlobal, m_elementCount);
          m_queIn.EnQue(indexLocal);
      }
      __aicore__ inline void Compute()
      {
          LocalTensor<T> valueLocal = m_queIn.DeQue<T>();
          LocalTensor<uint32_t> indexLocal = m_queIn.DeQue<uint32_t>();
          LocalTensor<T> sortedLocal = m_queCalc.AllocTensor<T>();
          LocalTensor<T> concatTmpLocal = m_queTmp.AllocTensor<T>();
          LocalTensor<T> sortTmpLocal = m_queTmp.AllocTensor<T>();
          LocalTensor<T> dstValueLocal = m_queOut.AllocTensor<T>();
          LocalTensor<uint32_t> dstIndexLocal = m_queOut.AllocTensor<uint32_t>();
          LocalTensor<T> concatLocal;
  
          Concat(concatLocal, valueLocal, concatTmpLocal, concatRepeatTimes);
          Sort<T, false>(sortedLocal, concatLocal, indexLocal, sortTmpLocal, sortRepeatTimes);
          uint32_t singleMergeTmpElementCount = m_elementCount / 4;
          uint32_t baseOffset;
          baseOffset = GetSortOffset(singleMergeTmpElementCount);
          MrgSortSrcList sortList = MrgSortSrcList(sortedLocal[0], sortedLocal[baseOffset], sortedLocal[2 * baseOffset], sortedLocal[3 * baseOffset]);
          uint16_t singleDataSize = m_elementCount / 4;
          const uint16_t elementCountList[4] = {singleDataSize, singleDataSize, singleDataSize, singleDataSize};
          uint32_t sortedNum[4];
  
          MrgSort<T, false>(sortTmpLocal, sortList, elementCountList, sortedNum, 0b1111, 1);
          Extract(dstValueLocal, dstIndexLocal, sortTmpLocal, extractRepeatTimes);
  
          m_queTmp.FreeTensor(concatTmpLocal);
          m_queTmp.FreeTensor(sortTmpLocal);
          m_queIn.FreeTensor(valueLocal);
          m_queIn.FreeTensor(indexLocal);
          m_queCalc.FreeTensor(sortedLocal);
          m_queOut.EnQue(dstValueLocal);
          m_queOut.EnQue(dstIndexLocal);
  
      }
      __aicore__ inline void CopyOut()
      {
          LocalTensor<T> dstValueLocal = m_queOut.DeQue<T>();
          LocalTensor<uint32_t> dstIndexLocal = m_queOut.DeQue<uint32_t>();
          DataCopy(m_dstValueGlobal, dstValueLocal, m_elementCount);
          DataCopy(m_dstIndexGlobal, dstIndexLocal, m_elementCount);
          m_queOut.FreeTensor(dstValueLocal);
          m_queOut.FreeTensor(dstIndexLocal);
      }
  private:
      TPipe m_pipe;
      TQue<QuePosition::VECIN, 2> m_queIn;
      TQue<QuePosition::VECOUT, 2> m_queOut;
      TQue<QuePosition::VECIN, 2> m_queTmp;
      TQue<QuePosition::VECIN, 1> m_queCalc;
      GlobalTensor<T> m_valueGlobal;
      GlobalTensor<uint32_t> m_indexGlobal;
      GlobalTensor<T> m_dstValueGlobal;
      GlobalTensor<uint32_t> m_dstIndexGlobal;
      uint32_t m_elementCount = 64;
      uint32_t concatRepeatTimes;
      uint32_t inBufferSize;
      uint32_t outBufferSize;
      uint32_t calcBufferSize;
      uint32_t tmpBufferSize;
      uint32_t sortedLocalSize;
      uint32_t sortTmpLocalSize;
      uint32_t sortRepeatTimes;
      uint32_t extractRepeatTimes;
  
  }; // class FullSort
  } // namespace AscendC
  
  extern "C" __global__ __aicore__ void FullSort(__gm__ uint8_t* src0Gm, __gm__ uint8_t* src1Gm, __gm__ uint8_t* dst0Gm, __gm__ uint8_t* dst1Gm)
  {
      AscendC::FullSort<half> op;
      op.Init(src0Gm, src1Gm, dst0Gm, dst1Gm);
      op.Process();
  }

  示例结果
  输入数据(srcValueGm): 128个float类型数据
  [15 14 13 ... 2 1 0
   31 30 29 ... 18 17 16
   47 46 45 ... 34 33 32
   63 62 61 ... 50 49 48]
  输入数据(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]
  输出数据(dstValueGm):
  [63 62 61 ... 2 1 0]
  输出数据(dstIndexGm):
  [63 62 61 ... 2 1 0]

• 处理128个half类型数据。

  该样例适用于：

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

  #include "kernel_operator.h"
  
  namespace AscendC {
  template <typename T>
  class FullSort {
  public:
      __aicore__ inline FullSort() {}
      __aicore__ inline void Init(__gm__ uint8_t* srcValueGm, __gm__ uint8_t* srcIndexGm, __gm__ uint8_t* dstValueGm, __gm__ uint8_t* dstIndexGm)
      {
          concatRepeatTimes = m_elementCount / 16;
          inBufferSize = m_elementCount * sizeof(uint32_t);
          outBufferSize = m_elementCount * sizeof(uint32_t);
          calcBufferSize = m_elementCount * 8;
          tmpBufferSize = m_elementCount * 8;
          sortedLocalSize = m_elementCount * 4;
          sortRepeatTimes = m_elementCount / 32;
          extractRepeatTimes = m_elementCount / 32;
          sortTmpLocalSize = m_elementCount * 4;
          m_valueGlobal.SetGlobalBuffer((__gm__ T*)srcValueGm);
          m_indexGlobal.SetGlobalBuffer((__gm__ uint32_t*)srcIndexGm);
          m_dstValueGlobal.SetGlobalBuffer((__gm__ T*)dstValueGm);
          m_dstIndexGlobal.SetGlobalBuffer((__gm__ uint32_t*)dstIndexGm);
          m_pipe.InitBuffer(m_queIn, 2, inBufferSize);
          m_pipe.InitBuffer(m_queOut, 2, outBufferSize);
          m_pipe.InitBuffer(m_queCalc, 1, calcBufferSize*sizeof(T));
          m_pipe.InitBuffer(m_queTmp, 2, tmpBufferSize*sizeof(T));
      }
      __aicore__ inline void Process()
      {
          CopyIn();
          Compute();
          CopyOut();
      }
  private:
      __aicore__ inline void CopyIn()
      {
          LocalTensor<T> valueLocal = m_queIn.AllocTensor<T>();
          DataCopy(valueLocal, m_valueGlobal, m_elementCount);
          m_queIn.EnQue(valueLocal);
          LocalTensor<uint32_t> indexLocal = m_queIn.AllocTensor<uint32_t>();
          DataCopy(indexLocal, m_indexGlobal, m_elementCount);
          m_queIn.EnQue(indexLocal);
      }
      __aicore__ inline void Compute()
      {
          LocalTensor<T> valueLocal = m_queIn.DeQue<T>();
          LocalTensor<uint32_t> indexLocal = m_queIn.DeQue<uint32_t>();
          LocalTensor<T> sortedLocal = m_queCalc.AllocTensor<T>();
          LocalTensor<T> concatTmpLocal = m_queTmp.AllocTensor<T>();
          LocalTensor<T> sortTmpLocal = m_queTmp.AllocTensor<T>();
          LocalTensor<T> dstValueLocal = m_queOut.AllocTensor<T>();
          LocalTensor<uint32_t> dstIndexLocal = m_queOut.AllocTensor<uint32_t>();
          LocalTensor<T> concatLocal;
  
          Concat(concatLocal, valueLocal, concatTmpLocal, concatRepeatTimes);
          Sort<T, false>(sortedLocal, concatLocal, indexLocal, sortTmpLocal, sortRepeatTimes);
          uint32_t singleMergeTmpElementCount = m_elementCount / 4;
          uint32_t baseOffset;
          baseOffset = GetSortOffset(singleMergeTmpElementCount);
          MrgSortSrcList sortList = MrgSortSrcList(sortedLocal[0], sortedLocal[baseOffset], sortedLocal[2 * baseOffset], sortedLocal[3 * baseOffset]);
          uint16_t singleDataSize = m_elementCount / 4;
          const uint16_t elementCountList[4] = {singleDataSize, singleDataSize, singleDataSize, singleDataSize};
          uint32_t sortedNum[4];
          MrgSort<T, false>(sortTmpLocal, sortList, elementCountList, sortedNum, 0b1111, 1);
          Extract(dstValueLocal, dstIndexLocal, sortTmpLocal, extractRepeatTimes);
  
          m_queTmp.FreeTensor(concatTmpLocal);
          m_queTmp.FreeTensor(sortTmpLocal);
          m_queIn.FreeTensor(valueLocal);
          m_queIn.FreeTensor(indexLocal);
          m_queCalc.FreeTensor(sortedLocal);
          m_queOut.EnQue(dstValueLocal);
          m_queOut.EnQue(dstIndexLocal);
  
      }
      __aicore__ inline void CopyOut()
      {
          LocalTensor<T> dstValueLocal = m_queOut.DeQue<T>();
          LocalTensor<uint32_t> dstIndexLocal = m_queOut.DeQue<uint32_t>();
          DataCopy(m_dstValueGlobal, dstValueLocal, m_elementCount);
          DataCopy(m_dstIndexGlobal, dstIndexLocal, m_elementCount);
          m_queOut.FreeTensor(dstValueLocal);
          m_queOut.FreeTensor(dstIndexLocal);
      }
  private:
      TPipe m_pipe;
      TQue<QuePosition::VECIN, 2> m_queIn;
      TQue<QuePosition::VECOUT, 2> m_queOut;
      TQue<QuePosition::VECIN, 2> m_queTmp;
      TQue<QuePosition::VECIN, 1> m_queCalc;
      GlobalTensor<T> m_valueGlobal;
      GlobalTensor<uint32_t> m_indexGlobal;
      GlobalTensor<T> m_dstValueGlobal;
      GlobalTensor<uint32_t> m_dstIndexGlobal;
      uint32_t m_elementCount = 128;
      uint32_t concatRepeatTimes;
      uint32_t inBufferSize;
      uint32_t outBufferSize;
      uint32_t calcBufferSize;
      uint32_t tmpBufferSize;
      uint32_t sortedLocalSize;
      uint32_t sortTmpLocalSize;
      uint32_t sortRepeatTimes;
      uint32_t extractRepeatTimes;
  
  }; // class FullSort
  } // namespace AscendC
  
  extern "C" __global__ __aicore__ void FullSort(__gm__ uint8_t* src0Gm, __gm__ uint8_t* src1Gm, __gm__ uint8_t* dst0Gm, __gm__ uint8_t* dst1Gm)
  {
      AscendC::FullSort<half> op;
      op.Init(src0Gm, src1Gm, dst0Gm, dst1Gm);
      op.Process();
  }

  示例结果
  输入数据(srcValueGm): 128个float类型数据
  [31 30 29 ... 2 1 0
   63 62 61 ... 34 33 32
   95 94 93 ... 66 65 64
   127 126 125 ... 98 97 96]
  输入数据(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]
  输出数据(dstValueGm):
  [127 126 125 ... 2 1 0]
  输出数据(dstIndexGm):
  [127 126 125 ... 2 1 0]