in source/backend/cpu/compute/ConvInt8TiledExecutor.cpp [928:1498]
ErrorCode DenseConvInt8TiledExecutor::onExecute(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) {
const auto input = inputs[0];
auto output = outputs[0];
auto core = static_cast<CPUBackend*>(backend())->int8Functions();
auto gcore = static_cast<CPUBackend*>(backend())->functions();
auto dynamicOption = static_cast<CPUBackend*>(backend())->getRuntime()->hint().dynamicQuantOption;
#ifdef MNN_KLEIDIAI_ENABLED
KleidiAI& kai = KleidiAI::getInstance();
if(mResourceInt8->mDynamicQuant && mResourceInt8->mActBits == 4 && kai.canAccelerate(mAccelType)) {
MNN_ASSERT(kai.isLoaded(mAccelType));
const size_t m = input->batch(); //lhs vector number.
const size_t n = output->channel(); //rhs vector number.
const size_t k = input->channel(); //vector size.
const size_t blkSize = mBlockNum == 1 ? 0 : k / mBlockNum;
bool bHalf = kai.isHalf();
size_t elementSize = bHalf ? sizeof(__fp16) : sizeof(float);
size_t lhsPackedSize = kai.getLhsQuantedPackedSize(mAccelType, m, k, blkSize);
auto lhs = input->host<uint8_t>();
auto lhsPacked = mTempIm2ColBuffer->host<int8_t>();
auto rhsPacked = mResourceInt8->mWeightInt8->host<uint8_t>();
auto dst = output->host<uint8_t>();
uint8_t *linearLhs, *linearDst;
if(m > 1 && !kai.isLinear()) {
linearLhs = (uint8_t *)lhsPacked + lhsPackedSize;
linearDst = linearLhs;
} else {
linearLhs = lhs;
linearDst = dst;
}
int threadNum = static_cast<CPUBackend*>(backend())->threadNumber();
int threadNeed, vecPerThread;
//Dynamic quant pack lhs.
if(m == 1) {
kai.runLhsQuantPack(mAccelType, 1, k, blkSize, 1, linearLhs, lhsPacked);
} else {
if(!kai.isLinear()) {
if(bHalf) {
KleidiAIUtil::transferNC4HW4ToNCHW((__fp16 *)lhs, (__fp16 *)linearLhs, m, k);
} else {
KleidiAIUtil::transferNC4HW4ToNCHW((float *)lhs, (float *)linearLhs, m, k);
}
}
vecPerThread = kai.getVecNumPerThread(m, threadNum, kai.getMr(mAccelType, m));
threadNeed = m % vecPerThread == 0 ? m / vecPerThread : (m / vecPerThread + 1);
size_t srcStride = vecPerThread * k * elementSize;
auto BatchDynamicQuant = [=, &kai](int tId) {
auto threadSrc = linearLhs + tId * srcStride;
auto threadDst = lhsPacked + kai.getLhsQuantedPackedOffset(mAccelType, m, tId * vecPerThread, k, blkSize);
int vecNum = (tId == threadNeed - 1) ? (m - vecPerThread * tId) : vecPerThread; //Last threadN may less than vecPerThread.
kai.runLhsQuantPack(mAccelType, vecNum, k, blkSize, kai.getMr(mAccelType, m), threadSrc, threadDst);
};
MNN_CONCURRENCY_BEGIN(tId, threadNeed) {
BatchDynamicQuant((int)tId);
}
MNN_CONCURRENCY_END();
}
//Run matmul.
if(kai.bSupportSme2() && mAccelType == KleidiAI::AccelType::QI4_SYM_CHNLQT) {
//SME prefer running on single thread to obtain better performance/power consumption ratio.
threadNum = 1;
}
vecPerThread = kai.getVecNumPerThread(n, threadNum, kai.getNStep(mAccelType));
threadNeed = n % vecPerThread == 0 ? n / vecPerThread : (n / vecPerThread + 1);
auto ThreadFunction = [=, &kai](int tId) {
auto threadRhsPacked = rhsPacked + kai.getRhsPackedOffset(mAccelType, tId * vecPerThread, k, blkSize);
auto threadDst = linearDst + kai.getDstOffset(0, tId * vecPerThread, n, elementSize);
int vecNum = (tId == threadNeed - 1) ? (n - vecPerThread * tId) : vecPerThread; //Last threadN may less than vecPerThread.
float scalarMax = bHalf ? FLT16_MAX : FLT_MAX;
kai.runMatmul(mAccelType, m, vecNum, k, blkSize, lhsPacked, threadRhsPacked, threadDst, n * elementSize, elementSize, scalarMax, -scalarMax);
};
MNN_CONCURRENCY_BEGIN(tId, threadNeed) {
ThreadFunction((int)tId);
}
MNN_CONCURRENCY_END();
if(m > 1 && !kai.isLinear()) {
if(bHalf) {
KleidiAIUtil::transferNCHWToNC4HW4((__fp16 *)linearDst, (__fp16 *)dst, m, n);
} else {
KleidiAIUtil::transferNCHWToNC4HW4((float *)linearDst, (float *)dst, m, n);
}
}
return NO_ERROR;
}
#endif
int UNIT, SRC_UNIT, DST_XUNIT;
core->MNNGetGemmUnit(&UNIT, &SRC_UNIT, &DST_XUNIT);
auto blitProc = core->MNNPackC4Int8ForMatMul_A;
const int plane = output->batch() * mIm2ColParamter.oh * mIm2ColParamter.ow;
const int batch = input->batch();
const int PackUnit = gcore->pack;
const int dstZStep = plane * PackUnit;
const int ocDiv4 = UP_DIV(output->channel(), PackUnit);
const int ocUp4 = ROUND_UP(output->channel(), PackUnit);
const auto kernelCountUnit = mIm2ColParamter.kernelCountUnit;
const auto unitColBufferSize = kernelCountUnit * DST_XUNIT * SRC_UNIT * sizeof(int8_t);
const auto colBufferSize = unitColBufferSize * mIm2ColCount;
const int dstBytes = static_cast<CPUBackend*>(backend())->getBytes(backend(), output);
const int blockL = kernelCountUnit / mBlockNum; // source depthQuad for each block.
const int kxky = mIm2ColParamter.kernelX * mIm2ColParamter.kernelY;
const int blocklu = blockL / kxky; // UP_DIV(ic,src_unit) per block
float weightBytes = 1.f;
int weightStepY = weightBytes * (UNIT * SRC_UNIT);
int inputPlane = batch * input->width() * input->height();
auto im2colPtr = mTempIm2ColBuffer->host<int8_t>();
if (SRC_UNIT > PackUnit) {
memset(im2colPtr, 0, mTempIm2ColBuffer->size());
}
const auto weightDataPtr = mResourceInt8->mWeightInt8->host<int8_t>();
auto srcKernelSumPtr = (int8_t*)mTempSrcSum.ptr();
auto im2colSrc = input->host<uint8_t>();
auto outputDataPtr = output->host<int8_t>();
uint8_t* biasPtr = nullptr;
int32_t inputZeroPoint = 0;
int im2colBytes = mIm2ColBasedInt8 == true ? 1 : gcore->bytes;
if (nullptr != mMutableResource.get()) {
biasPtr = mMutableResource->mBiasFloat->host<uint8_t>();
inputZeroPoint = mMutableResource->mInputZeroPoint;
if (mBatchQuantInfo.get()) {
float scalein = TensorUtils::getQuantInfo(inputs[0])[0];
float scaleou = TensorUtils::getQuantInfo(outputs[0])[0];
auto scaleX = scalein / scaleou;
for (int i = 0; i < DST_XUNIT; ++i) {
mBatchQuantInfo->host<float>()[i] = scaleX;
}
}
}
#ifdef MNN_LOW_MEMORY
auto BatchAsyDynamicQuant = [&](uint8_t* floatPtr, int32_t& inputZero, uint8_t* inputDequantScale, int LDiv4, int eCount, int innerSide, int32_t availableThreads, int8_t* dstInt8, uint8_t* inputDequantBias, int tId) {
// if mIm2ColBasedInt8=false, input shape: [kernelsize,mBlockNum,blocklu,EP,LP]
// if mIm2ColBasedInt8=true, input shape: [ic/pack,EP,pack]
auto scalePtr = (float*)inputDequantScale;
auto zeroPtr = (float*)inputDequantBias;
int scaleCount = mSizeInputBlockQuant;
int kernelsize = 1;
if (!mIm2ColBasedInt8) {
kernelsize = kxky;
}
auto minPtr = mTempMaxMinValueBuffer.ptr() + tId * scaleCount * gcore->bytes;
auto maxPtr = mTempMaxMinValueBuffer.ptr() + tId * scaleCount * gcore->bytes + (scaleCount / 2) * gcore->bytes;
auto qscale = (float*)(mQScaleZero.ptr() + tId * scaleCount * QUANT_INFO_BYTES);
auto qbias = (float*)(mQScaleZero.ptr() + tId * scaleCount * QUANT_INFO_BYTES + (scaleCount / 2) * QUANT_INFO_BYTES);
size_t info[9] = {(size_t)mInputBlockNum, (size_t)eCount, (size_t)innerSide, (size_t)DST_XUNIT, (size_t)SRC_UNIT, (size_t)kernelsize, (size_t)blocklu, 0, 0};
if (mIm2ColBasedInt8) {
info[6] = LDiv4 / mInputBlockNum;
}
if (mToFuseInputbias2Bias) {
info[7] = 1;
}
if (mIm2ColParamter.padX > 0 || mIm2ColParamter.padY > 0) {
info[8] = 1;
}
// scale&bias:float32
gcore->MNNAsyQuantInfo(scalePtr, zeroPtr, qscale, qbias, (float*)minPtr, (float*)maxPtr, (float*)floatPtr, info);
// quant: float->int8_t
if (!mToFuseInputbias2Bias) {
gcore->MNNAsyQuantFunc(dstInt8, (float*)floatPtr, qscale, qbias, info);
} else {
auto sizeDiv4 = UP_DIV(eCount * LDiv4 * innerSide, PackUnit);
mQuantFunc((float*)floatPtr, dstInt8, sizeDiv4, qscale, -128, 127, qbias, 0);
}
if (mToFuseInputbias2Bias) { // Decode
inputZero = qbias[0];
auto updatedBiasPtr = (float*)(mBiasBufferFusedInputzero.ptr() + tId * ocUp4 * QUANT_INFO_BYTES);
auto matmulBiasPtr = mResourceInt8->mOriginBias->host<float>();
auto weightKernelSum = mResourceInt8->mWeightKernelSum->host<float>();
auto zero_ = -inputZero * scalePtr[0];
gcore->MNNDynamicUpdateConvBiasScale(updatedBiasPtr, matmulBiasPtr, weightKernelSum, &zero_, UP_DIV(ocUp4, 4));
biasPtr = (uint8_t*)updatedBiasPtr;
auto unitsize = mBatchQuantInfo->length(1) / (2 * QUANT_INFO_BYTES);
auto inputScale = scalePtr[0];
for (int i = 0; i < unitsize; ++i) {
((float*)inputDequantScale)[i] = inputScale;
}
}
};
auto BatchSymDynamicQuant = [&](uint8_t* floatPtr, int32_t& inputZero, uint8_t* inputDequantScale, int LU, int EP, int LP, int32_t availableThreads, int8_t* dstInt8, int tId) {
auto quantPtr = mQScaleZero.ptr() + tId * mSizeInputBlockQuant * QUANT_INFO_BYTES;
auto maxPtr = mTempMaxMinValueBuffer.ptr() + tId * mSizeInputBlockQuant * gcore->bytes;
// compute sum and absmax
int divlu = UP_DIV(LU, availableThreads);
MNN_CONCURRENCY_BEGIN (tIdx, ALIMIN(availableThreads, UP_DIV(LU, divlu))) {
auto exeLu = ALIMIN(divlu, LU - tIdx * divlu);
auto batchMax = reinterpret_cast<float*>(maxPtr + tIdx * EP * gcore->bytes);
auto ptr_ = reinterpret_cast<float*>(floatPtr + tIdx * divlu * gcore->bytes * EP * LP);
gcore->MNNAbsMax((float*)ptr_, batchMax, exeLu, EP, LP);
} MNN_CONCURRENCY_END();
// Compute quant scale
gcore->MNNQuantScale((float*)maxPtr, (float*)quantPtr, (float*)inputDequantScale, availableThreads, EP);
// quant
auto scale_ptr = reinterpret_cast<float*>(quantPtr);
gcore->MNNDynamicQuant((float*)floatPtr, dstInt8, scale_ptr, LU, EP, LP, nullptr);
inputZero = 0;
};
if (mResourceInt8->mDynamicQuant) {
biasPtr = mResourceInt8->mOriginBias->host<uint8_t>();
}
if (mIm2ColBasedInt8 && mResourceInt8->mDynamicQuant) {
int icDiv4 = UP_DIV(input->channel(), PackUnit);
if (mUseBatchQuan) {
int availthreads = (icDiv4 > mThreadNums && inputPlane > 255 ) ? mThreadNums : 1;
if (dynamicOption != 2) {
BatchSymDynamicQuant(input->host<uint8_t>(), inputZeroPoint, mBatchQuantInfo->host<uint8_t>(), icDiv4, inputPlane, PackUnit, availthreads, mQuantInput->host<int8_t>(), 0);
} else {
BatchAsyDynamicQuant(input->host<uint8_t>(), inputZeroPoint, mBatchQuantInfo->host<uint8_t>(), icDiv4, inputPlane, PackUnit, availthreads, mQuantInput->host<int8_t>(), mBatchQuantInfo->host<uint8_t>() + mBatchQuantInfo->stride(0) / 2, 0);
}
} else {
BatchAsyDynamicQuant(input->host<uint8_t>(), inputZeroPoint, mBatchQuantInfo->host<uint8_t>(), icDiv4, inputPlane, PackUnit, 1, mQuantInput->host<int8_t>(), mBatchQuantInfo->host<uint8_t>() + mBatchQuantInfo->stride(0) / 2, 0);
}
im2colSrc = mQuantInput->host<uint8_t>();
}
#endif
if (mResourceInt8->mActBits == 4) {
weightBytes = 0.5;
weightStepY /= 2;
}
int blockunit = ocUp4 * 2 * QUANT_INFO_BYTES + blockL * weightStepY * UP_DIV(output->channel(), UNIT);
auto inputchannel = input->channel();
SumByAxisParams sumParams;
sumParams.oneScale = (mUseBatchQuan || dynamicOption == 2) ? 0 : 1;
sumParams.SRC_UNIT = SRC_UNIT;
sumParams.blockNum = mBlockNum;
sumParams.DST_XUNIT = DST_XUNIT;
sumParams.unitColBufferSize = unitColBufferSize;
sumParams.kernelCountUnitDouble = kernelCountUnit;
sumParams.valid = inputchannel % SRC_UNIT;
sumParams.kernelxy = kxky;
sumParams.LU = UP_DIV(inputchannel, SRC_UNIT);
sumParams.inputBlock = (mInputBlockNum > 1) ? 1 : 0;
auto tileSplitFunction = [&](int tId, int eStartIndex, int eEndIndex, int estep) {
auto ocDivThread = ocDiv4;
float* reluPtr = mResourceInt8->mReluThreshold.data();
float* accumbuff = nullptr;
uint8_t* inputScale = nullptr;
uint8_t* inputBias = nullptr;
uint8_t* ptrInputScale = nullptr;
uint8_t* ptrInputBias = nullptr;
if (mBatchQuantInfo.get()) {
if (mIm2ColBasedInt8) {
inputScale = mBatchQuantInfo->host<uint8_t>();
ptrInputScale = inputScale;
}
if (dynamicOption == 2 && mUseBatchQuan && mIm2ColBasedInt8) {
inputBias = inputScale + mBatchQuantInfo->stride(0) / 2;
ptrInputBias = inputBias;
}
}
if (mBlockNum > 1) {
accumbuff = reinterpret_cast<float*>(mAccumBuffer->host<int8_t>() + tId * mAccumBuffer->stride(0) * sizeof(int32_t));
}
float* ptrY = nullptr;
if ((dstBytes != 1)) {
ptrY = mResourceInt8->mWeightKernelSum->host<float>();
}
QuanPostTreatParameters quanParam;
quanParam.blockNum = mBlockNum;
quanParam.weightKernelSum = ptrY;
if (dstBytes != 1) {
quanParam.useInt8 = 0;
quanParam.fp32minmax = reluPtr;
} else {
quanParam.maxValue = mMutableResource->mClampMax;
if (mResourceInt8->mRelu) {
quanParam.minValue = mMutableResource->mOutputZeroPoint;
} else {
quanParam.minValue = mMutableResource->mClampMin;
}
}
auto weightPtrTid = weightDataPtr;
quanParam.biasFloat = reinterpret_cast<float*>(biasPtr);
auto im2colDstThread = im2colPtr + tId * mTempIm2ColBuffer->stride(0);
auto srcPtr = (int8_t const **)(mBlitInfo.ptr() + tId * mBlitInfoStride.first);
auto el = (int32_t *)(srcPtr + mBlitInfoStride.second);
auto xKernelSumPtrTid = reinterpret_cast<float*>(srcKernelSumPtr + tId * mBlockNum * DST_XUNIT * mIm2ColCount * QUANT_INFO_BYTES);
int32_t info[5];
info[1] = mIm2ColParamter.iw * mIm2ColParamter.ih * batch;
info[2] = static_cast<int32_t>(unitColBufferSize);
info[3] = mIm2ColParamter.strideX;
for (int tIndex = eStartIndex; tIndex < eEndIndex; tIndex += estep) {
const int xIndexStart = tIndex * DST_XUNIT * mIm2ColCount;
auto outputInTilePtr = outputDataPtr + xIndexStart * PackUnit * dstBytes;
int realDstCount = ALIMIN(plane - xIndexStart, DST_XUNIT * mIm2ColCount);
ptrInputScale = (mUseBatchQuan && mIm2ColBasedInt8) ? (inputScale + xIndexStart * mInputBlockNum * QUANT_INFO_BYTES) : inputScale;
ptrInputBias = (inputBias != nullptr) ? (inputBias + xIndexStart * mInputBlockNum * QUANT_INFO_BYTES) : inputBias;
// im2col
auto im2colDst = im2colDstThread;
auto res = ConvolutionTiledExecutor::turnIm2ColToBlitInfo((const float**)srcPtr, el, xIndexStart, realDstCount, mIm2ColParamter, (uint8_t*)im2colSrc, im2colBytes);
int number = res.first;
bool needZero = res.second;
if (needZero && mIm2ColBasedInt8) {
#ifdef MNN_USE_SSE
::memset(im2colDst, inputZeroPoint + 128, colBufferSize);
#else
::memset(im2colDst, inputZeroPoint, colBufferSize);
#endif
}
info[0] = number;
info[4] = realDstCount;
if (mIm2ColBasedInt8 && number > 0) {
blitProc(im2colDst, srcPtr, info, el);
}
#ifdef MNN_LOW_MEMORY
if (!mIm2ColBasedInt8) {
if (needZero) {
::memset(im2colDst, 0, mTempIm2ColBuffer->stride(0));
}
if (number > 0) {
if (SRC_UNIT > PackUnit && !needZero) {
memset(im2colDst, 0, mTempIm2ColBuffer->stride(0));
}
info[2] = realDstCount;
gcore->MNNGeneralIm2Col((float*)im2colDst, (float const**)srcPtr, info, el, SRC_UNIT, PackUnit); // im2colDst: [lu, realDstCount, lp]
}
ptrInputScale = mBatchQuantInfo->host<uint8_t>() + tId * mBatchQuantInfo->stride(0);
if (dynamicOption == 2) {
ptrInputBias = ptrInputScale + mBatchQuantInfo->stride(0) / 2;
BatchAsyDynamicQuant((uint8_t*)im2colDst, inputZeroPoint, ptrInputScale, kernelCountUnit, realDstCount, SRC_UNIT, 1, mQuantInput->host<int8_t>() + tId * mQuantInput->stride(0), ptrInputBias, tId);
} else if (mUseBatchQuan) {
BatchSymDynamicQuant((uint8_t*)im2colDst, inputZeroPoint, ptrInputScale, kernelCountUnit, realDstCount, SRC_UNIT, 1, mQuantInput->host<int8_t>() + tId * mQuantInput->stride(0), tId);
} else {
auto maxMinPtr = mTempMaxMinValueBuffer.ptr() + tId * 2 * gcore->bytes;
ptrInputBias = ptrInputScale + mBatchQuantInfo->stride(0) / 2;
BatchAsyDynamicQuant((uint8_t*)im2colDst, inputZeroPoint, ptrInputScale, kernelCountUnit, realDstCount, SRC_UNIT, 1, mQuantInput->host<int8_t>() + tId * mQuantInput->stride(0), ptrInputBias, tId);
quanParam.biasFloat = (float*)(mBiasBufferFusedInputzero.ptr() + tId * ocUp4 * QUANT_INFO_BYTES);
}
im2colDst = mQuantInput->host<int8_t>() + tId * mQuantInput->stride(0);
}
if (mBlockNum > 1 && kxky > 1) {
auto eU = UP_DIV(realDstCount, DST_XUNIT); // eU <= mIm2ColCount
auto reorderBuffer = mReorderBuffer.ptr() + tId * colBufferSize;
for (int k = 0; k < eU; ++k) {
int inside = blocklu * SRC_UNIT * ALIMIN(realDstCount - k * DST_XUNIT, DST_XUNIT);
auto dstbuffer = reorderBuffer + k * unitColBufferSize;
auto srcbuffer = im2colDst + k * unitColBufferSize;
for (int i = 0; i < mBlockNum; ++i) {
for (int j = 0; j < kxky; ++j) {
memcpy(dstbuffer + i * kxky * inside + j * inside, srcbuffer + i * inside + j * mBlockNum * inside, inside);
}
}
}
im2colDst = (int8_t*)reorderBuffer;
}
#endif
if (mResourceInt8->mWeightAsymmetricQuant) {
MNN_ASSERT(mBatchQuantInfo.get() && mBatchQuantInfo->host<float>());
gcore->MNNSumByAxisLForMatmul_A(xKernelSumPtrTid, im2colDst, (float*)ptrInputScale, realDstCount, sumParams);
} else {
memset(xKernelSumPtrTid, 0, mBlockNum * DST_XUNIT * mIm2ColCount * QUANT_INFO_BYTES);
}
auto ptrX = xKernelSumPtrTid;
do {
int step = ALIMIN(DST_XUNIT, realDstCount);
quanParam.inputScale = (float*)ptrInputScale;
quanParam.inputBias = (float*)ptrInputBias;
if (mBlockNum > 1) {
memset(accumbuff, 0, UNIT * 4 * DST_XUNIT);
quanParam.accumBuffer = accumbuff;
}
quanParam.srcKernelSum = ptrX;
mGemmKernel(outputInTilePtr, im2colDst, weightPtrTid, blockL, dstZStep * dstBytes, ocDivThread, &quanParam, step);
ptrX += (step * mBlockNum);
realDstCount-=step;
outputInTilePtr += DST_XUNIT * PackUnit * dstBytes;
im2colDst += unitColBufferSize;
ptrInputScale = mUseBatchQuan ? (ptrInputScale + step * mInputBlockNum * QUANT_INFO_BYTES) : ptrInputScale;
ptrInputBias = (ptrInputBias != nullptr) ? (ptrInputBias + step * mInputBlockNum * QUANT_INFO_BYTES) : ptrInputBias;
} while(realDstCount > 0);
}
};
auto ocSplitFunction = [&](int threads) { // Thread split by OC
auto im2colDst = mTempIm2ColBuffer->host<int8_t>();
auto srcPtr = (int8_t const **)(mBlitInfo.ptr());
auto el = (int32_t *)(srcPtr + mBlitInfoStride.second);
auto xKernelSumPtr = reinterpret_cast<float*>(mTempSrcSum.ptr());
auto eU = UP_DIV(plane, DST_XUNIT);
int32_t info[5];
info[1] = mIm2ColParamter.iw * mIm2ColParamter.ih * batch;
info[2] = static_cast<int32_t>(unitColBufferSize);
info[3] = mIm2ColParamter.strideX;
float* reluPtr = mResourceInt8->mReluThreshold.data();
if (mIm2ColBasedInt8) { // im2col
auto res = ConvolutionTiledExecutor::turnIm2ColToBlitInfo((const float**)srcPtr, el, 0, plane, mIm2ColParamter, (uint8_t*)im2colSrc, im2colBytes);
int number = res.first;
bool needZero = res.second;
if (needZero) {
#ifdef MNN_USE_SSE
::memset(im2colDst, inputZeroPoint + 128, mTempIm2ColBuffer->size());
#else
::memset(im2colDst, inputZeroPoint, mTempIm2ColBuffer->size());
#endif
}
info[0] = number;
info[4] = plane;
if (number > 0) {
blitProc(im2colDst, srcPtr, info, el);
}
}
#ifdef MNN_LOW_MEMORY
if (false == mIm2ColBasedInt8) {
int realDstCount = plane;
int start = 0;
auto ptrInputscale = mBatchQuantInfo->host<uint8_t>();
auto ptrInputbias = ptrInputscale + mBatchQuantInfo->stride(0) / 2;
auto int8Ptr = mQuantInput->host<int8_t>();
int sizePacked = 0;
auto im2colDstTmp = im2colDst;
while (realDstCount > 0) {
int work = std::min(realDstCount, DST_XUNIT);
sizePacked += (work * SRC_UNIT * kernelCountUnit);
auto res = ConvolutionTiledExecutor::turnIm2ColToBlitInfo((const float**)srcPtr, el, start, work, mIm2ColParamter, (uint8_t*)im2colSrc, im2colBytes);
int number = res.first;
bool needZero = res.second;
if (needZero) {
::memset(im2colDstTmp, 0, unitColBufferSize * gcore->bytes);
}
info[0] = number;
info[2] = work;
if (number > 0) { // im2col
gcore->MNNGeneralIm2Col((float*)im2colDstTmp, (float const**)srcPtr, info, el, SRC_UNIT, PackUnit); // im2colDst: [lu, realDstCount, lp]
}
if (mUseBatchQuan || dynamicOption == 2) {
if (dynamicOption == 2) {
BatchAsyDynamicQuant((uint8_t*)im2colDstTmp, inputZeroPoint, ptrInputscale, kernelCountUnit, work, SRC_UNIT, 1, int8Ptr, ptrInputbias, 0);
ptrInputbias += (mInputBlockNum * work * sizeof(int32_t));
} else {
BatchSymDynamicQuant((uint8_t*)im2colDstTmp, inputZeroPoint, ptrInputscale, kernelCountUnit, work, SRC_UNIT, 1, int8Ptr, 0);
}
ptrInputscale += (mInputBlockNum * work * sizeof(int32_t));
int8Ptr += unitColBufferSize;
}
realDstCount -= work;
start += work;
im2colDstTmp += (unitColBufferSize * gcore->bytes);
}
if (!mUseBatchQuan && dynamicOption != 2) {
BatchAsyDynamicQuant((uint8_t*)im2colDst, inputZeroPoint, ptrInputscale, kernelCountUnit, plane, SRC_UNIT, 1, mQuantInput->host<int8_t>(), ptrInputscale + plane * mInputBlockNum* QUANT_INFO_BYTES, 0);
}
im2colDst = mQuantInput->host<int8_t>();
}
if (mBlockNum > 1 && kxky > 1) {
for (int k = 0; k < eU; ++k) {
int inside = blocklu * SRC_UNIT * ALIMIN(DST_XUNIT, plane - k * DST_XUNIT);
auto dstbuffer = mReorderBuffer.ptr() + k * unitColBufferSize;
auto srcbuffer = im2colDst + k * unitColBufferSize;
for (int i = 0; i < mBlockNum; ++i) {
for (int j = 0; j < kxky; ++j) {
memcpy(dstbuffer + i * kxky * inside + j * inside, srcbuffer + i * inside + j * mBlockNum * inside, inside);
}
}
}
im2colDst = (int8_t*)mReorderBuffer.ptr();
}
#endif
if (mResourceInt8->mWeightAsymmetricQuant) {
MNN_ASSERT(mBatchQuantInfo.get() && mBatchQuantInfo->host<float>());
gcore->MNNSumByAxisLForMatmul_A(xKernelSumPtr, im2colDst, mBatchQuantInfo->host<float>(), plane, sumParams);
} else {
memset(xKernelSumPtr, 0, mTileCount * mBlockNum * DST_XUNIT * mIm2ColCount * QUANT_INFO_BYTES);
}
MNN_CONCURRENCY_BEGIN(tId, threads) {
int ocIndex = PackUnit * mDivides[tId];
auto ocDivThread = ALIMIN(mDivides[tId + 1] - mDivides[tId], ocDiv4 - mDivides[tId]);
if (ocIndex < ocUp4) {
auto im2colDstThread = im2colDst;
float* ptrY = nullptr;
if (dstBytes != 1) {
ptrY = mResourceInt8->mWeightKernelSum->host<float>() + (ocIndex / UNIT) * UNIT * mInputBlockNum;
}
QuanPostTreatParameters quanParam;
quanParam.blockNum = mBlockNum;
quanParam.weightKernelSum = ptrY;
quanParam.biasFloat = reinterpret_cast<float*>(biasPtr + ocIndex * 4);
if (dstBytes != 1) {
quanParam.useInt8 = 0;
quanParam.fp32minmax = reluPtr;
} else {
quanParam.maxValue = mMutableResource->mClampMax;
if (mResourceInt8->mRelu) {
quanParam.minValue = mMutableResource->mOutputZeroPoint;
} else {
quanParam.minValue = mMutableResource->mClampMin;
}
}
uint8_t* inputScale = nullptr; // input scale for batch dynamic quant.
uint8_t* inputBias = nullptr;
float* accumbuff = nullptr;
if (mBatchQuantInfo.get()) {
inputScale = mBatchQuantInfo->host<uint8_t>();
if (dynamicOption == 2) {
inputBias = inputScale + mInputBlockNum * plane * QUANT_INFO_BYTES;
}
}
if (mBlockNum > 1) {
accumbuff = reinterpret_cast<float*>(mAccumBuffer->host<int8_t>() + tId * mAccumBuffer->stride(0) * sizeof(int32_t));
}
auto outputInTilePtr = outputDataPtr + ocIndex * plane * dstBytes;
const auto weightPtrTid = weightDataPtr + static_cast<int32_t>(ocIndex * mBlockNum * blockL * SRC_UNIT * weightBytes + ocIndex * 2 * mBlockNum * QUANT_INFO_BYTES);
int realDstCount = plane;
auto ptrX = xKernelSumPtr;
do {
int step = ALIMIN(DST_XUNIT, realDstCount);
quanParam.inputScale = (float*)inputScale;
quanParam.inputBias = (float*)inputBias;
quanParam.srcKernelSum = ptrX;
if (mBlockNum > 1) {
memset(accumbuff, 0, UNIT * 4 * DST_XUNIT);
quanParam.accumBuffer = accumbuff;
}
mGemmKernel(outputInTilePtr, im2colDstThread, weightPtrTid, blockL, dstZStep * dstBytes, ocDivThread, &quanParam, step);
ptrX += (step * mBlockNum);
realDstCount-=step;
outputInTilePtr += DST_XUNIT * PackUnit * dstBytes;
im2colDstThread += unitColBufferSize;
inputScale = mUseBatchQuan ? (inputScale + mInputBlockNum * step * QUANT_INFO_BYTES) : inputScale;
inputBias = (inputBias != nullptr) ? (inputBias + mInputBlockNum * step * QUANT_INFO_BYTES) : inputBias;
} while(realDstCount > 0);
}
}
MNN_CONCURRENCY_END();
};
const int threads = static_cast<CPUBackend*>(backend())->threadNumber();
if (!mSplitByOc) {
MNN_CONCURRENCY_BEGIN(tId, threads) {
if (mDivides[tId + 1] - mDivides[tId] > 0) {
tileSplitFunction((int)tId, mDivides[tId], mDivides[tId + 1], 1);
}
}
MNN_CONCURRENCY_END();
} else {
ocSplitFunction(threads);
}
return NO_ERROR;
}