通过Unified Buffer融合实现连续vector计算
【优先级】高
【描述】算子实现中涉及多次vector计算,且前一次计算输出是后一次计算输入的情况下,可将前一次计算输出暂存在UB(Unified Buffer)上直接作为下一次计算的输入,不需要将前一次的计算输出从UB搬运到GM后再从GM搬运到UB。这种UB Buffer融合的方式可以减少搬入搬出次数,实现连续vector计算,提升内存使用效率。数据流图对比如下:
图1 数据流图对比
【反例】
该算子的计算逻辑为进行Exp计算后再进行Abs计算。计算过程中先把源操作数从GM搬运到UB进行Exp计算,Exp计算完成后将Exp的结果从UB搬运到GM;再从GM中把Exp的结果搬运到UB上作为Abs计算的输入,Abs计算完成后将目的操作数结果从UB搬运到GM。整个过程从GM搬进搬出共4次。当需要进行的vector计算为n次时,从GM搬进搬出共需要2n次。
class KernelSample {
public:
__aicore__ inline KernelSample() {}
__aicore__ inline void Init(__gm__ uint8_t* src0Gm, __gm__ uint8_t* dstGm)
{
src0Global.SetGlobalBuffer((__gm__ float*)src0Gm);
dstGlobal.SetGlobalBuffer((__gm__ float*)dstGm);
pipe.InitBuffer(inQueueSrc0, 1, 1024 * sizeof(float));
pipe.InitBuffer(outQueueDst, 1, 1024 * sizeof(float));
}
__aicore__ inline void Process()
{
CopyIn();
Compute();
CopyOut();
CopyIn1();
Compute1();
CopyOut1();
}
private:
__aicore__ inline void CopyIn()
{
LocalTensor<float> src0Local = inQueueSrc0.AllocTensor<float>();
DataCopy(src0Local, src0Global, 1024);
inQueueSrc0.EnQue(src0Local);
}
__aicore__ inline void Compute()
{
LocalTensor<float> src0Local = inQueueSrc0.DeQue<float>();
LocalTensor<float> dstLocal = outQueueDst.AllocTensor<float>();
Exp(dstLocal, src0Local, 1024);
outQueueDst.EnQue<float>(dstLocal);
inQueueSrc0.FreeTensor(src0Local);
}
__aicore__ inline void CopyOut()
{
LocalTensor<float> dstLocal = outQueueDst.DeQue<float>();
DataCopy(dstGlobal, dstLocal, 1024);
outQueueDst.FreeTensor(dstLocal);
}
__aicore__ inline void CopyIn1()
{
PipeBarrier<PIPE_ALL>();
LocalTensor<float> src0Local = inQueueSrc0.AllocTensor<float>();
DataCopy(src0Local, dstGlobal, 1024);
inQueueSrc0.EnQue(src0Local);
}
__aicore__ inline void Compute1()
{
LocalTensor<float> src0Local = inQueueSrc0.DeQue<float>();
LocalTensor<float> dstLocal = outQueueDst.AllocTensor<float>();
Abs(dstLocal, src0Local, 1024);
outQueueDst.EnQue<float>(dstLocal);
inQueueSrc0.FreeTensor(src0Local);
}
__aicore__ inline void CopyOut1()
{
LocalTensor<float> dstLocal = outQueueDst.DeQue<float>();
DataCopy(dstGlobal, dstLocal, 1024);
outQueueDst.FreeTensor(dstLocal);
}
private:
TPipe pipe;
TQue<QuePosition::VECIN, 1> inQueueSrc0;
TQue<QuePosition::VECOUT, 1> outQueueDst;
GlobalTensor<float> src0Global, dstGlobal;
};
【正例】
使用UB Buffer融合方式后,在UB上进行连续vector计算时,前一次的结果可直接作为后一次计算的输入,继续在UB上进行计算,不需要中间的搬进搬出,只需在开始计算时将源操作数搬运到UB,以及全部计算结束后将最终结果从UB搬运到GM,共2次搬进搬出。
class KernelSample {
public:
__aicore__ inline KernelSample() {}
__aicore__ inline void Init(__gm__ uint8_t* src0Gm, __gm__ uint8_t* dstGm)
{
src0Global.SetGlobalBuffer((__gm__ float*)src0Gm);
dstGlobal.SetGlobalBuffer((__gm__ float*)dstGm);
pipe.InitBuffer(inQueueSrc0, 1, 1024 * sizeof(float));
pipe.InitBuffer(outQueueDst, 1, 1024 * sizeof(float));
}
__aicore__ inline void Process()
{
CopyIn();
Compute();
CopyOut();
}
private:
__aicore__ inline void CopyIn()
{
LocalTensor<float> src0Local = inQueueSrc0.AllocTensor<float>();
DataCopy(src0Local, src0Global, 1024);
inQueueSrc0.EnQue(src0Local);
}
__aicore__ inline void Compute()
{
LocalTensor<float> src0Local = inQueueSrc0.DeQue<float>();
LocalTensor<float> dstLocal = outQueueDst.AllocTensor<float>();
Exp(dstLocal, src0Local, 1024);
Abs(dstLocal, dstLocal, 1024);
outQueueDst.EnQue<float>(dstLocal);
inQueueSrc0.FreeTensor(src0Local);
}
__aicore__ inline void CopyOut()
{
LocalTensor<float> dstLocal = outQueueDst.DeQue<float>();
DataCopy(dstGlobal, dstLocal, 1024);
outQueueDst.FreeTensor(dstLocal);
}
private:
TPipe pipe;
TQue<QuePosition::VECIN, 1> inQueueSrc0;
TQue<QuePosition::VECOUT, 1> outQueueDst;
GlobalTensor<float> src0Global, dstGlobal;
};
父主题: 内存优化