#include "VulkanConv1x1General.hpp"
#include "VulkanBackend.hpp"
#include "core/Macro.h"
#include <algorithm>
#include <cstdint>
#include <cstring>
#include <vector>

namespace MNN {

namespace {

struct QuantWeightPrepareParams {
    uint32_t ci;
    uint32_t co;
    uint32_t padN;
    uint32_t weightStride;
    uint32_t srcBytes;
};

struct QuantMetaPrepareParams {
    uint32_t co;
    uint32_t padN;
    uint32_t blockCount;
    uint32_t blockStride;
    uint32_t soSize;
    uint32_t alphaSize;
};

static size_t _alignUp4(size_t size) {
    return (size + 3u) & ~size_t(3u);
}

static bool _prepareQuantBuffersGPU(VulkanBackend* vkBn, const ConvolutionCommon::Int8Common* quantCommon,
                                    bool useFP16, int ci, int co, uint32_t padN, uint32_t blockStride,
                                    uint32_t decodeWeightStrideWords, int quantBits,
                                    std::shared_ptr<VulkanBuffer>& quantWeightBuffer,
                                    std::shared_ptr<VulkanBuffer>& quantMetaBuffer) {
    if (nullptr == vkBn || nullptr == quantCommon || nullptr == quantCommon->weight.get()) {
        return false;
    }

    const int soSize = quantCommon->asymmetric ? 2 : 1;
    const int alphaSize = quantCommon->alpha.size();
    const int alphaDenominator = std::max(1, co * soSize);
    const int blockCount = std::max(1, alphaSize / alphaDenominator);
    const int8_t* qWeight = quantCommon->weight.get();
    const size_t rawWeightBytes = static_cast<size_t>(quantCommon->weight.size());
    const size_t alignedWeightBytes = std::max<size_t>(4u, _alignUp4(rawWeightBytes));
    // int3 stores 2 uints per group of 16 weights; others store 1 uint per group.
    const uint32_t wordsPerGroup = (quantBits == 3) ? 2u : 1u;
    const size_t decodeWeightBytes =
        static_cast<size_t>(padN) * static_cast<size_t>(decodeWeightStrideWords) *
        static_cast<size_t>(wordsPerGroup) * sizeof(uint32_t);
    const size_t metaElem = static_cast<size_t>(padN) * static_cast<size_t>(blockStride) * 2u;
    const size_t metaBytes = metaElem * (useFP16 ? sizeof(int16_t) : sizeof(float));

    const void* rawWeightSrc = qWeight;
    std::vector<uint8_t> weightAlignedHost;
    if (alignedWeightBytes != rawWeightBytes) {
        weightAlignedHost.resize(alignedWeightBytes, 0);
        if (rawWeightBytes > 0u) {
            ::memcpy(weightAlignedHost.data(), qWeight, rawWeightBytes);
        }
        rawWeightSrc = weightAlignedHost.data();
    }

    std::shared_ptr<VulkanBuffer> rawWeightBuffer(new VulkanBuffer(
        vkBn->getMemoryPool(), false, alignedWeightBytes, nullptr,
        VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT, VK_SHARING_MODE_EXCLUSIVE, 0));
    vkBn->copyToGPUBuffer(rawWeightSrc, rawWeightBuffer->buffer(), alignedWeightBytes, 0);

    quantWeightBuffer.reset(new VulkanBuffer(vkBn->getMemoryPool(), false, decodeWeightBytes, nullptr,
                                             VK_BUFFER_USAGE_STORAGE_BUFFER_BIT, VK_SHARING_MODE_EXCLUSIVE, 0));

    const float* alphaPtr = quantCommon->alpha.get();
    const size_t rawAlphaBytes = static_cast<size_t>(std::max(alphaSize, 1)) * sizeof(float);
    std::shared_ptr<VulkanBuffer> rawAlphaBuffer(new VulkanBuffer(
        vkBn->getMemoryPool(), false, rawAlphaBytes, nullptr,
        VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT, VK_SHARING_MODE_EXCLUSIVE, 0));
    if (alphaSize > 0 && nullptr != alphaPtr) {
        vkBn->copyToGPUBuffer(alphaPtr, rawAlphaBuffer->buffer(), static_cast<size_t>(alphaSize) * sizeof(float), 0);
    } else {
        const float zero = 0.0f;
        vkBn->copyToGPUBuffer(&zero, rawAlphaBuffer->buffer(), sizeof(float), 0);
    }

    quantMetaBuffer.reset(new VulkanBuffer(vkBn->getMemoryPool(), false, metaBytes, nullptr,
                                           VK_BUFFER_USAGE_STORAGE_BUFFER_BIT, VK_SHARING_MODE_EXCLUSIVE, 0));

    const char* weightShader = nullptr;
    switch (quantBits) {
        case 2: weightShader = "glsl_conv1x1_int2_weight_prepare_comp"; break;
        case 3: weightShader = "glsl_conv1x1_int3_weight_prepare_comp"; break;
        case 4: weightShader = "glsl_conv1x1_int4_weight_prepare_comp"; break;
        default: weightShader = "glsl_conv1x1_int8_weight_prepare_comp"; break;
    }
    const char* metaShader = useFP16 ? "glsl_conv1x1_quant_meta_prepare_FP16_comp"
                                     : "glsl_conv1x1_quant_meta_prepare_comp";

    std::vector<VkDescriptorType> prepareTypes = {
        VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
        VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
    };
    auto weightPipeline = vkBn->getPipeline(weightShader, prepareTypes);
    auto metaPipeline = vkBn->getPipeline(metaShader, prepareTypes);
    if (nullptr == weightPipeline || nullptr == metaPipeline) {
        return false;
    }

    std::shared_ptr<VulkanLayout::DescriptorSet> weightSet(weightPipeline->createSet());
    std::shared_ptr<VulkanLayout::DescriptorSet> metaSet(metaPipeline->createSet());
    if (nullptr == weightSet.get() || nullptr == metaSet.get()) {
        return false;
    }

    std::shared_ptr<VulkanCommandPool::Buffer> prepareCmd(vkBn->getPool().allocBuffer());
    prepareCmd->begin(0);

    {
        QuantWeightPrepareParams pc;
        pc.ci = static_cast<uint32_t>(ci);
        pc.co = static_cast<uint32_t>(co);
        pc.padN = padN;
        pc.weightStride = decodeWeightStrideWords;
        pc.srcBytes = static_cast<uint32_t>(rawWeightBytes);

        weightSet->writeBuffer(rawWeightBuffer->buffer(), 0, rawWeightBuffer->size());
        weightSet->writeBuffer(quantWeightBuffer->buffer(), 1, quantWeightBuffer->size());
        weightPipeline->bind(prepareCmd->get(), weightSet->get());
        vkCmdPushConstants(prepareCmd->get(), weightPipeline->layout(), VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(pc), &pc);
        vkCmdDispatch(prepareCmd->get(), UP_DIV(decodeWeightStrideWords, 16u), UP_DIV(padN, 16u), 1);
        prepareCmd->barrierSource(quantWeightBuffer->buffer(), 0, quantWeightBuffer->size());
    }

    {
        QuantMetaPrepareParams pc;
        pc.co = static_cast<uint32_t>(co);
        pc.padN = padN;
        pc.blockCount = static_cast<uint32_t>(blockCount);
        pc.blockStride = blockStride;
        pc.soSize = static_cast<uint32_t>(soSize);
        pc.alphaSize = static_cast<uint32_t>(alphaSize);

        metaSet->writeBuffer(rawAlphaBuffer->buffer(), 0, rawAlphaBuffer->size());
        metaSet->writeBuffer(quantMetaBuffer->buffer(), 1, quantMetaBuffer->size());
        metaPipeline->bind(prepareCmd->get(), metaSet->get());
        vkCmdPushConstants(prepareCmd->get(), metaPipeline->layout(), VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(pc), &pc);
        vkCmdDispatch(prepareCmd->get(), UP_DIV(blockStride, 16u), UP_DIV(padN, 16u), 1);
        prepareCmd->barrierSource(quantMetaBuffer->buffer(), 0, quantMetaBuffer->size());
    }

    prepareCmd->end();
    vkBn->getPool().submitAndWait(prepareCmd->get());
    return true;
}

} // namespace

VulkanConv1x1General::VulkanConv1x1General(VulkanBackend* backend, const Convolution2DCommon* convOption,
                                           const float* biasPtr, int ci, int co,
                                           std::shared_ptr<ConvolutionCommon::Int8Common> quantInfo)
    : VulkanBasicExecution(backend), mCommon(convOption), mCi(ci), mCo(co), mQuantCommon(std::move(quantInfo)) {
    if (!_init(biasPtr, true)) {
        MNN_ERROR("VulkanConv1x1General init failed\n");
    }
}

VulkanConv1x1General::VulkanConv1x1General(VulkanBackend* backend, const Convolution2DCommon* convOption, int ci,
                                           int co, std::shared_ptr<ConvolutionCommon::Int8Common> quantInfo,
                                           bool initStaticResource)
    : VulkanBasicExecution(backend), mCommon(convOption), mCi(ci), mCo(co), mQuantCommon(std::move(quantInfo)) {
    if (!_init(nullptr, initStaticResource)) {
        MNN_ERROR("VulkanConv1x1General clone init failed\n");
    }
}

VulkanConv1x1General::~VulkanConv1x1General() {
}

bool VulkanConv1x1General::_init(const float* biasPtr, bool initStaticResource) {
    auto vkBn = static_cast<VulkanBackend*>(backend());
    if (nullptr == vkBn || nullptr == mQuantCommon.get() || nullptr == mQuantCommon->weight.get()) {
        return false;
    }

    const bool useFP16 = vkBn->useFP16();
    // Precedence: int2 > int3 > int4 > int8 (default).
    if (mQuantCommon->canUseInt2) {
        mQuantBits = 2;
    } else if (mQuantCommon->canUseInt3) {
        mQuantBits = 3;
    } else if (mQuantCommon->canUseInt4) {
        mQuantBits = 4;
    } else {
        mQuantBits = 8;
    }
    mPadK = ROUND_UP(static_cast<uint32_t>(mCi), 4u);
    mPadN = ROUND_UP(static_cast<uint32_t>(mCo), 32u);

    if (mPadK == 0u || mPadN == 0u) {
        MNN_ERROR("VulkanConv1x1General invalid shape, ci=%d, co=%d\n", mCi, mCo);
        return false;
    }

    const int soSize = mQuantCommon->asymmetric ? 2 : 1;
    const int alphaSize = mQuantCommon->alpha.size();
    const int alphaDenominator = std::max(1, mCo * soSize);
    const int blockCount = std::max(1, alphaSize / alphaDenominator);
    mBlockSize = std::max<uint32_t>(1u, static_cast<uint32_t>(UP_DIV(mCi, blockCount)));

    if ((mBlockSize & 3u) != 0u) {
        MNN_ERROR("VulkanConv1x1General requires blockSize %% 4 == 0, blockSize=%u\n", mBlockSize);
        return false;
    }
    if (mCi % static_cast<int>(mBlockSize) != 0) {
        MNN_ERROR("VulkanConv1x1General requires K %% blockSize == 0, K=%d, blockSize=%u\n", mCi, mBlockSize);
        return false;
    }
    if ((mPadK % mBlockSize) != 0u) {
        MNN_ERROR("VulkanConv1x1General requires padK %% blockSize == 0, padK=%u, blockSize=%u\n", mPadK, mBlockSize);
        return false;
    }
    mBlockStride = mPadK / mBlockSize;
    // For int2/int3, weightStride counts groups-of-16; int3 uses 2 uints per group, int2 uses 1.
    mDecodeWeightStrideWords = (mQuantBits == 2 || mQuantBits == 3) ? UP_DIV(mPadK, 16u)
                              : (mQuantBits == 4)                   ? UP_DIV(mPadK, 8u)
                                                                    : (mPadK / 4u);

    if (initStaticResource) {
        if (!_prepareQuantBuffersGPU(vkBn, mQuantCommon.get(), useFP16, mCi, mCo, mPadN, mBlockStride,
                                     mDecodeWeightStrideWords, mQuantBits,
                                     mQuantWeightBuffer, mQuantMetaBuffer)) {
            return false;
        }

        std::vector<float> biasHost(mPadN, 0.0f);
        if (nullptr != biasPtr) {
            ::memcpy(biasHost.data(), biasPtr, static_cast<size_t>(mCo) * sizeof(float));
        }
        if (useFP16) {
            std::vector<int16_t> biasHalf(mPadN);
            FLOAT_TO_HALF(biasHost.data(), biasHalf.data(), static_cast<int>(mPadN));
            mBiasBuffer.reset(new VulkanBuffer(vkBn->getMemoryPool(), false, mPadN * sizeof(int16_t), nullptr,
                                               VK_BUFFER_USAGE_STORAGE_BUFFER_BIT |
                                                   VK_BUFFER_USAGE_TRANSFER_DST_BIT,
                                               VK_SHARING_MODE_EXCLUSIVE, 0));
            vkBn->copyToGPUBuffer(biasHalf.data(), mBiasBuffer->buffer(), mPadN * sizeof(int16_t), 0);
        } else {
            mBiasBuffer.reset(new VulkanBuffer(vkBn->getMemoryPool(), false, mPadN * sizeof(float), nullptr,
                                               VK_BUFFER_USAGE_STORAGE_BUFFER_BIT |
                                                   VK_BUFFER_USAGE_TRANSFER_DST_BIT,
                                               VK_SHARING_MODE_EXCLUSIVE, 0));
            vkBn->copyToGPUBuffer(biasHost.data(), mBiasBuffer->buffer(), mPadN * sizeof(float), 0);
        }
    }

    int activation = 0;
    if (mCommon->relu()) {
        activation = 1;
    }
    if (mCommon->relu6()) {
        activation = 2;
    }

    {
        const auto& subgroup = vkBn->getDevice().getSubgroupInfo();
        const VkSubgroupFeatureFlags requiredOps = VK_SUBGROUP_FEATURE_BASIC_BIT | VK_SUBGROUP_FEATURE_ARITHMETIC_BIT;
        mUseSubgroup = subgroup.size > 0 &&
                       (subgroup.stages & VK_SHADER_STAGE_COMPUTE_BIT) &&
                       ((subgroup.ops & requiredOps) == requiredOps);
    }

    mDecodeSubgroupSize = vkBn->getDevice().getSubgroupSize();
    if (mDecodeSubgroupSize == 0u) {
        mDecodeSubgroupSize = 64u;
    }

    {
        std::vector<VkDescriptorType> types = {
            VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
            VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
            VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
            VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
            VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
        };
        const char* shader = nullptr;
        if (mUseSubgroup) {
            std::vector<uint32_t> spec = {static_cast<uint32_t>(activation)};
            switch (mQuantBits) {
                case 2: shader = useFP16 ? "glsl_gemv_dequant_int2_FP16_comp" : "glsl_gemv_dequant_int2_comp"; break;
                case 3: shader = useFP16 ? "glsl_gemv_dequant_int3_FP16_comp" : "glsl_gemv_dequant_int3_comp"; break;
                case 4: shader = useFP16 ? "glsl_gemv_dequant_int4_FP16_comp" : "glsl_gemv_dequant_int4_comp"; break;
                default: shader = useFP16 ? "glsl_gemv_dequant_int8_FP16_comp" : "glsl_gemv_dequant_int8_comp"; break;
            }
            mDecodePipeline = vkBn->getPipeline(shader, types, {mDecodeSubgroupSize, 1, 1}, spec);
        } else {
            uint32_t localSize = 64u;
            std::vector<uint32_t> spec = {static_cast<uint32_t>(activation), localSize};
            switch (mQuantBits) {
                case 2: shader = useFP16 ? "glsl_gemv_dequant_int2_nosubgroup_FP16_comp" : "glsl_gemv_dequant_int2_nosubgroup_comp"; break;
                case 3: shader = useFP16 ? "glsl_gemv_dequant_int3_nosubgroup_FP16_comp" : "glsl_gemv_dequant_int3_nosubgroup_comp"; break;
                case 4: shader = useFP16 ? "glsl_gemv_dequant_int4_nosubgroup_FP16_comp" : "glsl_gemv_dequant_int4_nosubgroup_comp"; break;
                default: shader = useFP16 ? "glsl_gemv_dequant_int8_nosubgroup_FP16_comp" : "glsl_gemv_dequant_int8_nosubgroup_comp"; break;
            }
            mDecodePipeline = vkBn->getPipeline(shader, types, {localSize, 1, 1}, spec);
        }
        if (nullptr == mDecodePipeline) {
            return false;
        }
        mDecodeSet.reset(mDecodePipeline->createSet());
    }

    {
        std::vector<VkDescriptorType> types = {
            VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
            VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
        };
        const char* shader = useFP16 ? "glsl_pack_a_k4m4_to_m64k4_FP16_comp" : "glsl_pack_a_k4m4_to_m64k4_comp";
        mPackAPipeline = vkBn->getPipeline(shader, types);
        if (nullptr == mPackAPipeline) {
            return false;
        }
        mPackASet.reset(mPackAPipeline->createSet());
    }

    {
        std::vector<VkDescriptorType> types = {
            VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
            VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
            VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
        };
        const char* shader = nullptr;
        switch (mQuantBits) {
            case 2: shader = useFP16 ? "glsl_int2_weight_to_pack_FP16_comp" : "glsl_int2_weight_to_pack_comp"; break;
            case 3: shader = useFP16 ? "glsl_int3_weight_to_pack_FP16_comp" : "glsl_int3_weight_to_pack_comp"; break;
            case 4: shader = useFP16 ? "glsl_int4_weight_to_pack_FP16_comp" : "glsl_int4_weight_to_pack_comp"; break;
            default: shader = useFP16 ? "glsl_int8_weight_to_pack_FP16_comp" : "glsl_int8_weight_to_pack_comp"; break;
        }
        mWeightToPackPipeline = vkBn->getPipeline(shader, types);
        if (nullptr == mWeightToPackPipeline) {
            return false;
        }
        mWeightToPackSet.reset(mWeightToPackPipeline->createSet());
    }

    {
        std::vector<VkDescriptorType> types = {
            VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
            VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
            VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
            VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
        };
        std::vector<uint32_t> spec = {static_cast<uint32_t>(activation)};
        const char* shader = useFP16 ? "glsl_gemm_m8n4_FP16_comp" : "glsl_gemm_m8n4_comp";
        mGemmPipeline = vkBn->getPipeline(shader, types, {}, spec);
        if (nullptr == mGemmPipeline) {
            return false;
        }
        mGemmSet.reset(mGemmPipeline->createSet());
    }

    return mDecodeSet != nullptr && mPackASet != nullptr && mWeightToPackSet != nullptr && mGemmSet != nullptr;
}

bool VulkanConv1x1General::onClone(Backend* bn, const Op* op, VulkanBasicExecution** dst) {
    if (nullptr == dst) {
        return true;
    }
    auto vkBn = static_cast<VulkanBackend*>(bn);
    auto conv2D = op->main_as_Convolution2D();
    if (nullptr == conv2D || nullptr == conv2D->common()) {
        return false;
    }
    auto res = new VulkanConv1x1General(vkBn, conv2D->common(), mCi, mCo, mQuantCommon, false);
    res->mQuantBits = mQuantBits;
    res->mPadK = mPadK;
    res->mPadN = mPadN;
    res->mBlockSize = mBlockSize;
    res->mBlockStride = mBlockStride;
    res->mDecodeWeightStrideWords = mDecodeWeightStrideWords;
    res->mDecodeSubgroupSize = mDecodeSubgroupSize;
    res->mUseSubgroup = mUseSubgroup;
    res->mQuantWeightBuffer = mQuantWeightBuffer;
    res->mQuantMetaBuffer = mQuantMetaBuffer;
    res->mBiasBuffer = mBiasBuffer;
    *dst = res;
    return true;
}

ErrorCode VulkanConv1x1General::onEncode(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs,
                                         const VulkanCommandPool::Buffer* cmdBuffer) {
    if (inputs.empty() || outputs.empty() || nullptr == mQuantWeightBuffer.get() || nullptr == mQuantMetaBuffer.get() ||
        nullptr == mBiasBuffer.get()) {
        return INVALID_VALUE;
    }

    auto input = inputs[0];
    auto output = outputs[0];
    const int M = output->batch() * output->height() * output->width();
    if (M <= 0 || mCi <= 0 || mCo <= 0) {
        return NO_ERROR;
    }

    auto vkBn = static_cast<VulkanBackend*>(backend());
    auto srcBuffer = vkBn->getTensorBuffer(input);
    auto dstBuffer = vkBn->getTensorBuffer(output);

    if (M == 1) {
        struct DecodeParams {
            uint32_t K;
            uint32_t N;
            uint32_t blockSize;
            uint32_t blockStride;
            uint32_t weightStride;
        } pc;
        pc.K = static_cast<uint32_t>(mCi);
        pc.N = static_cast<uint32_t>(mCo);
        pc.blockSize = mBlockSize;
        pc.blockStride = mBlockStride;
        pc.weightStride = mDecodeWeightStrideWords;

        mDecodeSet->writeBuffer(srcBuffer.first->buffer(), 0, vkBn->getTensorSize(input), srcBuffer.second);
        mDecodeSet->writeBuffer(mQuantWeightBuffer->buffer(), 1, mQuantWeightBuffer->size());
        mDecodeSet->writeBuffer(mQuantMetaBuffer->buffer(), 2, mQuantMetaBuffer->size());
        mDecodeSet->writeBuffer(mBiasBuffer->buffer(), 3, mBiasBuffer->size());
        mDecodeSet->writeBuffer(dstBuffer.first->buffer(), 4, vkBn->getTensorSize(output), dstBuffer.second);
        mDecodePipeline->bind(cmdBuffer->get(), mDecodeSet->get());
        vkCmdPushConstants(cmdBuffer->get(), mDecodePipeline->layout(), VK_SHADER_STAGE_COMPUTE_BIT, 0,
                           sizeof(DecodeParams), &pc);
        vkCmdDispatch(cmdBuffer->get(), static_cast<uint32_t>(mCo), 1, 1);
        return NO_ERROR;
    }

    const uint32_t padM = ROUND_UP(static_cast<uint32_t>(M), 64u);
    if ((mPadK & 3u) != 0u || (mPadN & 31u) != 0u || padM == 0u) {
        return INVALID_VALUE;
    }

    if (vkBn->useFP16()) {
        mTempInputPacked.reset(Tensor::createDevice<int16_t>({static_cast<int>(padM), static_cast<int>(mPadK)}));
        mTempWeightPacked.reset(Tensor::createDevice<int16_t>({static_cast<int>(mPadK), static_cast<int>(mPadN)}));
    } else {
        mTempInputPacked.reset(Tensor::createDevice<float>({static_cast<int>(padM), static_cast<int>(mPadK)}));
        mTempWeightPacked.reset(Tensor::createDevice<float>({static_cast<int>(mPadK), static_cast<int>(mPadN)}));
    }

    bool acquiredTempA = false;
    bool acquiredTempB = false;
    auto releaseTemp = [&]() {
        if (acquiredTempA && nullptr != mTempInputPacked.get()) {
            vkBn->onReleaseBuffer(mTempInputPacked.get(), Backend::DYNAMIC);
            acquiredTempA = false;
        }
        if (acquiredTempB && nullptr != mTempWeightPacked.get()) {
            vkBn->onReleaseBuffer(mTempWeightPacked.get(), Backend::DYNAMIC);
            acquiredTempB = false;
        }
    };

    if (!vkBn->onAcquireBuffer(mTempInputPacked.get(), Backend::DYNAMIC)) {
        return OUT_OF_MEMORY;
    }
    acquiredTempA = true;

    if (!vkBn->onAcquireBuffer(mTempWeightPacked.get(), Backend::DYNAMIC)) {
        releaseTemp();
        return OUT_OF_MEMORY;
    }
    acquiredTempB = true;

    auto packedABuffer = vkBn->getTensorBuffer(mTempInputPacked.get());
    auto packedBBuffer = vkBn->getTensorBuffer(mTempWeightPacked.get());
    const size_t packedASize = vkBn->getTensorSize(mTempInputPacked.get());
    const size_t packedBSize = vkBn->getTensorSize(mTempWeightPacked.get());

    {
        struct PackAParams {
            uint32_t M;
            uint32_t K;
        } pc;
        pc.M = static_cast<uint32_t>(M);
        pc.K = mPadK;

        mPackASet->writeBuffer(srcBuffer.first->buffer(), 0, vkBn->getTensorSize(input), srcBuffer.second);
        mPackASet->writeBuffer(packedABuffer.first->buffer(), 1, packedASize, packedABuffer.second);
        mPackAPipeline->bind(cmdBuffer->get(), mPackASet->get());
        vkCmdPushConstants(cmdBuffer->get(), mPackAPipeline->layout(), VK_SHADER_STAGE_COMPUTE_BIT, 0,
                           sizeof(PackAParams), &pc);
        vkCmdDispatch(cmdBuffer->get(), mPadK / 4u, padM / 64u, 1);
        cmdBuffer->barrierSource(packedABuffer.first->buffer(), packedABuffer.second, packedASize);
    }

    {
        struct WeightToPackParams {
            uint32_t N;
            uint32_t K;
            uint32_t blockSize;
            uint32_t KBlocks;
        } pc;
        pc.N = mPadN;
        pc.K = mPadK;
        pc.blockSize = mBlockSize;
        pc.KBlocks = mBlockStride;

        mWeightToPackSet->writeBuffer(mQuantWeightBuffer->buffer(), 0, mQuantWeightBuffer->size());
        mWeightToPackSet->writeBuffer(mQuantMetaBuffer->buffer(), 1, mQuantMetaBuffer->size());
        mWeightToPackSet->writeBuffer(packedBBuffer.first->buffer(), 2, packedBSize, packedBBuffer.second);
        mWeightToPackPipeline->bind(cmdBuffer->get(), mWeightToPackSet->get());
        vkCmdPushConstants(cmdBuffer->get(), mWeightToPackPipeline->layout(), VK_SHADER_STAGE_COMPUTE_BIT, 0,
                           sizeof(WeightToPackParams), &pc);
        vkCmdDispatch(cmdBuffer->get(), UP_DIV(mPadN / 4u, 16u), UP_DIV(mPadK / 4u, 8u), 1);
        cmdBuffer->barrierSource(packedBBuffer.first->buffer(), packedBBuffer.second, packedBSize);
    }

    {
        struct GemmParams {
            uint32_t M;
            uint32_t N;
            uint32_t K;
            uint32_t padN;
        } pc;
        pc.M = static_cast<uint32_t>(M);
        pc.N = static_cast<uint32_t>(mCo);
        pc.K = mPadK;
        pc.padN = mPadN;

        mGemmSet->writeBuffer(packedABuffer.first->buffer(), 0, packedASize, packedABuffer.second);
        mGemmSet->writeBuffer(packedBBuffer.first->buffer(), 1, packedBSize, packedBBuffer.second);
        mGemmSet->writeBuffer(mBiasBuffer->buffer(), 2, mBiasBuffer->size());
        mGemmSet->writeBuffer(dstBuffer.first->buffer(), 3, vkBn->getTensorSize(output), dstBuffer.second);
        mGemmPipeline->bind(cmdBuffer->get(), mGemmSet->get());
        vkCmdPushConstants(cmdBuffer->get(), mGemmPipeline->layout(), VK_SHADER_STAGE_COMPUTE_BIT, 0,
                           sizeof(GemmParams), &pc);
        vkCmdDispatch(cmdBuffer->get(), mPadN / 32u, padM / 64u, 1);
    }

    releaseTemp();
    return NO_ERROR;
}

} // namespace MNN