Scatter
功能说明
给定一个连续的输入张量和一个目的地址偏移张量，Scatter指令根据偏移地址生成新的结果张量后将输入张量分散到结果张量中。
将源操作数src中的element按照指定的位置（由dst_offset和base_addr共同作用）分散到目的操作数dst中。
函数原型
- tensor前n个数据计算
__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模式
__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[2], const uint8_t repeatTimes, const uint16_t srcRepStride);
- mask连续模式
__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 uint16_t srcRepStride);
- mask逐bit模式
参数说明
参数名称 |
输入/输出 |
含义 |
|---|---|---|
dstLocal |
输出 |
目的操作数，类型为LocalTensor。支持数据类型(uint16_t/uint32_t/float/half), 地址需要32bytes对齐。 |
srcLocal |
输入 |
源操作数，类型为LocalTensor。数据类型和dstLocal保持一致。 |
dstOffsetLocal |
输入 |
每个元素在dst中对应的地址偏移，地址偏移要大于等于0。该偏移量是相对于dst的基地址而言，支持数据类型(uint32_t)。每个数值的单位为Bytes。 |
dstBaseAddr |
输入 |
dstLocal的起始偏移地址，单位是字节。 |
count |
输入 |
执行处理的数据个数，不得超过srcLocal和dstOffsetLocal的元素个数。 |
mask |
输入 |
mask用于控制每次迭代内参与计算的元素。
|
repeatTimes |
输入 |
指令迭代次数，每次迭代完成8个datablock的数据收集，数据范围：repeatTimes∈[0,255]。 |
srcRepStride |
输入 |
相邻迭代间的地址步长，单位是datablock。 |
支持的型号
Atlas推理系列产品AI Core
约束说明
- 操作数地址偏移对齐要求请参见通用约束。
调用示例
#include "kernel_operator.h"
namespace AscendC {
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()
{
LocalTensor<T> srcLocal = m_queIn.AllocTensor<T>();
DataCopy(srcLocal, m_srcGlobal, m_elementCount);
m_queIn.EnQue(srcLocal);
LocalTensor<uint32_t> dstOffsetLocal = m_queIn.AllocTensor<uint32_t>();
DataCopy(dstOffsetLocal, m_dstOffsetGlobal, m_elementCount);
m_queIn.EnQue(dstOffsetLocal);
}
__aicore__ inline void Compute()
{
LocalTensor<T> srcLocal = m_queIn.DeQue<T>();
LocalTensor<uint32_t> dstOffsetLocal = m_queIn.DeQue<uint32_t>();
LocalTensor<T> dstLocal = m_queOut.AllocTensor<T>();
dstLocal.SetSize(m_elementCount);
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()
{
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_dstOffsetGlobal;
GlobalTensor<T> m_srcGlobal;
GlobalTensor<T> m_dstGlobal;
TQue<QuePosition::VECIN, 2> m_queIn;
TQue<QuePosition::VECOUT, 1> m_queOut;
}; // class ScatterTest
} // namespace AscendC
using namespace AscendC;
#define KERNEL_SCATTER(T, count) \
extern "C" __global__ __aicore__ void kernel_scatter_##T##_##count(GM_ADDR dstGm, GM_ADDR srcGm,\
GM_ADDR dstOffsetGm) \
{ \
AscendC::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]
