Scatter
功能说明
给定一个连续的输入张量和一个目的地址偏移张量,Scatter指令根据偏移地址生成新的结果张量后将输入张量分散到结果张量中。
将源操作数src中的element按照指定的位置(由dst_offset和base_addr共同作用)分散到目的操作数dst中。
函数原型
- tensor前n个数据计算
1 2
template <typename T> __aicore__ inline void Scatter(const LocalTensor<T>& dstLocal, const LocalTensor<T>& srcLocal, const LocalTensor<uint32_t>& dstOffsetLocal, const uint32_t dstBaseAddr, const uint32_t count)
- tensor高维切分计算
- mask逐bit模式
1 2
template <typename T> __aicore__ inline void Scatter(const LocalTensor<T>& dstLocal, const LocalTensor<T>& srcLocal, const LocalTensor<uint32_t>& dstOffsetLocal, const uint32_t dstBaseAddr, const uint64_t mask[], const uint8_t repeatTimes, const uint8_t srcRepStride)
- mask连续模式
1 2
template <typename T> __aicore__ inline void Scatter(const LocalTensor<T>& dstLocal, const LocalTensor<T>& srcLocal, const LocalTensor<uint32_t>& dstOffsetLocal, const uint32_t dstBaseAddr, const uint64_t mask, const uint8_t repeatTimes, const uint8_t srcRepStride)
- mask逐bit模式
参数说明
参数名 |
描述 |
---|---|
T |
操作数数据类型。 Atlas推理系列产品AI Core,支持的数据类型为:uint16_t/uint32_t/float/half Atlas 200/500 A2推理产品,支持的数据类型为:uint8_t/int8_t/uint16_t/int16_t/half/uint32_t/int32_t/float |
参数名称 |
输入/输出 |
含义 |
---|---|---|
dstLocal |
输出 |
目的操作数,类型为LocalTensor。地址需要32bytes对齐。 |
srcLocal |
输入 |
源操作数,类型为LocalTensor。数据类型和dstLocal保持一致。 |
dstOffsetLocal |
输入 |
每个元素在dst中对应的地址偏移,地址偏移要大于等于0。该偏移量是相对于dst的基地址而言,支持数据类型(uint32_t)。每个数值的单位为Bytes。 针对以下型号,地址偏移的取值范围不超出uint32_t的范围即可。 Atlas推理系列产品AI Core 针对以下型号,地址偏移的取值范围如下:当操作数为8位时,取值范围为[0, 216-1];当操作数为16位时,取值范围为[0, 217-1],当操作数为32位或者64位时,不超过uint32_t的范围即可,超出取值范围可能导致非预期输出。 Atlas 200/500 A2推理产品 |
dstBaseAddr |
输入 |
dstLocal的起始偏移地址,单位是字节。 |
count |
输入 |
执行处理的数据个数,不得超过srcLocal和dstOffsetLocal的元素个数。 |
mask |
输入 |
|
repeatTimes |
输入 |
指令迭代次数,每次迭代完成8个datablock的数据收集,数据范围:repeatTimes∈[0,255]。 |
srcRepStride |
输入 |
相邻迭代间的地址步长,单位是datablock。 |
支持的型号
Atlas推理系列产品AI Core
Atlas 200/500 A2推理产品
约束说明
- 操作数地址偏移对齐要求请参见通用约束。
调用示例
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 |
#include "kernel_operator.h" template <typename T> class ScatterTest { public: __aicore__ inline ScatterTest() {} __aicore__ inline void Init(__gm__ uint8_t* dstGm, __gm__ uint8_t* srcGm, __gm__ uint8_t* dstOffsetGm, const uint32_t count) { m_elementCount = count; m_dstGlobal.SetGlobalBuffer((__gm__ T*)dstGm); m_srcGlobal.SetGlobalBuffer((__gm__ T*)srcGm); m_dstOffsetGlobal.SetGlobalBuffer((__gm__ uint32_t*)dstOffsetGm); m_pipe.InitBuffer(m_queIn, 2, m_elementCount * sizeof(uint32_t)); m_pipe.InitBuffer(m_queOut, 1, m_elementCount * sizeof(uint32_t)); } __aicore__ inline void Process() { CopyIn(); Compute(); CopyOut(); } private: __aicore__ inline void CopyIn() { AscendC::LocalTensor<T> srcLocal = m_queIn.AllocTensor<T>(); AscendC::DataCopy(srcLocal, m_srcGlobal, m_elementCount); m_queIn.EnQue(srcLocal); AscendC::LocalTensor<uint32_t> dstOffsetLocal = m_queIn.AllocTensor<uint32_t>(); AscendC::DataCopy(dstOffsetLocal, m_dstOffsetGlobal, m_elementCount); m_queIn.EnQue(dstOffsetLocal); } __aicore__ inline void Compute() { AscendC::LocalTensor<T> srcLocal = m_queIn.DeQue<T>(); AscendC::LocalTensor<uint32_t> dstOffsetLocal = m_queIn.DeQue<uint32_t>(); AscendC::LocalTensor<T> dstLocal = m_queOut.AllocTensor<T>(); dstLocal.SetSize(m_elementCount); AscendC::Scatter(dstLocal, srcLocal, dstOffsetLocal, (uint32_t)0, m_elementCount); m_queIn.FreeTensor(srcLocal); m_queIn.FreeTensor(dstOffsetLocal); m_queOut.EnQue(dstLocal); } __aicore__ inline void CopyOut() { AscendC::LocalTensor<T> dstLocal = m_queOut.DeQue<T>(); AscendC::DataCopy(m_dstGlobal, dstLocal, m_elementCount); m_queOut.FreeTensor(dstLocal); } private: AscendC::TPipe m_pipe; AscendC::TQue<AscendC::QuePosition::VECIN, 1> m_queCalc; AscendC::GlobalTensor<T> m_valueGlobal; uint32_t m_concatRepeatTimes; uint32_t m_sortRepeatTimes; uint32_t m_extractRepeatTimes; uint32_t m_elementCount; AscendC::GlobalTensor<uint32_t> m_dstOffsetGlobal; AscendC::GlobalTensor<T> m_srcGlobal; AscendC::GlobalTensor<T> m_dstGlobal; AscendC::TQue<AscendC::QuePosition::VECIN, 2> m_queIn; AscendC::TQue<AscendC::QuePosition::VECOUT, 1> m_queOut; }; // class ScatterTest #define KERNEL_SCATTER(T, count) \ extern "C" __global__ __aicore__ void kernel_scatter_##T##_##count(GM_ADDR dstGm, GM_ADDR srcGm,\ GM_ADDR dstOffsetGm) \ { \ ScatterTest<T> op; \ op.Init(dstGm, srcGm, dstOffsetGm, count); \ op.Process(); \ } |
输入数据dstOffsetLocal: [254 252 250 ... 4 2 0] 输入数据srcLocal(128个half类型数据): [0 1 2 ... 125 126 127] 输出数据dstGlobal: [127 126 125 ... 2 1 0]