wehub/alibaba--mnn
1
0
Fork 0
Code Issues Pull Requests Actions 1 Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
main
alibaba--mnn/source/backend/cpu/compute/ConvInt8TiledExecutor.cpp
T
wehub-resource-sync 5b57521aa1 chore: import upstream snapshot with attribution
2026-07-13 13:33:03 +08:00

2527 lines
124 KiB
C++
Raw Permalink Blame History

// ConvInt8TiledExecutor.cpp
// MNN
//
// Created by MNN on 2019/5/17.
// Copyright © 2018, Alibaba Group Holding Limited
//
#include "ConvInt8TiledExecutor.hpp"
#include "ConvolutionTiledExecutor.hpp"
#include "core/Macro.h"
#include "core/BufferAllocator.hpp"
#include "SharedGather.hpp"
#include <math.h>
#include "backend/cpu/CPUBackend.hpp"
#include "core/Concurrency.h"
#include "core/TensorUtils.hpp"
#define QUANT_INFO_BYTES 4
#define WEIGHT_ONLINE_REORDER 8
namespace MNN {
ConvInt8TiledExecutor::ConvInt8TiledExecutor(Backend* backend, const Op* op)
: CPUConvolution(op->main_as_Convolution2D()->common(), backend) {}
ConvInt8TiledExecutor::ConvInt8TiledExecutor(Backend* backend, const Op* op, std::shared_ptr<ResourceInt8> res)
: CPUConvolution(op->main_as_Convolution2D()->common(), backend), mResourceInt8(res) {
if (!res->mDynamicQuant) {
mMutableResource.reset(new MutableResourceInt8(res, backend));
mValid = mMutableResource->mValid;
}
}
ConvInt8TiledExecutor::~ConvInt8TiledExecutor() {
// Do nothing
}
bool ConvInt8TiledExecutor::onClone(Backend* bn, const Op* op, Execution** dst) {
return false;
}
ErrorCode ConvInt8TiledExecutor::onResize(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) {
if (nullptr != mMutableResource) {
mMutableResource->updateInputOutputScale(TensorUtils::getQuantInfo(inputs[0]),
TensorUtils::getQuantInfo(outputs[0]));
}
CPUConvolution::onResize(inputs, outputs);
ConvolutionTiledExecutor::setIm2ColParameter(mIm2ColParamter, mCommon, inputs[0], outputs[0], mPadX, mPadY,
static_cast<CPUBackend*>(backend())->functions(),
static_cast<CPUBackend*>(backend())->int8Functions());
return NO_ERROR;
}
void ConvInt8TiledExecutor::initializeConvInt8QuantInfo(std::shared_ptr<CPUConvolution::ResourceInt8>& resourceInt8,
const Convolution2D* conv2D,
std::shared_ptr<ConvolutionCommon::Int8Common> quanCommon) {
// input/output scale&zeorpoint
if (conv2D->symmetricQuan()) {
resourceInt8->mWeightBits = conv2D->symmetricQuan()->nbits();
}
if (conv2D->bias() && (conv2D->quanParameter()->alpha() || quanCommon->alpha.get())) {
resourceInt8->mUseConvQuan = false;
}
resourceInt8->mInputZeroPoint = 0;
resourceInt8->mOutputZeroPoint = 0;
resourceInt8->mClampMin = -128;
resourceInt8->mClampMax = 127;
if (conv2D->symmetricQuan()) {
resourceInt8->mInputZeroPoint = conv2D->symmetricQuan()->zeroPoint();
resourceInt8->mOutputZeroPoint = conv2D->symmetricQuan()->outputZeroPoint();
resourceInt8->mClampMin = conv2D->symmetricQuan()->clampMin();
resourceInt8->mClampMax = conv2D->symmetricQuan()->clampMax();
}
if (conv2D->quanParameter() != nullptr) {
resourceInt8->mInputScale = conv2D->quanParameter()->scaleIn();
resourceInt8->mOutputScale = conv2D->quanParameter()->scaleOut();
}
resourceInt8->mRelu = conv2D->common()->relu() || conv2D->common()->relu6();
if (conv2D->symmetricQuan() && conv2D->symmetricQuan()->outputDataType() == MNN::DataType_DT_FLOAT) {
resourceInt8->mOutputZeroPoint = 0;
resourceInt8->mOutputScale = 1.0f;
}
}
void ConvInt8TiledExecutor::reorderWeight(uint8_t* dst, const uint8_t* src, int32_t* info, int32_t initval,
float* kernelsum, weightSummerFuncion summerFunc) {
// weight shape = {UP_DIV(oc, UNIT), blockNum, kernelCount* UP_DIV(ic / blockNum, SRC_UNIT), UNIT, SRC_UNIT};
MNN_ASSERT(dst != nullptr && src != nullptr);
int blockNum = info[0];
int oc = info[1];
int ic = info[2];
int kernelCount = info[3];
int UNIT = info[4];
int SRC_UNIT = info[5];
int blockL = UP_DIV(ic / blockNum, SRC_UNIT) * kernelCount;
int stride0 = blockNum * SRC_UNIT * blockL * UNIT; // weight->stride(0)
int stride1 = blockL * SRC_UNIT * UNIT; // weight->stride(1)
int stride2 = UNIT * SRC_UNIT; // weight->stride(2)
int weightlen = stride0 * UP_DIV(oc, UNIT);
memset(dst, initval, weightlen);
auto hU = UP_DIV(oc, UNIT);
auto lU = UP_DIV(ic / blockNum, SRC_UNIT) * kernelCount;
bool fast = (kernelCount == 1 && ROUND_UP(oc, UNIT) == oc && (ic % (blockNum * SRC_UNIT)) == 0);
if (fast) {
for (int i = 0; i < hU; ++i) {
for (int k = 0; k < UNIT; ++k) {
for (int bl = 0; bl < blockNum; ++bl) {
for (int j = 0; j < blockL; ++j) {
int srcindex = (i * UNIT + k) * ic + bl * (lU * SRC_UNIT) + j * SRC_UNIT;
int dstindex = i * stride0 + bl * stride1 + j * stride2 + k * SRC_UNIT;
memcpy(dst + dstindex, src + srcindex, SRC_UNIT);
}
}
}
}
} else {
auto blockic = ic / blockNum;
auto bklU = UP_DIV(blockic, SRC_UNIT);
for (int i_hU = 0; i_hU < hU; ++i_hU) {
int dst_i_hU = i_hU * stride0;
for (int bk = 0; bk < blockNum; ++bk) {
int dst_bk = dst_i_hU + bk * stride1;
int src_bk = bk * blockic * kernelCount;
for (int k2 = 0; k2 < kernelCount; ++k2) {
int dst_k2 = dst_bk + k2 * bklU * stride2;
int src_k2 = src_bk + k2;
for (int blu = 0; blu < bklU; ++blu) {
int dst_blu = dst_k2 + blu * stride2;
int src_blu = src_k2 + blu * SRC_UNIT * kernelCount;
for (int inId = 0; inId < UNIT; ++inId) {
int i = i_hU * UNIT + inId;
if (i >= oc)
continue;
int dst_inId = dst_blu + inId * SRC_UNIT;
int src_inId = src_blu + i * ic * kernelCount;
for (int blp = 0; blp < SRC_UNIT; ++blp) {
int j_in_block = blu * SRC_UNIT + blp;
if (j_in_block >= blockic)
continue;
int dstindex = dst_inId + blp;
int srcindex = src_inId + blp * kernelCount;
dst[dstindex] = src[srcindex];
}
}
}
}
}
}
} // not fast
if (summerFunc != nullptr && kernelsum != nullptr) {
summerFunc(kernelsum, (int8_t*)dst, blockNum * hU, blockL, UNIT, SRC_UNIT);
}
}
void ConvInt8TiledExecutor::packWeightAndQuantInfo(int8_t* dstbuffer, const int8_t* weight, const int8_t* quantInfo,
int32_t* info, int infoBytes) {
int blockNum = info[0];
int ocDiv = info[1];
int blockL = info[2];
int UNIT = info[3];
int SRC_UNIT = info[4];
auto ocUp4 = info[5];
auto src0 = weight; // int8 weight: [oc/hp, blocknum, ic/lp*(kx*ky)/blocknum, hp, lp]
auto src1 = quantInfo; // dequant scale: [blocknum, ocUp4]
auto src2 = src1 + infoBytes * ocUp4 * blockNum; // dequant bias: [blocknum, ocUp4]
int stride0 = info[0] * info[2] * info[3] * info[4];
int stride1 = info[2] * info[3] * info[4];
// dst: [oc/hp, blocknum, packedUnit]
// packedUnit: [ic/lp*(kx*ky)/blocknum, hp, lp] + [hp] + [hp]
for (int hU = 0; hU < ocDiv; ++hU) {
auto huPtr = dstbuffer + hU * blockNum * (stride1 + 2 * UNIT * infoBytes);
int scaleCount = ALIMIN(ocUp4 - hU * UNIT, UNIT);
for (int bl = 0; bl < blockNum; ++bl) {
auto blockPtr = huPtr + bl * (stride1 + 2 * UNIT * infoBytes);
memcpy(blockPtr, src0 + bl * stride1 + hU * stride0, stride1);
memcpy(blockPtr + stride1, src1 + (bl * ocUp4 + hU * UNIT) * infoBytes, scaleCount * infoBytes);
memcpy(blockPtr + stride1 + UNIT * infoBytes, src2 + (bl * ocUp4 + hU * UNIT) * infoBytes,
scaleCount * infoBytes);
}
}
}
static void _computeReorderQuantInfo(float* weightKernelSum, int32_t* paramsKernelSum, bool blockQuantInput,
int weightBits, bool asyQuantWeight, float* quanInfoPtr, int outputCount,
int kernelCount, int pack, AutoStorage<int8_t>& reorderedQuantInfo,
float* ikernelSum, int HP, bool realInt4OrInt8) {
// Only used for dynamic quant:
// copy gemm bias
// copy/compute real dequant bias/scale
// dequant weight kernel sum
int ocUp4 = ROUND_UP(outputCount, pack);
int ocUpHp = ROUND_UP(outputCount, HP);
int blockNum = paramsKernelSum[0];
int kernelSumSize = paramsKernelSum[1];
int scaleSize = blockNum * ocUp4; // pack size.
int blockSize = kernelCount / blockNum;
int originOffset = 0;
if (weightBits == 4) {
originOffset = -8;
} else if (weightBits == 3) {
// w3 kernels produce unsigned 0..7; -4 re-centers to signed [-4, 3].
originOffset = -4;
} else if (weightBits == 2) {
originOffset = -2;
}
// Save weight quant alpha and zero: wf=alpha*wi+zero
auto alphaPtr = reinterpret_cast<float*>(reorderedQuantInfo.get());
auto biasPtr = reinterpret_cast<float*>(reinterpret_cast<uint8_t*>(alphaPtr) + scaleSize * QUANT_INFO_BYTES);
if (outputCount % pack != 0) {
::memset(alphaPtr, 0, scaleSize * QUANT_INFO_BYTES);
::memset(biasPtr, 0, scaleSize * QUANT_INFO_BYTES);
}
::memset(weightKernelSum, 0, kernelSumSize * QUANT_INFO_BYTES);
int ocDiv4 = UP_DIV(outputCount, pack);
// resource->mWeightKernelSum: [hU,blocknum,hP]
if (asyQuantWeight) {
for (int i = 0; i < outputCount; ++i) {
float accum = 0.f;
auto ocOutside = i / HP;
auto ocInside = i % HP;
for (int j = 0; j < blockNum; ++j) {
int index = i * blockNum + j;
int srcSumIndex = ocOutside * blockNum * HP + j * HP + ocInside; // ikernelsum: [hU,blocknum,hP]
alphaPtr[j * ocUp4 + i] = quanInfoPtr[2 * index + 1];
biasPtr[j * ocUp4 + i] = quanInfoPtr[2 * index] + (float)originOffset * quanInfoPtr[2 * index + 1];
if (realInt4OrInt8) {
accum +=
(ikernelSum[srcSumIndex] * quanInfoPtr[2 * index + 1] + blockSize * biasPtr[j * ocUp4 + i]);
} else {
accum += ((ikernelSum[srcSumIndex] - blockSize * 8) * quanInfoPtr[2 * index + 1] +
blockSize * quanInfoPtr[2 * index]);
}
if (blockQuantInput) {
int dstSumIndex = ocOutside * blockNum * HP + j * HP + ocInside;
weightKernelSum[dstSumIndex] = accum;
accum = 0;
}
}
if (!blockQuantInput) {
weightKernelSum[i] = accum;
}
}
} else {
for (int i = 0; i < outputCount; ++i) {
float accum = 0.f;
auto ocOutside = i / HP;
auto ocInside = i % HP;
for (int j = 0; j < blockNum; ++j) {
int index = i * blockNum + j;
int srcSumIndex = ocOutside * blockNum * HP + j * HP + ocInside; // ikernelsum: [hU,blocknum,hP]
alphaPtr[j * ocUp4 + i] = quanInfoPtr[index];
biasPtr[j * ocUp4 + i] = (float)originOffset * quanInfoPtr[index];
if (realInt4OrInt8) {
accum += (ikernelSum[srcSumIndex] * quanInfoPtr[index] + blockSize * biasPtr[j * ocUp4 + i]);
} else {
accum += ((ikernelSum[srcSumIndex] - blockSize * 8) * quanInfoPtr[index]);
}
if (blockQuantInput) {
int dstSumIndex = ocOutside * blockNum * HP + j * HP + ocInside;
weightKernelSum[dstSumIndex] = accum;
accum = 0;
}
}
if (!blockQuantInput) {
weightKernelSum[i] = accum;
}
}
}
}
static inline void calculateSmeNeonWorkDivision(int& ocMain, int& ocBranch, std::vector<int>& divides, int oc,
int threads, int pack, int planeSize, int divisionRatio, int smeCores) {
// workload
auto ocDivPack = UP_DIV(oc, pack);
auto workUnit = UP_DIV(ocDivPack, divisionRatio * smeCores + 1 * (threads - smeCores));
int calOcMain = ALIMIN(ROUND_UP(workUnit * pack * smeCores * divisionRatio, GEMM_INT8_UNIT_SME2_128), oc);
if (calOcMain <= ocMain) { // The purpose of this function is to increase the value of ocMain.
return;
}
ocMain = calOcMain;
ocBranch = oc - ocMain;
divides.assign(threads + 1, ocDivPack);
divides[0] = 0;
// runtime UNIT for different core and different process(prefill or decode)
auto rtUnit4Sme = planeSize == 1 ? GEMM_INT8_UNIT_SME2_128 : GEMM_INT8_UNIT_SME2;
// mOcMain
auto ocPerSmeCore = ALIMIN(UP_DIV(UP_DIV(ROUND_UP(ocMain, pack), rtUnit4Sme), smeCores) * (rtUnit4Sme / pack),
UP_DIV(ocMain, pack));
for (int i = 0; i < smeCores; ++i) {
divides[i + 1] = ALIMIN(divides[i] + ocPerSmeCore, UP_DIV(ocMain, pack));
}
// ocRemain
if (ocBranch > 0) {
auto ocPerNeonCore = UP_DIV(UP_DIV(ROUND_UP(ocBranch, pack), GEMM_INT8_UNIT_ARM82), threads - smeCores) *
(GEMM_INT8_UNIT_ARM82 / pack);
for (int i = smeCores + 1; i < threads + 1; ++i) {
divides[i] = ALIMIN(divides[i - 1] + ocPerNeonCore, ocDivPack);
}
}
}
static inline void _getProportions(int totalProp, int& intensiveProp, int& lightProp) {
// compute the proportions of different kernels
lightProp = totalProp % 8;
intensiveProp = totalProp / 8 % 8;
if (lightProp == 0 && intensiveProp == 0) {
// pass
// Don't use mixed kernels
} else if (lightProp == 0) {
lightProp = 1;
} else if (intensiveProp == 0) {
intensiveProp = 6;
} else if (lightProp > intensiveProp) {
lightProp = 1;
}
}
static inline void _computeDivides4Sme(std::vector<int>& divides, int threads, int smeCoreNums, int size) {
divides.resize(threads + 1);
divides[0] = 0;
auto length = UP_DIV(size, smeCoreNums);
auto cur = length;
for (int i = 1; i < smeCoreNums + 1; ++i) {
divides[i] = cur;
cur = ALIMIN(cur + length, size);
}
}
static inline void _updateMixedKernelFlag(bool& mixedKernel, bool& onlineReorderWeightSme, int threads, int eP,
bool isDynamciQuant, bool postiveBothProp) {
mixedKernel = false;
if (threads >= 4 && eP == GEMM_INT8_DST_XUNIT_SME2 && isDynamciQuant && postiveBothProp) {
mixedKernel = true;
onlineReorderWeightSme = true;
}
}
DenseConvInt8TiledExecutor::DenseConvInt8TiledExecutor(Backend* backend, const Op* op,
std::shared_ptr<ConvolutionCommon::Int8Common> quanCommon,
bool isDynamicQuant)
: ConvInt8TiledExecutor(backend, op) {
// convolution info
auto convOp = op->main_as_Convolution2D();
int kernelCount = mCommon->kernelX() * mCommon->kernelY();
int oc = convOp->common()->outputCount();
int ic = convOp->common()->inputCount();
bool asyWeight = quanCommon ? quanCommon->asymmetric : false;
mOcBranch = 0;
mOcMain = oc;
int blockNum = 1;
int inputBlockNum = 1;
if (quanCommon) {
int dequantCnt = quanCommon->alphaSize;
if (quanCommon->asymmetric) {
dequantCnt /= 2;
}
blockNum = dequantCnt / oc;
}
mBlockNum = blockNum;
// backend info
auto core = static_cast<CPUBackend*>(backend)->int8Functions();
auto gcore = static_cast<CPUBackend*>(backend)->functions();
const int threads = static_cast<CPUBackend*>(backend)->threadNumber();
const int pack = gcore->pack;
// runtime hint
auto option = static_cast<CPUBackend*>(backend)->getRuntime()->hint().dynamicQuantOption;
auto weightOnlineReorderOption = WEIGHT_ONLINE_REORDER & option;
auto inputBlockQuantOption = option % WEIGHT_ONLINE_REORDER;
if (inputBlockQuantOption == 2) {
inputBlockNum = blockNum;
}
_getProportions(static_cast<CPUBackend*>(backend)->getRuntime()->hint().divisionRatio, mRatioPrefill, mRatioDecode);
mSmeCores = gcore->smeCoreNumber;
mRelatedFunctions = *(static_cast<CPUBackend*>(backend)->int8GemmFunctions());
mArm82Functions = gcore->arm82MatmulRelatedFunctions;
int UNITMain, SRC_UNITMain, DST_XUNITMain;
int UNITBranch = 0;
int SRC_UNITBranch = 0, DST_XUNITBranch = 0;
mRelatedFunctions.MNNGetGemmUnit(&UNITMain, &SRC_UNITMain, &DST_XUNITMain);
if (mArm82Functions.MNNGetGemmUnit != nullptr) { // exclude cpu does not support arm82
mArm82Functions.MNNGetGemmUnit(&UNITBranch, &SRC_UNITBranch, &DST_XUNITBranch);
}
// prefer to maximum decode performance & the machine supports 'sme2' & the runtime backend is 'sme2' ->
// mOnlineReorderWeightSme=true
mOnlineReorderWeightSme = (weightOnlineReorderOption > 0 && DST_XUNITMain == GEMM_INT8_DST_XUNIT_SME2);
if (isDynamicQuant == false) {
mOnlineReorderWeightSme = false;
}
_updateMixedKernelFlag(mMixedKernel, mOnlineReorderWeightSme, threads, DST_XUNITMain, isDynamicQuant,
mRatioDecode && mRatioPrefill);
if (mMixedKernel) {
// total work: UP_DIV(oc, pack)
// (sme's work / neon's work) = divisionRatio
auto workUnit = UP_DIV(UP_DIV(oc, pack), mRatioDecode * mSmeCores + 1 * (threads - mSmeCores));
mOcMain = ALIMIN(ROUND_UP(workUnit * pack * mSmeCores * mRatioDecode, GEMM_INT8_UNIT_SME2_128), oc);
;
mOcBranch = oc - mOcMain;
}
if (mOnlineReorderWeightSme) {
UNITMain = GEMM_INT8_UNIT_SME2_128;
}
// compute info
int ocUp4Main = ROUND_UP(mOcMain, pack);
int ocUpHpMain = ROUND_UP(mOcMain, UNITMain);
int lUMain = UP_DIV(ic / blockNum, SRC_UNITMain) * kernelCount;
int scaleSizeMain = ocUp4Main * blockNum;
int ocUp4Branch = ROUND_UP(mOcBranch, pack);
int ocUpHpBranch = UNITBranch != 0 ? ROUND_UP(mOcBranch, UNITBranch) : 0;
int ocDivHpBranch = UNITBranch != 0 ? UP_DIV(mOcBranch, UNITBranch) : 0;
int lUBranch = UNITBranch != 0 ? UP_DIV(ic / blockNum, SRC_UNITBranch) * kernelCount : 0;
int scaleSizeBranch = ocUp4Branch * blockNum;
std::vector<int> shapeMain = {blockNum, UP_DIV(mOcMain, UNITMain), lUMain, UNITMain, SRC_UNITMain};
std::vector<int> shapeBranch = {blockNum, ocDivHpBranch, lUBranch, UNITBranch, SRC_UNITBranch};
mResourceInt8.reset(new CPUConvolution::ResourceInt8);
mResourceInt8->mWeightAsymmetricQuant = asyWeight;
mResourceInt8->mWeightBits = 8;
mResourceInt8->mBlockNum = blockNum;
mResourceInt8->mHp = UNITMain;
mResourceInt8->mLp = SRC_UNITMain;
if (DST_XUNITMain == GEMM_INT8_DST_XUNIT_SME2) {
mResourceInt8->mPackMode = 1;
}
if (quanCommon && quanCommon->canUseInt4) {
shapeMain[4] = SRC_UNITMain / 2;
shapeBranch[4] = SRC_UNITBranch / 2;
mResourceInt8->mWeightBits = 4;
mResourceInt8->mWeightAsymmetricQuant = true; // offset: 8 from uint8_t
} else if (quanCommon && quanCommon->canUseInt2 &&
((gcore->bytes == 2 && gcore->pack == 8 &&
mRelatedFunctions.MNNGemmInt8AddBiasScale_w2_Unit_FP16 != nullptr) ||
mRelatedFunctions.Int8GemmKernel_W2 != nullptr)) {
// FP16 path uses MNNGemmInt8AddBiasScale_w2_Unit_FP16; FP32 path uses Int8GemmKernel_W2.
// Backends without either kernel fall through to the int8 (loader-expanded) path.
shapeMain[4] = SRC_UNITMain / 4;
shapeBranch[4] = SRC_UNITBranch / 4;
mResourceInt8->mWeightBits = 2;
mResourceInt8->mWeightAsymmetricQuant = true; // offset: 2 from uint8_t
} else if (quanCommon && quanCommon->canUseInt3 &&
((gcore->bytes == 2 && gcore->pack == 8 &&
mRelatedFunctions.MNNGemmInt8AddBiasScale_w3_Unit_FP16 != nullptr) ||
mRelatedFunctions.Int8GemmKernel_W3 != nullptr)) {
// bit-plane (2 + 1) split: per-cell layout is [main | aux] + optional padding.
// i8mm (SRC_UNIT=8): 16B main + 8B aux = 24 bytes (shape[4] = 3, exact).
// sdot (SRC_UNIT=4): 8B main + 4B aux = 12 bytes; 12/8 isn't integer, so we
// ceil-round shape[4] to 2 (= 16-byte cell) and pad the trailing 4 bytes with
// zero. The kernel skips those bytes after the aux load. Without ceiling,
// truncating to 1 undersizes the buffer by 33% and the packer/kernel write
// past the end → SIGSEGV in a later malloc/memset.
shapeMain[4] = (SRC_UNITMain * 3 + 7) / 8;
shapeBranch[4] = (SRC_UNITBranch * 3 + 7) / 8;
mResourceInt8->mWeightBits = 3;
mResourceInt8->mWeightAsymmetricQuant = true; // offset: 4 from uint8_t
}
mResourceInt8->mDynamicQuant = isDynamicQuant ? true : false;
// Relu/Relu6 post parameters
auto postPtr = getPostParameters();
mResourceInt8->mReluThreshold.resize(2);
mResourceInt8->mReluThreshold[0] = postPtr[2];
mResourceInt8->mReluThreshold[1] = postPtr[3];
if (gcore->bytes == 2) {
gcore->MNNFp32ToLowp(mResourceInt8->mReluThreshold.data(),
reinterpret_cast<int16_t*>(mResourceInt8->mReluThreshold.data()), 2);
}
// buffer allocate
auto quantlenMain = 2 * blockNum * ROUND_UP(mOcMain, UNITMain) * QUANT_INFO_BYTES;
auto weightlenMain = shapeMain[0] * shapeMain[1] * shapeMain[2] * shapeMain[3] * shapeMain[4];
auto quantlenBranch = 2 * blockNum * ocUpHpBranch * QUANT_INFO_BYTES;
auto weightlenBranch = shapeBranch[0] * shapeBranch[1] * shapeBranch[2] * shapeBranch[3] * shapeBranch[4];
mResourceInt8->mWeightInt8.reset(
Tensor::createDevice<uint8_t>({weightlenMain + quantlenMain + weightlenBranch + quantlenBranch}));
mResourceInt8->mOriginBias.reset(Tensor::createDevice<int32_t>({ocUp4Main + ocUpHpBranch})); // float
mResourceInt8->mWeightKernelSum.reset(
Tensor::createDevice<uint8_t>({inputBlockNum * QUANT_INFO_BYTES * (ocUpHpMain + ocUpHpBranch)}));
auto res = backend->onAcquireBuffer(mResourceInt8->mOriginBias.get(), Backend::STATIC);
res &= backend->onAcquireBuffer(mResourceInt8->mWeightKernelSum.get(), Backend::STATIC);
res &= backend->onAcquireBuffer(mResourceInt8->mWeightInt8.get(), Backend::STATIC);
if (!res) {
MNN_ERROR("weight acquire buffer error\n");
return;
}
bool useCachedMmap = backend->getRuntime()->hint().useCachedMmap > 1;
if (useCachedMmap) {
return;
}
// read weight, weight's scale&bias, convolution bias
::memset(mResourceInt8->mOriginBias->host<float>(), 0, mResourceInt8->mOriginBias->size());
// dynamic quant
bool directReadInt4weight = (kernelCount == 1 && ROUND_UP(mOcMain, UNITMain) == mOcMain &&
ROUND_UP(ic, SRC_UNITMain) == ic); // TODO:fix this
auto ocMain = mOcMain;
auto ocBranch = mOcBranch;
auto target = mResourceInt8;
auto funcsMain = mRelatedFunctions;
auto funcsBranch = mArm82Functions;
auto needToReorderWeightOnline4Sme = mOnlineReorderWeightSme;
// Save bias
if (convOp->bias()) {
::memcpy(mResourceInt8->mOriginBias->host<float>(), convOp->bias()->data(),
convOp->bias()->size() * sizeof(float));
}
auto coreFuncs = static_cast<CPUBackend*>(backend)->functions();
auto reorderFunc = [=](decltype(mRelatedFunctions) funcs, std::vector<int> shape, int UNIT, int SRC_UNIT,
int DST_XUNIT, int weightlen, int scaleSize, int oc, int offsetTg, bool fastReadWeight,
int8_t** addressPtr, weightSummerFuncion sumFunc) -> int {
auto sh = shape;
AutoStorage<int8_t> weightReordered(weightlen);
AutoStorage<int8_t> reorderedQuantInfo(2 * scaleSize * QUANT_INFO_BYTES);
AutoStorage<int8_t> kernelsum(blockNum * ROUND_UP(oc, UNIT) * QUANT_INFO_BYTES);
if (weightReordered.get() == nullptr || reorderedQuantInfo.get() == nullptr || kernelsum.get() == nullptr) {
MNN_ERROR("Memory not enough\n");
return -1;
}
memset(kernelsum.get(), 0, blockNum * ROUND_UP(oc, UNIT) * QUANT_INFO_BYTES);
/* 1. reorder weight */
auto srcPtr = (uint8_t*)addressPtr[0];
if (target->mWeightBits == 4 && fastReadWeight) {
auto dstPtr = (uint8_t*)weightReordered.get();
::memset(dstPtr, 0, weightlen);
funcs.MNNReorderWeightInt4(dstPtr, srcPtr, sh.data(), sh.size(), (float*)kernelsum.get());
} else { // int4 weight but oc/ic not packed
int blocksize = ic * kernelCount / blockNum;
int originOffset = 0;
int32_t info[6] = {blockNum, oc, ic, kernelCount, UNIT, SRC_UNIT};
if (target->mWeightBits == 4) {
if (kernelCount == 1 && ROUND_UP(ic, SRC_UNIT) == ic) {
int ocUp = ROUND_UP(oc, UNIT);
int rowBytes = ic * kernelCount / 2;
AutoStorage<uint8_t> paddedWeight(ocUp * rowBytes);
if (paddedWeight.get() == nullptr) {
MNN_ERROR("Weight reorder memory not enough!\n");
return -1;
}
::memset(paddedWeight.get(), 0, ocUp * rowBytes);
for (int o = 0; o < oc; ++o) {
::memcpy(paddedWeight.get() + o * rowBytes, srcPtr + o * rowBytes, rowBytes);
}
funcs.MNNReorderWeightInt4(reinterpret_cast<uint8_t*>(weightReordered.get()), paddedWeight.get(),
sh.data(), sh.size(), reinterpret_cast<float*>(kernelsum.get()));
} else {
originOffset = -8;
std::vector<uint8_t> tmpWeight(oc * ic * kernelCount);
for (int j = 0; j < oc; ++j) {
for (int k = 0; k < blockNum; ++k) {
for (int i = 0; i < blocksize; ++i) {
int index = j * blockNum * blocksize + k * blocksize + i;
uint8_t w_ = srcPtr[index / 2];
int truew = index % 2 ? (w_ & 0x0f) : (w_ >> 4);
tmpWeight[index] = truew;
}
}
}
AutoStorage<uint8_t> packedInt8weight(weightlen * 2);
if (packedInt8weight.get() == nullptr) {
MNN_ERROR("Weight reorder memory not enough!\n");
return -1;
}
reorderWeight(packedInt8weight.get(), (uint8_t*)tmpWeight.data(), info, 0, (float*)kernelsum.get(),
sumFunc);
// pack two int4 to int8
int leng = weightlen * 2;
auto srcint4Ptr = (uint8_t*)packedInt8weight.get();
auto dstint4Ptr = (uint8_t*)weightReordered.get();
int permuteUnit = UNIT * SRC_UNIT;
int halfPermuteStride = static_cast<int32_t>(permuteUnit / 2);
for (int i = 0; i < leng / permuteUnit; ++i) {
auto src0 = srcint4Ptr + i * permuteUnit;
auto dst0 = dstint4Ptr + i * halfPermuteStride;
for (int j = 0; j < halfPermuteStride; ++j) {
int s0, s1, d;
if (DST_XUNIT == GEMM_INT8_DST_XUNIT_SME2) { // SME2
s0 = src0[2 * j + 0];
s1 = src0[2 * j + 1];
d = s0 + (s1) * 16;
} else {
s0 = src0[j];
s1 = src0[j + halfPermuteStride];
d = (s0) * 16 + (s1);
}
dst0[j] = d;
}
}
}
} else if (target->mWeightBits == 2) {
// Loader gave us signed int8 in [-2, 1]; convert to unsigned [0, 3], reorder via int8
// path, then pack 4 weights/byte into the kernel-expected layout.
originOffset = -2;
std::vector<uint8_t> tmpWeight(oc * ic * kernelCount);
for (int idx = 0; idx < oc * ic * kernelCount; ++idx) {
tmpWeight[idx] = (uint8_t)((int)((int8_t*)srcPtr)[idx] + 2);
}
AutoStorage<uint8_t> packedInt8weight(weightlen * 4);
if (packedInt8weight.get() == nullptr) {
MNN_ERROR("Weight reorder memory not enough!\n");
return -1;
}
reorderWeight(packedInt8weight.get(), (uint8_t*)tmpWeight.data(), info, 0, (float*)kernelsum.get(),
sumFunc);
int permuteUnit = UNIT * SRC_UNIT; // 64 for ARMV86 (UNIT=8, SRC_UNIT=8); 32 for ARMV82 sdot (8x4)
int quarterPermuteStride = permuteUnit / 4; // 16 bytes for ARMV86, 8 bytes for sdot
int leng = weightlen * 4; // unsigned int8 stage size
auto srcw2Ptr = (uint8_t*)packedInt8weight.get();
auto dstw2Ptr = (uint8_t*)weightReordered.get();
MNN_ASSERT(UNIT == 8 && (SRC_UNIT == 8 || SRC_UNIT == 4));
if (SRC_UNIT == 8) {
// ARMV86 i8mm layout. 16 bytes per cell:
// bytes 0..7 : bits[6:7]=oc0, [4:5]=oc2, [2:3]=oc4, [0:1]=oc6 (at IC = b)
// bytes 8..15 : bits[6:7]=oc1, [4:5]=oc3, [2:3]=oc5, [0:1]=oc7 (at IC = b-8)
for (int i = 0; i < leng / permuteUnit; ++i) {
auto src0 = srcw2Ptr + i * permuteUnit;
auto dst0 = dstw2Ptr + i * quarterPermuteStride;
for (int b = 0; b < quarterPermuteStride; ++b) {
int ic_in_cell = b % 8;
int oc_offset = b / 8;
uint8_t out = 0;
out |= ((uint8_t)(src0[(0 + oc_offset) * 8 + ic_in_cell] & 0x3) << 6);
out |= ((uint8_t)(src0[(2 + oc_offset) * 8 + ic_in_cell] & 0x3) << 4);
out |= ((uint8_t)(src0[(4 + oc_offset) * 8 + ic_in_cell] & 0x3) << 2);
out |= ((uint8_t)(src0[(6 + oc_offset) * 8 + ic_in_cell] & 0x3) << 0);
dst0[b] = out;
}
}
} else { // SRC_UNIT == 4 (ARMV82 sdot)
// 8 bytes per cell:
// bytes 0..3 : bits[6:7]=oc0, [4:5]=oc1, [2:3]=oc2, [0:1]=oc3 (at IC = b)
// bytes 4..7 : bits[6:7]=oc4, [4:5]=oc5, [2:3]=oc6, [0:1]=oc7 (at IC = b-4)
for (int i = 0; i < leng / permuteUnit; ++i) {
auto src0 = srcw2Ptr + i * permuteUnit;
auto dst0 = dstw2Ptr + i * quarterPermuteStride;
for (int b = 0; b < quarterPermuteStride; ++b) {
int ic_in_cell = b % 4;
int oc_offset = (b / 4) * 4;
uint8_t out = 0;
out |= ((uint8_t)(src0[(oc_offset + 0) * 4 + ic_in_cell] & 0x3) << 6);
out |= ((uint8_t)(src0[(oc_offset + 1) * 4 + ic_in_cell] & 0x3) << 4);
out |= ((uint8_t)(src0[(oc_offset + 2) * 4 + ic_in_cell] & 0x3) << 2);
out |= ((uint8_t)(src0[(oc_offset + 3) * 4 + ic_in_cell] & 0x3) << 0);
dst0[b] = out;
}
}
}
} else if (target->mWeightBits == 3) {
// Loader gave us signed int8 in [-4, 3]; convert to unsigned [0, 7].
// Reorder via int8 path, then split into 2-bit main plane (16B) + 1-bit aux plane (8B)
// per (UNIT * SRC_UNIT) cell, contiguously laid out as [main | aux].
// The kernel produces unsigned [0, 7]; _computeReorderQuantInfo applies
// originOffset = -4 in the post-process bias term to re-center to signed [-4, 3].
originOffset = -4;
std::vector<uint8_t> tmpWeight(oc * ic * kernelCount);
for (int idx = 0; idx < oc * ic * kernelCount; ++idx) {
tmpWeight[idx] = (uint8_t)((int)((int8_t*)srcPtr)[idx] + 4);
}
// Stage int8 buffer holds reorderWeight output (1 byte per weight).
// Cell stride in dst depends on SRC_UNIT: i8mm packs 24B exactly per cell;
// sdot's 12B doesn't divide UNIT, so shape[4] is rounded up to 2 -> 16B cell
// with 4 trailing pad bytes (already zeroed by reorderWeight's memset).
int permuteUnit = UNIT * SRC_UNIT; // 64 i8mm / 32 sdot
int cellStride = UNIT * shape[4]; // 24 i8mm / 16 sdot
int cellCount = weightlen / cellStride;
int int8WeightLen = cellCount * permuteUnit;
AutoStorage<uint8_t> packedInt8weight(int8WeightLen);
if (packedInt8weight.get() == nullptr) {
MNN_ERROR("Weight reorder memory not enough!\n");
return -1;
}
reorderWeight(packedInt8weight.get(), (uint8_t*)tmpWeight.data(), info, 0, (float*)kernelsum.get(),
sumFunc);
// Per cell layout (useful bytes; remainder up to cellStride is zero-pad):
// i8mm (8x8 cell, 64 weights): 16B main + 8B aux = 24 useful, cellStride=24
// sdot (8x4 cell, 32 weights): 8B main + 4B aux = 12 useful, cellStride=16
int cellMain = (SRC_UNIT == 8) ? 16 : 8;
int cellAux = (SRC_UNIT == 8) ? 8 : 4;
auto srcStgPtr = (uint8_t*)packedInt8weight.get();
auto dstPtr = (uint8_t*)weightReordered.get();
MNN_ASSERT(UNIT == 8 && (SRC_UNIT == 8 || SRC_UNIT == 4));
if (SRC_UNIT == 8) {
for (int i = 0; i < cellCount; ++i) {
auto src0 = srcStgPtr + i * permuteUnit; // 64 unsigned int8 in [0, 7]
auto dst0 = dstPtr + i * cellStride; // main, aux, [pad]
// Main plane (2-bit), bytes 0..7 OC even, 8..15 OC odd, IC interleaved
for (int b = 0; b < cellMain; ++b) {
int ic_in_cell = b % 8;
int oc_offset = b / 8;
uint8_t out = 0;
out |= ((uint8_t)(src0[(0 + oc_offset) * 8 + ic_in_cell] & 0x3) << 6);
out |= ((uint8_t)(src0[(2 + oc_offset) * 8 + ic_in_cell] & 0x3) << 4);
out |= ((uint8_t)(src0[(4 + oc_offset) * 8 + ic_in_cell] & 0x3) << 2);
out |= ((uint8_t)(src0[(6 + oc_offset) * 8 + ic_in_cell] & 0x3) << 0);
dst0[b] = out;
}
// Aux plane (1-bit, IC-major): byte ic, bit (7 - oc) is the high bit.
// This lets the i8mm w3 kernel unpack all OC pairs with vector shifts instead of tbl/ext.
for (int icIdx = 0; icIdx < 8; ++icIdx) {
uint8_t out = 0;
for (int oc = 0; oc < 8; ++oc) {
uint8_t hi = (uint8_t)((src0[oc * 8 + icIdx] >> 2) & 1);
out |= (hi << (7 - oc));
}
dst0[cellMain + icIdx] = out;
}
}
} else { // SRC_UNIT == 4 (ARMV82 sdot)
// Main bytes 0..3: bits[6:7]=oc0, [4:5]=oc1, [2:3]=oc2, [0:1]=oc3 at IC=b
// Main bytes 4..7: bits[6:7]=oc4, ..., [0:1]=oc7 at IC=b-4
// Aux byte 0: bits[3:0]=oc0_aux at IC[0..3], bits[7:4]=oc1_aux
// Aux byte 1: bits[3:0]=oc2_aux, bits[7:4]=oc3_aux
// Aux byte 2..3: oc4..oc7 (same pattern)
// Bytes 12..15: zero pad (cellStride - useful = 4 bytes).
for (int i = 0; i < cellCount; ++i) {
auto src0 = srcStgPtr + i * permuteUnit; // 32 unsigned int8 in [0, 7]
auto dst0 = dstPtr + i * cellStride; // main, aux, [pad]
// Main plane
for (int b = 0; b < cellMain; ++b) {
int ic_in_cell = b % 4;
int oc_offset = (b / 4) * 4;
uint8_t out = 0;
out |= ((uint8_t)(src0[(oc_offset + 0) * 4 + ic_in_cell] & 0x3) << 6);
out |= ((uint8_t)(src0[(oc_offset + 1) * 4 + ic_in_cell] & 0x3) << 4);
out |= ((uint8_t)(src0[(oc_offset + 2) * 4 + ic_in_cell] & 0x3) << 2);
out |= ((uint8_t)(src0[(oc_offset + 3) * 4 + ic_in_cell] & 0x3) << 0);
dst0[b] = out;
}
// Aux plane: each output byte holds 2 OCs * 4 IC = 8 bits.
for (int b = 0; b < cellAux; ++b) {
int ocPair = b * 2; // OC pair at this aux byte: ocPair / ocPair+1
uint8_t out = 0;
for (int icIdx = 0; icIdx < 4; ++icIdx) {
uint8_t hi0 = (uint8_t)((src0[(ocPair + 0) * 4 + icIdx] >> 2) & 1);
uint8_t hi1 = (uint8_t)((src0[(ocPair + 1) * 4 + icIdx] >> 2) & 1);
out |= (hi0 << icIdx); // bits[3:0] = OC[ocPair]
out |= (hi1 << (icIdx + 4)); // bits[7:4] = OC[ocPair+1]
}
dst0[cellMain + b] = out;
}
}
}
} else { // int8 weight
reorderWeight((uint8_t*)weightReordered.get(), srcPtr, info, 0, (float*)kernelsum.get(), sumFunc);
}
}
if (convOp->symmetricQuan() && convOp->symmetricQuan()->bias()) {
// Compability for old model
::memcpy(target->mOriginBias->host<float>(), convOp->symmetricQuan()->bias()->data(), oc * sizeof(int32_t));
#ifdef MNN_USE_SSE
if (target->mUseConvQuan) {
for (int ks = 0; ks < oc; ++ks) {
target->mOriginBias->host<int32_t>()[ks] -= 128 * ((float*)kernelsum.get())[ks];
}
}
#endif
}
/* 2. compute and order dequant scale&bias */
bool notConvertInt4ToInt8 = true;
if (target->mWeightBits == 4 && !fastReadWeight) {
notConvertInt4ToInt8 = false;
}
int32_t paramsKernelSum[2] = {blockNum, inputBlockNum * ROUND_UP(oc, UNIT)};
float* weightKernelSum = (float*)addressPtr[2];
float* quanScalePtr = (float*)addressPtr[3];
_computeReorderQuantInfo(weightKernelSum, paramsKernelSum, (inputBlockQuantOption == 2), target->mWeightBits,
asyWeight, quanScalePtr, oc, kernelCount * ic, pack, reorderedQuantInfo,
(float*)kernelsum.get(), UNIT, notConvertInt4ToInt8);
/* 3. put weight and quantInfo together */
int32_t params[6] = {shape[0], shape[1], shape[2], shape[3], shape[4], ROUND_UP(oc, pack)};
int8_t* weightInt8 = addressPtr[1];
ConvInt8TiledExecutor::packWeightAndQuantInfo(weightInt8, (int8_t*)weightReordered.get(),
reorderedQuantInfo.get(), params, QUANT_INFO_BYTES);
return 0;
};
auto function = [=]() -> int {
bool fastReadWeight =
(kernelCount == 1 && ROUND_UP(ocMain, UNITMain) == ocMain && ROUND_UP(ic, SRC_UNITMain) == ic);
weightSummerFuncion sumFunc = funcsMain.MNNSumWeightInt8;
if (mOnlineReorderWeightSme) {
sumFunc = funcsMain.MNNSumWeightInt8SmeHp128;
}
int8_t* addressPtr[4];
addressPtr[0] = quanCommon ? quanCommon->weight.get() : (int8_t*)convOp->symmetricQuan()->weight()->data();
addressPtr[1] = target->mWeightInt8->host<int8_t>();
addressPtr[2] = target->mWeightKernelSum->host<int8_t>();
addressPtr[3] =
quanCommon ? (int8_t*)quanCommon->alpha.get() : (int8_t*)convOp->symmetricQuan()->scale()->data();
reorderFunc(funcsMain, shapeMain, UNITMain, SRC_UNITMain, DST_XUNITMain, weightlenMain, scaleSizeMain, ocMain,
0, fastReadWeight, addressPtr, sumFunc);
if (ocBranch > 0) {
// update the address of weight source, weight destination, weight kernel sum and weight scale
addressPtr[0] +=
(target->mWeightBits == 4 ? ocMain * ic * kernelCount / 2
: ocMain * ic * kernelCount); // ocMain%2==0, so divides 2 directly
addressPtr[1] += (weightlenMain + quantlenMain);
addressPtr[2] += ROUND_UP(ocMain, UNITMain) * inputBlockNum * QUANT_INFO_BYTES;
addressPtr[3] += (quanCommon->asymmetric ? 2 * ocMain * blockNum * QUANT_INFO_BYTES
: ocMain * blockNum * QUANT_INFO_BYTES);
sumFunc = funcsBranch.MNNSumWeightInt8;
fastReadWeight =
(kernelCount == 1 && ROUND_UP(ocBranch, UNITMain) == ocBranch && ROUND_UP(ic, SRC_UNITMain) == ic);
reorderFunc(funcsBranch, shapeBranch, UNITBranch, SRC_UNITBranch, DST_XUNITBranch, weightlenBranch,
scaleSizeBranch, ocBranch, 1, fastReadWeight, addressPtr, sumFunc);
}
return 0;
};
static_cast<CPUBackend*>(backend)->enqueueTask(std::move(function));
if (!isDynamicQuant) {
mResourceInt8->mDynamicQuant = false;
std::shared_ptr<float> scaleAndBias(new float[ocUpHpMain * 2 * mBlockNum],
[](void* ptr) { delete[] (float*)ptr; });
memset(scaleAndBias.get(), 0, ocUpHpMain * 2 * mBlockNum * sizeof(float));
int weightSize;
bool weightAsy = false;
if (quanCommon && quanCommon->asymmetric) {
weightAsy = true;
}
if (convOp->symmetricQuan() && convOp->symmetricQuan()->bias() && convOp->symmetricQuan()->scale()) {
// Compability for old model
MNN_ASSERT(convOp->symmetricQuan()->bias()->size() == oc && convOp->symmetricQuan()->scale()->size() == oc);
::memcpy(scaleAndBias.get(), convOp->symmetricQuan()->scale()->data(), oc * sizeof(float));
}
if ((convOp->quanParameter() && convOp->quanParameter()->alpha()) || (quanCommon && quanCommon->alpha.get())) {
int quantCount;
if (convOp->quanParameter() && convOp->quanParameter()->alpha()) {
quantCount = convOp->quanParameter()->alpha()->size();
} else {
quantCount = quanCommon->alpha.size();
}
if (false == weightAsy) { // symmetric quant
if (convOp->quanParameter() && convOp->quanParameter()->alpha()) {
::memcpy(scaleAndBias.get(), convOp->quanParameter()->alpha()->data(), quantCount * sizeof(float));
} else {
::memcpy(scaleAndBias.get(), quanCommon->alpha.get(), quanCommon->alpha.size() * sizeof(float));
}
} else if (true == weightAsy) { // asymmetric
int scaleSize = quantCount / 2;
for (int i = 0; i < scaleSize; ++i) {
((float*)scaleAndBias.get())[i] = quanCommon->alpha.get()[2 * i + 1];
((float*)scaleAndBias.get())[i + ocUpHpMain] = quanCommon->alpha.get()[2 * i];
}
}
}
initializeConvInt8QuantInfo(mResourceInt8, convOp, quanCommon);
mMutableResource.reset(new MutableResourceInt8(mResourceInt8, backend, scaleAndBias.get()));
// gemmInt8 kernel
mGemmKernel = mRelatedFunctions.Int8GemmKernel;
#ifdef MNN_USE_SSE
if (convOp->symmetricQuan()) {
int actBits = convOp->symmetricQuan()->nbits();
if (actBits <= 7) {
mGemmKernel = mRelatedFunctions.Int8GemmKernelFast;
}
}
#else
if (convOp->symmetricQuan() && convOp->symmetricQuan()->method() == QuantizeAlgo_OVERFLOW_AWARE) {
mGemmKernel = mRelatedFunctions.Int8GemmKernelFast;
}
if (mResourceInt8->mWeightBits == 4) {
mGemmKernel = mRelatedFunctions.Int8GemmKernel_W4;
}
#endif
}
}
DenseConvInt8TiledExecutor::DenseConvInt8TiledExecutor(Backend* backend, const Op* op,
const DenseConvInt8TiledExecutor& exe)
: ConvInt8TiledExecutor(backend, op, exe.mResourceInt8), mGemmKernel(exe.mGemmKernel) {}
DenseConvInt8TiledExecutor::~DenseConvInt8TiledExecutor() {
// Do nothing
}
bool DenseConvInt8TiledExecutor::onClone(Backend* bn, const Op* op, Execution** dst) {
if (nullptr == dst) {
return true;
}
if (op->type() == OpType_GatherV2) {
*dst = new SharedGather(bn, mResourceInt8);
return true;
}
auto exe = new DenseConvInt8TiledExecutor(bn, op, *this);
if (!exe->valid()) {
return false;
}
*dst = exe;
return true;
}
ErrorCode DenseConvInt8TiledExecutor::onResize(const std::vector<Tensor*>& inputs,
const std::vector<Tensor*>& outputs) {
// Initialize.
mUseBatchQuan = false;
mIm2ColBasedInt8 = true;
m4BitPtq = false;
if (mResourceInt8->mDynamicQuant == false && mResourceInt8->mWeightBits == 4) {
m4BitPtq = true;
}
// backend info
auto core = static_cast<CPUBackend*>(backend())->int8Functions();
auto gcore = static_cast<CPUBackend*>(backend())->functions();
const int threads = static_cast<CPUBackend*>(backend())->threadNumber();
mRelatedFunctions = *(static_cast<CPUBackend*>(backend())->int8GemmFunctions());
mArm82Functions = gcore->arm82MatmulRelatedFunctions;
// runtime hint
auto option = static_cast<CPUBackend*>(backend())->getRuntime()->hint().dynamicQuantOption;
mSmeCores = gcore->smeCoreNumber;
auto inputBlockQuantOption = option % WEIGHT_ONLINE_REORDER;
auto weightOnlineReorderOption = WEIGHT_ONLINE_REORDER & option;
_getProportions(static_cast<CPUBackend*>(backend())->getRuntime()->hint().divisionRatio, mRatioPrefill,
mRatioDecode);
// feature map info
int batch = inputs[0]->batch();
int inC = inputs[0]->channel();
auto output = outputs[0];
int kernelCount = mCommon->kernelY() * mCommon->kernelX();
int inputPlane = batch * inputs[0]->width() * inputs[0]->height();
auto planeSize = output->width() * output->height() * output->batch();
int UNIT, SRC_UNIT, DST_XUNIT;
mRelatedFunctions.MNNGetGemmUnit(&UNIT, &SRC_UNIT, &DST_XUNIT);
mOnlineReorderWeightSme = (weightOnlineReorderOption > 0 && DST_XUNIT == GEMM_INT8_DST_XUNIT_SME2);
if (mResourceInt8->mDynamicQuant == false) {
mOnlineReorderWeightSme = false;
}
_updateMixedKernelFlag(mMixedKernel, mOnlineReorderWeightSme, threads, DST_XUNIT, mResourceInt8->mDynamicQuant,
mRatioDecode && mRatioPrefill);
if (mOnlineReorderWeightSme && planeSize == 1) { // Decode, set runtime unit
UNIT = GEMM_INT8_UNIT_SME2_128;
}
mGemmUnits[0] = UNIT;
mGemmUnits[1] = SRC_UNIT;
mGemmUnits[2] = DST_XUNIT;
bool fastway = (kernelCount == 1) && (output->width() == inputs[0]->width()) &&
(output->height() == inputs[0]->height()) && (mCommon->strideX() * mCommon->strideY()) == 1;
if (inputPlane > 1) {
mUseBatchQuan = true;
}
if (!fastway) { // general conv
mIm2ColBasedInt8 = false;
if (planeSize > 1) {
mUseBatchQuan = true;
}
if (inputBlockQuantOption == 1) { // lowest level.
mIm2ColBasedInt8 = true;
mUseBatchQuan = false;
}
}
float weightBytes = 1.0f;
if (mResourceInt8->mWeightBits == 4) {
weightBytes = 0.5f;
} else if (mResourceInt8->mWeightBits == 3) {
auto packedBytesPerOc = (SRC_UNIT * 3 + 7) / 8;
weightBytes = static_cast<float>(packedBytesPerOc) / SRC_UNIT;
} else if (mResourceInt8->mWeightBits == 2) {
weightBytes = 0.25f;
}
mBlockNum = mResourceInt8->mBlockNum;
CPUConvolution::onResize(inputs, outputs);
if (mResourceInt8->mDynamicQuant == false) {
mMutableResource->updateInputOutputScale(TensorUtils::getQuantInfo(inputs[0]),
TensorUtils::getQuantInfo(outputs[0]));
if (!mMutableResource->mResource->mUseConvQuan) {
// In some previous quantized models, input's scale already fused with weight's scale and output's scale.
// So there is no need to read input's scale additionally.
mBatchQuantInfo.reset(Tensor::createDevice<int8_t>({1, DST_XUNIT * QUANT_INFO_BYTES}));
auto success = backend()->onAcquireBuffer(mBatchQuantInfo.get(), Backend::DYNAMIC);
if (!success) {
return OUT_OF_MEMORY;
}
}
mIm2ColBasedInt8 = true;
mUseBatchQuan = false;
}
int matmulUnits[3] = {UNIT, SRC_UNIT, DST_XUNIT};
ConvolutionTiledExecutor::setIm2ColParameter(mIm2ColParamter, mCommon, inputs[0], outputs[0], mPadX, mPadY, gcore,
core, gcore->pack, matmulUnits);
// Im2col info
int im2colBytes = 1;
const int L2Size = 2048;
int tileLimitByC = UP_DIV(L2Size, mIm2ColParamter.kernelCountUnit * SRC_UNIT);
if (mIm2ColBasedInt8 == false) {
im2colBytes = gcore->bytes;
tileLimitByC = 1;
}
int ic = inputs[0]->channel();
int tileLimit = 0;
int outC = output->channel();
int outC4 = UP_DIV(outC, gcore->pack);
mOcMain = outC;
mOcBranch = 0;
const int pack = gcore->pack;
auto kernelCountUnit = mIm2ColParamter.kernelCountUnit;
mSplitByOc = true;
// flop and io
float flop = gcore->bytes * planeSize *
(ROUND_UP(output->channel(), gcore->pack) * kernelCountUnit * SRC_UNIT / 1024.0 / 1024.0 / 1024.0);
float ios = (((CPUBackend*)backend())->getTensorSize(outputs[0], true) +
((CPUBackend*)backend())->getTensorSize(inputs[0], true) +
((CPUBackend*)backend())->getTensorSize(mResourceInt8->mWeightInt8.get()) * weightBytes) /
(1024.0 * 1024.0 * 1024.0);
if ((threads < planeSize || mOnlineReorderWeightSme) && !mMixedKernel) { // Thread split by output nhw.
tileLimit = ALIMIN(tileLimitByC, UP_DIV(planeSize, threads));
mIm2ColCount = UP_DIV(tileLimit, DST_XUNIT);
auto DynamicDestUnit = DST_XUNIT * mIm2ColCount;
mTileCount = UP_DIV(planeSize, DynamicDestUnit);
if (mTileCount > threads || (mOnlineReorderWeightSme && planeSize > 1)) {
mSplitByOc = false;
}
}
if (mSplitByOc) {
tileLimit = ALIMIN(tileLimitByC, planeSize);
mIm2ColCount = UP_DIV(tileLimit, DST_XUNIT);
auto DynamicDestUnit = DST_XUNIT * mIm2ColCount;
mTileCount = UP_DIV(planeSize, DynamicDestUnit);
mDivides.resize(threads + 1);
mDivides[0] = 0;
// output channel divided by threads
if (!mMixedKernel) {
auto ocPerThread = UP_DIV(outC4, threads);
auto threadNeed = UP_DIV(outC4, ocPerThread);
int totalWork = outC4;
int part = 1;
if (UNIT > gcore->pack) { // AVX512:UNIT=64,pack=16
MNN_ASSERT(UNIT % gcore->pack == 0);
int ocDivUnit = UP_DIV(outC4 * gcore->pack, UNIT);
ocPerThread = UP_DIV(ocDivUnit, threads);
threadNeed = UP_DIV(ocDivUnit, ocPerThread);
totalWork = ocDivUnit;
part = UNIT / gcore->pack;
}
mThreadNums = ALIMIN(threads, threadNeed);
if (threads >= 4 && DST_XUNIT == GEMM_INT8_DST_XUNIT_SME2 && mResourceInt8->mDynamicQuant) {
_computeDivides4Sme(mDivides, threads, mSmeCores, totalWork);
} else {
mDivides.resize(threads + 1);
mDivides[0] = 0;
static_cast<CPUBackend*>(backend())->computeDivideSizes(totalWork, mDivides.data() + 1, flop / ios);
}
for (int i = 0; i < mDivides.size(); ++i) {
mDivides[i] *= part;
}
} else {
// workload
mOcMain = 0; // initialize for mixed kernel, before calculate
calculateSmeNeonWorkDivision(mOcMain, mOcBranch, mDivides, outC, threads, pack, planeSize, mRatioDecode,
mSmeCores);
mThreadNums = threads;
}
}
if (!mSplitByOc) {
mThreadNums = ALIMIN(threads, mTileCount);
if (threads >= 4 && DST_XUNIT == GEMM_INT8_DST_XUNIT_SME2 && mResourceInt8->mDynamicQuant && !mMixedKernel) {
_computeDivides4Sme(mDivides, threads, mSmeCores, mTileCount);
} else {
mDivides.resize(threads + 1);
mDivides[0] = 0;
static_cast<CPUBackend*>(backend())->computeDivideSizes(mTileCount, mDivides.data() + 1, flop / ios);
}
}
mDividesTmp.resize(threads + 1);
if (mMixedKernel) {
mOriginSmeWork = mDivides[mSmeCores];
}
int ocUp4 = ROUND_UP(outC, gcore->pack);
int k = mThreadNums;
int workPT = DST_XUNIT * mIm2ColCount;
if (mSplitByOc) {
k = 1; // Use one thread to finish im2col.
workPT = mTileCount * DST_XUNIT * mIm2ColCount;
}
auto bufferAlloc = static_cast<CPUBackend*>(backend())->getBufferAllocator();
auto blitInfoSize = ConvolutionTiledExecutor::computeBlitInfoSize(
workPT, mIm2ColParamter.ow, mIm2ColParamter.kernelX * mIm2ColParamter.kernelY, k);
mBlitInfoStride = blitInfoSize.second;
mBlitInfo = bufferAlloc->alloc(blitInfoSize.first);
const int unitColBufferSize = kernelCountUnit * DST_XUNIT * SRC_UNIT * sizeof(int8_t);
const int colBufferSize = unitColBufferSize * mIm2ColCount;
if (!mSplitByOc) {
mTempIm2ColBuffer.reset(Tensor::createDevice<int8_t>({threads, colBufferSize * im2colBytes}));
mTempSrcSum = bufferAlloc->alloc(threads * mBlockNum * DST_XUNIT * mIm2ColCount * QUANT_INFO_BYTES);
} else {
mTempIm2ColBuffer.reset(Tensor::createDevice<int8_t>({mTileCount, colBufferSize * im2colBytes}));
mTempSrcSum = bufferAlloc->alloc(mTileCount * mBlockNum * DST_XUNIT * mIm2ColCount * QUANT_INFO_BYTES);
}
mAccumBuffer.reset(Tensor::createDevice<int32_t>({threads, DST_XUNIT * ALIMAX(UNIT, gcore->pack)}));
auto success = backend()->onAcquireBuffer(mTempIm2ColBuffer.get(), Backend::DYNAMIC);
success &= backend()->onAcquireBuffer(mAccumBuffer.get(), Backend::DYNAMIC);
if (!success || mBlitInfo.invalid() || mTempSrcSum.invalid()) {
return OUT_OF_MEMORY;
}
if (false == mResourceInt8->mDynamicQuant && false == m4BitPtq) {
bufferAlloc->free(mBlitInfo);
bufferAlloc->free(mTempSrcSum);
backend()->onReleaseBuffer(mTempIm2ColBuffer.get(), Backend::DYNAMIC);
if (mBatchQuantInfo.get()) {
backend()->onReleaseBuffer(mBatchQuantInfo.get(), Backend::DYNAMIC);
}
backend()->onReleaseBuffer(mAccumBuffer.get(), Backend::DYNAMIC);
return NO_ERROR;
}
#ifdef MNN_LOW_MEMORY
if (!mMixedKernel) { // Dynamic Quant kernels, use single gemm kernel.
mGemmKernel = mRelatedFunctions.Int8GemmKernel;
if (mOnlineReorderWeightSme && planeSize == 1) {
mGemmKernel = mRelatedFunctions.MNNGemmInt8AddBiasScale_Unit_FP32_DecodeMax;
}
if (mResourceInt8->mWeightBits == 4) {
mGemmKernel = mRelatedFunctions.Int8GemmKernel_W4;
if (mOnlineReorderWeightSme && planeSize == 1) {
mGemmKernel = mRelatedFunctions.MNNGemmInt8AddBiasScale_w4_Unit_FP32_DecodeMax;
}
} else if (mResourceInt8->mWeightBits == 2 && mRelatedFunctions.Int8GemmKernel_W2 != nullptr) {
mGemmKernel = mRelatedFunctions.Int8GemmKernel_W2;
} else if (mResourceInt8->mWeightBits == 3 && mRelatedFunctions.Int8GemmKernel_W3 != nullptr) {
mGemmKernel = mRelatedFunctions.Int8GemmKernel_W3;
}
mQuantFunc = core->MNNFloat2Int8;
if (gcore->bytes == 2 && gcore->pack == 8) {
mGemmKernel = mRelatedFunctions.MNNGemmInt8AddBiasScale_Unit_FP16;
if (mOnlineReorderWeightSme && planeSize == 1) {
mGemmKernel = mRelatedFunctions.MNNGemmInt8AddBiasScale_Unit_FP16_DecodeMax;
}
if (mResourceInt8->mWeightBits == 4) {
mGemmKernel = mRelatedFunctions.MNNGemmInt8AddBiasScale_w4_Unit_FP16;
if (mOnlineReorderWeightSme && planeSize == 1) {
mGemmKernel = mRelatedFunctions.MNNGemmInt8AddBiasScale_w4_Unit_FP16_DecodeMax;
}
} else if (mResourceInt8->mWeightBits == 2 &&
mRelatedFunctions.MNNGemmInt8AddBiasScale_w2_Unit_FP16 != nullptr) {
mGemmKernel = mRelatedFunctions.MNNGemmInt8AddBiasScale_w2_Unit_FP16;
} else if (mResourceInt8->mWeightBits == 3 &&
mRelatedFunctions.MNNGemmInt8AddBiasScale_w3_Unit_FP16 != nullptr) {
mGemmKernel = mRelatedFunctions.MNNGemmInt8AddBiasScale_w3_Unit_FP16;
}
mQuantFunc = core->DynamicQuanInput_ARM82;
mQuantAndReorderFunc = core->DynamicQuanInputAndReorder_ARM82;
}
// A axisSum kernel
} else { // use sme and neon gemmInt8
// Fp32
if (planeSize == 1) { // Decode
mGemmKernels.push_back(mRelatedFunctions.MNNGemmInt8AddBiasScale_Unit_FP32_DecodeMax);
mGemmKernels.push_back(mArm82Functions.Int8GemmKernel);
if (mResourceInt8->mWeightBits == 4) {
mGemmKernels[0] = mRelatedFunctions.MNNGemmInt8AddBiasScale_w4_Unit_FP32_DecodeMax;
mGemmKernels[1] = mArm82Functions.Int8GemmKernel_W4;
} else if (mResourceInt8->mWeightBits == 2 && mRelatedFunctions.Int8GemmKernel_W2 != nullptr) {
// No SME2 DecodeMax for w2 yet; reuse plain w2 for both branches.
mGemmKernels[0] = mRelatedFunctions.Int8GemmKernel_W2;
mGemmKernels[1] = mArm82Functions.Int8GemmKernel_W2 != nullptr ? mArm82Functions.Int8GemmKernel_W2
: mRelatedFunctions.Int8GemmKernel_W2;
} else if (mResourceInt8->mWeightBits == 3 && mRelatedFunctions.Int8GemmKernel_W3 != nullptr) {
mGemmKernels[0] = mRelatedFunctions.Int8GemmKernel_W3;
mGemmKernels[1] = mArm82Functions.Int8GemmKernel_W3 != nullptr ? mArm82Functions.Int8GemmKernel_W3
: mRelatedFunctions.Int8GemmKernel_W3;
}
} else { // Prefill
mGemmKernels.push_back(mRelatedFunctions.Int8GemmKernel);
mGemmKernels.push_back(mArm82Functions.Int8GemmKernel);
if (mResourceInt8->mWeightBits == 4) {
mGemmKernels[0] = mRelatedFunctions.Int8GemmKernel_W4;
mGemmKernels[1] = mArm82Functions.Int8GemmKernel_W4;
} else if (mResourceInt8->mWeightBits == 2 && mRelatedFunctions.Int8GemmKernel_W2 != nullptr) {
mGemmKernels[0] = mRelatedFunctions.Int8GemmKernel_W2;
mGemmKernels[1] = mArm82Functions.Int8GemmKernel_W2 != nullptr ? mArm82Functions.Int8GemmKernel_W2
: mRelatedFunctions.Int8GemmKernel_W2;
} else if (mResourceInt8->mWeightBits == 3 && mRelatedFunctions.Int8GemmKernel_W3 != nullptr) {
mGemmKernels[0] = mRelatedFunctions.Int8GemmKernel_W3;
mGemmKernels[1] = mArm82Functions.Int8GemmKernel_W3 != nullptr ? mArm82Functions.Int8GemmKernel_W3
: mRelatedFunctions.Int8GemmKernel_W3;
}
}
mQuantFunc = core->MNNFloat2Int8;
// fp16
if (gcore->bytes == 2 && gcore->pack == 8) {
if (planeSize == 1) { // Decode
mGemmKernels[0] = mRelatedFunctions.MNNGemmInt8AddBiasScale_Unit_FP16_DecodeMax;
mGemmKernels[1] = mArm82Functions.MNNGemmInt8AddBiasScale_Unit_FP16;
if (mResourceInt8->mWeightBits == 4) {
mGemmKernels[0] = mRelatedFunctions.MNNGemmInt8AddBiasScale_w4_Unit_FP16_DecodeMax;
mGemmKernels[1] = mArm82Functions.MNNGemmInt8AddBiasScale_w4_Unit_FP16;
} else if (mResourceInt8->mWeightBits == 2 &&
mArm82Functions.MNNGemmInt8AddBiasScale_w2_Unit_FP16 != nullptr) {
// No SME2 DecodeMax for w2 yet; reuse plain ARMV86 for both branches.
mGemmKernels[0] = mArm82Functions.MNNGemmInt8AddBiasScale_w2_Unit_FP16;
mGemmKernels[1] = mArm82Functions.MNNGemmInt8AddBiasScale_w2_Unit_FP16;
} else if (mResourceInt8->mWeightBits == 3 &&
mArm82Functions.MNNGemmInt8AddBiasScale_w3_Unit_FP16 != nullptr) {
mGemmKernels[0] = mArm82Functions.MNNGemmInt8AddBiasScale_w3_Unit_FP16;
mGemmKernels[1] = mArm82Functions.MNNGemmInt8AddBiasScale_w3_Unit_FP16;
}
} else { // Prefill
mGemmKernels[0] = mRelatedFunctions.MNNGemmInt8AddBiasScale_Unit_FP16;
mGemmKernels[1] = mArm82Functions.MNNGemmInt8AddBiasScale_Unit_FP16;
if (mResourceInt8->mWeightBits == 4) {
mGemmKernels[0] = mRelatedFunctions.MNNGemmInt8AddBiasScale_w4_Unit_FP16;
mGemmKernels[1] = mArm82Functions.MNNGemmInt8AddBiasScale_w4_Unit_FP16;
} else if (mResourceInt8->mWeightBits == 2 &&
mArm82Functions.MNNGemmInt8AddBiasScale_w2_Unit_FP16 != nullptr) {
mGemmKernels[0] = mRelatedFunctions.MNNGemmInt8AddBiasScale_w2_Unit_FP16;
mGemmKernels[1] = mArm82Functions.MNNGemmInt8AddBiasScale_w2_Unit_FP16;
} else if (mResourceInt8->mWeightBits == 3 &&
mArm82Functions.MNNGemmInt8AddBiasScale_w3_Unit_FP16 != nullptr) {
mGemmKernels[0] = mRelatedFunctions.MNNGemmInt8AddBiasScale_w3_Unit_FP16;
mGemmKernels[1] = mArm82Functions.MNNGemmInt8AddBiasScale_w3_Unit_FP16;
}
}
mQuantFunc = core->DynamicQuanInput_ARM82;
mQuantAndReorderFunc = core->DynamicQuanInputAndReorder_ARM82;
}
// A axisSum kernel
}
mInputBlockNum = (inputBlockQuantOption == 2) ? mBlockNum : 1;
bool symmetricQuant = (inputBlockQuantOption != 2 && mUseBatchQuan) ? true : false;
int size = 0;
if (!mUseBatchQuan) { // single quant
if (mSplitByOc) {
size = 2 * mInputBlockNum * ALIMIN(DST_XUNIT, planeSize) * QUANT_INFO_BYTES;
} else {
size = 2 * mInputBlockNum * mIm2ColCount * DST_XUNIT * QUANT_INFO_BYTES;
}
}
if (mUseBatchQuan) {
if (mIm2ColBasedInt8) {
size = 2 * mInputBlockNum * inputPlane * QUANT_INFO_BYTES;
} else if (!mSplitByOc) { // only threads buffer needed by this case
size = 2 * mInputBlockNum * mIm2ColCount * DST_XUNIT * QUANT_INFO_BYTES;
} else {
size = 2 * mInputBlockNum * planeSize * QUANT_INFO_BYTES;
}
}
if (symmetricQuant) { // symmetric quant
size /= 2;
}
if (false == m4BitPtq) {
if (!mIm2ColBasedInt8 && !mSplitByOc) {
mBatchQuantInfo.reset(Tensor::createDevice<int8_t>({threads, size}));
} else {
mBatchQuantInfo.reset(Tensor::createDevice<int8_t>({1, size})); // keep dimensions=2!
}
success &= backend()->onAcquireBuffer(mBatchQuantInfo.get(), Backend::DYNAMIC);
}
// Dynamic quant.
// set im2col tensor info
if (mIm2ColBasedInt8) {
mQuantInput.reset((
Tensor::createDevice<int8_t>({batch, mIm2ColParamter.ih, mIm2ColParamter.iw, ROUND_UP(inC, gcore->pack)})));
} else if (!mSplitByOc) {
mQuantInput.reset((Tensor::createDevice<int8_t>({threads, colBufferSize * 1})));
} else {
mQuantInput.reset((Tensor::createDevice<int8_t>({mTileCount, colBufferSize * 1})));
}
success &= backend()->onAcquireBuffer(mQuantInput.get(), Backend::DYNAMIC);
// set compute buffer
int tempSize = threads * 2 * mInputBlockNum * inputPlane;
if (!mIm2ColBasedInt8) {
if (!mSplitByOc) {
tempSize = threads * 2 * mInputBlockNum * DST_XUNIT * mIm2ColCount;
} else {
tempSize = threads * 2 * mInputBlockNum * ROUND_UP(planeSize, DST_XUNIT);
}
}
if (symmetricQuant) { // symmetric batch quant.
tempSize /= 2;
}
mSizeInputBlockQuant = tempSize / threads;
mTempMaxMinValueBuffer = bufferAlloc->alloc(tempSize * gcore->bytes);
mQScaleZero = bufferAlloc->alloc(tempSize * QUANT_INFO_BYTES);
if (mQScaleZero.invalid()) {
return OUT_OF_MEMORY;
}
if (mOnlineReorderWeightSme && planeSize > 1) { // only prefill need
int ocProcessedBySme = mOcMain;
int ocProcessedByNeon = 0;
if (mMixedKernel && mRatioDecode != mRatioPrefill) {
auto workUnit = UP_DIV(outC4, mRatioPrefill * mSmeCores + 1 * (threads - mSmeCores));
ocProcessedBySme =
ALIMIN(ROUND_UP(workUnit * pack * mSmeCores * mRatioPrefill, GEMM_INT8_UNIT_SME2_128), outC);
ocProcessedBySme = ALIMAX(ocProcessedBySme, mOcMain);
ocProcessedByNeon = outC - ocProcessedBySme;
}
int weightlenSme = ROUND_UP(ocProcessedBySme, GEMM_INT8_UNIT_SME2_128) * mBlockNum *
ROUND_UP(ic / mBlockNum, SRC_UNIT) * kernelCount;
int weightlenNeon =
ROUND_UP(ocProcessedByNeon, 8) * mBlockNum * ROUND_UP(ic / mBlockNum, SRC_UNIT) * kernelCount;
if (mResourceInt8->mWeightBits == 4) {
weightlenSme /= 2;
weightlenNeon /= 2;
}
int scalebiasLenSme = 2 * mBlockNum * ROUND_UP(ocProcessedBySme, GEMM_INT8_UNIT_SME2_128) * QUANT_INFO_BYTES;
int scalebiasLenNeon = 2 * mBlockNum * ROUND_UP(ocProcessedByNeon, 8) * QUANT_INFO_BYTES;
mWeight4Prefill = bufferAlloc->alloc(weightlenSme + scalebiasLenSme + weightlenNeon + scalebiasLenNeon);
if (mWeight4Prefill.invalid()) {
return OUT_OF_MEMORY;
}
if (mInputBlockNum > 1) { // only in this case, need to use weight_kernel_sum
mWeightKernelSum4Prefill =
bufferAlloc->alloc(ROUND_UP(outC, GEMM_INT8_UNIT_SME2_128) * mBlockNum * sizeof(float));
if (mWeightKernelSum4Prefill.invalid()) {
return OUT_OF_MEMORY;
}
}
}
mToFuseInputbias2Bias = (!mUseBatchQuan && inputBlockQuantOption != 2) ? true : false;
if (mToFuseInputbias2Bias) { // input data has only one bias&scale
if (mIm2ColBasedInt8) {
mBiasBufferFusedInputzero = bufferAlloc->alloc(
ROUND_UP(outC, UNIT) * QUANT_INFO_BYTES); // should be UP_DIV(oc, UNIT),not UP_DIV(oc, pack)
} else {
mBiasBufferFusedInputzero = bufferAlloc->alloc(threads * ROUND_UP(outC, UNIT) * QUANT_INFO_BYTES);
}
if (mBiasBufferFusedInputzero.invalid()) {
return OUT_OF_MEMORY;
}
}
if (mBlockNum > 1 && kernelCount > 1) {
if (mSplitByOc) {
mReorderBuffer = bufferAlloc->alloc(UP_DIV(planeSize, DST_XUNIT) * unitColBufferSize);
} else {
mReorderBuffer = bufferAlloc->alloc(threads * colBufferSize);
}
if (mReorderBuffer.invalid()) {
return OUT_OF_MEMORY;
}
}
if (!success || mTempMaxMinValueBuffer.invalid()) {
return OUT_OF_MEMORY;
}
bufferAlloc->free(mBlitInfo);
bufferAlloc->free(mTempSrcSum);
bufferAlloc->free(mTempMaxMinValueBuffer);
bufferAlloc->free(mQScaleZero);
if (mOnlineReorderWeightSme && planeSize > 1) {
bufferAlloc->free(mWeight4Prefill);
if (mInputBlockNum > 1) {
bufferAlloc->free(mWeightKernelSum4Prefill);
}
}
if (mBlockNum > 1 && kernelCount > 1) {
bufferAlloc->free(mReorderBuffer);
}
if (mToFuseInputbias2Bias) {
bufferAlloc->free(mBiasBufferFusedInputzero);
}
// Additional Adjustments
if (m4BitPtq) {
mTempOutput = bufferAlloc->alloc(ocUp4 * planeSize * gcore->bytes);
if (mTempOutput.invalid()) {
return OUT_OF_MEMORY;
}
bufferAlloc->free(mTempOutput);
}
backend()->onReleaseBuffer(mTempIm2ColBuffer.get(), Backend::DYNAMIC);
backend()->onReleaseBuffer(mBatchQuantInfo.get(), Backend::DYNAMIC);
backend()->onReleaseBuffer(mQuantInput.get(), Backend::DYNAMIC);
backend()->onReleaseBuffer(mAccumBuffer.get(), Backend::DYNAMIC);
return NO_ERROR;
#else
return NO_ERROR;
#endif
}
static void _onlineReorderWeightPackH128ToH32(int8_t* dst, int8_t* src, int hPSrc, int hPDst, int hU, int blockNum,
int blockLu, int lp, bool int4weight) {
// hPSrc = 4 * hPDst
int unitsize_ = hPDst * lp;
if (int4weight) {
lp /= 2;
unitsize_ /= 2;
}
int unitsize4 = unitsize_ * 4;
// Calculate strides based on source and destination h-pack sizes
int srcStride1 = blockLu * hPSrc * lp + 2 * hPSrc * sizeof(float);
int srcStride0 = blockNum * srcStride1;
int dstStride1 = blockLu * hPDst * lp + 2 * hPDst * sizeof(float);
int dstStride0 = blockNum * dstStride1;
for (int i = 0; i < hU; ++i) {
for (int k = 0; k < blockNum; ++k) {
auto weightsrc = (int8_t*)(src + i * srcStride0 + k * srcStride1);
auto weightdst0 = (int8_t*)(dst + (4 * i) * dstStride0 + k * dstStride1);
auto weightdst1 = (int8_t*)(dst + (4 * i + 1) * dstStride0 + k * dstStride1);
auto weightdst2 = (int8_t*)(dst + (4 * i + 2) * dstStride0 + k * dstStride1);
auto weightdst3 = (int8_t*)(dst + (4 * i + 3) * dstStride0 + k * dstStride1);
auto lu = blockLu;
while (lu > 7) {
for (int j = 0; j < 8; ++j) {
memcpy(weightdst0 + j * unitsize_, weightsrc + j * unitsize4 + 0 * unitsize_, unitsize_);
memcpy(weightdst1 + j * unitsize_, weightsrc + j * unitsize4 + 1 * unitsize_, unitsize_);
memcpy(weightdst2 + j * unitsize_, weightsrc + j * unitsize4 + 2 * unitsize_, unitsize_);
memcpy(weightdst3 + j * unitsize_, weightsrc + j * unitsize4 + 3 * unitsize_, unitsize_);
}
weightsrc += unitsize4 * 8;
weightdst0 += unitsize_ * 8;
weightdst1 += unitsize_ * 8;
weightdst2 += unitsize_ * 8;
weightdst3 += unitsize_ * 8;
lu -= 8;
}
if (lu > 3) {
for (int j = 0; j < 4; ++j) {
memcpy(weightdst0 + j * unitsize_, weightsrc + j * unitsize4 + 0 * unitsize_, unitsize_);
memcpy(weightdst1 + j * unitsize_, weightsrc + j * unitsize4 + 1 * unitsize_, unitsize_);
memcpy(weightdst2 + j * unitsize_, weightsrc + j * unitsize4 + 2 * unitsize_, unitsize_);
memcpy(weightdst3 + j * unitsize_, weightsrc + j * unitsize4 + 3 * unitsize_, unitsize_);
}
weightsrc += unitsize4 * 4;
weightdst0 += unitsize_ * 4;
weightdst1 += unitsize_ * 4;
weightdst2 += unitsize_ * 4;
weightdst3 += unitsize_ * 4;
lu -= 4;
}
if (lu > 1) {
memcpy(weightdst0, weightsrc, unitsize_);
memcpy(weightdst0 + unitsize_, weightsrc + unitsize4, unitsize_);
memcpy(weightdst1, weightsrc + unitsize_, unitsize_);
memcpy(weightdst1 + unitsize_, weightsrc + unitsize4 + unitsize_, unitsize_);
memcpy(weightdst2, weightsrc + unitsize_ * 2, unitsize_);
memcpy(weightdst2 + unitsize_, weightsrc + unitsize4 + unitsize_ * 2, unitsize_);
memcpy(weightdst3, weightsrc + unitsize_ * 3, unitsize_);
memcpy(weightdst3 + unitsize_, weightsrc + unitsize4 + unitsize_ * 3, unitsize_);
weightsrc += unitsize4 * 2;
weightdst0 += unitsize_ * 2;
weightdst1 += unitsize_ * 2;
weightdst2 += unitsize_ * 2;
weightdst3 += unitsize_ * 2;
lu -= 2;
}
if (lu > 0) {
memcpy(weightdst0, weightsrc, unitsize_);
memcpy(weightdst1, weightsrc + unitsize_, unitsize_);
memcpy(weightdst2, weightsrc + unitsize_ * 2, unitsize_);
memcpy(weightdst3, weightsrc + unitsize_ * 3, unitsize_);
}
// Reorder scale and bias
auto scaleSrc = src + i * srcStride0 + k * srcStride1 + blockLu * hPSrc * lp;
auto scaleDst0 = dst + (4 * i) * dstStride0 + k * dstStride1 + blockLu * hPDst * lp;
auto scaleDst1 = dst + (4 * i + 1) * dstStride0 + k * dstStride1 + blockLu * hPDst * lp;
auto scaleDst2 = dst + (4 * i + 2) * dstStride0 + k * dstStride1 + blockLu * hPDst * lp;
auto scaleDst3 = dst + (4 * i + 3) * dstStride0 + k * dstStride1 + blockLu * hPDst * lp;
// Copy scales (first part of the scale/bias region)
int scaleSize = hPDst * sizeof(float);
memcpy(scaleDst0, scaleSrc, scaleSize);
memcpy(scaleDst1, scaleSrc + scaleSize, scaleSize);
memcpy(scaleDst2, scaleSrc + scaleSize * 2, scaleSize);
memcpy(scaleDst3, scaleSrc + scaleSize * 3, scaleSize);
// Copy biases (second part of the scale/bias region)
auto biasSrcOffset = hPSrc * sizeof(float);
memcpy(scaleDst0 + scaleSize, scaleSrc + biasSrcOffset, scaleSize);
memcpy(scaleDst1 + scaleSize, scaleSrc + biasSrcOffset + scaleSize, scaleSize);
memcpy(scaleDst2 + scaleSize, scaleSrc + biasSrcOffset + scaleSize * 2, scaleSize);
memcpy(scaleDst3 + scaleSize, scaleSrc + biasSrcOffset + scaleSize * 3, scaleSize);
}
}
}
static void _onlineReorderWeightPackH8ToH32(int8_t* dst, const int8_t* src, int blockLu, int lp, bool isInt4Weight,
int srcH, int blockNum, int resOcBranch) {
constexpr int hPSrc = 8;
constexpr int hPDst = 32;
int srcUnitLp = isInt4Weight ? lp / 2 : lp;
const size_t srcUnitSize = (size_t)hPSrc * srcUnitLp;
const size_t dstUnitSize = (size_t)hPDst * srcUnitLp;
const size_t srcStride1 = (size_t)blockLu * srcUnitSize + 2 * hPSrc * sizeof(float);
const size_t srcStride0 = (size_t)blockNum * srcStride1;
const size_t dstStride1 = (size_t)blockLu * dstUnitSize + 2 * hPDst * sizeof(float);
const size_t dstStride0 = (size_t)blockNum * dstStride1;
const int hUDst = srcH / 4;
const int hTail = srcH % 4;
for (int i = 0; i < hUDst; ++i) {
for (int k = 0; k < blockNum; ++k) {
auto weightSrcBase0 = src + (4 * i + 0) * srcStride0 + k * srcStride1;
auto weightSrcBase1 = src + (4 * i + 1) * srcStride0 + k * srcStride1;
auto weightSrcBase2 = src + (4 * i + 2) * srcStride0 + k * srcStride1;
auto weightSrcBase3 = src + (4 * i + 3) * srcStride0 + k * srcStride1;
auto weightDstBase = dst + i * dstStride0 + k * dstStride1;
int lu = blockLu;
// --- Reorder Weights ---
if (isInt4Weight) {
auto process_int4_block = [](uint8_t* dst_b, const uint8_t* src_b, size_t size) {
auto half_size = size / 2;
for (int s = 0; s < half_size; ++s) {
uint8_t p0 = src_b[2 * s];
uint8_t p1 = src_b[2 * s + 1];
dst_b[s] = (p1 & 0xF0) | (p0 >> 4);
dst_b[s + half_size] = (p1 << 4) | (p0 & 0x0F);
}
};
while (lu >= 4) {
for (int j = 0; j < 4; ++j) {
const auto* srcPtr0 = (const uint8_t*)(weightSrcBase0 + j * srcUnitSize);
const auto* srcPtr1 = (const uint8_t*)(weightSrcBase1 + j * srcUnitSize);
const auto* srcPtr2 = (const uint8_t*)(weightSrcBase2 + j * srcUnitSize);
const auto* srcPtr3 = (const uint8_t*)(weightSrcBase3 + j * srcUnitSize);
auto* dstPtr = (uint8_t*)(weightDstBase + j * dstUnitSize);
process_int4_block(dstPtr + 0 * srcUnitSize, srcPtr0, srcUnitSize);
process_int4_block(dstPtr + 1 * srcUnitSize, srcPtr1, srcUnitSize);
process_int4_block(dstPtr + 2 * srcUnitSize, srcPtr2, srcUnitSize);
process_int4_block(dstPtr + 3 * srcUnitSize, srcPtr3, srcUnitSize);
}
weightSrcBase0 += 4 * srcUnitSize;
weightSrcBase1 += 4 * srcUnitSize;
weightSrcBase2 += 4 * srcUnitSize;
weightSrcBase3 += 4 * srcUnitSize;
weightDstBase += 4 * dstUnitSize;
lu -= 4;
}
for (int j = 0; j < lu; ++j) {
const auto* srcPtr0 = (const uint8_t*)(weightSrcBase0);
const auto* srcPtr1 = (const uint8_t*)(weightSrcBase1);
const auto* srcPtr2 = (const uint8_t*)(weightSrcBase2);
const auto* srcPtr3 = (const uint8_t*)(weightSrcBase3);
auto* dstPtr = (uint8_t*)(weightDstBase);
process_int4_block(dstPtr + 0 * srcUnitSize, srcPtr0, srcUnitSize);
process_int4_block(dstPtr + 1 * srcUnitSize, srcPtr1, srcUnitSize);
process_int4_block(dstPtr + 2 * srcUnitSize, srcPtr2, srcUnitSize);
process_int4_block(dstPtr + 3 * srcUnitSize, srcPtr3, srcUnitSize);
weightSrcBase0 += srcUnitSize;
weightSrcBase1 += srcUnitSize;
weightSrcBase2 += srcUnitSize;
weightSrcBase3 += srcUnitSize;
weightDstBase += dstUnitSize;
}
} else {
while (lu >= 4) {
// j = 0
memcpy(weightDstBase + 0 * dstUnitSize + 0 * srcUnitSize, weightSrcBase0 + 0 * srcUnitSize,
srcUnitSize);
memcpy(weightDstBase + 0 * dstUnitSize + 1 * srcUnitSize, weightSrcBase1 + 0 * srcUnitSize,
srcUnitSize);
memcpy(weightDstBase + 0 * dstUnitSize + 2 * srcUnitSize, weightSrcBase2 + 0 * srcUnitSize,
srcUnitSize);
memcpy(weightDstBase + 0 * dstUnitSize + 3 * srcUnitSize, weightSrcBase3 + 0 * srcUnitSize,
srcUnitSize);
// j = 1
memcpy(weightDstBase + 1 * dstUnitSize + 0 * srcUnitSize, weightSrcBase0 + 1 * srcUnitSize,
srcUnitSize);
memcpy(weightDstBase + 1 * dstUnitSize + 1 * srcUnitSize, weightSrcBase1 + 1 * srcUnitSize,
srcUnitSize);
memcpy(weightDstBase + 1 * dstUnitSize + 2 * srcUnitSize, weightSrcBase2 + 1 * srcUnitSize,
srcUnitSize);
memcpy(weightDstBase + 1 * dstUnitSize + 3 * srcUnitSize, weightSrcBase3 + 1 * srcUnitSize,
srcUnitSize);
// j = 2
memcpy(weightDstBase + 2 * dstUnitSize + 0 * srcUnitSize, weightSrcBase0 + 2 * srcUnitSize,
srcUnitSize);
memcpy(weightDstBase + 2 * dstUnitSize + 1 * srcUnitSize, weightSrcBase1 + 2 * srcUnitSize,
srcUnitSize);
memcpy(weightDstBase + 2 * dstUnitSize + 2 * srcUnitSize, weightSrcBase2 + 2 * srcUnitSize,
srcUnitSize);
memcpy(weightDstBase + 2 * dstUnitSize + 3 * srcUnitSize, weightSrcBase3 + 2 * srcUnitSize,
srcUnitSize);
// j = 3
memcpy(weightDstBase + 3 * dstUnitSize + 0 * srcUnitSize, weightSrcBase0 + 3 * srcUnitSize,
srcUnitSize);
memcpy(weightDstBase + 3 * dstUnitSize + 1 * srcUnitSize, weightSrcBase1 + 3 * srcUnitSize,
srcUnitSize);
memcpy(weightDstBase + 3 * dstUnitSize + 2 * srcUnitSize, weightSrcBase2 + 3 * srcUnitSize,
srcUnitSize);
memcpy(weightDstBase + 3 * dstUnitSize + 3 * srcUnitSize, weightSrcBase3 + 3 * srcUnitSize,
srcUnitSize);
weightSrcBase0 += 4 * srcUnitSize;
weightSrcBase1 += 4 * srcUnitSize;
weightSrcBase2 += 4 * srcUnitSize;
weightSrcBase3 += 4 * srcUnitSize;
weightDstBase += 4 * dstUnitSize;
lu -= 4;
}
for (int j = 0; j < lu; ++j) {
memcpy(weightDstBase + 0 * srcUnitSize, weightSrcBase0, srcUnitSize);
memcpy(weightDstBase + 1 * srcUnitSize, weightSrcBase1, srcUnitSize);
memcpy(weightDstBase + 2 * srcUnitSize, weightSrcBase2, srcUnitSize);
memcpy(weightDstBase + 3 * srcUnitSize, weightSrcBase3, srcUnitSize);
weightSrcBase0 += srcUnitSize;
weightSrcBase1 += srcUnitSize;
weightSrcBase2 += srcUnitSize;
weightSrcBase3 += srcUnitSize;
weightDstBase += dstUnitSize;
}
}
// --- Reorder scale and bias ---
const int scaleSrcSize = hPSrc * sizeof(float);
const int8_t* scaleSrcBase = src + (4 * i) * srcStride0 + k * srcStride1 + (size_t)blockLu * srcUnitSize;
int8_t* scaleDstBase = dst + i * dstStride0 + k * dstStride1 + (size_t)blockLu * dstUnitSize;
memcpy(scaleDstBase + 0 * scaleSrcSize, scaleSrcBase + 0 * srcStride0, scaleSrcSize);
memcpy(scaleDstBase + 1 * scaleSrcSize, scaleSrcBase + 1 * srcStride0, scaleSrcSize);
memcpy(scaleDstBase + 2 * scaleSrcSize, scaleSrcBase + 2 * srcStride0, scaleSrcSize);
memcpy(scaleDstBase + 3 * scaleSrcSize, scaleSrcBase + 3 * srcStride0, scaleSrcSize);
const int8_t* biasSrcBase = scaleSrcBase + scaleSrcSize;
int8_t* biasDstBase = scaleDstBase + hPDst * sizeof(float);
memcpy(biasDstBase + 0 * scaleSrcSize, biasSrcBase + 0 * srcStride0, scaleSrcSize);
memcpy(biasDstBase + 1 * scaleSrcSize, biasSrcBase + 1 * srcStride0, scaleSrcSize);
memcpy(biasDstBase + 2 * scaleSrcSize, biasSrcBase + 2 * srcStride0, scaleSrcSize);
memcpy(biasDstBase + 3 * scaleSrcSize, biasSrcBase + 3 * srcStride0, scaleSrcSize);
}
}
// --- 2. Process the tail ---
if (hTail > 0) {
// The last block starts at index hUDst.
const int i = hUDst;
for (int k = 0; k < blockNum; ++k) {
const int8_t* srcBases[4] = {nullptr, nullptr, nullptr, nullptr};
for (int j = 0; j < hTail; ++j) {
srcBases[j] = src + (4 * i + j) * srcStride0 + k * srcStride1;
}
auto weightDstBase = dst + i * dstStride0 + k * dstStride1;
int lu = blockLu;
if (isInt4Weight) {
auto process_int4_block = [](uint8_t* dst_b, const uint8_t* src_b, size_t size) {
auto half_size = size / 2;
for (int s = 0; s < half_size; ++s) {
uint8_t p0 = src_b[2 * s];
uint8_t p1 = src_b[2 * s + 1];
dst_b[s] = (p1 & 0xF0) | (p0 >> 4);
dst_b[s + half_size] = (p1 << 4) | (p0 & 0x0F);
}
};
while (lu-- > 0) {
for (int j = 0; j < hTail; ++j) {
process_int4_block((uint8_t*)(weightDstBase + j * srcUnitSize), (const uint8_t*)(srcBases[j]),
srcUnitSize);
}
// For the remaining part of the destination block, set 0
if (hTail < 4) {
memset(weightDstBase + hTail * srcUnitSize, 0, (4 - hTail) * srcUnitSize);
}
for (int j = 0; j < hTail; ++j) {
srcBases[j] += srcUnitSize;
}
weightDstBase += dstUnitSize;
}
} else { // int8 weight
while (lu-- > 0) {
for (int j = 0; j < hTail; ++j) {
memcpy(weightDstBase + j * srcUnitSize, srcBases[j], srcUnitSize);
}
// Zero out the rest of the destination block
if (hTail < 4) {
memset(weightDstBase + hTail * srcUnitSize, 0, (4 - hTail) * srcUnitSize);
}
for (int j = 0; j < hTail; ++j) {
srcBases[j] += srcUnitSize;
}
weightDstBase += dstUnitSize;
}
}
// --- Reorder scale and bias for tail ---
const int scaleSrcSize = hPSrc * sizeof(float);
const int8_t* scaleSrcBase = src + (4 * i) * srcStride0 + k * srcStride1 + (size_t)blockLu * srcUnitSize;
int8_t* scaleDstBase = dst + i * dstStride0 + k * dstStride1 + (size_t)blockLu * dstUnitSize;
for (int j = 0; j < hTail; ++j) {
memcpy(scaleDstBase + j * scaleSrcSize, scaleSrcBase + j * srcStride0, scaleSrcSize);
}
if (hTail < 4) {
memset(scaleDstBase + hTail * scaleSrcSize, 0, (4 - hTail) * scaleSrcSize);
}
const int8_t* biasSrcBase = scaleSrcBase + scaleSrcSize;
int8_t* biasDstBase = scaleDstBase + hPDst * sizeof(float);
for (int j = 0; j < hTail; ++j) {
memcpy(biasDstBase + j * scaleSrcSize, biasSrcBase + j * srcStride0, scaleSrcSize);
}
if (hTail < 4) {
memset(biasDstBase + hTail * scaleSrcSize, 0, (4 - hTail) * scaleSrcSize);
}
}
}
// --- 3. Copy the residual part ---
if (resOcBranch > 0) {
size_t resLp = isInt4Weight ? lp / 2 : lp;
size_t resChannels = ROUND_UP(resOcBranch, hPSrc);
size_t resDataLen =
(size_t)blockNum * ((size_t)blockLu * resChannels * resLp + 2 * resChannels * sizeof(float));
// The source for residual data starts after ALL processed srcH blocks.
memcpy(dst + (size_t)hUDst * dstStride0 + (hTail > 0 ? dstStride0 : 0), src + (size_t)srcH * srcStride0,
resDataLen);
}
}
static void _onlineReorderWeightKernelSumH128ToH32(float* dst, float* src, int blockNum, int hpSrc, int hpDst, int oc) {
// hpSrc = 4 * hpDst
// src shape: [huSrc, blockNum, hpSrc]
// dst shape: [huDst, blockNum, hpDst], where huDst = huSrc * 4
auto huSrc = UP_DIV(oc, hpSrc);
auto strideSrc = blockNum * hpSrc;
auto strideDst = blockNum * hpDst;
for (int i = 0; i < huSrc; ++i) {
for (int k = 0; k < blockNum; ++k) {
auto srcBase = src + i * strideSrc + k * hpSrc;
auto dst0 = dst + (4 * i + 0) * strideDst + k * hpDst;
auto dst1 = dst + (4 * i + 1) * strideDst + k * hpDst;
auto dst2 = dst + (4 * i + 2) * strideDst + k * hpDst;
auto dst3 = dst + (4 * i + 3) * strideDst + k * hpDst;
memcpy(dst0, srcBase, hpDst * sizeof(float));
memcpy(dst1, srcBase + hpDst, hpDst * sizeof(float));
memcpy(dst2, srcBase + 2 * hpDst, hpDst * sizeof(float));
memcpy(dst3, srcBase + 3 * hpDst, hpDst * sizeof(float));
}
}
}
static void _onlineReorderWeightKernelSumH8ToH32(float* dst, float* src, int blockNum, int hpSrc, int hpDst,
int ocNeedReorder, int ocPreserve) {
// hpDst = 4 * hpSrc
// src shape: [huSrc, blockNum, hpSrc], where huSrc = huDst * 4
// dst shape: [huDst, blockNum, hpDst]
auto huDst = UP_DIV(ocNeedReorder, hpDst);
auto strideSrc = blockNum * hpSrc;
auto strideDst = blockNum * hpDst;
for (int i = 0; i < huDst; ++i) {
for (int k = 0; k < blockNum; ++k) {
auto dstBase = dst + i * strideDst + k * hpDst;
auto src0 = src + (4 * i + 0) * strideSrc + k * hpSrc;
auto src1 = src + (4 * i + 1) * strideSrc + k * hpSrc;
auto src2 = src + (4 * i + 2) * strideSrc + k * hpSrc;
auto src3 = src + (4 * i + 3) * strideSrc + k * hpSrc;
memcpy(dstBase, src0, hpSrc * sizeof(float));
memcpy(dstBase + hpSrc, src1, hpSrc * sizeof(float));
memcpy(dstBase + 2 * hpSrc, src2, hpSrc * sizeof(float));
memcpy(dstBase + 3 * hpSrc, src3, hpSrc * sizeof(float));
}
}
if (ocPreserve) {
memcpy(dst + huDst * strideDst, src + 4 * huDst * strideSrc,
ROUND_UP(ocPreserve, hpSrc) * blockNum * sizeof(float));
}
}
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 % 8;
int UNIT = mGemmUnits[0];
int SRC_UNIT = mGemmUnits[1];
int DST_XUNIT = mGemmUnits[2];
auto blitProc = mRelatedFunctions.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 int ocUpHp = ROUND_UP(output->channel(), UNIT);
const auto kernelCountUnit = mIm2ColParamter.kernelCountUnit;
const auto unitColBufferSize = kernelCountUnit * DST_XUNIT * SRC_UNIT * sizeof(int8_t);
const auto colBufferSize = unitColBufferSize * mIm2ColCount;
auto 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
const int oc = output->channel();
const int ic = input->channel();
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());
}
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;
// Additional Adjustments for 4Bit Ptq model
if (m4BitPtq) {
outputDataPtr = (int8_t*)mTempOutput.ptr();
dstBytes = 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];
if (true == m4BitPtq) {
scaleou = 1;
}
auto scaleX = scalein / scaleou;
for (int i = 0; i < DST_XUNIT; ++i) {
mBatchQuantInfo->host<float>()[i] = scaleX;
}
}
}
// Declare variables used in dynamic quantization
const int threads = static_cast<CPUBackend*>(backend())->threadNumber();
int dropBranch = 0;
#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 = static_cast<int32_t>(roundf(qbias[0]));
auto updatedBiasPtr = (float*)(mBiasBufferFusedInputzero.ptr() + tId * ocUpHp * QUANT_INFO_BYTES);
auto matmulBiasPtr = mResourceInt8->mOriginBias->host<float>();
auto weightKernelSum = mResourceInt8->mWeightKernelSum->host<float>();
auto inputZeroF = -qbias[0] * scalePtr[0];
gcore->MNNDynamicUpdateConvBiasScale(updatedBiasPtr, matmulBiasPtr, weightKernelSum, &inputZeroF,
UP_DIV(ocUpHp, 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>();
}
if (mOnlineReorderWeightSme && plane > 1) {
_onlineReorderWeightPackH128ToH32((int8_t*)mWeight4Prefill.ptr(), weightDataPtr, GEMM_INT8_UNIT_SME2_128, UNIT,
UP_DIV(mOcMain, GEMM_INT8_UNIT_SME2_128), mBlockNum, blockL, SRC_UNIT,
mResourceInt8->mWeightBits == 4);
int kernelSumMainSize = 0;
int kernelSumBranchSize = 0;
if (dstBytes > 1 && mInputBlockNum > 1) {
_onlineReorderWeightKernelSumH128ToH32((float*)mWeightKernelSum4Prefill.ptr(),
mResourceInt8->mWeightKernelSum->host<float>(), mBlockNum,
GEMM_INT8_UNIT_SME2_128, UNIT, mOcMain);
kernelSumMainSize = ROUND_UP(mOcMain, UNIT) * mBlockNum * QUANT_INFO_BYTES;
kernelSumBranchSize = ROUND_UP(mOcBranch, 8) * mBlockNum * QUANT_INFO_BYTES;
}
// If change the workload distribution among SME and NEON cores.
if (mMixedKernel && mRatioDecode != mRatioPrefill) {
auto offsetWeight =
UP_DIV(mOcMain, GEMM_INT8_UNIT_SME2_128) * mBlockNum * blockL * SRC_UNIT * GEMM_INT8_UNIT_SME2_128;
if (mResourceInt8->mWeightBits == 4) {
offsetWeight /= 2;
}
offsetWeight += (ROUND_UP(mOcMain, GEMM_INT8_UNIT_SME2_128) * mBlockNum * 2 * sizeof(float));
// Don't change mOcMain&mOcBranch here.
int tmpMain = mOcMain;
int tmpBranch = mOcBranch;
calculateSmeNeonWorkDivision(tmpMain, tmpBranch, mDividesTmp, oc, threads, PackUnit, plane, mRatioPrefill,
mSmeCores);
auto updatedSmeWork = mDividesTmp[mSmeCores];
if (updatedSmeWork - mOriginSmeWork > 0 &&
((updatedSmeWork - mOriginSmeWork) * 4 % 8 == 0)) { // To ensure pack=4, dropBranch % 2 == 0
dropBranch = updatedSmeWork - mOriginSmeWork; // Ensure update "dropBranch" inner the loop.
memcpy(mDivides.data(), mDividesTmp.data(), (threads + 1) * sizeof(float));
dropBranch = mDivides[mSmeCores] - mOriginSmeWork;
_onlineReorderWeightPackH8ToH32((int8_t*)(mWeight4Prefill.ptr() + offsetWeight),
weightDataPtr + offsetWeight, blockL, SRC_UNIT,
mResourceInt8->mWeightBits == 4, (int)(dropBranch * PackUnit / 8),
mBlockNum, (mDivides[threads] - mDivides[mSmeCores]) * PackUnit);
}
if (dstBytes > 1 && mInputBlockNum > 1) {
if (dropBranch > 0) {
// reorder
_onlineReorderWeightKernelSumH8ToH32(
(float*)(mWeightKernelSum4Prefill.ptr() + kernelSumMainSize),
(float*)(mResourceInt8->mWeightKernelSum->host<int8_t>() + kernelSumMainSize), mBlockNum, 8,
UNIT, dropBranch * PackUnit, (mDivides[threads] - mDivides[mSmeCores]) * PackUnit);
}
}
}
if (dropBranch == 0) { // If dropBranch == 0, it means that the arrangement of the weights processed by the
// Arm82 architecture remains unchanged.
// copy
memcpy(mWeightKernelSum4Prefill.ptr() + kernelSumMainSize,
mResourceInt8->mWeightKernelSum->host<uint8_t>() + kernelSumMainSize, kernelSumBranchSize);
}
weightDataPtr = (int8_t*)mWeight4Prefill.ptr();
}
#endif
if (mResourceInt8->mWeightBits == 4) {
weightBytes = 0.5;
weightStepY /= 2;
} else if (mResourceInt8->mWeightBits == 3) {
auto packedBytesPerOc = (SRC_UNIT * 3 + 7) / 8;
weightBytes = static_cast<float>(packedBytesPerOc) / SRC_UNIT;
weightStepY = UNIT * packedBytesPerOc;
} else if (mResourceInt8->mWeightBits == 2) {
weightBytes = 0.25f;
weightStepY /= 4;
}
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;
std::vector<float> fakeInputScales(DST_XUNIT, 1.f);
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;
}
} else {
inputScale = (uint8_t*)fakeInputScales.data();
ptrInputScale = inputScale;
}
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 = (mOnlineReorderWeightSme && mInputBlockNum > 1) ? (float*)mWeightKernelSum4Prefill.ptr()
: mResourceInt8->mWeightKernelSum->host<float>();
}
QuanPostTreatParameters quanParam;
quanParam.blockNum = mBlockNum;
int32_t indices[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15};
quanParam.indices = indices;
if (dstBytes != 1) {
quanParam.useInt8 = 0;
quanParam.fp32minmax = reluPtr;
#ifdef MNN_USE_SSE
if (!mBatchQuantInfo.get()) {
quanParam.weightKernelSum = nullptr;
}
#endif
} else {
quanParam.maxValue = mMutableResource->mClampMax;
if (mResourceInt8->mRelu) {
quanParam.minValue = mMutableResource->mOutputZeroPoint;
} else {
quanParam.minValue = mMutableResource->mClampMin;
}
}
auto weightPtrTid = weightDataPtr;
quanParam.weightKernelSum = ptrY;
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;
mRelatedFunctions.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 * ocUpHp * 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>());
mRelatedFunctions.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
mRelatedFunctions.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>());
mRelatedFunctions.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 && ocDivThread > 0) {
decltype(mGemmKernel) gemmInt8;
if (mMixedKernel) {
gemmInt8 = tId < mSmeCores ? mGemmKernels[0] : mGemmKernels[1];
} else {
gemmInt8 = mGemmKernel;
}
auto im2colDstThread = im2colDst;
float* ptrY = nullptr;
if (dstBytes != 1) {
float* wkernelSum = (mOnlineReorderWeightSme && mInputBlockNum > 1 && plane > 1)
? (float*)mWeightKernelSum4Prefill.ptr()
: mResourceInt8->mWeightKernelSum->host<float>();
ptrY = wkernelSum + ocIndex * mInputBlockNum;
}
QuanPostTreatParameters quanParam;
quanParam.blockNum = mBlockNum;
quanParam.weightKernelSum = ptrY;
quanParam.biasFloat = reinterpret_cast<float*>(biasPtr + ocIndex * 4);
int32_t indices[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15};
quanParam.indices = indices;
if (dstBytes != 1) {
quanParam.useInt8 = 0;
quanParam.fp32minmax = reluPtr;
#ifdef MNN_USE_SSE
if (!mBatchQuantInfo.get()) {
quanParam.weightKernelSum = nullptr;
}
#endif
} 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;
}
} else {
inputScale = (uint8_t*)fakeInputScales.data();
}
if (mBlockNum > 1) {
accumbuff = reinterpret_cast<float*>(mAccumBuffer->host<int8_t>() +
tId * mAccumBuffer->stride(0) * sizeof(int32_t));
}
auto outputInTilePtr = outputDataPtr + ocIndex * plane * dstBytes;
auto weightSrc = weightDataPtr;
if (tId >= mSmeCores && dropBranch == 0 && mMixedKernel) {
weightSrc = mResourceInt8->mWeightInt8->host<int8_t>();
}
auto weightPtrTid =
weightSrc + 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;
}
gemmInt8(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();
};
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);
}
if (m4BitPtq) {
std::vector<float> outputQuantScale(PackUnit);
float s = TensorUtils::getQuantInfo(outputs[0])[0] == 0 ? 0 : 1.f / TensorUtils::getQuantInfo(outputs[0])[0];
for (int i = 0; i < PackUnit; ++i) {
outputQuantScale[i] = s;
}
float zero_ = TensorUtils::getQuantInfo(outputs[0])[1];
mQuantFunc((float*)mTempOutput.ptr(), output->host<int8_t>(), plane * ocDiv4, outputQuantScale.data(),
mResourceInt8->mClampMin, mResourceInt8->mClampMax, &zero_, 0);
}
return NO_ERROR;
}
} // namespace MNN
Reference in New Issue View Git Blame Copy Permalink