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)
参数说明
参数名称 |
输入/输出 |
含义 |
---|---|---|
dstLocal |
输出 |
目的操作数。 类型为LocalTensor,支持的TPosition为VECIN/VECCALC/VECOUT。 LocalTensor的起始地址需要32字节对齐。 Atlas推理系列产品AI Core,支持的数据类型为:half/uint16_t/int16_t/float/uint32_t/int32_t Atlas A2训练系列产品/Atlas 800I A2推理产品,支持的数据类型为:half/bfloat16_t/uint16_t/int16_t/float/uint32_t/int32_t Atlas 200/500 A2推理产品,支持的数据类型为:half/uint16_t/int16_t/float/uint32_t/int32_t |
src0Local |
输入 |
源操作数。 类型为LocalTensor,支持的TPosition为VECIN/VECCALC/VECOUT。 LocalTensor的起始地址需要32字节对齐。 数据类型需要与目的操作数保持一致。 Atlas推理系列产品AI Core,支持的数据类型为:half/uint16_t/int16_t/float/uint32_t/int32_t Atlas A2训练系列产品/Atlas 800I A2推理产品,支持的数据类型为:half/uint16_t/int16_t/float/uint32_t/int32_t Atlas 200/500 A2推理产品,支持的数据类型为:half/uint16_t/int16_t/float/uint32_t/int32_t |
src1Pattern |
输入 |
gather mask(数据收集的掩码),分为内置固定模式和用户自定义模式两种,根据内置固定模式对应的二进制或者用户自定义输入的Tensor数值对应的二进制从源操作数中选取元素写入目的操作数中。1为选取,0为不选取。
|
reduceMode |
输入 |
mask模式选择参数,支持的数据类型:bool,取值为:
|
mask |
输入 |
mask值仅在reduceMode为true时生效。 根据reduceMode,分为两种模式:
|
gatherMaskParams |
输入 |
控制操作数地址步长的数据结构。结构体内包含操作数相邻迭代间相同datablock的地址步长,操作数同一迭代内不同datablock的地址步长等参数。数据结构的定义如下: struct GatherMaskParams{ uint8_t src0BlockStride; uint16_t repeatTimes; uint16_t src0RepeatStride; uint8_t src1RepeatStride; }; 相邻迭代间的地址步长参数说明请参考repeatStride(相邻迭代间相同datablock的地址步长);同一迭代内datablock的地址步长参数说明请参考dataBlockStride(同一迭代内不同datablock的地址步长)。 |
rsvdCnt |
输出 |
该条指令筛选后保留下来的元素计数,对应dstLocal中有效元素个数,数据类型为uint64_t。 |
参数名称 |
输入/输出 |
含义 |
---|---|---|
mode |
输入 |
预留参数,为后续功能做预留,当前提供默认值,用户无需设置该参数。 |
返回值
无
支持的型号
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