//
//  LayerNormBufExecution.cpp
//  MNN
//
//  Created by MNN on 2023/07/05.
//  Copyright © 2018, Alibaba Group Holding Limited
//

#ifndef MNN_OPENCL_BUFFER_CLOSED
#include "backend/opencl/execution/buffer/LayerNormBufExecution.hpp"

namespace MNN {
namespace OpenCL {

LayerNormBufExecution::LayerNormBufExecution(const std::vector<Tensor*>& inputs, const MNN::Op* op, Backend* backend)
    : CommonExecution(backend, op) {
    mOpenCLBackend = static_cast<OpenCLBackend*>(backend);
    auto runtime = mOpenCLBackend->getOpenCLRuntime();
    const auto* layer_norm_param = op->main_as_LayerNorm();
    mResource.reset(new LayernormResource);
    if (nullptr != layer_norm_param->axis()) {
        mResource->axis_size = layer_norm_param->axis()->size();
    }
    mResource->epsilon_ = layer_norm_param->epsilon();
    mResource->group_ = layer_norm_param->group();
    mResource->RMSNorm = layer_norm_param->useRMSNorm();
    auto bufferUnitSize =
        mOpenCLBackend->getPrecision() != BackendConfig::Precision_High ? sizeof(half_float::half) : sizeof(float);
    auto kernel =
        runtime->buildKernel("layernorm_buf", "layernorm_buf", {"-DLOCAL_SIZE=512"}, mOpenCLBackend->getPrecision());
    OPENCL_CHECK_KERNEL_CTOR(kernel);
    mResource->mMaxWorkGroupSize = static_cast<uint32_t>(runtime->getMaxWorkGroupSize(kernel));

    mResource->has_gamma_beta_ = (layer_norm_param->gamma() && layer_norm_param->beta());
    int gammasize = 0;
    if (mResource->has_gamma_beta_) {
        MNN_ASSERT(layer_norm_param->gamma()->size() == layer_norm_param->beta()->size());
        gammasize = layer_norm_param->gamma()->size();
    }
    mResource->has_gamma_beta_ =
        mResource->has_gamma_beta_ || (layer_norm_param->external() && layer_norm_param->external()->size() > 1 &&
                                       layer_norm_param->external()->data()[1] > 0);
    if (mResource->has_gamma_beta_ && gammasize == 0) {
        gammasize = layer_norm_param->external()->data()[1] / sizeof(float);
    }

    auto staticMapAlloc = mOpenCLBackend->getStaticAllocatorMMap();
    if (mResource->has_gamma_beta_) {
        {
            auto error = CL_SUCCESS;
            int size = gammasize;
            if (mOpenCLBackend->getRuntime()->hint().useCachedMmap && staticMapAlloc != nullptr) {
                mResource->mGammaBuffer = staticMapAlloc.get()->allocBuffer(ALIGN_UP4(size) * bufferUnitSize);
            } else {
                mResource->mGammaBuffer.reset(new cl::Buffer(mOpenCLBackend->getOpenCLRuntime()->context(),
                                                             CL_MEM_READ_WRITE | CL_MEM_ALLOC_HOST_PTR,
                                                             ALIGN_UP4(size) * bufferUnitSize));
            }
            if (mOpenCLBackend->getRuntime()->hint().useCachedMmap <= 1) {
                auto GammaPtrCL = mOpenCLBackend->getOpenCLRuntime()->commandQueue().enqueueMapBuffer(
                    *(mResource->mGammaBuffer.get()), true, CL_MAP_WRITE, 0, ALIGN_UP4(size) * bufferUnitSize, nullptr,
                    nullptr, &error);
                const float* gamma_data = layer_norm_param->gamma()->data();
                if (GammaPtrCL != nullptr && error == CL_SUCCESS) {
                    if (mOpenCLBackend->getPrecision() != BackendConfig::Precision_High) {
                        for (int i = 0; i < size; i++) {
                            ((half_float::half*)GammaPtrCL)[i] = (half_float::half)(gamma_data[i]);
                        }
                        for (int i = size; i < ALIGN_UP4(size); i++) {
                            ((half_float::half*)GammaPtrCL)[i] = (half_float::half)(0.0f);
                        }
                    } else {
                        ::memset(GammaPtrCL, 0, ALIGN_UP4(size) * sizeof(float));
                        ::memcpy(GammaPtrCL, gamma_data, size * sizeof(float));
                    }
                } else {
                    MNN_ERROR("Map error GammaPtrCL == nullptr \n");
                }
                mOpenCLBackend->getOpenCLRuntime()->commandQueue().enqueueUnmapMemObject(*mResource->mGammaBuffer.get(),
                                                                                         GammaPtrCL);
            }
        }
        {
            auto error = CL_SUCCESS;
            int size = gammasize;
            if (mOpenCLBackend->getRuntime()->hint().useCachedMmap && staticMapAlloc != nullptr) {
                mResource->mBetaBuffer = staticMapAlloc.get()->allocBuffer(ALIGN_UP4(size) * bufferUnitSize);
            } else {
                mResource->mBetaBuffer.reset(new cl::Buffer(mOpenCLBackend->getOpenCLRuntime()->context(),
                                                            CL_MEM_READ_WRITE | CL_MEM_ALLOC_HOST_PTR,
                                                            ALIGN_UP4(size) * bufferUnitSize));
            }
            if (mOpenCLBackend->getRuntime()->hint().useCachedMmap <= 1) {
                auto BetaPtrCL = mOpenCLBackend->getOpenCLRuntime()->commandQueue().enqueueMapBuffer(
                    *(mResource->mBetaBuffer.get()), true, CL_MAP_WRITE, 0, ALIGN_UP4(size) * bufferUnitSize, nullptr,
                    nullptr, &error);
                const float* beta_data = layer_norm_param->beta()->data();
                if (BetaPtrCL != nullptr && error == CL_SUCCESS) {
                    if (mOpenCLBackend->getPrecision() != BackendConfig::Precision_High) {
                        for (int i = 0; i < size; i++) {
                            ((half_float::half*)BetaPtrCL)[i] = (half_float::half)(beta_data[i]);
                        }
                        for (int i = size; i < ALIGN_UP4(size); i++) {
                            ((half_float::half*)BetaPtrCL)[i] = (half_float::half)(0.0f);
                        }
                    } else {
                        ::memset(BetaPtrCL, 0, ALIGN_UP4(size) * sizeof(float));
                        ::memcpy(BetaPtrCL, beta_data, size * sizeof(float));
                    }
                } else {
                    MNN_ERROR("Map error BetaPtrCL == nullptr \n");
                }
                mOpenCLBackend->getOpenCLRuntime()->commandQueue().enqueueUnmapMemObject(*mResource->mBetaBuffer.get(),
                                                                                         BetaPtrCL);
            }
        }
    }
}

LayerNormBufExecution::LayerNormBufExecution(std::shared_ptr<LayernormResource> resource, const Op* op,
                                             Backend* backend)
    : CommonExecution(backend, op) {
    mResource = resource;
    mOpenCLBackend = (OpenCLBackend*)backend;
}

bool LayerNormBufExecution::onClone(Backend* bn, const Op* op, Execution** dst) {
    if (!mValid) {
        return false;
    }
    if (nullptr == dst) {
        return true;
    }
    *dst = new LayerNormBufExecution(mResource, op, bn);
    return true;
}

int LayerNormBufExecution::getLocalSize(int size, int maxGroupSize) {
    int local_size = 1;
    while (local_size * 2 <= maxGroupSize && local_size * 2 <= size) {
        local_size *= 2;
    }
    return local_size;
}

ErrorCode LayerNormBufExecution::onEncode(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) {
    Tensor* input = inputs[0];
    Tensor* output = outputs[0];
    auto runtime = ((OpenCLBackend*)backend())->getOpenCLRuntime();
    auto MaxLocalSize =
        std::min(std::min(runtime->getMaxWorkItemSizes()[0], mResource->mMaxWorkGroupSize), (uint32_t)256);

    const auto layout = TensorUtils::getDescribe(input)->dimensionFormat;
    bool isNC4HW4 = layout == MNN_DATA_FORMAT_NC4HW4;

    std::vector<int> inputShape = tensorShapeFormat(input);
    std::vector<int> outputShape = tensorShapeFormat(output);

    int rank = inputs.at(0)->dimensions();
    int outter_size = 1;
    int inner_size = 1;
    for (int i = 0; i < rank - mResource->axis_size; ++i) {
        outter_size *= inputs.at(0)->length(i);
    }
    for (int i = rank - mResource->axis_size; i < rank; ++i) {
        inner_size *= inputs.at(0)->length(i);
    }

    if (mResource->group_ > 1) {
        outter_size = inputs[0]->length(0) * mResource->group_;
        inner_size = 1;
        for (int i = 1; i < rank; i++) {
            inner_size *= inputs[0]->length(i);
        }
        inner_size /= mResource->group_;
    }

    if (isNC4HW4) {
        inner_size = inputs.at(0)->length(1);
        outter_size = 1;
        for (int i = 0; i < rank; i++) {
            if (i != 1) {
                outter_size *= inputs.at(0)->length(i);
            }
        }
    }
    bool splitBinaryLN = (isNC4HW4 && inputs.size() == 2 && outputs.size() == 2);

    int local_size;
    std::string kernelName;
    if (isNC4HW4) {
        int channelUnit = UP_DIV(inner_size, 4);
        local_size = getLocalSize(channelUnit, MaxLocalSize);
        if (splitBinaryLN) {
            // The second stage kernel (LayerNorm only)
            kernelName = "layernorm_c4_buf";
        } else {
            kernelName = "layernorm_c4_buf";
        }
    } else {
        local_size = getLocalSize(inner_size / 4, MaxLocalSize);
        kernelName = "layernorm_buf";
    }

    std::set<std::string> buildOptions;
    buildOptions.emplace("-DLOCAL_SIZE=" + std::to_string(local_size));
    if (mResource->RMSNorm) {
        buildOptions.emplace("-DRMSNORM");
    }
    if (mResource->has_gamma_beta_) {
        buildOptions.emplace("-DGAMMA_BETA");
    }
    if (!isNC4HW4 && inner_size % 4 != 0) {
        buildOptions.emplace("-DPACK_LEAVE");
    }

    if (splitBinaryLN) {
        // ---------- Two-kernel SPLIT path ----------
        int total_size_float = outter_size * inner_size;
        int gws_x = UP_DIV(total_size_float, 4);
        mUnits.resize(2);

        // unit[0]: binary_add_c4_buf (1D + vload4, following binary_buf.cl pattern)
        {
            auto& u0 = mUnits[0];
            std::set<std::string> addOpts;
            if (total_size_float % 4 != 0) {
                addOpts.emplace("-DPACK_LEAVE");
            }
            u0.kernel =
                runtime->buildKernel("layernorm_buf", "binary_add_c4_buf", addOpts, mOpenCLBackend->getPrecision());
            OPENCL_CHECK_KERNEL(u0.kernel);
            std::vector<uint32_t> gwsVec = {(uint32_t)gws_x, 1u};
            uint32_t maxWGS = static_cast<uint32_t>(runtime->getMaxWorkGroupSize(u0.kernel));
            uint32_t aidx = 0;
            cl_int aret = CL_SUCCESS;
            aret |= u0.kernel->get().setArg(aidx++, gws_x);
            aret |= u0.kernel->get().setArg(aidx++, openCLBuffer(inputs[0]));
            aret |= u0.kernel->get().setArg(aidx++, openCLBuffer(inputs[1]));
            aret |= u0.kernel->get().setArg(aidx++, openCLBuffer(outputs[0]));
            aret |= u0.kernel->get().setArg(aidx++, total_size_float);
            MNN_CHECK_CL_SUCCESS(aret, "setArg binary_add_c4_buf");
            std::vector<uint32_t> lwsVec = localWS2DDefault(gwsVec, maxWGS, runtime, "binary_add_c4_buf", u0.kernel,
                                                            mOpenCLBackend->getCLTuneLevel(), "layernorm_buf")
                                               .first;
            mOpenCLBackend->recordKernel2d(u0.kernel, gwsVec, lwsVec);
            u0.globalWorkSize = {gwsVec[0], gwsVec[1]};
            u0.localWorkSize = {lwsVec[0], lwsVec[1]};
        }

        // unit[1]: layernorm_c4_buf (reads output0 residual, writes output1 norm)
        {
            auto& u1 = mUnits[1];
            u1.kernel =
                runtime->buildKernel("layernorm_buf", "layernorm_c4_buf", buildOptions, mOpenCLBackend->getPrecision());
            OPENCL_CHECK_KERNEL(u1.kernel);
            mGWS = {(uint32_t)local_size, (uint32_t)outter_size};
            mLWS = {(uint32_t)local_size, 1};
            uint32_t lidx = 0;
            cl_int lret = CL_SUCCESS;
            lret |= u1.kernel->get().setArg(lidx++, mGWS[0]);
            lret |= u1.kernel->get().setArg(lidx++, mGWS[1]);
            lret |= u1.kernel->get().setArg(lidx++, openCLBuffer(outputs[0])); // input = residual
            lret |= u1.kernel->get().setArg(lidx++, openCLBuffer(outputs[1])); // output = norm
            lret |= u1.kernel->get().setArg(lidx++, (int32_t)inner_size);
            if (mResource->has_gamma_beta_) {
                lret |= u1.kernel->get().setArg(lidx++, *mResource->mGammaBuffer.get());
                lret |= u1.kernel->get().setArg(lidx++, *mResource->mBetaBuffer.get());
            }
            lret |= u1.kernel->get().setArg(lidx++, mResource->epsilon_);
            MNN_CHECK_CL_SUCCESS(lret, "setArg layernorm_c4_buf (SPLIT)");
            mOpenCLBackend->recordKernel2d(u1.kernel, mGWS, mLWS);
            u1.globalWorkSize = {mGWS[0], mGWS[1]};
            u1.localWorkSize = {mLWS[0], mLWS[1]};
        }
        return NO_ERROR;
    }

    // ---------- Single-kernel path (no residual add, or non-NC4HW4) ----------
    mUnits.resize(1);
    auto& unit = mUnits[0];
    unit.kernel = runtime->buildKernel("layernorm_buf", kernelName, buildOptions, mOpenCLBackend->getPrecision());
    OPENCL_CHECK_KERNEL(unit.kernel);
    mGWS = {static_cast<uint32_t>(local_size), static_cast<uint32_t>(outter_size)};
    mLWS = {static_cast<uint32_t>(local_size), 1};

    uint32_t idx = 0;
    cl_int ret = CL_SUCCESS;
    ret |= unit.kernel->get().setArg(idx++, mGWS[0]);
    ret |= unit.kernel->get().setArg(idx++, mGWS[1]);
    ret |= unit.kernel->get().setArg(idx++, openCLBuffer(input));
    ret |= unit.kernel->get().setArg(idx++, openCLBuffer(output));
    ret |= unit.kernel->get().setArg(idx++, static_cast<int32_t>(inner_size));
    if (mResource->has_gamma_beta_) {
        ret |= unit.kernel->get().setArg(idx++, *mResource->mGammaBuffer.get());
        ret |= unit.kernel->get().setArg(idx++, *mResource->mBetaBuffer.get());
    }
    ret |= unit.kernel->get().setArg(idx++, mResource->epsilon_);
    MNN_CHECK_CL_SUCCESS(ret, "setArg LayerNormBufExecution");
    mOpenCLBackend->recordKernel2d(unit.kernel, mGWS, mLWS);
    unit.globalWorkSize = {mGWS[0], mGWS[1]};
    unit.localWorkSize = {mLWS[0], mLWS[1]};

    return NO_ERROR;
}

class LayerNormBufCreator : public OpenCLBackend::Creator {
public:
    virtual ~LayerNormBufCreator() = default;
    virtual Execution* onCreate(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs,
                                const MNN::Op* op, Backend* backend) const override {
        for (int i = 0; i < inputs.size(); ++i) {
            TensorUtils::setTensorSupportPack(inputs[i], false);
        }
        for (int i = 0; i < outputs.size(); ++i) {
            TensorUtils::setTensorSupportPack(outputs[i], false);
        }
        OPENCL_CREATOR_CHECK(new LayerNormBufExecution(inputs, op, backend));
    }
};

REGISTER_OPENCL_OP_CREATOR(LayerNormBufCreator, OpType_LayerNorm, BUFFER);

} // namespace OpenCL
} // namespace MNN

#endif /* MNN_OPENCL_BUFFER_CLOSED */