#include "kernel_operator.h"
constexpr int32_t BUFFER_NUM = 2; // tensor num for each queue

class KernelMatmulVec {

public:

__aicore__ inline KernelMatmulVec() {}

__aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR z, uint32_t m, uint32_t k, uint32_t smallCoreN, uint32_t bigCoreN, uint32_t bigCoreNum)

{

this->m = m;

this->k = k;

// this->smallCoreN;

// this->bigCoreN;

// this->bigCoreNum;

// this->coreNum;

uint32_t addrOffset = 0;

uint32_t addrOffsetZ = 0;

//判断是否最后一块是有有偏移量

if (AscendC::GetBlockIdx() >= bigCoreNum){

this->n = smallCoreN;

addrOffset = this->n * this->k * AscendC::GetBlockIdx() + (this->k * bigCoreNum *16);

addrOffsetZ = this->n * AscendC::GetBlockIdx() + bigCoreNum *16 ;

}else{

this->n = bigCoreN;

addrOffset = this->n * this->k * AscendC::GetBlockIdx();

addrOffsetZ = this->n * AscendC::GetBlockIdx();

}

//起始地址

xGm.SetGlobalBuffer((__gm__ DTYPE_X *)x , m * k );

yGm.SetGlobalBuffer((__gm__ DTYPE_Y *)y + addrOffset, this->n * k);

zGm.SetGlobalBuffer((__gm__ DTYPE_Z *)z + addrOffsetZ, this->n * 1);

//AscendC::printf("init-ok-----------------------------------------1\n");

//临时计算结果存储，不涉及搬运

pipe.InitBuffer(tmp1, this->m * this->k * sizeof(DTYPE_Z));

pipe.InitBuffer(tmp2, this->m * this->k * sizeof(DTYPE_Z));

pipe.InitBuffer(tmp3, this->n * sizeof(DTYPE_Z));

// pipe.InitBuffer(reduceRes, this->n * sizeof(DTYPE_Z));

//核内单次搬运空间大小

// AscendC::printf("this->n = %d\n",this->n);

pipe.InitBuffer(inQueueX, BUFFER_NUM, this->m * this->k * sizeof(DTYPE_X));

pipe.InitBuffer(inQueueY, BUFFER_NUM, this->k * sizeof(DTYPE_Y));

pipe.InitBuffer(outQueueZ, BUFFER_NUM, 16 * sizeof(DTYPE_Z));

}

__aicore__ inline void Process()

{

int32_t loopCount = this->n ;

this->vecLeft = inQueueX.AllocTensor<DTYPE_X>();

AscendC::LocalTensor<DTYPE_Z> zLocal;

AscendC::DataCopy(vecLeft, xGm, this->k);

inQueueX.EnQue(vecLeft);

vecLeft = inQueueX.DeQue<DTYPE_X>();

for (int32_t i = 0; i < loopCount; i++) {

if (i % 16 ==0) {

zLocal = outQueueZ.AllocTensor<DTYPE_Z>();

}

AscendC::LocalTensor<DTYPE_Y> yLocal = inQueueY.AllocTensor<DTYPE_Y>();

AscendC::DataCopy(yLocal, yGm[i * this->k], this->k);

inQueueY.EnQue(yLocal);

yLocal = inQueueY.DeQue<DTYPE_Y>();

auto p1 = tmp1.Get<DTYPE_Z>();

auto p2 = tmp2.Get<DTYPE_Z>();

auto p3 = tmp3.Get<DTYPE_Z>();

AscendC::Cast(p1, vecLeft, AscendC::RoundMode::CAST_NONE, this->k);

AscendC::Cast(p2, yLocal, AscendC::RoundMode::CAST_NONE, this->k);

AscendC::Mul(p1, p1, p2, this->k);

AscendC::ReduceSum<DTYPE_Z>(zLocal[i % 16], p1, p3, k);

// AscendC::WholeReduceSum<DTYPE_Z>(zLocal[i % 16], p1, k, 1, 1, 1, 1);

// AscendC::DumpTensor(zLocal, 1, 16);

if ((i+1) % 16 ==0) {

outQueueZ.EnQue<DTYPE_Z>(zLocal);

zLocal = outQueueZ.DeQue<DTYPE_Z>();

AscendC::DataCopy(zGm[i-15], zLocal, 16);

outQueueZ.FreeTensor(zLocal);

}

inQueueY.FreeTensor(yLocal);

}

inQueueX.FreeTensor(vecLeft);

}

private:

AscendC::TPipe pipe;

AscendC::TQue<AscendC::QuePosition::VECIN, BUFFER_NUM> inQueueX;

AscendC::TQue<AscendC::QuePosition::VECIN, BUFFER_NUM> inQueueY;

AscendC::TQue<AscendC::QuePosition::VECOUT, BUFFER_NUM> outQueueZ;

AscendC::GlobalTensor<DTYPE_X> xGm;

AscendC::GlobalTensor<DTYPE_Y> yGm;

AscendC::GlobalTensor<DTYPE_Z> zGm;

AscendC::LocalTensor<DTYPE_X> vecLeft;

// AscendC::LocalTensor<DTYPE_Z> zLocal;

// AscendC::LocalTensor<DTYPE_Z> reduceResult;

uint32_t m;

uint32_t k;

uint32_t smallCoreN;

uint32_t bigCoreN;

uint32_t bigCoreNum;

uint32_t n;

AscendC::TBuf<AscendC::QuePosition::VECCALC> tmp1, tmp2 ,tmp3;

};

extern "C" __global__ __aicore__ void matmul_vec(GM_ADDR x, GM_ADDR y, GM_ADDR z, GM_ADDR workspace, GM_ADDR tiling) {

GET_TILING_DATA(tiling_data, tiling);

// TODO: user kernel impl

//AscendC::printf("matmul_vec-ok-----------------------------------------1\n");

KernelMatmulVec kernel;

kernel.Init(x, y, z, tiling_data.m, tiling_data.k, tiling_data.smallCoreN, tiling_data.bigCoreN, tiling_data.bigCoreNum);

kernel.Process();

}