Function _createKleidiAIUnit

source/backend/cpu/compute/ConvolutionFloatFactory.cpp:36–118

#ifdef MNN_KLEIDIAI_ENABLED
static Execution* _createKleidiAIUnit(const Tensor* input, const Tensor* output, Backend* backend, const Op* op,
                                      const float* originWeight, size_t originWeightSize, const float* bias,
                                      size_t biasSize, std::shared_ptr<ConvolutionCommon::Int8Common> weightQuantInfo,
                                      bool supportSparse, bool lowMemory) {
    auto cpuBackend = (CPUBackend*)backend;
    auto conv2d     = op->main_as_Convolution2D();
    auto common     = conv2d->common();

    bool fastWay = common->kernelY() == 1 && common->kernelX() == 1 && output->width() == input->width() &&
                   output->height() == input->height() && common->strideX() == 1 && common->strideY() == 1;

#ifdef MNN_LOW_MEMORY
    if (lowMemory && nullptr != weightQuantInfo.get() && originWeightSize == 0) {
        if (cpuBackend->memoryMode() == BackendConfig::Memory_Low) {
            do {
                if (!weightQuantInfo->canUseInt4) {
                    break;
                }
                auto convOp = op->main_as_Convolution2D();
                auto core   = static_cast<CPUBackend*>(backend)->functions();
                int oc      = convOp->common()->outputCount();
                int ic      = convOp->common()->inputCount();

                int blockNum   = 1;
                int dequantCnt = weightQuantInfo->alphaSize;
                if (weightQuantInfo->asymmetric) {
                    dequantCnt /= 2;
                }
                blockNum = dequantCnt / oc;

                bool bAsym     = weightQuantInfo->asymmetric;
                size_t blkSize = blockNum == 1 ? 0 : ic / blockNum;

                KleidiAI::AccelType accelType = KleidiAI::getQIntAccelType(4, bAsym, blkSize, core->bytes);

                KleidiAI& kai = KleidiAI::getInstance(*MNNGetCPUInfo());
                if (!kai.canAccelerate(accelType, convOp->common())) {
                    break;
                }

                if (!kai.isLoaded(accelType)) {
                    kai.setLoaded(accelType);
                    kai.printInfo(accelType);
                }

                return new KleidiAIConvInt8(backend, op, weightQuantInfo, true, kai, accelType, blockNum);
            } while (0);
        }

        // Have not supported the quantized weight.
        return nullptr;
    }
#else
    if (cpuBackend->memoryMode() == BackendConfig::Memory_Low) {
        if (MNNGetCPUInfo()->sme2 && !weightQuantInfo) {
            return new KleidiAIDenseConvolution(common, backend, originWeight, originWeightSize, bias, biasSize,
                                                weightQuantInfo);
        }