GatherMask

函数功能

以内置固定模式对应的二进制或者用户自定义输入的Tensor数值对应的二进制为gather mask（数据收集的掩码），从源操作数中选取元素写入目的操作数中。

函数原型

• 用户自定义模式

  1 2

  template <typename T, typename U, GatherMaskMode mode = defaultGahterMaskMode> __aicore__ inline void GatherMask(const LocalTensor<T>& dstLocal, const LocalTensor<T>& src0Local, const LocalTensor<U>& src1Pattern, const bool reduceMode, const uint32_t mask, const GatherMaskParams& gatherMaskParams, uint64_t& rsvdCnt)

• 内置固定模式

  1 2	template <typename T, GatherMaskMode mode = defaultGahterMaskMode> __aicore__ inline void GatherMask(const LocalTensor<T>& dstLocal, const LocalTensor<T>& src0Local, const uint8_t src1Pattern, const bool reduceMode, const uint32_t mask, const GatherMaskParams& gatherMaskParams, uint64_t& rsvdCnt)

参数说明

struct GatherMaskParams{
    uint8_t src0BlockStride;
    uint16_t repeatTimes;
    uint16_t src0RepeatStride;
    uint8_t src1RepeatStride;

};

返回值

无

支持的型号

Atlas推理系列产品AI Core

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

Atlas 200/500 A2推理产品

注意事项

• 为了节省地址空间，开发者可以定义一个Tensor，供源操作数与目的操作数同时使用（即地址重叠），相关约束如下：
  • 单次迭代内，要求源操作数和目的操作数之间100%重叠，不支持部分重叠。
  • 多次迭代间，第N次目的操作数是第N+1次源操作数的情况下，不支持地址重叠。
• 操作数地址偏移对齐要求请参见通用约束。
• 若调用该接口前为counter模式，在调用该接口后需要显示设置回counter模式(接口内部结束后会设置为norm模式)。

调用示例

• 用户自定义Tensor样例

  #include "kernel_operator.h"
  class KernelGatherMask {
  public:
      __aicore__ inline KernelGatherMask () {}
      __aicore__ inline void Init(__gm__ uint8_t* src0Gm, __gm__ uint8_t* src1Gm, __gm__ uint8_t* dstGm)
      {
          src0Global.SetGlobalBuffer((__gm__ uint16_t*)src0Gm);
          src1Global.SetGlobalBuffer((__gm__ uint16_t*)src1Gm);
          dstGlobal.SetGlobalBuffer((__gm__ uint16_t*)dstGm);
          pipe.InitBuffer(inQueueSrc0, 1, 128 * sizeof(uint16_t));
          pipe.InitBuffer(inQueueSrc1, 1, 32 * sizeof(uint16_t));
          pipe.InitBuffer(outQueueDst, 1, 128 * sizeof(uint16_t));
      }
      __aicore__ inline void Process()
      {
          CopyIn();
          Compute();
          CopyOut();
      }
  private:
      __aicore__ inline void CopyIn()
      {
          AscendC::LocalTensor<uint16_t> src0Local = inQueueSrc0.AllocTensor<uint16_t>();
          AscendC::LocalTensor<uint16_t> src1Local = inQueueSrc1.AllocTensor<uint16_t>();
          AscendC::DataCopy(src0Local, src0Global, 128);
          AscendC::DataCopy(src1Local, src1Global, 32);
          inQueueSrc0.EnQue(src0Local);
          inQueueSrc1.EnQue(src1Local);
      }
      __aicore__ inline void Compute()
      {
          AscendC::LocalTensor<uint16_t> src0Local = inQueueSrc0.DeQue<uint16_t>();
          AscendC::LocalTensor<uint16_t> src1Local = inQueueSrc1.DeQue<uint16_t>();
          AscendC::LocalTensor<uint16_t> dstLocal = outQueueDst.AllocTensor<uint16_t>();
   
          uint32_t mask = 128;
          uint64_t rsvdCnt = 0;
          // reduceMode = true; 使用counter模式
          // src0BlockStride = 1; 单次迭代内数据间隔1个datablock，即数据连续读取和写入
          // repeatTimes = 1;该参数在counter模式下不生效
          // src0RepeatStride = 8;源操作数迭代间数据间隔8个datablock
          // src1RepeatStride = 8;源操作数迭代间数据间隔8个datablock
          AscendC::GatherMask (dstLocal, src0Local, src1Local, true, mask, { 1, 1, 8, 8 }, rsvdCnt);
   
          outQueueDst.EnQue<uint16_t>(dstLocal);
          inQueueSrc0.FreeTensor(src0Local);
          inQueueSrc1.FreeTensor(src1Local);
      }
      __aicore__ inline void CopyOut()
      {
          AscendC::LocalTensor<uint16_t> dstLocal = outQueueDst.DeQue<uint16_t>();
          
          AscendC::DataCopy(dstGlobal, dstLocal, 128);
          outQueueDst.FreeTensor(dstLocal);
      }
  private:
      AscendC::TPipe pipe;
      AscendC::TQue<AscendC::QuePosition::VECIN, 1> inQueueSrc0, inQueueSrc1;
      AscendC::TQue<AscendC::QuePosition::VECOUT, 1> outQueueDst;
      AscendC::GlobalTensor<uint16_t> src0Global, src1Global, dstGlobal;
  };
  extern "C" __global__ __aicore__ void gather_mask_simple_kernel(__gm__ uint8_t* src0Gm, __gm__ uint8_t* src1Gm, __gm__ uint8_t* dstGm)
  {
      KernelGatherMask op;
      op.Init(src0Gm, src1Gm, dstGm);
      op.Process();
  }

  结果示例如下：

  输入数据(src0Local): [1 2 3 ... 128]
                       // 43690对应的二进制：0b1010 1010 1010 1010
  输入数据(src1Local): [43690 43690 43690 43690 43690 43690 43690 43690 0 0 0 0 ...0]
  输出数据(dstLocal): [2 4 6 8 10 12 14 16 18 20 22 24 26 28 30 32 34 36 38 40 42 44 46 48 50 52 54 56 58 60 62 64 66 68 70 72 74 76 78 80 82 84 86 88 90 92 94 96 98 100 102 104 106 108 110 112 114 116 118 120 122 124 126 128 undefined ..undefined]
  输出数据(rsvdCnt): 64

• 内置固定模式

  #include "kernel_operator.h"
  class KernelGatherMask {
  public:
      __aicore__ inline KernelGatherMask () {}
      __aicore__ inline void Init(__gm__ uint8_t* src0Gm, __gm__ uint8_t* dstGm)
      {
          src0Global.SetGlobalBuffer((__gm__ uint16_t*)src0Gm);
          dstGlobal.SetGlobalBuffer((__gm__ uint16_t*)dstGm);
          pipe.InitBuffer(inQueueSrc0, 1, 128 * sizeof(uint16_t));
          pipe.InitBuffer(outQueueDst, 1, 128 * sizeof(uint16_t));
      }
      __aicore__ inline void Process()
      {
          CopyIn();
          Compute();
          CopyOut();
      }
  private:
      __aicore__ inline void CopyIn()
      {
          AscendC::LocalTensor<uint16_t> src0Local = inQueueSrc0.AllocTensor<uint16_t>();
          AscendC::DataCopy(src0Local, src0Global, 128);
          inQueueSrc0.EnQue(src0Local);
      }
      __aicore__ inline void Compute()
      {
          AscendC::LocalTensor<uint16_t> src0Local = inQueueSrc0.DeQue<uint16_t>();
          AscendC::LocalTensor<uint16_t> dstLocal = outQueueDst.AllocTensor<uint16_t>();
   
          uint32_t mask = 0; // normal模式下mask需要设置为0
          uint64_t rsvdCnt = 0; // 用于保存筛选后保留下来的元素个数
          uint8_t src1Pattern = 2; // 内置固定模式
          // reduceMode = false; 使用normal模式
          // src0BlockStride = 1; 单次迭代内数据间隔1个Block，即数据连续读取和写入
          // repeatTimes = 1;重复迭代一次
          // src0RepeatStride = 0;重复一次，故设置为0
          // src1RepeatStride = 0;重复一次，故设置为0
          AscendC::GatherMask(dstLocal, src0Local, src1Pattern, false, mask, { 1, 1, 0, 0 }, rsvdCnt);
   
          outQueueDst.EnQue<uint16_t>(dstLocal);
          inQueueSrc0.FreeTensor(src0Local);
      }
      __aicore__ inline void CopyOut()
      {
          AscendC::LocalTensor<uint16_t> dstLocal = outQueueDst.DeQue<uint16_t>();
          
          AscendC::DataCopy(dstGlobal, dstLocal, 128);
          outQueueDst.FreeTensor(dstLocal);
      }
  private:
      AscendC::TPipe pipe;
      AscendC::TQue<AscendC::QuePosition::VECIN, 1> inQueueSrc0;
      AscendC::TQue<AscendC::QuePosition::VECOUT, 1> outQueueDst;
      AscendC::GlobalTensor<uint16_t> src0Global, dstGlobal;
  };
   
  extern "C" __global__ __aicore__ void gather_mask_simple_kernel(__gm__ uint8_t* src0Gm, __gm__ uint8_t* dstGm)
  {
      KernelGatherMask op;
      op.Init(src0Gm, dstGm);
      op.Process();
  }

  结果示例如下：

  输入数据(src0Local): [1 2 3 ... 128]
  输入数据(src1Pattern): src1Pattern = 2;
  输出数据(dstLocal): [2 4 6 8 10 12 14 16 18 20 22 24 26 28 30 32 34 36 38 40 42 44 46 48 50 52 54 56 58 60 62 64 66 68 70 72 74 76 78 80 82 84 86 88 90 92 94 96 98 100 102 104 106 108 110 112 114 116 118 120 122 124 126 128 undefine ..undefined]
  输出数据(rsvdCnt): 64