Gather
功能说明
给定输入的张量和一个地址偏移张量,Gather指令根据偏移地址将输入张量按元素收集到结果张量中。
函数原型
- tensor前n个数据计算
__aicore__ inline void Gather(const LocalTensor<T>& dstLocal, const LocalTensor<T>& srcLocal, const LocalTensor<uint32_t>& srcOffsetLocal, const uint32_t srcBaseAddr, const uint32_t count);
- tensor高维切分计算
- mask逐bit模式
__aicore__ inline void Gather(const LocalTensor<T>& dstLocal, const LocalTensor<T>& srcLocal, const LocalTensor<uint32_t>& srcOffsetLocal, const uint32_t srcBaseAddr, const uint64_t mask[2], const uint8_t repeatTimes, const uint16_t dstRepStride);
- mask连续模式
__aicore__ inline void Gather(const LocalTensor<T>& dstLocal, const LocalTensor<T>& srcLocal, const LocalTensor<uint32_t>& srcOffsetLocal, const uint32_t srcBaseAddr, const uint64_t mask, const uint8_t repeatTimes, const uint16_t dstRepStride);
- mask逐bit模式
参数说明
参数名称 |
输入/输出 |
含义 |
---|---|---|
dstLocal |
输出 |
目的操作数。 类型为LocalTensor,支持的TPosition为VECIN/VECCALC/VECOUT。 Atlas A2训练系列产品/Atlas 800I A2推理产品,支持的数据类型为:int16_t/uint16_t/int32_t/uint32_t/float/half/bfloat16_t Atlas推理系列产品AI Core,支持的数据类型为:int16_t/uint16_t/int32_t/uint32_t/float/half |
srcLocal |
输入 |
源操作数。 类型为LocalTensor,支持的TPosition为VECIN/VECCALC/VECOUT。 数据类型和dstLocal保持一致。 |
srcOffsetLocal |
输入 |
每个元素在src中对应的地址偏移,地址偏移要大于等于0。该偏移量是相对于src的基地址而言,支持数据类型(uint32_t)。每个数值的单位为Bytes。取值应保证src元素类型位宽对齐。 |
srcBaseAddr |
输入 |
srcLocal的起始偏移地址,单位是字节。 |
count |
输入 |
执行处理的数据个数,不得超过srcLocal和srcOffsetLocal的元素个数。 |
mask |
输入 |
mask用于控制每次迭代内参与计算的元素。
|
repeatTimes |
输入 |
指令迭代次数,每次迭代完成8个datablock(32Bytes)的数据收集,数据范围:repeatTimes∈[0,255]。 |
dstRepStride |
输入 |
相邻迭代间的地址步长,单位是datablock(32Bytes)。 |
支持的型号
Atlas A2训练系列产品/Atlas 800I A2推理产品
Atlas推理系列产品AI Core
约束说明
- 操作数地址偏移对齐要求请参见通用约束。
调用示例
#include "kernel_operator.h" namespace AscendC { template <typename T> class GatherTest { public: __aicore__ inline GatherTest() {} __aicore__ inline void Init(__gm__ uint8_t* dstGm, __gm__ uint8_t* srcGm, __gm__ uint8_t* srcOffsetGm, const uint32_t count) { m_elementCount = count; m_dstGlobal.SetGlobalBuffer((__gm__ T*)dstGm); m_srcGlobal.SetGlobalBuffer((__gm__ T*)srcGm); m_srcOffsetGlobal.SetGlobalBuffer((__gm__ uint32_t*)srcOffsetGm); m_pipe.InitBuffer(m_queIn, 2, m_elementCount * sizeof(uint32_t)); m_pipe.InitBuffer(m_queOut, 2, m_elementCount * sizeof(uint32_t)); } __aicore__ inline void Process() { CopyIn(); Compute(); CopyOut(); } private: __aicore__ inline void CopyIn() { LocalTensor<T> srcLocal = m_queIn.AllocTensor<T>(); DataCopy(srcLocal, m_srcGlobal, m_elementCount); m_queIn.EnQue(srcLocal); LocalTensor<uint32_t> srcOffsetLocal = m_queIn.AllocTensor<uint32_t>(); DataCopy(srcOffsetLocal, m_srcOffsetGlobal, m_elementCount); m_queIn.EnQue(srcOffsetLocal); } __aicore__ inline void Compute() { LocalTensor<T> srcLocal = m_queIn.DeQue<T>(); LocalTensor<uint32_t> srcOffsetLocal = m_queIn.DeQue<uint32_t>(); LocalTensor<T> dstLocal = m_queOut.AllocTensor<T>(); srcLocal.SetSize(m_elementCount); Gather(dstLocal, srcLocal, srcOffsetLocal, (uint32_t)0, m_elementCount); m_queIn.FreeTensor(srcLocal); m_queIn.FreeTensor(srcOffsetLocal); m_queOut.EnQue(dstLocal); } __aicore__ inline void CopyOut() { LocalTensor<T> dstLocal = m_queOut.DeQue<T>(); DataCopy(m_dstGlobal, dstLocal, m_elementCount); m_queOut.FreeTensor(dstLocal); } private: TPipe m_pipe; TQue<QuePosition::VECIN, 1> m_queCalc; GlobalTensor<T> m_valueGlobal; uint32_t m_concatRepeatTimes; uint32_t m_sortRepeatTimes; uint32_t m_extractRepeatTimes; uint32_t m_elementCount; GlobalTensor<uint32_t> m_srcOffsetGlobal; GlobalTensor<T> m_srcGlobal; GlobalTensor<T> m_dstGlobal; TQue<QuePosition::VECIN, 2> m_queIn; TQue<QuePosition::VECOUT, 2> m_queOut; }; // class GatherTest } // namespace AscendC using namespace AscendC; #define KERNEL_GATHER(T, count) \ extern "C" __global__ __aicore__ void kernel_gather_##T##_##count(GM_ADDR dstGm, GM_ADDR srcGm, \ GM_ADDR srcOffsetGm) \ { \ AscendC::GatherTest<T> op; \ op.Init(dstGm, srcGm, srcOffsetGm, count); \ op.Process(); \ }
输入数据srcOffsetLocal: [254 252 250 ... 4 2 0] 输入数据srcLocal(128个half类型数据): [0 1 2 ... 125 126 127] 输出数据dstGlobal: [127 126 125 ... 2 1 0]